saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Decouple use_sequence_as_input_ids from has_hidden_states (#13130)

Browse source 
### Description
<!-- Describe your changes. -->

A fix for parity issue in huggingface bart model with beam search
https://github.com/microsoft/onnxruntime/pull/12779

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
Ye Wang 2022-09-29 22:45:52 -07:00 committed by GitHub
parent 32395e2e16
commit c8781b77f6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 25 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h
View file

@ -289,7 +289,7 @@ Status BeamSearchT5<T>::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()));

9
onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc
View file

@ -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<int32_t>(), 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<Tensor>()->MutableDataAsSpan<int32_t>());
} else {
int32_t* input_ids_data = input_ids.GetMutable<Tensor>()->MutableData<int32_t>();
@ -865,7 +864,7 @@ template Status UpdateDecoderFeeds<float>(
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);

4
onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.h
View file

@ -146,7 +146,7 @@ using UpdateDecoderFeedsFunc = std::function<Status(
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)>;
@ -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);

7
onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc
View file

@ -77,6 +77,13 @@ Status T5DecoderSubgraph::Validate(const std::vector<const NodeArg*>& subgraph_i
ORT_RETURN_IF_ERROR(GetParameters(past_shape, logits_shape, false));
num_layers = (static_cast<int>(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;

8
onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h
View file

@ -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

10
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc
View file

@ -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<float>(
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<MLFloat16>(
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);

2
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h
View file

@ -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);