diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 22607e543a..c6f0f4e406 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -371,7 +371,7 @@ This version of the operator has been available since version 1 of the 'com.micr
The id of the padding token
-#### Inputs (5 - 9) +#### Inputs (5 - 10)
input_ids : I
@@ -392,6 +392,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Mask of vocabulary. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (vacab_size)
prefix_vocab_mask (optional) : M
Mask of vocabulary for first step. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (batch_size, vocab_size)
+
attention_mask (optional) : I
+
Custom attention mask. Shape is (batch_size, sequence_length)
#### Outputs (1 - 3) diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index d4a9483d77..5869514a3b 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -389,7 +389,7 @@ Do not modify directly.* |**Operator Domain:** *com.microsoft*|||| |Attention|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* mask_index:**M**
*in* past:**T**
*in* extra_add:**T**
*out* output:**T**
*out* present:**T**|1+|**T** = tensor(float)| |AttnLSTM|*in* X:**T**
*in* W:**T**
*in* R:**T**
*in* B:**T**
*in* sequence_lens:**T1**
*in* initial_h:**T**
*in* initial_c:**T**
*in* P:**T**
*in* QW:**T**
*in* MW:**T**
*in* V:**T**
*in* M:**T**
*in* memory_seq_lens:**T1**
*in* AW:**T**
*out* Y:**T**
*out* Y_h:**T**
*out* Y_c:**T**|1+|**T** = tensor(double), tensor(float)
**T1** = tensor(int32)| -|BeamSearch|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* num_beams:**I**
*in* num_return_sequences:**I**
*in* length_penalty:**T**
*in* repetition_penalty:**T**
*in* vocab_mask:**M**
*in* prefix_vocab_mask:**M**
*out* sequences:**I**
*out* sequences_scores:**T**
*out* scores:**T**|1+|**T** = tensor(float)| +|BeamSearch|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* num_beams:**I**
*in* num_return_sequences:**I**
*in* length_penalty:**T**
*in* repetition_penalty:**T**
*in* vocab_mask:**M**
*in* prefix_vocab_mask:**M**
*in* attention_mask:**I**
*out* sequences:**I**
*out* sequences_scores:**T**
*out* scores:**T**|1+|**T** = tensor(float)| |BiasGelu|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float)| |BifurcationDetector|*in* src_tokens:**T**
*in* cur_tokens:**T**
*in* prev_suffix_match_idx:**T**
*in* pred_tokens:**T**
*out* tokens:**T**
*out* suffix_match_idx:**T**|1+|**T** = tensor(int64)| |CDist|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(double), tensor(float)| @@ -747,7 +747,7 @@ Do not modify directly.* | | |**Operator Domain:** *com.microsoft*|||| |Attention|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* mask_index:**M**
*in* past:**T**
*in* extra_add:**T**
*out* output:**T**
*out* present:**T**|1+|**T** = tensor(float), tensor(float16)| -|BeamSearch|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* num_beams:**I**
*in* num_return_sequences:**I**
*in* length_penalty:**T**
*in* repetition_penalty:**T**
*in* vocab_mask:**M**
*in* prefix_vocab_mask:**M**
*out* sequences:**I**
*out* sequences_scores:**T**
*out* scores:**T**|1+|**T** = tensor(float), tensor(float16)| +|BeamSearch|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* num_beams:**I**
*in* num_return_sequences:**I**
*in* length_penalty:**T**
*in* repetition_penalty:**T**
*in* vocab_mask:**M**
*in* prefix_vocab_mask:**M**
*in* attention_mask:**I**
*out* sequences:**I**
*out* sequences_scores:**T**
*out* scores:**T**|1+|**T** = tensor(float), tensor(float16)| |BiasDropout|*in* data:**T**
*in* bias:**T**
*in* residual:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T2**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)| |BiasGelu|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |BiasSoftmax|*in* data:**T**
*in* bias:**T**
*out* output:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search.cc b/onnxruntime/contrib_ops/cpu/transformers/beam_search.cc index 8683317834..4cc3533f45 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search.cc @@ -59,7 +59,7 @@ namespace transformers { void BeamSearch::Init(const OpKernelInfo& info) { parameters_.ParseFromAttributes(info); - // Model_type could be either 0 (GPT-2) or 1 (encoder-decoder like T5). + // Model_type could be either 0 (GPT-2) or 1 (encoder-decoder like T5) ORT_ENFORCE(parameters_.model_type == IBeamSearchParameters::kModelTypeGpt || parameters_.model_type == IBeamSearchParameters::kModelTypeT5); diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_device_helper.cc b/onnxruntime/contrib_ops/cpu/transformers/beam_search_device_helper.cc index 55fefe7460..2e6af56fc4 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_device_helper.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_device_helper.cc @@ -334,10 +334,12 @@ template void PickGptPastState(const std::vector& last_outputs, std::vector& next_inputs, gsl::span& beam_indices, + int gpt_subgraph_first_past_input_idx, + int gpt_subgraph_first_present_output_idx, AllocatorPtr allocator) { - int num_present_tensors = static_cast(last_outputs.size()) - transformers::GptSubgraph::kFirstPresentOutputIndex; + int num_present_tensors = static_cast(last_outputs.size()) - gpt_subgraph_first_present_output_idx; for (int i = 0; i < num_present_tensors; ++i) { - const OrtValue& present = last_outputs[transformers::GptSubgraph::kFirstPresentOutputIndex + i]; + const OrtValue& present = last_outputs[gpt_subgraph_first_present_output_idx + i]; // shape is like (2, batch_beam_size, 12, past_seq_len, 64) const TensorShape& past_shape = present.Get().Shape(); @@ -364,7 +366,7 @@ void PickGptPastState(const std::vector& last_outputs, gsl::copy(present_value, past_value); } - next_inputs[transformers::GptSubgraph::kFirstPastInputIndex + i] = past; + next_inputs[gpt_subgraph_first_past_input_idx + i] = past; } } @@ -379,6 +381,8 @@ Status UpdateGptFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int gpt_subgraph_first_past_input_idx, + int gpt_subgraph_first_present_output_idx, const transformers::IConsoleDumper* dumper) { // last_outputs: logits, present_0, present_1, ... // next_inputs: input_ids, position_id, attention_mask, past_0, past_1 @@ -434,12 +438,14 @@ Status UpdateGptFeeds( // Update past state if (num_beams == 1) { // feed present_* output to past_* inputs one by one - const int k = transformers::GptSubgraph::kFirstPastInputIndex - transformers::GptSubgraph::kFirstPresentOutputIndex; - for (size_t i = transformers::GptSubgraph::kFirstPresentOutputIndex; i < last_outputs.size(); ++i) { + const int k = gpt_subgraph_first_past_input_idx - gpt_subgraph_first_present_output_idx; + for (size_t i = gpt_subgraph_first_present_output_idx; i < last_outputs.size(); ++i) { next_inputs[i + k] = last_outputs[i]; } } else { - PickGptPastState(last_outputs, next_inputs, beam_indices, allocator); + PickGptPastState(last_outputs, next_inputs, beam_indices, + gpt_subgraph_first_past_input_idx, + gpt_subgraph_first_present_output_idx, allocator); } return Status::OK(); } @@ -449,6 +455,7 @@ Status UpdateGptFeeds( // --------------------------------------------------------------- Status CreateEncoderInputs( const Tensor* original_encoder_input_ids, + const OrtValue* attn_mask_value, int num_beams, int pad_token_id, int start_token_id, @@ -476,24 +483,30 @@ Status CreateEncoderInputs( encoder_input_ids); OrtValue encoder_attention_mask; - auto mask_type = DataTypeImpl::GetType(); - Tensor::InitOrtValue(mask_type, input_ids_shape, allocator, encoder_attention_mask); + if (attn_mask_value != nullptr) { + const Tensor& attention_mask = attn_mask_value->Get(); + Tensor::InitOrtValue(element_type, input_ids_shape, const_cast(&attention_mask)->MutableData(), + allocator->Info(), encoder_attention_mask); + } else { + auto mask_type = DataTypeImpl::GetType(); + Tensor::InitOrtValue(mask_type, input_ids_shape, allocator, encoder_attention_mask); - // Set attention mask to be 0 for pad tokens, and 1 for all other tokens. - int32_t* mask_data = encoder_attention_mask.GetMutable()->MutableData(); - const int32_t* word_id = original_encoder_input_ids->Data(); - int32_t* mask = mask_data; - for (int i = 0; i < batch_size; i++) { - int32_t abs_position = 0; - for (int j = 0; j < sequence_length; j++, word_id++, mask++) { - // T5Tokenizer might add one EOS pad token at the end. - // That EOS token shall have attention mask 1 even when EOS token is same as pad token. - // Here we only set attention mask to be 0 for left padding only, so as to be parity with huggingface. - if (*word_id == pad_token_id && abs_position == 0) { - *mask = 0; - } else { - *mask = 1; - abs_position++; + // Set attention mask to be 0 for pad tokens, and 1 for all other tokens. + int32_t* mask_data = encoder_attention_mask.GetMutable()->MutableData(); + const int32_t* word_id = original_encoder_input_ids->Data(); + int32_t* mask = mask_data; + for (int i = 0; i < batch_size; i++) { + int32_t abs_position = 0; + for (int j = 0; j < sequence_length; j++, word_id++, mask++) { + // T5Tokenizer might add one EOS pad token at the end. + // That EOS token shall have attention mask 1 even when EOS token is same as pad token. + // Here we only set attention mask to be 0 for left padding only, so as to be parity with huggingface. + if (*word_id == pad_token_id && abs_position == 0) { + *mask = 0; + } else { + *mask = 1; + abs_position++; + } } } } @@ -525,9 +538,11 @@ void PickT5PastState(const std::vector& last_outputs, std::vector& next_inputs, int num_present_tensors, gsl::span& beam_indices, + int t5_decoder_first_past_input_idx, + int t5_decoder_first_present_output_idx, AllocatorPtr allocator) { for (int i = 0; i < num_present_tensors; ++i) { - const OrtValue& present = last_outputs[transformers::T5DecoderSubgraph::kFirstPresentOutputIndex + i]; + const OrtValue& present = last_outputs[t5_decoder_first_present_output_idx + i]; // shape is like (batch_beam_size, 12, past_seq_len, 64) const TensorShape& past_shape = present.Get().Shape(); @@ -547,7 +562,7 @@ void PickT5PastState(const std::vector& last_outputs, gsl::copy(present_beam, past_beam); } - next_inputs[transformers::T5DecoderSubgraph::kFirstPastInputIndex + i] = past; + next_inputs[t5_decoder_first_past_input_idx + i] = past; } } @@ -562,26 +577,44 @@ Status UpdateDecoderFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int t5_decoder_first_past_input_idx, + int t5_decoder_first_present_output_idx, + bool has_hidden_state, + int current_length, + transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper) { ORT_UNUSED_PARAMETER(stream); // last_outputs: logits, present_key_self_0, present_value_self_0, ... // next_inputs: input_ids, - // encoder_attention_mask, encoder_hidden_states, + // encoder_attention_mask, encoder_hidden_states(optional), // past_key_self_0, past_value_self_0, ... // past_key_cross_0, past_value_cross_0, ... // Only need copy beam next tokens to input_ids, and copy present_*_self_* to past_*_self_*, // Update input_ids with next tokens. int batch_beam_size = static_cast(beam_next_tokens.length()); - int64_t dims[] = {batch_beam_size, 1}; - TensorShape input_ids_shape(&dims[0], 2); // TODO(tianleiwu): Reuse buffer for input_ids to reduce memory allocation. OrtValue input_ids; + int sequence_length = has_hidden_state ? 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); - gsl::copy(beam_next_tokens, input_ids.GetMutable()->MutableDataAsSpan()); + // TODO: decouple has_hidden_state with full input_ids + if (has_hidden_state) { + gsl::copy(beam_next_tokens, input_ids.GetMutable()->MutableDataAsSpan()); + } else { + int32_t* input_ids_data = input_ids.GetMutable()->MutableData(); + for (int i = 0; i < batch_beam_size; i++) { + gsl::span sequence = sequences.GetSequence(i); + const int32_t* sequence_data = sequence.data(); + for (int j = 0; j < current_length; j++) { + input_ids_data[i * current_length + j] = sequence_data[j]; + } + } + } next_inputs[0] = input_ids; @@ -597,11 +630,12 @@ Status UpdateDecoderFeeds( if (num_beams == 1) { // feed present_* output to past_* inputs one by one for (int i = 0; i < num_present_tensors; ++i) { - next_inputs[transformers::T5DecoderSubgraph::kFirstPastInputIndex + i] = - last_outputs[transformers::T5DecoderSubgraph::kFirstPresentOutputIndex + i]; + next_inputs[t5_decoder_first_past_input_idx + i] = + last_outputs[t5_decoder_first_present_output_idx + i]; } } else { - PickT5PastState(last_outputs, next_inputs, num_present_tensors, beam_indices, allocator); + PickT5PastState(last_outputs, next_inputs, num_present_tensors, beam_indices, + t5_decoder_first_past_input_idx, t5_decoder_first_present_output_idx, allocator); } return Status::OK(); } @@ -653,6 +687,8 @@ template Status UpdateGptFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int gpt_subgraph_first_past_input_idx, + int gpt_subgraph_first_present_output_idx, const transformers::IConsoleDumper* dumper); template Status UpdateDecoderFeeds( @@ -664,6 +700,11 @@ template Status UpdateDecoderFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int t5_decoder_first_past_input_idx, + int t5_decoder_first_present_output_idx, + bool has_hidden_state, + int current_length, + transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper); template void ExpandInputs(const OrtValue& input, int num_beams, AllocatorPtr allocator, OrtValue& expanded); diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_device_helper.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_device_helper.h index b58a95d4d0..8cd7a0291a 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_device_helper.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_device_helper.h @@ -99,11 +99,14 @@ using UpdateGptFeedsFunc = std::function beam_next_tokens, gsl::span beam_indices, int num_beams, + int gpt_subgraph_first_past_input_idx, + int gpt_subgraph_first_present_output_idx, const transformers::IConsoleDumper* dumper)>; // Create encoder inputs (for encoder-decoder model like T5). using CreateEncoderInputsFunc = std::function beam_next_tokens, gsl::span beam_indices, int num_beams, + int t5_decoder_first_past_input_idx, + int t5_decoder_first_present_output_idx, + bool has_hidden_state, + int current_length, + transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper)>; } // namespace BeamSearchDeviceHelper @@ -194,6 +202,8 @@ Status UpdateGptFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int gpt_subgraph_first_past_input_idx, + int gpt_subgraph_first_present_output_idx, const transformers::IConsoleDumper* dumper); // --------------------------------------------------------------- @@ -201,6 +211,7 @@ Status UpdateGptFeeds( // --------------------------------------------------------------- Status CreateEncoderInputs( const Tensor* original_encoder_input_ids, + const OrtValue* attn_mask_value, int num_beams, int pad_token_id, int start_token_id, @@ -220,6 +231,11 @@ Status UpdateDecoderFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int t5_decoder_first_past_input_idx, + int t5_decoder_first_present_output_idx, + bool has_hidden_state, + int current_length, + transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper); // --------------------------------------------------------------- diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h index 9885b4ae54..03c88b5aa8 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h @@ -262,6 +262,19 @@ Status BeamSearchBase::CheckInputs(const OpKernelContextInternal& context) { parameters_->prefix_vocab_mask = prefix_vocab_mask->DataAsSpan(); } + const Tensor* attention_mask = context.Input(9); + if (attention_mask != nullptr) { + const auto& dims_attn = attention_mask->Shape().GetDims(); + if (dims_attn.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'attention_mask' is expected to have 2 dimensions, got ", dims_attn.size()); + } + if (dims_attn != dims) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'attention_mask' is expected to have same shape as input_ids"); + } + } + return Status::OK(); } @@ -288,9 +301,7 @@ Status BeamSearchBase::Initialize() { CHECK_SCALAR_INPUT(num_return_sequences, 4, true); - CHECK_SCALAR_INPUT(temperature, 5, true); - - CHECK_SCALAR_INPUT(length_penalty, 6, true); + CHECK_SCALAR_INPUT(length_penalty, 5, true); ORT_RETURN_IF(parameters_->num_return_sequences > parameters_->num_beams, "'num_return_sequences' has to be smaller or equal to 'num_beams'."); diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h index bd5d3fcc3a..9cf5daeba9 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h @@ -104,6 +104,8 @@ Status BeamSearchGpt::UpdateFeeds( beam_next_tokens, beam_indices, this->parameters_->num_beams, + gpt_subgraph_.GetFirstPastInputIndex(), + gpt_subgraph_.GetFirstPresentOutputIndex(), this->GetConsoleDumper()); } 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 491defe314..6acbe857e5 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h @@ -100,6 +100,8 @@ Status BeamSearchT5::Execute(const FeedsFetchesManager& encoder_feeds_fetches const OrtValue* encoder_input_ids_value = this->context_.GetInputOrtValue(0); const Tensor& encoder_input_ids = encoder_input_ids_value->Get(); + const OrtValue* encoder_attn_mask_value = this->context_.GetInputOrtValue(9); + BeamSearchCpuState cpu_state; cpu_state.Init(this->cpu_allocator_, static_cast(parameters->BatchBeamSize()), @@ -111,6 +113,7 @@ Status BeamSearchT5::Execute(const FeedsFetchesManager& encoder_feeds_fetches OrtValue expanded_decoder_input_ids; // Tensor in CPU, and it will be used to initialize sequence in cpu_state ORT_RETURN_IF_ERROR(this->encoder_subgraph_.CreateInitialFeeds( encoder_input_ids, + encoder_attn_mask_value, this->implicit_inputs_, parameters->num_beams, parameters->pad_token_id, @@ -137,7 +140,7 @@ Status BeamSearchT5::Execute(const FeedsFetchesManager& encoder_feeds_fetches dumper->Print("", encoder_feeds[i]); } - for (int i = 0; i <= T5EncoderSubgraph::kFirstPresentOutputIndex; i++) { + for (int i = 0; i <= encoder_subgraph_.GetFirstPresentOutputIndex(); i++) { dumper->Print("encoder_fetches", i, true); dumper->Print("", encoder_fetches[i]); } @@ -219,7 +222,7 @@ Status BeamSearchT5::Execute(const FeedsFetchesManager& encoder_feeds_fetches auto cur_len = std::to_string(current_length); dumper->Print("***CurrentLength", cur_len, true); - for (int i = 0; i <= T5DecoderSubgraph::kFirstPastInputIndex; i++) { + for (int i = 0; i <= decoder_subgraph_.GetFirstPastInputIndex(); i++) { dumper->Print("decoder_feeds", i, true); dumper->Print("", decoder_feeds[i]); } @@ -237,7 +240,7 @@ Status BeamSearchT5::Execute(const FeedsFetchesManager& encoder_feeds_fetches ORT_RETURN_IF_ERROR(status); #ifdef DEBUG_BEAM_SEARCH - for (int i = 0; i <= T5DecoderSubgraph::kFirstPresentOutputIndex; i++) { + for (int i = 0; i <= decoder_subgraph_.GetFirstPresentOutputIndex(); i++) { dumper->Print("decoder_fetches", i, true); dumper->Print("", decoder_fetches[i]); } @@ -271,6 +274,11 @@ Status BeamSearchT5::Execute(const FeedsFetchesManager& encoder_feeds_fetches beam_next_tokens.as_span(), beam_indices.as_span(), parameters->num_beams, + decoder_subgraph_.GetFirstPastInputIndex(), + decoder_subgraph_.GetFirstPresentOutputIndex(), + decoder_subgraph_.HasHiddenStates(), + current_length, + cpu_state.sequences, this->GetConsoleDumper())); } decoder_fetches.clear(); diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc b/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc index c69e78e223..8241b3afb5 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc @@ -73,7 +73,7 @@ Status GptSubgraph::CreateInitialFeeds( buffer)); // The remaining inputs are past state. - for (int i = kFirstPastInputIndex; i < num_subgraph_inputs; ++i) { + for (int i = first_past_input_index_; i < num_subgraph_inputs; ++i) { feeds.push_back(empty_past); } @@ -87,7 +87,7 @@ Status GptSubgraph::CreateInitialFeeds( Status GptSubgraph::Validate(const std::vector& subgraph_inputs, const std::vector& subgraph_outputs) { - ORT_RETURN_IF(num_subgraph_outputs <= kFirstPresentOutputIndex, + ORT_RETURN_IF(num_subgraph_outputs <= first_present_output_index_, "Invalid GPT-2 subgraph: number of outputs shall be larger than 1 (Need past state in outputs)."); ORT_RETURN_IF(num_subgraph_inputs != num_subgraph_outputs + 2, @@ -152,9 +152,9 @@ Status GptSubgraph::Validate(const std::vector& subgraph_inputs, ORT_RETURN_IF(output_type != float32_type && output_type != float16_type, "subgraph output 0 (logits) shall be float or float16 data type"); - ORT_RETURN_IF(subgraph_inputs[kFirstPastInputIndex]->TypeAsProto()->tensor_type().elem_type() != output_type, + ORT_RETURN_IF(subgraph_inputs[first_past_input_index_]->TypeAsProto()->tensor_type().elem_type() != output_type, "subgraph input 3 (past_0) shall shall have same data type of logits output"); - ORT_RETURN_IF(subgraph_outputs[kFirstPresentOutputIndex]->TypeAsProto()->tensor_type().elem_type() != output_type, + ORT_RETURN_IF(subgraph_outputs[first_present_output_index_]->TypeAsProto()->tensor_type().elem_type() != output_type, "subgraph output 1 (present_0) shall shall have same data type of logits output"); is_output_float16_ = (output_type == float16_type); diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.h b/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.h index 728c446dd4..694baef788 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.h +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.h @@ -15,7 +15,10 @@ class GptSubgraph : public Subgraph { GptSubgraph( const onnxruntime::Node& node_in, const std::string& attribute_name, - const GraphViewer& subgraph_in) : Subgraph(node_in, attribute_name, subgraph_in) {} + const GraphViewer& subgraph_in) : Subgraph(node_in, attribute_name, subgraph_in) { + first_past_input_index_ = 3; + first_present_output_index_ = 1; + } // Create inputs for first inference of subgraph. Status CreateInitialFeeds( @@ -33,8 +36,17 @@ class GptSubgraph : public Subgraph { Status Validate(const std::vector& subgraph_inputs, const std::vector& subgraph_outputs) override; - constexpr static int kFirstPastInputIndex = 3; - constexpr static int kFirstPresentOutputIndex = 1; + int GetFirstPastInputIndex() const { + return first_past_input_index_; + } + + int GetFirstPresentOutputIndex() const { + return first_present_output_index_; + } + + private: + int first_past_input_index_; + int first_present_output_index_; }; } // namespace transformers diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc index 8a55cb3f2f..bb931255dd 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc @@ -44,28 +44,38 @@ namespace transformers { Status T5DecoderSubgraph::Validate(const std::vector& subgraph_inputs, const std::vector& subgraph_outputs) { - ORT_RETURN_IF(num_subgraph_inputs < 7 || (num_subgraph_inputs - kFirstPastInputIndex) % 4 != 0, - "number of outputs expected to be 3 + 4 * layers, got:", num_subgraph_inputs); - ORT_RETURN_IF(num_subgraph_outputs < 3 || (num_subgraph_outputs - kFirstPresentOutputIndex) % 2 != 0, + bool has_hidden_state = subgraph_inputs[2]->Name() == "encoder_hidden_states" ? true : false; + SetPastInputIndex(has_hidden_state); + + ORT_RETURN_IF(first_past_input_index_ != 2 && first_past_input_index_ != 3, + "kFirstPastInputIndex currently only supports 2 or 3"); + ORT_RETURN_IF(num_subgraph_inputs < 4 + first_past_input_index_ || + (num_subgraph_inputs - first_past_input_index_) % 4 != 0, + "number of outputs expected to be kFirstPastInputIndex + 4 * layers, got:", num_subgraph_inputs); + ORT_RETURN_IF(num_subgraph_outputs < 3 || (num_subgraph_outputs - first_present_output_index_) % 2 != 0, "number of outputs expected to be 1 + 2 * layers, got:", num_subgraph_outputs); ORT_RETURN_IF(subgraph_inputs[0]->Name() != "input_ids", "decoder subgraph input 0 shall be named as input_ids, got: ", subgraph_inputs[0]->Name()); ORT_RETURN_IF(subgraph_inputs[1]->Name() != "encoder_attention_mask", - "decoder subgraph input 1 shall be named as encoder_attention_mask, got: ", subgraph_inputs[1]->Name()); - ORT_RETURN_IF(subgraph_inputs[2]->Name() != "encoder_hidden_states", - "decoder subgraph input 2 shall be named as encoder_hidden_states, got: ", subgraph_inputs[2]->Name()); + "decoder subgraph input 1 shall be named as encoder_attention_mask, got: ", + subgraph_inputs[1]->Name()); + if (first_past_input_index_ == 3) { + ORT_RETURN_IF(subgraph_inputs[2]->Name() != "encoder_hidden_states", + "decoder subgraph input 2 shall be named as encoder_hidden_states, got: ", + subgraph_inputs[2]->Name()); + } // check subgraph outputs ORT_RETURN_IF(subgraph_outputs[0]->Name() != "logits", "decoder subgraph output 0 shall be named as logits, got: ", subgraph_outputs[0]->Name()); const ONNX_NAMESPACE::TensorShapeProto* logits_shape = subgraph_outputs[0]->Shape(); - const ONNX_NAMESPACE::TensorShapeProto* past_shape = subgraph_outputs[kFirstPresentOutputIndex]->Shape(); + const ONNX_NAMESPACE::TensorShapeProto* past_shape = subgraph_outputs[first_present_output_index_]->Shape(); // Save parameters related to the subgraph. ORT_RETURN_IF_ERROR(GetParameters(past_shape, logits_shape, false)); - num_layers = (static_cast(subgraph_outputs.size()) - kFirstPresentOutputIndex) / 2; + num_layers = (static_cast(subgraph_outputs.size()) - first_present_output_index_) / 2; constexpr auto int32_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT32; constexpr auto float32_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT; @@ -80,7 +90,7 @@ Status T5DecoderSubgraph::Validate(const std::vector& subgraph_i ORT_RETURN_IF(float_type != float32_type && float_type != float16_type, "decoder subgraph input 2 (encoder_hidden_states) shall have float or float16 type"); - for (int i = kFirstPastInputIndex; i < num_subgraph_inputs; i++) { + for (int i = first_past_input_index_; i < num_subgraph_inputs; i++) { ORT_RETURN_IF(subgraph_inputs[i]->TypeAsProto()->tensor_type().elem_type() != float_type, "decoder subgraph past inputs shall have same data type as that of encoder_hidden_states"); } @@ -136,8 +146,10 @@ Status T5DecoderSubgraph::CreateInitialFeeds( // The encoder_attention_mask is copied from the second input of encoder. decoder_feeds.push_back(encoder_feeds[1]); - // The encoder_hidden_states and past states are copied from the second output of encoder. - for (size_t j = 1; j < encoder_fetches.size(); j++) { + // When first_past_input_index_ == 3, the encoder_hidden_states and past states are copied from the second output + // of encoder. + // When first_past_input_index_ == 2, the past states are copied from the second output of encoder. + for (size_t j = 4 - first_past_input_index_; j < encoder_fetches.size(); j++) { decoder_feeds.push_back(encoder_fetches[j]); } diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h index 3e59181a65..108b1c298d 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h @@ -15,7 +15,10 @@ class T5DecoderSubgraph : public Subgraph { T5DecoderSubgraph( const onnxruntime::Node& node_in, const std::string& attribute_name, - const GraphViewer& subgraph_in) : Subgraph(node_in, attribute_name, subgraph_in) {} + const GraphViewer& subgraph_in) : Subgraph(node_in, attribute_name, subgraph_in), + has_hidden_state_(false) { + first_present_output_index_ = 1; + } // Create inputs for first inference of decoder subgraph. Status CreateInitialFeeds( @@ -30,8 +33,31 @@ class T5DecoderSubgraph : public Subgraph { Status Validate(const std::vector& subgraph_inputs, const std::vector& subgraph_outputs) override; - constexpr static int kFirstPastInputIndex = 3; - constexpr static int kFirstPresentOutputIndex = 1; + void SetPastInputIndex(bool has_hidden_state) { + has_hidden_state_ = has_hidden_state; + if (!has_hidden_state_) { + first_past_input_index_ = 2; + } else { + first_past_input_index_ = 3; + } + } + + int GetFirstPastInputIndex() const { + return first_past_input_index_; + } + + int GetFirstPresentOutputIndex() const { + return first_present_output_index_; + } + + int HasHiddenStates() const { + return has_hidden_state_; + } + + private: + int first_past_input_index_; + int first_present_output_index_; + bool has_hidden_state_; }; } // namespace transformers diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.cc b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.cc index 9b8292ed6f..7574c31ec5 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.cc @@ -43,7 +43,7 @@ Status T5EncoderSubgraph::Validate(const std::vector& subgraph_i ORT_RETURN_IF(num_subgraph_inputs != 3, "expect 3 inputs, got:", num_subgraph_inputs); ORT_RETURN_IF(num_subgraph_outputs < 6, "expect >=6 outputs, got:", num_subgraph_outputs); - ORT_RETURN_IF((static_cast(subgraph_outputs.size()) - kFirstPresentOutputIndex) % 4 != 0, + ORT_RETURN_IF((static_cast(subgraph_outputs.size()) - first_present_output_index_) % 4 != 0, "number of outputs expected to be 2 + 4 * layers, got:", num_subgraph_outputs); ORT_RETURN_IF(subgraph_inputs[0]->Name() != "encoder_input_ids", @@ -67,7 +67,7 @@ Status T5EncoderSubgraph::Validate(const std::vector& subgraph_i // Save parameters related to the subgraph. ORT_RETURN_IF_ERROR(GetParameters(past_shape, logits_shape, false)); - num_layers = (static_cast(subgraph_outputs.size()) - kFirstPresentOutputIndex) / 4; + num_layers = (static_cast(subgraph_outputs.size()) - first_present_output_index_) / 4; constexpr auto int32_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT32; constexpr auto float32_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT; @@ -97,6 +97,7 @@ Status T5EncoderSubgraph::Validate(const std::vector& subgraph_i // Create inputs for first inference of subgraph. Status T5EncoderSubgraph::CreateInitialFeeds( const Tensor& encoder_input_ids, + const OrtValue* attn_mask_value, const std::vector& implicit_inputs, int num_beams, int pad_token_id, @@ -113,11 +114,17 @@ Status T5EncoderSubgraph::CreateInitialFeeds( // Allocate subgraph inputs to be same device as encoder_input_ids. AllocatorPtr cpu_allocator = session_state_->GetAllocator(encoder_input_ids.Location()); + if (cpu_allocator == nullptr) { + const IExecutionProvider* provider = GetProvider(); + cpu_allocator = provider->GetAllocator(0, OrtMemTypeDefault); + } + ORT_RETURN_IF(cpu_allocator == nullptr, "cpu_allocator shouldn't be nullptr"); // TODO(tianleiwu): expand the outputs instead of inputs to save computation. OrtValue expanded_encoder_input_ids; OrtValue expanded_encoder_attention_mask; ORT_RETURN_IF_ERROR(create_encoder_inputs_func(&encoder_input_ids, + attn_mask_value, num_beams, pad_token_id, start_token_id, diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.h b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.h index 30ff3840a7..83c9cb22c6 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.h +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.h @@ -15,11 +15,14 @@ class T5EncoderSubgraph : public Subgraph { T5EncoderSubgraph( const onnxruntime::Node& node_in, const std::string& attribute_name, - const GraphViewer& subgraph_in) : Subgraph(node_in, attribute_name, subgraph_in) {} + const GraphViewer& subgraph_in) : Subgraph(node_in, attribute_name, subgraph_in) { + first_present_output_index_ = 2; + } // Create inputs for first inference of subgraph. Status CreateInitialFeeds( const Tensor& encoder_input_ids, + const OrtValue* attn_mask_value, const std::vector& implicit_inputs, int num_beams, int pad_token_id, @@ -33,7 +36,12 @@ class T5EncoderSubgraph : public Subgraph { Status Validate(const std::vector& subgraph_inputs, const std::vector& subgraph_outputs) override; - constexpr static int kFirstPresentOutputIndex = 2; + int GetFirstPresentOutputIndex() const { + return first_present_output_index_; + } + + private: + int first_present_output_index_; }; } // namespace transformers diff --git a/onnxruntime/contrib_ops/cuda/transformers/beam_search.cc b/onnxruntime/contrib_ops/cuda/transformers/beam_search.cc index b8afb13919..91b660c197 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/beam_search.cc +++ b/onnxruntime/contrib_ops/cuda/transformers/beam_search.cc @@ -22,9 +22,8 @@ ONNX_OPERATOR_KERNEL_EX( .InputMemoryType(OrtMemTypeCPUInput, 2) // 'min_length' needs to be on CPU .InputMemoryType(OrtMemTypeCPUInput, 3) // 'num_beams' needs to be on CPU .InputMemoryType(OrtMemTypeCPUInput, 4) // 'num_return_sequences' needs to be on CPU - .InputMemoryType(OrtMemTypeCPUInput, 5) // 'temperature' needs to be on CPU - .InputMemoryType(OrtMemTypeCPUInput, 6) // 'length_penalty' needs to be on CPU - .InputMemoryType(OrtMemTypeCPUInput, 7) // 'repetition_penalty' needs to be on CPU + .InputMemoryType(OrtMemTypeCPUInput, 5) // 'length_penalty' needs to be on CPU + .InputMemoryType(OrtMemTypeCPUInput, 6) // 'repetition_penalty' needs to be on CPU .OutputMemoryType(OrtMemTypeCPUOutput, 0) // 'sequences' output on CPU .OutputMemoryType(OrtMemTypeCPUOutput, 1) // 'sequences_scores' output on CPU .TypeConstraint("T", {DataTypeImpl::GetTensorType(), diff --git a/onnxruntime/contrib_ops/cuda/transformers/beam_search_device_helper.cc b/onnxruntime/contrib_ops/cuda/transformers/beam_search_device_helper.cc index 3134d3a020..3c45c2cc60 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/beam_search_device_helper.cc +++ b/onnxruntime/contrib_ops/cuda/transformers/beam_search_device_helper.cc @@ -404,10 +404,12 @@ Status PickGptPastState(const std::vector& last_outputs, std::vector& next_inputs, gsl::span& beam_indices, AllocatorPtr allocator, + int gpt_subgraph_first_past_input_idx, + int gpt_subgraph_first_present_output_idx, void* stream) { - int num_present_tensors = static_cast(last_outputs.size()) - transformers::GptSubgraph::kFirstPresentOutputIndex; + int num_present_tensors = static_cast(last_outputs.size()) - gpt_subgraph_first_present_output_idx; for (int i = 0; i < num_present_tensors; ++i) { - const OrtValue& present = last_outputs[transformers::GptSubgraph::kFirstPresentOutputIndex + i]; + const OrtValue& present = last_outputs[gpt_subgraph_first_present_output_idx + i]; // shape is like (2, batch_beam_size, 12, past_seq_len, 64) const TensorShape& past_shape = present.Get().Shape(); @@ -436,7 +438,7 @@ Status PickGptPastState(const std::vector& last_outputs, cudaMemcpyDeviceToDevice, reinterpret_cast(stream))); } - next_inputs[transformers::GptSubgraph::kFirstPastInputIndex + i] = past; + next_inputs[gpt_subgraph_first_past_input_idx + i] = past; } return Status::OK(); @@ -449,9 +451,11 @@ Status PickT5PastState(const std::vector& last_outputs, int num_present_tensors, gsl::span& beam_indices, AllocatorPtr allocator, + int t5_decoder_first_past_input_idx, + int t5_decoder_first_present_output_idx, void* stream) { for (int i = 0; i < num_present_tensors; ++i) { - const OrtValue& present = last_outputs[transformers::T5DecoderSubgraph::kFirstPresentOutputIndex + i]; + const OrtValue& present = last_outputs[t5_decoder_first_present_output_idx + i]; // shape is like (batch_beam_size, 12, past_seq_len, 64) const TensorShape& past_shape = present.Get().Shape(); @@ -472,7 +476,7 @@ Status PickT5PastState(const std::vector& last_outputs, cudaMemcpyDeviceToDevice, reinterpret_cast(stream))); } - next_inputs[transformers::T5DecoderSubgraph::kFirstPastInputIndex + i] = past; + next_inputs[t5_decoder_first_past_input_idx + i] = past; } return Status::OK(); @@ -489,6 +493,8 @@ Status UpdateGptFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int gpt_subgraph_first_past_input_idx, + int gpt_subgraph_first_present_output_idx, const transformers::IConsoleDumper* dumper) { // Update input_ids with next tokens. int batch_beam_size = static_cast(beam_next_tokens.length()); @@ -532,13 +538,15 @@ Status UpdateGptFeeds( // Update past state if (num_beams == 1) { - const int k = transformers::GptSubgraph::kFirstPastInputIndex - transformers::GptSubgraph::kFirstPresentOutputIndex; + const int k = gpt_subgraph_first_past_input_idx - gpt_subgraph_first_present_output_idx; // feed present_* output to past_* inputs one by one - for (size_t i = transformers::GptSubgraph::kFirstPresentOutputIndex; i < last_outputs.size(); ++i) { + for (size_t i = gpt_subgraph_first_present_output_idx; i < last_outputs.size(); ++i) { next_inputs[i + k] = last_outputs[i]; } } else { - ORT_RETURN_IF_ERROR(PickGptPastState(last_outputs, next_inputs, beam_indices, allocator, stream)); + ORT_RETURN_IF_ERROR(PickGptPastState(last_outputs, next_inputs, beam_indices, allocator, + gpt_subgraph_first_past_input_idx, + gpt_subgraph_first_present_output_idx, stream)); } // Make sure data is ready before next subgraph execution. @@ -557,6 +565,11 @@ Status UpdateDecoderFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int t5_decoder_first_past_input_idx, + int t5_decoder_first_present_output_idx, + bool has_hidden_state, + int current_length, + transformers::Sequences&, const transformers::IConsoleDumper* dumper) { // last_outputs: logits, present_key_self_0, present_value_self_0, ... // next_inputs: input_ids, @@ -565,6 +578,12 @@ 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) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "BeamSearch CUDA Op does not support no hidden state senario in decoder input"); + } + ORT_UNUSED_PARAMETER(current_length); + // Update input_ids with next tokens. int batch_beam_size = static_cast(beam_next_tokens.length()); int64_t dims[] = {batch_beam_size, 1}; @@ -589,13 +608,14 @@ Status UpdateDecoderFeeds( if (num_beams == 1) { // feed present_* output to past_* inputs one by one for (int i = 0; i < num_present_tensors; ++i) { - next_inputs[transformers::T5DecoderSubgraph::kFirstPastInputIndex + i] = - last_outputs[transformers::T5DecoderSubgraph::kFirstPresentOutputIndex + i]; + next_inputs[t5_decoder_first_past_input_idx + i] = + last_outputs[t5_decoder_first_present_output_idx + i]; return Status::OK(); } } - return PickT5PastState(last_outputs, next_inputs, num_present_tensors, beam_indices, allocator, stream); + return PickT5PastState(last_outputs, next_inputs, num_present_tensors, beam_indices, allocator, + t5_decoder_first_past_input_idx, t5_decoder_first_present_output_idx, stream); } // Explicit template instantiations of functions @@ -640,6 +660,8 @@ template Status UpdateGptFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int gpt_subgraph_first_past_input_idx, + int gpt_subgraph_first_present_output_idx, const transformers::IConsoleDumper* dumper); // Float16 @@ -672,6 +694,8 @@ template Status UpdateGptFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int gpt_subgraph_first_past_input_idx, + int gpt_subgraph_first_present_output_idx, const transformers::IConsoleDumper* dumper); template Status UpdateDecoderFeeds( @@ -683,6 +707,11 @@ template Status UpdateDecoderFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int t5_decoder_first_past_input_idx, + int t5_decoder_first_present_output_idx, + bool has_hidden_state, + int current_length, + transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper); template Status UpdateDecoderFeeds( @@ -694,6 +723,11 @@ template Status UpdateDecoderFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int t5_decoder_first_past_input_idx, + int t5_decoder_first_present_output_idx, + bool has_hidden_state, + int current_length, + transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper); } // namespace BeamSearchCudaDeviceHelper diff --git a/onnxruntime/contrib_ops/cuda/transformers/beam_search_device_helper.h b/onnxruntime/contrib_ops/cuda/transformers/beam_search_device_helper.h index dde2b9e3f5..a90d0c7ee8 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/beam_search_device_helper.h +++ b/onnxruntime/contrib_ops/cuda/transformers/beam_search_device_helper.h @@ -71,6 +71,8 @@ Status UpdateGptFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int gpt_subgraph_first_past_input_idx, + int gpt_subgraph_first_present_output_idx, const transformers::IConsoleDumper* dumper); // --------------------------------------------------------------- @@ -88,6 +90,11 @@ Status UpdateDecoderFeeds( gsl::span beam_next_tokens, gsl::span beam_indices, int num_beams, + int t5_decoder_first_past_input_idx, + int t5_decoder_first_present_output_idx, + bool has_hidden_state, + int current_length, + transformers::Sequences& sequences, const transformers::IConsoleDumper* dumper); } // namespace BeamSearchCudaDeviceHelper diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 469e0b7d5a..8cf2d278e0 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -976,6 +976,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA(BeamSearch, 1, .Input(6, "repetition_penalty", "The parameter for repetition penalty. Default value 1.0 means no penalty. Accepts value > 0.0. Shape is (1)", "T", OpSchema::Optional) .Input(7, "vocab_mask", "Mask of vocabulary. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (vacab_size)", "M", OpSchema::Optional) .Input(8, "prefix_vocab_mask", "Mask of vocabulary for first step. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (batch_size, vocab_size)", "M", OpSchema::Optional) + .Input(9, "attention_mask", "Custom attention mask. Shape is (batch_size, sequence_length)", "I", OpSchema::Optional) .Output(0, "sequences", "Word IDs of generated sequences. Shape is (batch_size, num_return_sequences, max_sequence_length)", "I") .Output(1, "sequences_scores", "Final beam score of the generated sequences. Shape is (batch_size, num_return_sequences)", "T", OpSchema::Optional) .Output(2, "scores",