mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
apply zcode changes to the beam search op (#11880)
* apply zcode changes to the beam search op * fix pipeline failure * add doc * workaround for C# * update * update * use name zcode * review comment * review comments * fix cpplint * review coments
This commit is contained in:
Ye Wang
GitHub
parent
cefceff5c9
commit
859ef277a0
18 changed files with 267 additions and 81 deletions
|
|
@ -371,7 +371,7 @@ This version of the operator has been available since version 1 of the 'com.micr
|
|||
<dd>The id of the padding token</dd>
|
||||
</dl>
|
||||
|
||||
#### Inputs (5 - 9)
|
||||
#### Inputs (5 - 10)
|
||||
|
||||
<dl>
|
||||
<dt><tt>input_ids</tt> : I</dt>
|
||||
|
|
@ -392,6 +392,8 @@ This version of the operator has been available since version 1 of the 'com.micr
|
|||
<dd>Mask of vocabulary. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (vacab_size)</dd>
|
||||
<dt><tt>prefix_vocab_mask</tt> (optional) : M</dt>
|
||||
<dd>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)</dd>
|
||||
<dt><tt>attention_mask</tt> (optional) : I</dt>
|
||||
<dd>Custom attention mask. Shape is (batch_size, sequence_length)</dd>
|
||||
</dl>
|
||||
|
||||
#### Outputs (1 - 3)
|
||||
|
|
|
|||
|
|
@ -389,7 +389,7 @@ Do not modify directly.*
|
|||
|**Operator Domain:** *com.microsoft*||||
|
||||
|Attention|*in* input:**T**<br> *in* weight:**T**<br> *in* bias:**T**<br> *in* mask_index:**M**<br> *in* past:**T**<br> *in* extra_add:**T**<br> *out* output:**T**<br> *out* present:**T**|1+|**T** = tensor(float)|
|
||||
|AttnLSTM|*in* X:**T**<br> *in* W:**T**<br> *in* R:**T**<br> *in* B:**T**<br> *in* sequence_lens:**T1**<br> *in* initial_h:**T**<br> *in* initial_c:**T**<br> *in* P:**T**<br> *in* QW:**T**<br> *in* MW:**T**<br> *in* V:**T**<br> *in* M:**T**<br> *in* memory_seq_lens:**T1**<br> *in* AW:**T**<br> *out* Y:**T**<br> *out* Y_h:**T**<br> *out* Y_c:**T**|1+|**T** = tensor(double), tensor(float)<br/> **T1** = tensor(int32)|
|
||||
|BeamSearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* num_beams:**I**<br> *in* num_return_sequences:**I**<br> *in* length_penalty:**T**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**M**<br> *in* prefix_vocab_mask:**M**<br> *out* sequences:**I**<br> *out* sequences_scores:**T**<br> *out* scores:**T**|1+|**T** = tensor(float)|
|
||||
|BeamSearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* num_beams:**I**<br> *in* num_return_sequences:**I**<br> *in* length_penalty:**T**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**M**<br> *in* prefix_vocab_mask:**M**<br> *in* attention_mask:**I**<br> *out* sequences:**I**<br> *out* sequences_scores:**T**<br> *out* scores:**T**|1+|**T** = tensor(float)|
|
||||
|BiasGelu|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T**|1+|**T** = tensor(float)|
|
||||
|BifurcationDetector|*in* src_tokens:**T**<br> *in* cur_tokens:**T**<br> *in* prev_suffix_match_idx:**T**<br> *in* pred_tokens:**T**<br> *out* tokens:**T**<br> *out* suffix_match_idx:**T**|1+|**T** = tensor(int64)|
|
||||
|CDist|*in* A:**T**<br> *in* B:**T**<br> *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**<br> *in* weight:**T**<br> *in* bias:**T**<br> *in* mask_index:**M**<br> *in* past:**T**<br> *in* extra_add:**T**<br> *out* output:**T**<br> *out* present:**T**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|BeamSearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* num_beams:**I**<br> *in* num_return_sequences:**I**<br> *in* length_penalty:**T**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**M**<br> *in* prefix_vocab_mask:**M**<br> *out* sequences:**I**<br> *out* sequences_scores:**T**<br> *out* scores:**T**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|BeamSearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* num_beams:**I**<br> *in* num_return_sequences:**I**<br> *in* length_penalty:**T**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**M**<br> *in* prefix_vocab_mask:**M**<br> *in* attention_mask:**I**<br> *out* sequences:**I**<br> *out* sequences_scores:**T**<br> *out* scores:**T**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|BiasDropout|*in* data:**T**<br> *in* bias:**T**<br> *in* residual:**T**<br> *in* ratio:**T1**<br> *in* training_mode:**T2**<br> *out* output:**T**<br> *out* mask:**T2**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)<br/> **T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)<br/> **T2** = tensor(bool)|
|
||||
|BiasGelu|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|
||||
|BiasSoftmax|*in* data:**T**<br> *in* bias:**T**<br> *out* output:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
|
|
@ -334,10 +334,12 @@ template <typename T>
|
|||
void PickGptPastState(const std::vector<OrtValue>& last_outputs,
|
||||
std::vector<OrtValue>& next_inputs,
|
||||
gsl::span<const int32_t>& beam_indices,
|
||||
int gpt_subgraph_first_past_input_idx,
|
||||
int gpt_subgraph_first_present_output_idx,
|
||||
AllocatorPtr allocator) {
|
||||
int num_present_tensors = static_cast<int>(last_outputs.size()) - transformers::GptSubgraph::kFirstPresentOutputIndex;
|
||||
int num_present_tensors = static_cast<int>(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<Tensor>().Shape();
|
||||
|
|
@ -364,7 +366,7 @@ void PickGptPastState(const std::vector<OrtValue>& 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<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<T>(last_outputs, next_inputs, beam_indices, allocator);
|
||||
PickGptPastState<T>(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<int32_t>();
|
||||
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>();
|
||||
Tensor::InitOrtValue(element_type, input_ids_shape, const_cast<Tensor*>(&attention_mask)->MutableData<int32_t>(),
|
||||
allocator->Info(), encoder_attention_mask);
|
||||
} else {
|
||||
auto mask_type = DataTypeImpl::GetType<int32_t>();
|
||||
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<Tensor>()->MutableData<int32_t>();
|
||||
const int32_t* word_id = original_encoder_input_ids->Data<int32_t>();
|
||||
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<Tensor>()->MutableData<int32_t>();
|
||||
const int32_t* word_id = original_encoder_input_ids->Data<int32_t>();
|
||||
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<OrtValue>& last_outputs,
|
|||
std::vector<OrtValue>& next_inputs,
|
||||
int num_present_tensors,
|
||||
gsl::span<const int32_t>& 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<Tensor>().Shape();
|
||||
|
|
@ -547,7 +562,7 @@ void PickT5PastState(const std::vector<OrtValue>& 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<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<int>(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<int32_t>(), input_ids_shape, allocator, input_ids);
|
||||
|
||||
gsl::copy(beam_next_tokens, input_ids.GetMutable<Tensor>()->MutableDataAsSpan<int32_t>());
|
||||
// TODO: decouple has_hidden_state with full input_ids
|
||||
if (has_hidden_state) {
|
||||
gsl::copy(beam_next_tokens, input_ids.GetMutable<Tensor>()->MutableDataAsSpan<int32_t>());
|
||||
} else {
|
||||
int32_t* input_ids_data = input_ids.GetMutable<Tensor>()->MutableData<int32_t>();
|
||||
for (int i = 0; i < batch_beam_size; i++) {
|
||||
gsl::span<const int32_t> 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<T>(last_outputs, next_inputs, num_present_tensors, beam_indices, allocator);
|
||||
PickT5PastState<T>(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<float>(
|
|||
gsl::span<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<float>(
|
||||
|
|
@ -664,6 +700,11 @@ template Status UpdateDecoderFeeds<float>(
|
|||
gsl::span<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<int32_t>(const OrtValue& input, int num_beams, AllocatorPtr allocator, OrtValue& expanded);
|
||||
|
|
|
|||
|
|
@ -99,11 +99,14 @@ using UpdateGptFeedsFunc = std::function<Status(
|
|||
gsl::span<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<Status(
|
||||
const Tensor* original_encoder_input_ids,
|
||||
const OrtValue* attn_mask_value,
|
||||
int num_beams,
|
||||
int pad_token_id,
|
||||
int start_token_id,
|
||||
|
|
@ -123,6 +126,11 @@ using UpdateDecoderFeedsFunc = std::function<Status(
|
|||
gsl::span<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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);
|
||||
|
||||
// ---------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -262,6 +262,19 @@ Status BeamSearchBase<T>::CheckInputs(const OpKernelContextInternal& context) {
|
|||
parameters_->prefix_vocab_mask = prefix_vocab_mask->DataAsSpan<int32_t>();
|
||||
}
|
||||
|
||||
const Tensor* attention_mask = context.Input<Tensor>(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<T>::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'.");
|
||||
|
|
|
|||
|
|
@ -104,6 +104,8 @@ Status BeamSearchGpt<T>::UpdateFeeds(
|
|||
beam_next_tokens,
|
||||
beam_indices,
|
||||
this->parameters_->num_beams,
|
||||
gpt_subgraph_.GetFirstPastInputIndex(),
|
||||
gpt_subgraph_.GetFirstPresentOutputIndex(),
|
||||
this->GetConsoleDumper());
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -100,6 +100,8 @@ Status BeamSearchT5<T>::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<Tensor>();
|
||||
|
||||
const OrtValue* encoder_attn_mask_value = this->context_.GetInputOrtValue(9);
|
||||
|
||||
BeamSearchCpuState cpu_state;
|
||||
cpu_state.Init(this->cpu_allocator_,
|
||||
static_cast<size_t>(parameters->BatchBeamSize()),
|
||||
|
|
@ -111,6 +113,7 @@ Status BeamSearchT5<T>::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<T>::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<T>::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<T>::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<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
|
|||
beam_next_tokens.as_span<const int32_t>(),
|
||||
beam_indices.as_span<const int32_t>(),
|
||||
parameters->num_beams,
|
||||
decoder_subgraph_.GetFirstPastInputIndex(),
|
||||
decoder_subgraph_.GetFirstPresentOutputIndex(),
|
||||
decoder_subgraph_.HasHiddenStates(),
|
||||
current_length,
|
||||
cpu_state.sequences,
|
||||
this->GetConsoleDumper()));
|
||||
}
|
||||
decoder_fetches.clear();
|
||||
|
|
|
|||
|
|
@ -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<const NodeArg*>& subgraph_inputs,
|
||||
const std::vector<const NodeArg*>& 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<const NodeArg*>& 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);
|
||||
|
|
|
|||
|
|
@ -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<const NodeArg*>& subgraph_inputs,
|
||||
const std::vector<const NodeArg*>& 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
|
||||
|
|
|
|||
|
|
@ -44,28 +44,38 @@ namespace transformers {
|
|||
|
||||
Status T5DecoderSubgraph::Validate(const std::vector<const NodeArg*>& subgraph_inputs,
|
||||
const std::vector<const NodeArg*>& 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<int>(subgraph_outputs.size()) - kFirstPresentOutputIndex) / 2;
|
||||
num_layers = (static_cast<int>(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<const NodeArg*>& 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]);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<const NodeArg*>& subgraph_inputs,
|
||||
const std::vector<const NodeArg*>& 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
|
||||
|
|
|
|||
|
|
@ -43,7 +43,7 @@ Status T5EncoderSubgraph::Validate(const std::vector<const NodeArg*>& 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<int>(subgraph_outputs.size()) - kFirstPresentOutputIndex) % 4 != 0,
|
||||
ORT_RETURN_IF((static_cast<int>(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<const NodeArg*>& subgraph_i
|
|||
|
||||
// Save parameters related to the subgraph.
|
||||
ORT_RETURN_IF_ERROR(GetParameters(past_shape, logits_shape, false));
|
||||
num_layers = (static_cast<int>(subgraph_outputs.size()) - kFirstPresentOutputIndex) / 4;
|
||||
num_layers = (static_cast<int>(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<const NodeArg*>& 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<const OrtValue*>& 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,
|
||||
|
|
|
|||
|
|
@ -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<const OrtValue*>& implicit_inputs,
|
||||
int num_beams,
|
||||
int pad_token_id,
|
||||
|
|
@ -33,7 +36,12 @@ class T5EncoderSubgraph : public Subgraph {
|
|||
Status Validate(const std::vector<const NodeArg*>& subgraph_inputs,
|
||||
const std::vector<const NodeArg*>& subgraph_outputs) override;
|
||||
|
||||
constexpr static int kFirstPresentOutputIndex = 2;
|
||||
int GetFirstPresentOutputIndex() const {
|
||||
return first_present_output_index_;
|
||||
}
|
||||
|
||||
private:
|
||||
int first_present_output_index_;
|
||||
};
|
||||
|
||||
} // namespace transformers
|
||||
|
|
|
|||
|
|
@ -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<float>(),
|
||||
|
|
|
|||
|
|
@ -404,10 +404,12 @@ Status PickGptPastState(const std::vector<OrtValue>& last_outputs,
|
|||
std::vector<OrtValue>& next_inputs,
|
||||
gsl::span<const int32_t>& 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<int>(last_outputs.size()) - transformers::GptSubgraph::kFirstPresentOutputIndex;
|
||||
int num_present_tensors = static_cast<int>(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<Tensor>().Shape();
|
||||
|
|
@ -436,7 +438,7 @@ Status PickGptPastState(const std::vector<OrtValue>& last_outputs,
|
|||
cudaMemcpyDeviceToDevice, reinterpret_cast<cudaStream_t>(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<OrtValue>& last_outputs,
|
|||
int num_present_tensors,
|
||||
gsl::span<const int32_t>& 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<Tensor>().Shape();
|
||||
|
|
@ -472,7 +476,7 @@ Status PickT5PastState(const std::vector<OrtValue>& last_outputs,
|
|||
cudaMemcpyDeviceToDevice, reinterpret_cast<cudaStream_t>(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<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<int>(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<T>(last_outputs, next_inputs, beam_indices, allocator, stream));
|
||||
ORT_RETURN_IF_ERROR(PickGptPastState<T>(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<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<int>(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<T>(last_outputs, next_inputs, num_present_tensors, beam_indices, allocator, stream);
|
||||
return PickT5PastState<T>(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<float>(
|
|||
gsl::span<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<MLFloat16>(
|
|||
gsl::span<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<float>(
|
||||
|
|
@ -683,6 +707,11 @@ template Status UpdateDecoderFeeds<float>(
|
|||
gsl::span<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<MLFloat16>(
|
||||
|
|
@ -694,6 +723,11 @@ template Status UpdateDecoderFeeds<MLFloat16>(
|
|||
gsl::span<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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
|
||||
|
|
|
|||
|
|
@ -71,6 +71,8 @@ Status UpdateGptFeeds(
|
|||
gsl::span<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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<const int32_t> beam_next_tokens,
|
||||
gsl::span<const int32_t> 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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Reference in a new issue