From c8781b77f62661ac9765e8d0b6f6adb539fcf4ca Mon Sep 17 00:00:00 2001 From: Ye Wang <52801275+wangyems@users.noreply.github.com> Date: Thu, 29 Sep 2022 22:45:52 -0700 Subject: [PATCH] Decouple use_sequence_as_input_ids from has_hidden_states (#13130) ### Description A fix for parity issue in huggingface bart model with beam search https://github.com/microsoft/onnxruntime/pull/12779 ### Motivation and Context --- .../contrib_ops/cpu/transformers/beam_search_impl_t5.h | 2 +- .../cpu/transformers/generation_device_helper.cc | 9 ++++----- .../cpu/transformers/generation_device_helper.h | 4 ++-- .../cpu/transformers/subgraph_t5_decoder.cc | 7 +++++++ .../contrib_ops/cpu/transformers/subgraph_t5_decoder.h | 8 +++++--- .../cuda/transformers/generation_device_helper.cc | 10 +++++----- .../cuda/transformers/generation_device_helper.h | 2 +- 7 files changed, 25 insertions(+), 17 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h index e95af66849..304bb64c1b 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h @@ -289,7 +289,7 @@ Status BeamSearchT5::Execute(const FeedsFetchesManager& encoder_feeds_fetches parameters->num_beams, decoder_subgraph_.GetFirstPastInputIndex(), decoder_subgraph_.GetFirstPresentOutputIndex(), - decoder_subgraph_.HasHiddenStates(), + decoder_subgraph_.UseSequenceAsInputIds(), current_length, cpu_state.sequences, this->GetConsoleDumper())); diff --git a/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc b/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc index eb29dea924..6631b20e87 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc @@ -725,7 +725,7 @@ Status UpdateDecoderFeeds( int num_beams, int t5_decoder_first_past_input_idx, int t5_decoder_first_present_output_idx, - bool has_hidden_state, + bool use_sequence_as_input_ids, int current_length, transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper) { @@ -743,13 +743,12 @@ Status UpdateDecoderFeeds( // TODO(tianleiwu): Reuse buffer for input_ids to reduce memory allocation. OrtValue input_ids; - int sequence_length = has_hidden_state ? 1 : current_length; + int sequence_length = !use_sequence_as_input_ids ? 1 : current_length; int64_t dims[] = {batch_beam_size, sequence_length}; TensorShape input_ids_shape(&dims[0], 2); Tensor::InitOrtValue(DataTypeImpl::GetType(), input_ids_shape, allocator, input_ids); - // TODO(wy): decouple has_hidden_state with full input_ids - if (has_hidden_state) { + if (!use_sequence_as_input_ids) { gsl::copy(beam_next_tokens, input_ids.GetMutable()->MutableDataAsSpan()); } else { int32_t* input_ids_data = input_ids.GetMutable()->MutableData(); @@ -865,7 +864,7 @@ template Status UpdateDecoderFeeds( int num_beams, int t5_decoder_first_past_input_idx, int t5_decoder_first_present_output_idx, - bool has_hidden_state, + bool use_sequence_as_input_ids, int current_length, transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper); diff --git a/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.h b/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.h index aa445e7c62..68ac60d058 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.h +++ b/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.h @@ -146,7 +146,7 @@ using UpdateDecoderFeedsFunc = std::function; @@ -276,7 +276,7 @@ Status UpdateDecoderFeeds( int num_beams, int t5_decoder_first_past_input_idx, int t5_decoder_first_present_output_idx, - bool has_hidden_state, + bool use_sequence_as_input_ids, int current_length, transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper); diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc index d9ee9a3ba9..6702a1ab6b 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc @@ -77,6 +77,13 @@ Status T5DecoderSubgraph::Validate(const std::vector& subgraph_i ORT_RETURN_IF_ERROR(GetParameters(past_shape, logits_shape, false)); num_layers = (static_cast(subgraph_outputs.size()) - first_present_output_index_) / 2; + // If input_ids's shape is ['batch_size', 1] then use next token as input_ids. + // Otherwise in the case of shape ['batch_size', 'sequence'], use sequence as input_ids. + const ONNX_NAMESPACE::TensorShapeProto* input_ids_shape = subgraph_inputs[0]->Shape(); + if (input_ids_shape->dim(1).has_dim_value() && input_ids_shape->dim(1).dim_value() == 1) { + use_sequence_as_input_ids_ = false; + } + constexpr auto int32_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT32; constexpr auto float32_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT; constexpr auto float16_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16; diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h index cdda69d04d..66437ee98f 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h @@ -16,7 +16,8 @@ class T5DecoderSubgraph : public Subgraph { const onnxruntime::Node& node_in, const std::string& attribute_name, const GraphViewer& subgraph_in) : Subgraph(node_in, attribute_name, subgraph_in), - has_hidden_state_(false) { + has_hidden_state_(false), + use_sequence_as_input_ids_(true) { first_present_output_index_ = 1; } @@ -54,14 +55,15 @@ class T5DecoderSubgraph : public Subgraph { return first_present_output_index_; } - int HasHiddenStates() const { - return has_hidden_state_; + bool UseSequenceAsInputIds() const { + return use_sequence_as_input_ids_; } private: int first_past_input_index_; int first_present_output_index_; bool has_hidden_state_; + bool use_sequence_as_input_ids_; }; } // namespace transformers diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc index 85c571c189..ccb1b017a5 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc +++ b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc @@ -714,7 +714,7 @@ Status UpdateDecoderFeeds( int num_beams, int t5_decoder_first_past_input_idx, int t5_decoder_first_present_output_idx, - bool has_hidden_state, + bool use_sequence_as_input_ids, int current_length, transformers::Sequences&, const transformers::IConsoleDumper* dumper) { @@ -725,9 +725,9 @@ Status UpdateDecoderFeeds( // past_key_cross_0, past_value_cross_0, ... // Only need copy beam next tokens to input_ids, and copy present_*_self_* to past_*_self_*, - if (!has_hidden_state) { + if (use_sequence_as_input_ids) { return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, - "BeamSearch CUDA Op does not support no hidden state senario in decoder input"); + "BeamSearch CUDA Op does not support using sequence as input_ids in decoder input"); } ORT_UNUSED_PARAMETER(current_length); @@ -941,7 +941,7 @@ template Status UpdateDecoderFeeds( int num_beams, int t5_decoder_first_past_input_idx, int t5_decoder_first_present_output_idx, - bool has_hidden_state, + bool use_sequence_as_input_ids, int current_length, transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper); @@ -957,7 +957,7 @@ template Status UpdateDecoderFeeds( int num_beams, int t5_decoder_first_past_input_idx, int t5_decoder_first_present_output_idx, - bool has_hidden_state, + bool use_sequence_as_input_ids, int current_length, transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper); diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h index ebd2642554..1f97060b94 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h +++ b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h @@ -109,7 +109,7 @@ Status UpdateDecoderFeeds( int num_beams, int t5_decoder_first_past_input_idx, int t5_decoder_first_present_output_idx, - bool has_hidden_state, + bool use_sequence_as_input_ids, int current_length, transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper);