diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 3a0171a064..a982999f0a 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -5,6 +5,7 @@ Do not modify directly.* * com.microsoft * com.microsoft.Attention * com.microsoft.AttnLSTM + * com.microsoft.BeamSearch * com.microsoft.BiasDropout * com.microsoft.BiasGelu * com.microsoft.BiasSoftmax @@ -337,6 +338,75 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.BeamSearch** + + Beam Search for text generation. Supports GPT-2 decoder. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
body : graph (required)
+
The GPT-2 subgraph with input_ids, position_ids, attention_mask, past_0, past_1, ... as inputs, and logits, present_0, present_1, ... as output
+
early_stopping : int
+
early stop or not
+
eos_token_id : int (required)
+
The id of the end-of-sequence token
+
no_repeat_ngram_size : int
+
no repeat ngrams size
+
pad_token_id : int (required)
+
The id of the padding token
+
+ +#### Inputs (6 - 9) + +
+
input_ids : I
+
The sequence used as a prompt for the generation. Shape is (batch_size, sequence_length)
+
max_length : I
+
The maximum length of the sequence to be generated. Shape is (1)
+
min_length (optional) : I
+
The minimum length below which the score of eos_token_id is set to -Inf. Shape is (1)
+
num_beams : I
+
Number of beams for beam search. 1 means no beam search. Shape is (1)
+
num_return_sequences : I
+
The number of returned sequences in the batch. Shape is (1)
+
temperature : T
+
The value used to module the next token probabilities. Accepts value > 0.0. Shape is (1)
+
length_penalty (optional) : T
+
Exponential penalty to the length. Default value 1.0 means no penalty.Value > 1.0 encourages longer sequences, while values < 1.0 produces shorter sequences.Shape is (1,)
+
repetition_penalty (optional) : T
+
The parameter for repetition penalty. Default value 1.0 means no penalty. Accepts value > 0.0. Shape is (1)
+
vocab_mask (optional) : M
+
Mask of vocabulary. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (vacab_size)
+
+ +#### Outputs (1 - 3) + +
+
sequences : I
+
Word IDs of generated sequences. Shape is (batch_size, num_return_sequences, max_sequence_length)
+
sequences_scores (optional) : T
+
Final beam score of the generated sequences. Shape is (batch_size, num_return_sequences)
+
scores (optional) : T
+
Processed beam scores for each vocabulary token at each generation step.Beam scores consisting of log softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this beam.Shape is (max_length - sequence_length, batch_size, num_beams, vocab_size)
+
+ +#### Type Constraints + +
+
T : tensor(float), tensor(float16)
+
Constrain input and output types to float tensors.
+
I : tensor(int32)
+
Constrain to integer types
+
M : tensor(int32)
+
Constrain mask to integer types
+
+ + ### **com.microsoft.BiasDropout** output, dropout_mask = Dropout(data + bias, ratio) + residual, Intended to specialize the dropout pattern commonly found in transformer models. diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 052367130e..6f09fb5026 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -377,6 +377,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* temperature:**T**
*in* length_penalty:**T**
*in* repetition_penalty:**T**
*in* vocab_mask:**M**
*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)| diff --git a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc index 77669f74ea..1f0ee6d17e 100644 --- a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc @@ -22,6 +22,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, GridSample); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, Attention); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, BeamSearch); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, EmbedLayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ExpandDims); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedConv); @@ -199,6 +200,7 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) { // add more kernels here BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search.cc b/onnxruntime/contrib_ops/cpu/transformers/beam_search.cc new file mode 100644 index 0000000000..1225442192 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search.cc @@ -0,0 +1,650 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// there's no way to use a raw pointer as the copy destination with std::copy_n +// (which gsl::copy uses with span::data() which returns a raw pointer) with the 14.11 toolset +// without generating a 4996 warning. going through an iterator is way too much overhead so turn off the warning. +#ifdef _MSC_VER +#pragma warning(push) +#pragma warning(disable : 4996) +#endif + +#include +#include "core/providers/cpu/controlflow/utils.h" +#include "core/providers/cpu/math/top_k.h" +#include "core/framework/allocator.h" +#include "core/framework/framework_common.h" +#include "core/framework/op_kernel_context_internal.h" +#include "core/framework/session_state.h" +#include "core/framework/tensorprotoutils.h" +#include "core/framework/utils.h" +#include "core/providers/cpu/tensor/utils.h" +#include "core/framework/session_options.h" +#include "core/framework/TensorSeq.h" +#include "gsl/gsl" +#include "core/providers/cpu/math/softmax_shared.h" +#include "beam_search.h" +#include "logits_processor.h" +#include "sequences.h" +#include "dump_tensor.h" + +#ifdef _MSC_VER +#pragma warning(pop) +#endif + +using namespace ONNX_NAMESPACE; +using namespace onnxruntime::common; + +namespace onnxruntime { +namespace contrib { + +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + BeamSearch, \ + kMSDomain, \ + 1, \ + T, \ + kCpuExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + transformers::BeamSearch); + +REGISTER_KERNEL_TYPED(float) + +namespace transformers { + +template +struct BeamSearchState { + gsl::span beam_scores; // shape (batch_size, num_beams) + gsl::span next_token_logits; // shape (batch_size * num_beams, vocab_size) + gsl::span next_token_scores; // shape (batch_size, num_beams * vocab_size) + gsl::span next_tokens; // shape (batch_size, 2 * num_beams) + gsl::span next_indices; // shape (batch_size, 2 * num_beams) + gsl::span next_positions; // shape (batch_size, num_beams). Next position value for position_ids. + + gsl::span scores; // shape (max_length - sequence_length + 1, batch_size, num_beams * vocab_size) + gsl::span remaining_scores; // subspan that is avaiable for appending next token scores. + + Sequences sequences; + + void Init(AllocatorPtr allocator, + int batch_size, + int num_beams, + int vocab_size, + int sequence_length, + int max_length, + bool output_scores) { + size_t batch_beam_size = SafeInt(batch_size) * num_beams; + beam_scores = AllocateBuffer(allocator, beam_scores_buffer_, batch_beam_size, true, static_cast(0)); + + // Initialize score of first beam of each group with 0 and the rest with -1e9. + // This ensures that the beams in the same group don't produce same tokens every time. + for (int i = 0; i < batch_size; i++) { + for (int j = 1; j < num_beams; j++) { + beam_scores[i * num_beams + j] = -1e9; + } + } + + size_t next_token_size = SafeInt(batch_beam_size) * vocab_size; + next_token_logits = AllocateBuffer(allocator, next_token_logits_buffer_, next_token_size, true, static_cast(0)); + next_token_scores = AllocateBuffer(allocator, next_token_scores_buffer_, next_token_size, true, static_cast(0)); + + next_tokens = AllocateBuffer(allocator, next_tokens_buffer_, SafeInt(2) * batch_beam_size, true, static_cast(0)); + + next_indices = AllocateBuffer(allocator, next_indices_buffer_, SafeInt(2) * batch_beam_size, true, static_cast(0)); + + next_positions = AllocateBuffer(allocator, next_positions_buffer_, batch_beam_size, true, static_cast(0)); + + if (output_scores) { + size_t elements = SafeInt(max_length - sequence_length) * batch_size * num_beams * vocab_size; + scores = AllocateBuffer(allocator, scores_buffer_, elements); + remaining_scores = scores; + } + + // sequences will be initialized later since it has dependency on input_ids + } + + private: + BufferUniquePtr beam_scores_buffer_; + BufferUniquePtr next_token_logits_buffer_; + BufferUniquePtr next_token_scores_buffer_; + BufferUniquePtr next_tokens_buffer_; + BufferUniquePtr next_indices_buffer_; + BufferUniquePtr next_positions_buffer_; + BufferUniquePtr scores_buffer_; +}; + +template +class BeamSearchImpl { + public: + BeamSearchImpl(OpKernelContextInternal& context, + const SessionState& session_state, + GptSubgraph& gpt_subgraph, + concurrency::ThreadPool* thread_pool, + void* stream, + BeamSearchParameters& params); + + // Initialize by validating all the inputs, and allocating the output tensors. + Status Initialize(); + + // Execute beam search in iterations util stopping criteria is reached. + // In each iteration, GPT subgraph is called, and next token for each sequence is generated. + Status Execute(const FeedsFetchesManager& cached_ffm); + + private: + // Validate inputs. + Status CheckInputs(const OpKernelContextInternal& context); + + // Prepare the inputs for first inference of subgraph + void CreateInitialFeeds(gsl::span& next_positions, std::vector& feeds); + + // Update the input for next iteration. + Status UpdateFeeds( + const std::vector& last_outputs, + std::vector& next_inputs, + int current_length, + gsl::span& next_positions, + gsl::span beam_next_tokens, + gsl::span beam_indices); + + // Process logits and append next tokens to sequences. + Status GenerateNextToken(const OrtValue& logits, + gsl::span& beam_next_tokens, + gsl::span& beam_indices, + BeamSearchState& beam_state); + + // Calculate scores from logits, then apply filtering and select next token for each beam. + Status ProcessLogits(const OrtValue& logits, // logits output of subgraph + BeamSearchState& beam_state, + AllocatorPtr& allocator); + + OpKernelContextInternal& context_; + + const SessionState& session_state_; + + GptSubgraph& gpt_subgraph_; + + concurrency::ThreadPool* thread_pool_; + + const std::vector& implicit_inputs_; + + // Not used in CPU. Stream is for CUDA only. + void* stream_; + + BeamSearchParameters* parameters_; + + LogitsProcessorList logits_processors_; + + std::unique_ptr> beam_scorer_; + + AllocatorPtr allocator_; +}; + +template +void BeamSearch::Init(const OpKernelInfo& info) { + // Make sure the body attribute was present even though we don't need it here. + ONNX_NAMESPACE::GraphProto proto; + ORT_ENFORCE(info.GetAttr("body", &proto).IsOK()); + ORT_IGNORE_RETURN_VALUE(proto); + + parameters_.ParseFromAttributes(info); + + stream_ = nullptr; +} + +template +std::unique_ptr BeamSearch::Create(const OpKernelInfo& info, + void* stream) { + auto result = std::make_unique(info); + result->SetComputeStream(stream); + return result; +} + +template +common::Status BeamSearch::SetupSubgraphExecutionInfo(const SessionState& session_state, + const std::string& attribute_name, + const SessionState& subgraph_session_state) { + ORT_ENFORCE(gpt_subgraph_ == nullptr, "SetupSubgraphExecutionInfo should only be called once for each subgraph."); + const auto& node = Node(); + gpt_subgraph_ = std::make_unique(node, attribute_name, subgraph_session_state.GetGraphViewer()); + ORT_RETURN_IF_ERROR(gpt_subgraph_->Setup(session_state, subgraph_session_state)); + feeds_fetches_manager_ = gpt_subgraph_->GetFeedsFetchesManager(); + parameters_.SetSubgraphParameters(gpt_subgraph_->vocab_size, + gpt_subgraph_->num_heads, + gpt_subgraph_->head_size, + gpt_subgraph_->num_layers); + return Status::OK(); +} + +template +Status BeamSearch::Compute(OpKernelContext* ctx) const { + auto* ctx_internal = static_cast(ctx); + auto* session_state = ctx_internal->SubgraphSessionState("body"); + ORT_ENFORCE(session_state, "Subgraph SessionState was not found for 'body' attribute."); + ORT_ENFORCE(feeds_fetches_manager_, "CreateFeedsFetchesManager must be called prior to execution of graph."); + + concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool(); + + BeamSearchParameters parameters = parameters_; // make a copy since we will update the parameters based on inputs later + + BeamSearchImpl impl{*ctx_internal, *session_state, *gpt_subgraph_, thread_pool, stream_, parameters}; + + auto status = impl.Initialize(); + ORT_RETURN_IF_ERROR(status); + + status = impl.Execute(*feeds_fetches_manager_); + + return status; +} + +template +BeamSearchImpl::BeamSearchImpl(OpKernelContextInternal& context, + const SessionState& session_state, + GptSubgraph& gpt_subgraph, + concurrency::ThreadPool* thread_pool, + void* stream, + BeamSearchParameters& params) + : context_(context), + session_state_(session_state), + gpt_subgraph_(gpt_subgraph), + thread_pool_(thread_pool), + implicit_inputs_(context_.GetImplicitInputs()), + stream_(stream), + parameters_(¶ms), + allocator_(nullptr) { + parameters_->ParseFromInputs(&context); + + allocator_ = session_state.GetExecutionProviders() + .Get(onnxruntime::kCpuExecutionProvider) + ->GetAllocator(0, OrtMemTypeDefault); +} + +template +Status BeamSearchImpl::CheckInputs(const OpKernelContextInternal& context) { + // Input shapes: + // input_ids : (batch_size, sequence_length) + // vocab_mask : (vocab_size) or nullptr + + const Tensor* input_ids = context.Input(0); + const auto& dims = input_ids->Shape().GetDims(); + if (dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'input_ids' is expected to have 2 dimensions, got ", + dims.size()); + } + + const Tensor* vocab_mask = context.Input(8); + if (vocab_mask != nullptr) { // vocab_mask is optional + const auto& vocab_mask_dims = vocab_mask->Shape().GetDims(); + if (vocab_mask_dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'vocab_mask' is expected to have 1 dimension, got ", + vocab_mask_dims.size()); + } + + // There is dependency on vocab_size parameter, which shall be set before calling this function. + if (static_cast(vocab_mask_dims[0]) != parameters_->vocab_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'vocab_mask' shape does not match with vocab_size, got ", + vocab_mask_dims[0]); + } + + // store vocab mask in parameters. + parameters_->vocab_mask = vocab_mask->DataAsSpan(); + } + + return Status::OK(); +} + +template +Status BeamSearchImpl::Initialize() { + auto status = Status::OK(); + +#define CHECK_SCALAR_INPUT(name, index, required) \ + auto* name##_tensor = context_.Input(index); \ + if (name##_tensor) { \ + if (!name##_tensor->Shape().IsScalar()) { \ + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "'BeamSearch' input " #name " should be a scalar. Got shape of ", \ + name##_tensor->Shape()); \ + } \ + } else if (required) { \ + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "'BeamSearch' input " #name " is required"); \ + } + + CHECK_SCALAR_INPUT(min_length, 1, false); + + CHECK_SCALAR_INPUT(max_length, 2, true); + + CHECK_SCALAR_INPUT(num_beams, 3, true); + + CHECK_SCALAR_INPUT(num_return_sequences, 4, true); + + CHECK_SCALAR_INPUT(temperature, 5, true); + + CHECK_SCALAR_INPUT(length_penalty, 6, true); + + ORT_RETURN_IF(parameters_->num_return_sequences > parameters_->num_beams, "'num_return_sequences' has to be smaller or equal to 'num_beams'."); + + ORT_RETURN_IF_ERROR(CheckInputs(context_)); + + // This flag will be updated later when the scores output exists. + parameters_->output_scores = false; + + // Initialize processsors after CheckInputs so that parameters_->vocab_mask is ready. + logits_processors_.Init(*parameters_); + + return status; +} + +template +void BeamSearchImpl::CreateInitialFeeds(gsl::span& next_positions, std::vector& feeds) { + const OrtValue* input_ids_value = context_.GetInputOrtValue(0); + const Tensor& input_ids = input_ids_value->Get(); + gpt_subgraph_.CreateInitialFeeds(input_ids, implicit_inputs_, parameters_->num_beams, parameters_->pad_token_id, next_positions, feeds); +} + +template +Status BeamSearchImpl::ProcessLogits( + const OrtValue& logits, + BeamSearchState& beam_state, + AllocatorPtr& allocator) { + const int64_t batch_beam_size = static_cast(parameters_->BatchBeamSize()); + const int& vocab_size = parameters_->vocab_size; + + const T* logits_data = logits.Get().Data(); + + // Logits has shape (batch_size * num_beams, input_length, vocab_size), + // where input_length equals to parameters_->sequence_length for first subgraph call, and 1 for the remaining calls. + const TensorShape& logits_shape = logits.Get().Shape(); + ORT_ENFORCE(logits_shape.NumDimensions() == 3); + auto input_length = logits_shape[1]; + + // Get logits for the last token: + // next_token_logits = logits[:, -1, :], and the result shape is (batch_size * num_beams, vocab_size) + // When input_length == 1, use logits directly in SoftmaxCPU below so it only need for input_length > 1. + gsl::span& next_token_logits = beam_state.next_token_logits; + if (input_length > 1) { + const T* current_logits = logits_data + (input_length - 1) * vocab_size; + for (int i = 0; i < batch_beam_size; i++) { + gsl::span source(current_logits, vocab_size); + gsl::span target = next_token_logits.subspan(i * vocab_size, vocab_size); + gsl::copy(source, target); + current_logits += input_length * vocab_size; + } + } + +#ifdef DEBUG_BEAM_SEARCH + //DumpOrtValue("logits", logits); + DumpTensor("next_token_logits", next_token_logits.data(), parameters_->batch_size, parameters_->num_beams, vocab_size); +#endif + + // Get scores for candidates of next token: next_token_scores = log_softmax(next_token_logits, dim=-1) + gsl::span& next_token_scores = beam_state.next_token_scores; + Status status = SoftmaxCPU(batch_beam_size, // rows + vocab_size, // elements per row + input_length > 1 ? next_token_logits.data() : logits_data, + next_token_scores.data(), + true, + thread_pool_); + if (!status.IsOK()) { + return status; + } + +#ifdef DEBUG_BEAM_SEARCH + DumpTensor("next_token_scores after softmax", next_token_scores.data(), parameters_->batch_size, parameters_->num_beams, vocab_size); +#endif + + // Apply all score processors that updates scores + logits_processors_.Process(&(beam_state.sequences), next_token_scores); + +#ifdef DEBUG_BEAM_SEARCH + DumpTensor("next_token_scores after logits processor", next_token_scores.data(), parameters_->batch_size, parameters_->num_beams, vocab_size); +#endif + + // Add beam score to next token scores. Corresponding python code is like: + // next_token_scores = next_token_scores + beam_scores[:, None].expand_as(next_token_scores) + // TODO: use thread pool to parrellel + int offset = 0; + int batch_beam_index = 0; + for (int i = 0; i < parameters_->batch_size; i++) { + for (int j = 0; j < parameters_->num_beams; j++, batch_beam_index++) { + for (int k = 0; k < parameters_->vocab_size; k++, offset++) { + next_token_scores[offset] += beam_state.beam_scores[batch_beam_index]; + } + } + } + +#ifdef DEBUG_BEAM_SEARCH + DumpTensor("next_token_scores after adding beam_scores", next_token_scores.data(), parameters_->batch_size, parameters_->num_beams, vocab_size); +#endif + + if (parameters_->output_scores) { + // Append next token scores to the scores output. + gsl::copy(next_token_scores, beam_state.remaining_scores); + beam_state.remaining_scores = beam_state.remaining_scores.subspan(next_token_scores.size()); + } + + // Apply top-k selection like the following: + // next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size) + // next_token_scores, next_tokens = torch.topk(next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True) + int64_t next_token_scores_dims[] = {parameters_->batch_size, parameters_->num_beams * vocab_size}; + TensorShape next_token_scores_shape(&next_token_scores_dims[0], 2); + auto element_type = DataTypeImpl::GetType(); + OrtValue next_token_scores_value; + Tensor::InitOrtValue(element_type, next_token_scores_shape, next_token_scores.data(), allocator->Info(), next_token_scores_value); + const Tensor& input = next_token_scores_value.Get(); + + const int axis = 1; + const unsigned top_k = static_cast(2 * parameters_->num_beams); + const bool largest = true; + const bool sorted = true; // results returned in sorted order. + + std::unique_ptr topk_scores; + std::unique_ptr topk_indices; + status = GetTopK(&input, axis, top_k, largest, sorted, allocator, thread_pool_, topk_scores, topk_indices); + if (!status.IsOK()) { + return status; + } + +#ifdef DEBUG_BEAM_SEARCH + DumpTensor("topk_scores", *(topk_scores.get())); + DumpTensor("topk_indices", *(topk_indices.get())); +#endif + + // Convert indices in range [0, num_beams * vocab_size) to token ID of range [0, vocab_size) like the following: + // next_indices = (next_tokens / vocab_size).long() + // next_tokens = next_tokens % vocab_size + gsl::span next_token_indices = topk_indices->DataAsSpan(); + offset = 0; + for (int i = 0; i < parameters_->batch_size; i++) { + for (unsigned int j = 0; j < top_k; j++, offset++) { + beam_state.next_indices[offset] = next_token_indices[offset] / vocab_size; + beam_state.next_tokens[offset] = next_token_indices[offset] % vocab_size; + } + } + + gsl::span next_scores = topk_scores->DataAsSpan(); + gsl::span next_tokens(beam_state.next_tokens.data(), beam_state.next_tokens.size()); + gsl::span next_indices(beam_state.next_indices.data(), beam_state.next_indices.size()); + +#ifdef DEBUG_BEAM_SEARCH + DumpTensor("next_scores before scorer", next_scores.data(), parameters_->batch_size, top_k); + DumpTensor("next_tokens before scorer", next_tokens.data(), parameters_->batch_size, top_k); + DumpTensor("next_indices before scorer", next_indices.data(), parameters_->batch_size, top_k); +#endif + + beam_scorer_->Process( + &(beam_state.sequences), + next_scores, + next_tokens, + next_indices); + + return Status::OK(); +} + +template +Status BeamSearchImpl::GenerateNextToken( + const OrtValue& logits, + gsl::span& beam_next_tokens, + gsl::span& beam_indices, + BeamSearchState& beam_state) { + // Process logits to get next token scores + ORT_RETURN_IF_ERROR(ProcessLogits(logits, beam_state, allocator_)); + + gsl::span& beam_scores = beam_scorer_->GetNextScores(); + // It is optional to clone beam_scores. Change it to use same buffer also works: + // beam_state.beam_scores = beam_scores + // Here we make a copy to reduce the coupling with little cost (the buffer size is small). + gsl::copy(beam_scores, beam_state.beam_scores); + + beam_next_tokens = beam_scorer_->GetNextTokens(); + beam_indices = beam_scorer_->GetNextIndices(); + +#ifdef DEBUG_BEAM_SEARCH + DumpTensor("beam_scores after scorer", beam_scores.data(), parameters_->batch_size, parameters_->num_beams); + DumpTensor("beam_next_tokens after scorer", beam_next_tokens.data(), parameters_->batch_size, parameters_->num_beams); + DumpTensor("beam_indices after scorer", beam_indices.data(), parameters_->batch_size, parameters_->num_beams); +#endif + + beam_state.sequences.AppendNextTokenToSequences(beam_indices, beam_next_tokens); + +#ifdef DEBUG_BEAM_SEARCH + beam_state.sequences.PrintSequences(); +#endif + return Status::OK(); +} + +template +Status BeamSearchImpl::UpdateFeeds( + const std::vector& last_outputs, + std::vector& next_inputs, + int current_length, + gsl::span& next_positions, + gsl::span beam_next_tokens, + gsl::span beam_indices) { + return gpt_subgraph_.UpdateFeeds(last_outputs, next_inputs, current_length, next_positions, + beam_next_tokens, beam_indices, parameters_->num_beams); +} + +template +Status BeamSearchImpl::Execute(const FeedsFetchesManager& ffm) { + auto status = Status::OK(); + + std::vector sequences_dims{parameters_->batch_size, parameters_->num_return_sequences, parameters_->max_length}; + TensorShape sequences_shape(sequences_dims); + Tensor* output_sequences = context_.Output(0, sequences_shape); + + std::vector sequences_scores_dims{parameters_->batch_size, parameters_->num_return_sequences}; + TensorShape sequences_scores_shape(sequences_scores_dims); + Tensor* output_sequences_scores = context_.Output(1, sequences_scores_shape); + + std::vector scores_dims{ + parameters_->max_length - parameters_->sequence_length, + parameters_->batch_size, parameters_->num_beams, parameters_->vocab_size}; + TensorShape scores_shape(scores_dims); + Tensor* output_scores = context_.Output(2, scores_shape); + + // Update the flag to indicate whether scores exists in output + parameters_->output_scores = (output_scores != nullptr); + + std::vector feeds; + std::vector fetches; + + // Initialize resources + AllocatorPtr temp_space_allocator; + ORT_RETURN_IF_ERROR(context_.GetTempSpaceAllocator(&temp_space_allocator)); + + BeamSearchState beam_state; + beam_state.Init(temp_space_allocator, + parameters_->batch_size, + parameters_->num_beams, + parameters_->vocab_size, + parameters_->sequence_length, + parameters_->max_length, + parameters_->output_scores); + + beam_scorer_ = std::make_unique>(parameters_->batch_size, + parameters_->num_beams, + parameters_->max_length, + parameters_->length_penalty, + parameters_->early_stopping, + parameters_->num_return_sequences, + parameters_->pad_token_id, + parameters_->eos_token_id); + beam_scorer_->Initialize(allocator_, parameters_->sequence_length); // TODO: use temp_space_allocator + + CreateInitialFeeds(beam_state.next_positions, feeds); + const OrtValue& input_ids = feeds[0]; + beam_state.sequences.Init(temp_space_allocator, + input_ids, + parameters_->BatchBeamSize(), + parameters_->sequence_length, + parameters_->max_length); + +#ifdef DEBUG_BEAM_SEARCH + DumpOrtValue("input_ids", input_ids); + DumpOrtValue("position_ids", feeds[1]); + DumpOrtValue("attention_mask", feeds[2]); +#endif + + int current_length = parameters_->sequence_length; + while (current_length < parameters_->max_length) { +#ifdef DEBUG_BEAM_SEARCH + DumpString("***CurrentLength", std::to_string(current_length), true); +#endif + + status = utils::ExecuteSubgraph(session_state_, ffm, feeds, fetches, {}, + ExecutionMode::ORT_SEQUENTIAL, context_.GetTerminateFlag(), context_.Logger()); + + ORT_RETURN_IF_ERROR(status); + + const OrtValue& logits = fetches[0]; + gsl::span beam_next_tokens; + gsl::span beam_indices; + ORT_RETURN_IF_ERROR(GenerateNextToken(logits, beam_next_tokens, beam_indices, beam_state)); + + // When all batches are finished, stop earlier to avoid wasting computation. + if (beam_scorer_->IsDone()) { + break; + } + + // Increase sequence length after a new token is generated. + ++current_length; + + // Prepare inputs for next round of subgraph call. + if (current_length < parameters_->max_length) { + ORT_RETURN_IF_ERROR(UpdateFeeds(fetches, feeds, current_length, + beam_state.next_positions, + beam_next_tokens.as_span(), + beam_indices.as_span())); + } + fetches.clear(); + +#ifdef DEBUG_BEAM_SEARCH + if (current_length - parameters_->sequence_length == 3) { // only dump a few steps. + DisableTensorDump(); + } +#endif + } + + gsl::span beam_scores(beam_state.beam_scores.data(), beam_state.beam_scores.size()); + beam_scorer_->Finalize(&(beam_state.sequences), + beam_scores, + output_sequences, + output_sequences_scores); + + // Output per token scores + if (output_scores != nullptr) { + gsl::span target = output_scores->MutableDataAsSpan(); + gsl::span source = gsl::span(beam_state.scores.data(), beam_state.scores.size()); + assert(target.length() == source.length()); + gsl::copy(source, target); + } + + return status; +} + +// Instantiation +template class BeamSearchImpl; +template class BeamSearch; + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search.h new file mode 100644 index 0000000000..9dc5cac408 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search.h @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include "gsl/gsl" +#include "core/common/common.h" +#include "core/framework/feeds_fetches_manager.h" +#include "core/framework/op_kernel.h" +#include "core/providers/cpu/controlflow/utils.h" +#include "beam_search_parameters.h" +#include "beam_search_scorer.h" +#include "gpt_subgraph.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +template +class BeamSearch : public controlflow::IControlFlowKernel { + public: + BeamSearch(const OpKernelInfo& info) : IControlFlowKernel(info) { Init(info); } + void Init(const OpKernelInfo& info); + + Status Compute(OpKernelContext* ctx) const override; + + Status SetupSubgraphExecutionInfo(const SessionState& session_state, + const std::string& attribute_name, + const SessionState& subgraph_session_state) override; + + static std::unique_ptr Create(const OpKernelInfo& info, void* stream); + + protected: + void SetComputeStream(void* stream) { stream_ = stream; } + + private: + // Subgraph and FeedsFetchesManager re-used for each subgraph execution. + std::unique_ptr gpt_subgraph_; + FeedsFetchesManager* feeds_fetches_manager_; + + void* stream_; + + BeamSearchParameters parameters_; +}; + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_parameters.cc b/onnxruntime/contrib_ops/cpu/transformers/beam_search_parameters.cc new file mode 100644 index 0000000000..fee3ec4753 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_parameters.cc @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include "beam_search_parameters.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +constexpr int kMaxSequenceLength = 4096; + +Status BeamSearchParameters::Validate() const { + ORT_RETURN_IF(eos_token_id < 0, "eos_token_id is invalid"); + ORT_RETURN_IF(pad_token_id < 0, "pad_token_id is invalid"); + ORT_RETURN_IF(min_length >= max_length, "min_length shall be smaller than max_length"); + return Status::OK(); +} + +void BeamSearchParameters::ParseFromAttributes(const OpKernelInfo& info) { + early_stopping = info.GetAttrOrDefault("early_stopping", 0) == 1; + eos_token_id = static_cast(info.GetAttrOrDefault("eos_token_id", -1)); + pad_token_id = static_cast(info.GetAttrOrDefault("pad_token_id", -1)); + no_repeat_ngram_size = static_cast(info.GetAttrOrDefault("no_repeat_ngram_size", 0)); +} + +void BeamSearchParameters::ParseFromInputs(OpKernelContext* context) { + ORT_ENFORCE(context != nullptr); + const Tensor* input_ids = context->Input(0); + const auto& dims = input_ids->Shape().GetDims(); + ORT_ENFORCE(dims.size() == 2, "input_ids shall have 2 dimensions. Got ", dims.size()); + batch_size = static_cast(dims[0]); + sequence_length = static_cast(dims[1]); + + auto* max_length_tensor = context->Input(1); + max_length = max_length_tensor ? static_cast(*max_length_tensor->Data()) : kMaxSequenceLength; + ORT_ENFORCE(max_length > sequence_length, "max_length (", max_length, ") shall be greater than input sequence length (", sequence_length, ")"); + ORT_ENFORCE(max_length <= kMaxSequenceLength, "max_length (", max_length, ") shall be no more than ", kMaxSequenceLength); + + auto* min_length_tensor = context->Input(2); + min_length = min_length_tensor ? static_cast(*min_length_tensor->Data()) : 0; + + auto* num_beams_tensor = context->Input(3); + num_beams = num_beams_tensor ? static_cast(*num_beams_tensor->Data()) : 1; + // TODO: limit num_beams > 1 when we can have another operator for greedy search. + ORT_ENFORCE(num_beams >= 1, "num_beams shall be a positive integer, got ", num_beams); + + auto* num_return_sequences_tensor = context->Input(4); + num_return_sequences = num_return_sequences_tensor ? static_cast(*num_return_sequences_tensor->Data()) : 1; + ORT_ENFORCE(num_return_sequences >= 1, "num_return_sequences shall be a positive integer, got ", num_return_sequences); + ORT_ENFORCE(num_beams >= num_return_sequences, "num_return_sequences (", num_return_sequences, ") shall be be no more than num_beams (", num_beams, ")"); + + auto* temperature_tensor = context->Input(5); + temperature = temperature_tensor ? static_cast(*temperature_tensor->Data()) : 1; + ORT_ENFORCE(temperature > 0.0f, "temperature shall be greater than 0, got ", temperature); + + auto* length_penalty_tensor = context->Input(6); + length_penalty = length_penalty_tensor ? static_cast(*length_penalty_tensor->Data()) : 1; + + auto* repetition_penalty_tensor = context->Input(7); + repetition_penalty = repetition_penalty_tensor ? static_cast(*repetition_penalty_tensor->Data()) : 1.0f; + ORT_ENFORCE(repetition_penalty > 0.0f, "repetition_penalty shall be greater than 0, got ", repetition_penalty); +} + +void BeamSearchParameters::SetSubgraphParameters(int vocabulary_size, int heads, int hidden_size_per_head, int layers) { + vocab_size = vocabulary_size; + num_heads = heads; + head_size = hidden_size_per_head; + num_layers = layers; +} + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_parameters.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_parameters.h new file mode 100644 index 0000000000..26de2a9840 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_parameters.h @@ -0,0 +1,55 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/common/common.h" +#include "core/framework/op_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +struct BeamSearchParameters { + // Parameters from node attributes + int eos_token_id; + int pad_token_id; + int no_repeat_ngram_size; + bool early_stopping; + + // Parameters from inputs + int min_length; + int max_length; + int num_beams; + int num_return_sequences; + float temperature; + float length_penalty; + float repetition_penalty; + int batch_size; // deduce from first dimension of input_ids + int sequence_length; // deduce from second dimension of input_ids + + gsl::span vocab_mask; + + // Parameters from outputs. + bool output_scores; // whether scores existed in output + + // Parameters from subgraph. + int vocab_size; + // Below are used in CPU, reserved for CUDA. + int num_heads; + int head_size; + int num_layers; + + Status Validate() const; + + int BatchBeamSize() const { return batch_size * num_beams; } + + void ParseFromAttributes(const OpKernelInfo& info); + + void ParseFromInputs(OpKernelContext* context); + + void SetSubgraphParameters(int vocab_size, int num_heads, int head_size, int num_layers); +}; + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.cc b/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.cc new file mode 100644 index 0000000000..bb7aeb989e --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.cc @@ -0,0 +1,285 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include "core/common/common.h" +#include "core/framework/allocator.h" +#include "core/framework/tensorprotoutils.h" +#include "core/framework/utils.h" +#include "core/providers/cpu/tensor/utils.h" +#include "core/providers/cpu/rnn/rnn_helpers.h" +#include "beam_search_scorer.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { +using ::onnxruntime::rnn::detail::Allocate; + +template +BeamHypotheses::BeamHypotheses(int num_beams, T length_penalty, bool early_stopping) + : num_beams_(num_beams), + length_penalty_(length_penalty), + early_stopping_(early_stopping), + worst_score_(1e9) {} + +template +void BeamHypotheses::Add(gsl::span& hypothesis, T sum_logprobs) { + auto length = hypothesis.size(); + // TODO: when T is FP16, compute in FP32, then cast result back to FP16. length_penalty_ might also be float. + T score = sum_logprobs / pow(static_cast(length), length_penalty_); + + if (this->Size() < num_beams_ || score > worst_score_) { + HypothesisScore item(hypothesis, score); + beams_.push(item); + if (this->Size() > num_beams_) { + beams_.pop(); + } + worst_score_ = beams_.top().score; + } +} + +template +bool BeamHypotheses::IsDone(T best_sum_logprobs, int current_length) { + // If there are enough hypotheses and that none of the hypotheses being generated can become better + // than the worst one in the heap, then we are done with this sentence. + + if (Size() < num_beams_) + return false; + + if (early_stopping_) + return true; + + T current_score = best_sum_logprobs / pow(static_cast(current_length), length_penalty_); + return worst_score_ >= current_score; +} + +template +void BeamHypotheses::Output( + int top_k, + int max_length, + gsl::span& sequences, // buffer filled with pad token ID, shape (num_return_sequences, max_length) + gsl::span& sequences_scores) // buffer of shape (num_return_sequences) or empty +{ + ORT_ENFORCE(top_k <= Size()); + int remove_count = Size() - top_k; + for (int i = 0; i < remove_count; i++) { + beams_.pop(); + } + + // Since pop get the worst sequence, so output it in the reverse order. + // The frist (worst) beam shall be put at the last position among top_k sequences. + int index = top_k - 1; + while (!beams_.empty()) { + auto item = beams_.top(); + gsl::span& source = item.hypothesis; + gsl::span target = sequences.subspan(index * max_length, max_length); + + // Note that word_ids might be less than max_length. + // Since the sequences has been filled with pad token ID, so padding is not needed here. + // Since data type need cast from int64_t to int32_t, we cannot use gsl::copy(word_ids, sequence) here. + for (size_t i = 0; i < source.length(); i++) { + target[i] = static_cast(source[i]); + } + + if (!sequences_scores.empty()) + sequences_scores[index] = item.score; + + beams_.pop(); + index--; + } +} + +template +BeamSearchScorer::BeamSearchScorer(int batch_size, + int num_beams, + int max_length, + T length_penalty, + bool early_stopping, + int num_return_sequences, + int pad_token_id, + int eos_token_id) + : batch_size_(batch_size), + num_beams_(num_beams), + max_length_(max_length), + num_beam_hyps_to_keep_(num_return_sequences), + pad_token_id_(pad_token_id), + eos_token_id_(eos_token_id), + hypothesis_buffer_length_(0), + hypothesis_buffer_offset_(0) { + for (int batch = 0; batch < batch_size; batch++) { + beam_hyps.push_back(BeamHypotheses(num_beams, length_penalty, early_stopping)); + } +} + +template +bool BeamSearchScorer::IsDone() { + for (int batch = 0; batch < batch_size_; batch++) { + if (!done_[batch]) + return false; + } + return true; +} + +template +void BeamSearchScorer::Initialize(AllocatorPtr& allocator, int sequence_length){ + ORT_ENFORCE(next_beam_scores_.empty()); // Make sure this is called only once. + + size_t batch_beam_size = static_cast(batch_size_ * num_beams_); + const bool no_fill = false; // do not fill values after allocation + next_beam_scores_ = Allocate(allocator, batch_beam_size, next_beam_scores_ptr_, no_fill); + next_beam_tokens_ = Allocate(allocator, batch_beam_size, next_beam_tokens_ptr_, no_fill); + next_beam_indices_ = Allocate(allocator, batch_beam_size, next_beam_indices_ptr_, no_fill); + + // Space to store intermediate sequence with length sequence_length, sequence_length + 1, ..., max_sequence_length. + int buffer_per_beam = (max_length_ * (max_length_ + 1) - (sequence_length - 1) * sequence_length) / 2; + hypothesis_buffer_length_ = batch_beam_size * static_cast(buffer_per_beam); + hypothesis_buffer_ = Allocate(allocator, hypothesis_buffer_length_, hypothesis_buffer_ptr_, no_fill); + + done_ = Allocate(allocator, static_cast(batch_size_), done_ptr_, no_fill); + std::fill_n(done_.data(), done_.size(), false); +} + +template +void BeamSearchScorer::Process(ISequences* sequences, + gsl::span& next_scores, + gsl::span& next_tokens, + gsl::span& next_indices) { + // Sequences shape is (batch_size * num_beams, total_sequence_length) + // It contains word ID of whole sequence generated so far. + // It is different from subgraph input_ids, which only need one word when past state is not empty. + + const int sequence_length = sequences->GetSequenceLength(); + + ORT_ENFORCE(next_scores.size() == next_tokens.size()); + ORT_ENFORCE(next_scores.size() == next_indices.size()); + + for (int batch = 0; batch < batch_size_; batch++) { + BeamHypotheses& beam_hyp = beam_hyps[batch]; + if (done_[batch]) { + ORT_ENFORCE(beam_hyp.Size() >= num_beams_, "Batch can only be done if all beams have been generated"); + + // Pad the batch. + for (int j = 0; j < num_beams_; j++) { + next_beam_scores_[batch * num_beams_ + j] = 0.0f; + next_beam_tokens_[batch * num_beams_ + j] = pad_token_id_; + next_beam_indices_[batch * num_beams_ + j] = 0; + } + continue; + } + + // Next tokens for this sentence. + int beam_idx = 0; + int top_k = 2 * num_beams_; + for (int j = 0; j < top_k; j++) { + int64_t next_token = next_tokens[batch * top_k + j]; + T next_score = next_scores[batch * top_k + j]; + int64_t next_index = next_indices[batch * top_k + j]; + + int batch_beam_idx = batch * num_beams_ + static_cast(next_index); + // Add to generated hypotheses if end of sentence. + if ((eos_token_id_ >= 0) && (next_token == eos_token_id_)) { + bool is_beam_token_worse_than_top_num_beams = (j >= num_beams_); + if (is_beam_token_worse_than_top_num_beams) { + continue; + } + + // Clone the sequence and append to buffer. + gsl::span src = sequences->GetSequence(batch_beam_idx); + auto clone = hypothesis_buffer_.subspan(hypothesis_buffer_offset_, sequence_length); + gsl::copy(src, clone); + hypothesis_buffer_offset_ += sequence_length; + auto sequence = clone.template as_span(); + beam_hyp.Add(sequence, next_score); + } else { + // Add next predicted token since it is not eos_token. + next_beam_scores_[batch * num_beams_ + beam_idx] = next_score; + next_beam_tokens_[batch * num_beams_ + beam_idx] = next_token; + next_beam_indices_[batch * num_beams_ + beam_idx] = batch_beam_idx; + ++beam_idx; + } + + // Once the beam for next step is full, don't add more tokens to it. + if (beam_idx == num_beams_) + break; + } + + ORT_ENFORCE(beam_idx == num_beams_); + ORT_ENFORCE(hypothesis_buffer_offset_ <= batch_size_ * num_beams_ * max_length_); + + // Check if we are done so that we can save a pad step if all(done) + if (!done_[batch]) { + gsl::span topk_scores = next_scores.subspan(batch * num_beams_, top_k); + const T* best_sum_logprobs = std::max_element(topk_scores.begin(), topk_scores.end()); + if (beam_hyp.IsDone(*best_sum_logprobs, sequence_length)) { + done_[batch] = true; + } + } + } +} + +template +void BeamSearchScorer::Finalize(ISequences* sequences, + gsl::span& final_beam_scores, + Tensor* output_sequences, + Tensor* output_sequence_scores) { + ORT_ENFORCE(sequences != nullptr); + ORT_ENFORCE(output_sequences != nullptr); + + // Finalize all open beam hypotheses and add to generated hypotheses. + for (int batch_index = 0; batch_index < batch_size_; batch_index++) { + BeamHypotheses& beam_hyp = beam_hyps[batch_index]; + if (done_[batch_index]) { + continue; + } + + for (int beam_index = 0; beam_index < num_beams_; beam_index++) { + int batch_beam_index = batch_index * num_beams_ + beam_index; + T final_score = final_beam_scores[batch_beam_index]; + auto final_tokens = sequences->GetSequence(batch_beam_index); + beam_hyp.Add(final_tokens, final_score); + } + } + + // Word IDs of each sequence, with shape (batch_size * num_return_sequences, max_sequence_length). + gsl::span output = output_sequences->MutableDataAsSpan(); + + // Fill output sequences with pad token ID so that we do not need append it later. + std::fill_n(output.data(), output.size(), pad_token_id_); + + // Score of each sequence, with shape (batch_size * num_return_sequences). + gsl::span sequence_scores; + if (output_sequence_scores != nullptr) { + sequence_scores = output_sequence_scores->MutableDataAsSpan(); + } + + // Span is empty when output_sequence_scores is NULL. + gsl::span batch_sequence_score; + + // Select the best hypotheses according to number of sequences to return. + for (int batch_index = 0; batch_index < batch_size_; batch_index++) { + BeamHypotheses& beam_hyp = beam_hyps[batch_index]; + + const int num_return_sequences = num_beam_hyps_to_keep_; + auto batch_output = output.subspan(batch_index * num_return_sequences * max_length_, num_return_sequences * max_length_); + + if (output_sequence_scores != nullptr) { + batch_sequence_score = sequence_scores.subspan(batch_index * num_return_sequences, num_return_sequences); + } + + beam_hyp.Output( + num_return_sequences, + max_length_, + batch_output, + batch_sequence_score); + } +} + +// Instantiation +template class HypothesisScoreCompare; +template class BeamHypotheses; +template class BeamSearchScorer; + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.h new file mode 100644 index 0000000000..2a15080236 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.h @@ -0,0 +1,144 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// The implementation is based on huggingface transformers generation_beam_search.py + +#pragma once +#include +#include +#include "core/common/common.h" +#include "core/framework/allocator.h" +#include "core/framework/tensorprotoutils.h" +#include "core/framework/utils.h" +#include "core/providers/cpu/tensor/utils.h" +#include "sequences.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +// Interface for all scorers for beam search or beam sample. +template +class IBeamScorer { + public: + virtual ~IBeamScorer() {} + + virtual void Initialize(AllocatorPtr& allocator, int sequence_length) = 0; + + virtual void Process(ISequences* sequences, + gsl::span& next_scores, + gsl::span& next_tokens, + gsl::span& next_indices) = 0; + + virtual void Finalize(ISequences* sequences, + gsl::span& final_beam_scores, + Tensor* output_sequences, + Tensor* output_sequence_scores) = 0; +}; + +template +struct HypothesisScore { + HypothesisScore(gsl::span& _hypothesis, T _score) + : hypothesis(_hypothesis), score(_score) {} + + gsl::span hypothesis; + T score; +}; + +template +class HypothesisScoreCompare { + public: + bool operator()(const HypothesisScore& a, const HypothesisScore& b) { + return a.score > b.score; + } +}; + +template +class BeamHypotheses { + public: + BeamHypotheses(int num_beams, T length_penalty, bool early_stopping); + + // Number of hypotheses + int Size() { return static_cast(beams_.size()); } + + // Add a new hypothesis + void Add(gsl::span& hypothesis, T sum_logprobs); + + bool IsDone(T best_sum_logprobs, int current_length); + + // Output results. Note that it will clear all beams. + void Output(int top_k, // number of sequences to return + int max_length, // max sequence length + gsl::span& sequences, // buffer filled with pad token ID, with shape (num_return_sequences, max_length) + gsl::span& sequences_scores); // buffer for sequence scores, with shape (num_return_sequences) + + private: + int num_beams_; + T length_penalty_; + bool early_stopping_; + T worst_score_; + std::priority_queue, std::vector>, HypothesisScoreCompare> beams_; // min-heap for top k +}; + +template +class BeamSearchScorer : public IBeamScorer { + public: + BeamSearchScorer(int batch_size, + int num_beams, + int max_length, + T length_penalty, + bool early_stopping, + int num_return_sequences, + int pad_token_id, + int eos_token_id); + + void Initialize(AllocatorPtr& allocator, int sequence_length) override; + + void Process(ISequences* sequences, + gsl::span& next_scores, + gsl::span& next_tokens, + gsl::span& next_indices) override; + + void Finalize(ISequences* sequences, + gsl::span& final_beam_scores, + Tensor* output_sequences, + Tensor* output_sequence_scores) override; + + bool IsDone(); + + gsl::span& GetNextScores() { return next_beam_scores_; } + gsl::span& GetNextTokens() { return next_beam_tokens_; } + gsl::span& GetNextIndices() { return next_beam_indices_; } + + private: + int batch_size_; + int num_beams_; + int max_length_; + int num_beam_hyps_to_keep_; + int pad_token_id_; + int eos_token_id_; + + // TODO: use ORT allocator to avoid allocating from heap directly + std::vector> beam_hyps; // List of batch result of beam search. Its shape is (batch_size) + + IAllocatorUniquePtr done_ptr_; // List of flags indicates whether each batch is finished or not. Its shape is (batch_size). + gsl::span done_; + + IAllocatorUniquePtr next_beam_scores_ptr_; + gsl::span next_beam_scores_; + + IAllocatorUniquePtr next_beam_tokens_ptr_; + gsl::span next_beam_tokens_; + + IAllocatorUniquePtr next_beam_indices_ptr_; + gsl::span next_beam_indices_; + + IAllocatorUniquePtr hypothesis_buffer_ptr_; // Allocated buffer to hold all hypotheses + gsl::span hypothesis_buffer_; // Span of the allocated buffer + size_t hypothesis_buffer_length_; // Total number of elements + int hypothesis_buffer_offset_; // Offset of avaiable buffer, or length of used buffer. +}; + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cpu/transformers/dump_tensor.cc b/onnxruntime/contrib_ops/cpu/transformers/dump_tensor.cc new file mode 100644 index 0000000000..8aa9e8907f --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/dump_tensor.cc @@ -0,0 +1,73 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "dump_tensor.h" +#include "core/platform/env.h" +#include "core/platform/env_var_utils.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +namespace dump_tensor_env_vars { +constexpr const char* kDumpBeamSearch = "ORT_DUMP_BEAM_SEARCH"; +} + +#ifdef NDEBUG +bool g_enable_tensor_dump = false; +#else +bool g_enable_tensor_dump = true; +#endif + +void DumpOrtValue(const char* name, const OrtValue& value) { + if (!g_enable_tensor_dump) + return; + std::cout << std::string(name) << "\n"; + const Tensor& tensor = value.Get(); + MLDataType dataType = tensor.DataType(); + if (dataType == DataTypeImpl::GetType()) { + DumpTensor(nullptr, tensor); + } else if (dataType == DataTypeImpl::GetType()) { + DumpTensor(nullptr, tensor); + } else if (dataType == DataTypeImpl::GetType()) { + DumpTensor(nullptr, tensor); + } else { + std::cout << "not float/int32/int64"; + } +} + +void ConfigureTensorDump() { + const auto parsed = ParseEnvironmentVariable(dump_tensor_env_vars::kDumpBeamSearch); + if (parsed.has_value()) { + g_enable_tensor_dump = *parsed; + } +} + +void DisableTensorDump() { + g_enable_tensor_dump = false; +} + +void DumpString(const char* name, int index, bool end_line) { + if (!g_enable_tensor_dump) + return; + std::cout << std::string(name) << "[" << index << "]"; + + if (end_line) { + std::cout << std::endl; + } +} + +void DumpString(const char* name, std::string value, bool end_line) { + if (!g_enable_tensor_dump) + return; + + std::cout << std::string(name) << "=" << value; + + if (end_line) { + std::cout << std::endl; + } +} + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cpu/transformers/dump_tensor.h b/onnxruntime/contrib_ops/cpu/transformers/dump_tensor.h new file mode 100644 index 0000000000..2f80410011 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/dump_tensor.h @@ -0,0 +1,148 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include +#include "core/framework/tensorprotoutils.h" + +#ifndef NDEBUG +//#define DEBUG_BEAM_SEARCH 1 // uncomment it for debugging beam search +#endif + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +#define MAX_ROW_OR_COLUMN 8 + +#define SKIP_IF_MORE_THAN(row_or_column_size, i, max_n, new_line) \ + if (row_or_column_size > max_n && i >= max_n / 2 && i + max_n / 2 < row_or_column_size) { \ + if (i == max_n / 2) { \ + std::cout << ", ..."; \ + if (new_line) \ + std::cout << std::endl; \ + } \ + continue; \ + } + +#define SKIP_IF_TOO_MANY(row_or_column_size, i, new_line) SKIP_IF_MORE_THAN(row_or_column_size, i, MAX_ROW_OR_COLUMN, new_line) + +extern bool g_enable_tensor_dump; // global variance to turn on/off dump + +template +void PrintValue(const T& value) { + if (std::is_floating_point::value) + std::cout << std::setprecision(8) << value; + else + std::cout << value; +} + +template +void DumpTensor(const char* name, const Tensor& tensor) { + if (!g_enable_tensor_dump) + return; + + if (nullptr != name) { + std::cout << std::string(name) << std::endl; + } + + const auto& shape = tensor.Shape(); + auto num_items = shape.Size(); + + if (num_items == 0) { + std::cout << "no data"; + return; + } + + size_t num_dims = shape.NumDimensions(); + size_t num_rows = 1; + if (num_dims > 1) { + num_rows = static_cast(shape[0]); + } + + size_t row_size = num_items / num_rows; + + auto data = tensor.DataAsSpan(); + + for (size_t row = 0; row < num_rows; ++row) { + SKIP_IF_TOO_MANY(num_rows, row, true); + std::cout << "[" << row << "]:"; + for (size_t i = 0; i < row_size; ++i) { + SKIP_IF_TOO_MANY(row_size, i, false); + + if (i > 0) + std::cout << ", "; + + PrintValue(data[row * row_size + i]); + } + std::cout << "\n"; + } + + std::cout << std::endl; +} + +void DumpOrtValue(const char* name, const OrtValue& value); + +template +void DumpTensor(const char* name, const T* tensor, int dim0, int dim1) { + if (!g_enable_tensor_dump) + return; + + if (nullptr != name) { + std::cout << std::string(name) << std::endl; + } + + for (int i = 0; i < dim0; i++) { + SKIP_IF_TOO_MANY(dim0, i, true); + std::cout << "[" << i << "]:"; + for (int j = 0; j < dim1; j++) { + SKIP_IF_TOO_MANY(dim1, j, false); + if (j > 0) + std::cout << ", "; + T value = tensor[i * dim1 + j]; + PrintValue(value); + } + std::cout << std::endl; + } +} + +void DumpString(const char* name, int index, bool end_line); + +void DumpString(const char* name, std::string value, bool end_line); + +template +void DumpTensor(const char* name, const T* tensor, int dim0, int dim1, int dim2) { + if (!g_enable_tensor_dump) + return; + + if (nullptr != name) { + std::cout << std::string(name) << std::endl; + } + + for (int i = 0; i < dim0; i++) { + SKIP_IF_TOO_MANY(dim0, i, true); + for (int j = 0; j < dim1; j++) { + SKIP_IF_TOO_MANY(dim1, j, true); + std::cout << "[" << i << "][" << j << "]:"; + for (int k = 0; k < dim2; k++) { + SKIP_IF_TOO_MANY(dim2, k, false); + if (k > 0) + std::cout << ", "; + T value = tensor[i * dim1 * dim2 + j * dim2 + k]; + PrintValue(value); + } + std::cout << std::endl; + } + std::cout << std::endl; + } + std::cout << std::endl; +} + +void ConfigureTensorDump(); + +void DisableTensorDump(); + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cpu/transformers/gpt_subgraph.cc b/onnxruntime/contrib_ops/cpu/transformers/gpt_subgraph.cc new file mode 100644 index 0000000000..80825c7e26 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/gpt_subgraph.cc @@ -0,0 +1,451 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// there's no way to use a raw pointer as the copy destination with std::copy_n +// (which gsl::copy uses with span::data() which returns a raw pointer) with the 14.11 toolset +// without generating a 4996 warning. going through an iterator is way too much overhead so turn off the warning. +#ifdef _MSC_VER +#pragma warning(push) +#pragma warning(disable : 4996) +#endif + +#include "core/framework/framework_common.h" +#include "core/framework/session_state.h" +#include "core/framework/tensorprotoutils.h" +#include "core/framework/utils.h" +#include "core/providers/cpu/tensor/utils.h" +#include "gsl/gsl" +#include "gpt_subgraph.h" +#include "dump_tensor.h" + +#ifdef _MSC_VER +#pragma warning(pop) +#endif + +using namespace ONNX_NAMESPACE; +using namespace onnxruntime::common; + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +GptSubgraph::GptSubgraph( + const onnxruntime::Node& node_in, + const std::string& attribute_name, + const GraphViewer& subgraph_in) + : node(node_in), attribute(attribute_name), subgraph(subgraph_in), allocator_(nullptr) { + num_implicit_inputs = static_cast(node.ImplicitInputDefs().size()); + + auto& subgraph_inputs = subgraph.GetInputs(); + auto& subgraph_outputs = subgraph.GetOutputs(); + + // inputs: input_ids, position_ids, attention_mask, past_0, past_1, ... + // outputs: logits, present_0, present_1, ... + num_subgraph_inputs = static_cast(subgraph_inputs.size()); + num_subgraph_outputs = static_cast(subgraph_outputs.size()); + + // CheckSubgraph will verify inputs and outputs later. + subgraph_input_names.reserve(num_subgraph_inputs); + for (int i = 0; i < num_subgraph_inputs; ++i) { + subgraph_input_names.push_back(subgraph_inputs[i]->Name()); + } + + subgraph_output_names.reserve(num_subgraph_outputs); + for (int i = 0; i < num_subgraph_outputs; ++i) { + subgraph_output_names.push_back(subgraph_outputs[i]->Name()); + } +} + +Status GptSubgraph::Validate(const std::vector& subgraph_inputs, + const std::vector& subgraph_outputs) { + ORT_RETURN_IF(num_subgraph_outputs <= 1, + "Invalid GPT-2 subgraph: number of outputs shall be larger than 1 (Need past state in inputs and outputs)."); + + ORT_RETURN_IF(num_subgraph_inputs != num_subgraph_outputs + 2, + "Invalid GPT-2 subgraph: number of inputs shall be number of outputs plus 2"); + + ORT_RETURN_IF(subgraph_inputs[0]->Name() != "input_ids", "subgraph input 0 shall be named as input_ids, got: ", + subgraph_inputs[0]->Name()); + ORT_RETURN_IF(subgraph_inputs[1]->Name() != "position_ids", "subgraph input 1 shall be named as position_ids, got: ", + subgraph_inputs[1]->Name()); + ORT_RETURN_IF(subgraph_inputs[2]->Name() != "attention_mask", "subgraph input 2 shall be named as attention_mask, got: ", + subgraph_inputs[2]->Name()); + ORT_RETURN_IF(subgraph_inputs[3]->Name() != "past_0", "subgraph input 3 shall be named as past_0, got: ", + subgraph_inputs[3]->Name()); + + // Past state shape is like (2, batch_size, 12, past_seq_len, 64). Here 12 and 64 are constants of num_heads and hidden_size/num_heads. + const ONNX_NAMESPACE::TensorShapeProto* past_shape = subgraph_inputs[3]->Shape(); + ORT_RETURN_IF(past_shape->dim_size() != 5, "subgraph past state is expected to have 5 dimension, got ", + past_shape->dim_size()); + + ORT_RETURN_IF(!past_shape->dim(0).has_dim_value() || past_shape->dim(0).dim_value() != 2, + "subgraph past state dimension 0 shall have length of 2"); + + ORT_RETURN_IF(!past_shape->dim(2).has_dim_value() || past_shape->dim(2).dim_value() <= 0, + "subgraph past state dimension 2 shall have a positive value for number of heads"); + + ORT_RETURN_IF(!past_shape->dim(4).has_dim_value() || past_shape->dim(4).dim_value() <= 0, + "subgraph past state dimension 4 shall have a positive value for hidden size per head"); + + // check subgraph outputs + ORT_RETURN_IF(subgraph_outputs[0]->Name() != "logits", "subgraph output 0 shall be named as logits, got: ", + subgraph_outputs[0]->Name()); + + ORT_RETURN_IF(subgraph_outputs[1]->Name() != "present_0", "subgraph input 1 shall be named as present_0, got: ", + subgraph_outputs[1]->Name()); + + // Logits shape is like (batch_size, seq_len, 50257). Here 50257 is the vocabulary size. + const ONNX_NAMESPACE::TensorShapeProto* logits_shape = subgraph_outputs[0]->Shape(); + ORT_RETURN_IF(logits_shape->dim_size() != 3, "subgraph logits output is expected to have 3 dimension, got ", + logits_shape->dim_size()); + + ORT_RETURN_IF(!logits_shape->dim(2).has_dim_value() || logits_shape->dim(2).dim_value() <= 0, + "subgraph past state dimension 2 shall have a positive value for vocabulary size"); + + // Save parameters related to the subgraph. + num_heads = static_cast(past_shape->dim(2).dim_value()); + head_size = static_cast(past_shape->dim(4).dim_value()); + vocab_size = static_cast(logits_shape->dim(2).dim_value()); + num_layers = static_cast(subgraph_outputs.size()) - 1; + + ORT_RETURN_IF(subgraph_inputs[0]->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT64, + "subgraph input 0 (input_ids) shall have int64 type"); + ORT_RETURN_IF(subgraph_inputs[1]->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT64, + "subgraph input 1 (position_ids) shall have int64 type"); + // TODO: support float16 + ORT_RETURN_IF(subgraph_inputs[2]->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT, + "subgraph input 2 (attention_mask) shall have float type"); + ORT_RETURN_IF(subgraph_inputs[3]->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT, + "subgraph input 3 (past_0) shall have float type"); + ORT_RETURN_IF(subgraph_outputs[0]->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT, + "subgraph output 0 (logits) shall have float type"); + ORT_RETURN_IF(subgraph_outputs[1]->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT, + "subgraph output 1 (present_0) shall have float type"); + + return Status::OK(); +} + +Status GptSubgraph::Setup(const SessionState& session_state, + const SessionState& subgraph_session_state) { + session_state_ = &session_state; + subgraph_session_state_ = &subgraph_session_state; + + std::vector feed_names; + feed_names.reserve(num_subgraph_inputs + num_implicit_inputs); + + // First, get the location of input_ids of current operator. + const auto& node_inputs = node.InputDefs(); + const OrtMemoryInfo& input_ids_location = utils::FindMemoryInfoForValue(session_state, node_inputs[0]->Name()); + + // position_ids, attention_mask, past_0, ... are created by this operator so the name doesn't matter. + // as we skip them when we call FindDevicesForValues, and default them to be in the same device as input_ids + feed_names.insert(feed_names.end(), subgraph_input_names.begin(), subgraph_input_names.end()); + + for (auto& entry : node.ImplicitInputDefs()) { + feed_names.push_back(entry->Name()); + } + + std::vector feed_locations; + feed_locations.resize(feed_names.size()); + + for (size_t i = 0, end = feed_names.size(); i < end; ++i) { + if (i >= subgraph_input_names.size()) { // implicit inputs + const auto& location = utils::FindMemoryInfoForValue(session_state, feed_names[i]); + feed_locations[i] = location.device; + } else { + feed_locations[i] = input_ids_location.device; + } + } + + std::unique_ptr ffm; + ORT_RETURN_IF_ERROR(FeedsFetchesManager::Create(feed_names, subgraph_output_names, + subgraph_session_state.GetOrtValueNameIdxMap(), ffm)); + ORT_RETURN_IF_ERROR(utils::InitializeFeedFetchCopyInfo(subgraph_session_state, *ffm)); + + // setup the locations where we want the subgraph output to end up on + std::vector fetch_locations; + fetch_locations.reserve(num_subgraph_outputs); + + // past state need to be where we can feed them in to the next iteration, so set the fetch location to match the feed location. + for (int i = 0; i < num_subgraph_outputs; ++i) { + fetch_locations.push_back(&input_ids_location); + } + + utils::FinalizeFeedFetchCopyInfo(*ffm, feed_locations, fetch_locations); + + feeds_fetches_manager_ = std::move(ffm); + + // Check subgraph only need once so put in Setup function. + auto& inputs = subgraph.GetInputs(); + auto& outputs = subgraph.GetOutputs(); + ORT_RETURN_IF_ERROR(Validate(inputs, outputs)); + + return Status::OK(); +} + +void GptSubgraph::CreateInitialFeeds( + const Tensor& input_ids, + const std::vector& implicit_inputs, + int num_beams, + int pad_token_id, + gsl::span& next_positions, + std::vector& feeds) { + ORT_ENFORCE(session_state_ != nullptr, "Setup must be called before CreateInitialFeeds"); + + // Subgraph inputs: + // input_ids: shape (B, S) wher B is batch size, and S is sequence length + // position_ids: shape (B, S) + // attention_mask: shape (B, P+S), where past_sequence_length (P) is 0 + // After expansion, their shapes will become (B, M*S), where M is num_beams. + + // Allocate subgraph inputs to be same device as input_ids + AllocatorPtr alloactor = session_state_->GetAllocator(input_ids.Location()); + + // Store allocator, which is needed in ExpandInputs. + allocator_ = alloactor; + + const TensorShape& input_ids_shape = input_ids.Shape(); + ORT_ENFORCE(input_ids_shape.NumDimensions() == 2); + const int64_t& batch_size = input_ids_shape[0]; + const int64_t& sequence_length = input_ids_shape[1]; + + // Allocate position_ids and attention_mask based on shape of input_ids + auto element_type = DataTypeImpl::GetType(); + + // input_ids for subgraph is int64, so we need Cast input_ids from int32 to int64. + OrtValue subgraph_input_ids; + // Current shape is (batch_size, sequence_length) + // Note that we will expand it to (batch_size * num_beams, sequence_length) later. + Tensor::InitOrtValue(element_type, input_ids_shape, alloactor, subgraph_input_ids); + + int64_t* subgraph_input_data = subgraph_input_ids.GetMutable()->MutableData(); + const int32_t* source = input_ids.Data(); + int64_t* target = subgraph_input_data; + for (int i = 0; i < batch_size; i++) { + for (int j = 0; j < sequence_length; j++, source++, target++) { + *target = static_cast(*source); + } + } + + OrtValue position_ids; + Tensor::InitOrtValue(element_type, input_ids_shape, alloactor, position_ids); + + OrtValue attention_mask; + auto mask_type = DataTypeImpl::GetType(); + Tensor::InitOrtValue(mask_type, input_ids_shape, alloactor, attention_mask); + + // Set attention mask to be 0 for pad tokens, and 1 for all other tokens. + // Set position id to be 0 for pad tokens, and cumulated sum of mask in a batch for other tokens + float* mask_data = attention_mask.GetMutable()->MutableData(); + int64_t* position_data = position_ids.GetMutable()->MutableData(); + source = input_ids.Data(); + float* mask = mask_data; + int64_t* position = position_data; + for (int i = 0; i < batch_size; i++) { + int64_t abs_position = 0; + for (int j = 0; j < sequence_length; j++, source++, mask++, position++) { + if (*source == pad_token_id) { + *mask = 0.0f; + *position = 0; + } else { + *mask = 1.0f; + *position = abs_position; + abs_position++; + } + } + for (int k = 0; k < num_beams; k++) { + next_positions[i * num_beams + k] = abs_position; + } + } + + // Initialize empty past state + auto past_type = DataTypeImpl::GetType(); + int64_t past_state_dims[] = {2, batch_size * num_beams, num_heads, 0, head_size}; + TensorShape past_shape(&past_state_dims[0], 5); + OrtValue empty_past; + Tensor::InitOrtValue(past_type, past_shape, allocator_, empty_past); + + // Expand (batch_size, sequence_length) to (batch_size * num_beams, sequence_length) for input_ids, position_ids and attention_mask + // TODO: Try expand inputs/outputs after first subgraph call instead. That may get better peroformance, but more complex to implement. + OrtValue expanded_input_ids = ExpandInputs(subgraph_input_ids, num_beams); + OrtValue expanded_position_ids = ExpandInputs(position_ids, num_beams); + OrtValue expanded_attention_mask = ExpandInputs(attention_mask, num_beams); + + // The ordering is the same as used in Setup + feeds.reserve(num_subgraph_inputs + num_implicit_inputs); + feeds.push_back(expanded_input_ids); + feeds.push_back(expanded_position_ids); + feeds.push_back(expanded_attention_mask); + + // The remaing inputs are past state. + for (int i = 3; i < num_subgraph_inputs; ++i) { + feeds.push_back(empty_past); + } + + // pass in implicit inputs + for (const auto* entry : implicit_inputs) { + feeds.push_back(*entry); + } +} + +OrtValue GptSubgraph::ExpandInputs(const OrtValue& input, int num_beams) const { + // Input shape (batch_size, sequence_length) + // Output shape (batch_size * num_beams, sequence_length) + if (num_beams == 1) + return input; + + const TensorShape& input_shape = input.Get().Shape(); + const int64_t& batch_size = input_shape[0]; + const int64_t& sequence_length = input_shape[1]; + + int64_t dims[] = {batch_size * num_beams, sequence_length}; + TensorShape expanded_shape(&dims[0], 2); + + OrtValue expanded; + MLDataType element_type = input.Get().DataType(); + Tensor::InitOrtValue(element_type, expanded_shape, allocator_, expanded); + + if (element_type == DataTypeImpl::GetType()) { + const int64_t* input_data = input.Get().Data(); + int64_t* expanded_data = expanded.GetMutable()->MutableData(); + int64_t* target = expanded_data; + for (int i = 0; i < batch_size; i++) { + for (int j = 0; j < num_beams; j++) { + memcpy(target, input_data + i * sequence_length, sizeof(int64_t) * sequence_length); + target += sequence_length; + } + } + } else if (element_type == DataTypeImpl::GetType()) { + const float* input_data = input.Get().Data(); + float* expanded_data = expanded.GetMutable()->MutableData(); + float* target = expanded_data; + for (int i = 0; i < batch_size; i++) { + for (int j = 0; j < num_beams; j++) { + memcpy(target, input_data + i * sequence_length, sizeof(float) * sequence_length); + target += sequence_length; + } + } + } + + return expanded; +} + +// TODO: support float16 +void GptSubgraph::PickPastState(const std::vector& last_outputs, + std::vector& next_inputs, + gsl::span& beam_indices) { + for (int i = 3; i < num_subgraph_inputs; ++i) { + const OrtValue& present = last_outputs[i - 2]; // shape is like (2, batch_beam_size, 12, past_seq_len, 64) + const TensorShape& past_shape = present.Get().Shape(); + + // Create a tensor with same shape. + OrtValue past; + auto past_type = DataTypeImpl::GetType(); + Tensor::InitOrtValue(past_type, past_shape, allocator_, past); + + auto block_size_per_beam = past_shape[2] * past_shape[3] * past_shape[4]; + auto past_key_size = past_shape[1] * past_shape[2] * past_shape[3] * past_shape[4]; + + gsl::span past_span = past.GetMutable()->MutableDataAsSpan(); + gsl::span present_span = present.Get().DataAsSpan(); + for (gsl::index j = 0; j < beam_indices.length(); j++) { + int64_t beam_index = beam_indices[j]; + gsl::span present_key = present_span.subspan(beam_index * block_size_per_beam, block_size_per_beam); + gsl::span present_value = present_span.subspan(past_key_size + beam_index * block_size_per_beam, block_size_per_beam); + + gsl::span past_key = past_span.subspan(j * block_size_per_beam, block_size_per_beam); + gsl::span past_value = past_span.subspan(past_key_size + j * block_size_per_beam, block_size_per_beam); + gsl::copy(present_key, past_key); + gsl::copy(present_value, past_value); +#ifdef DEBUG_BEAM_SEARCH + if (i == 3) // only dump past_0 + { + DumpString("past_key of beam", static_cast(j), true); + DumpTensor(nullptr, past_key.data(), 1, static_cast(block_size_per_beam)); + + DumpString("past_value of beam", static_cast(j), true); + DumpTensor(nullptr, past_value.data(), 1, static_cast(block_size_per_beam)); + } +#endif + } + + next_inputs[i] = past; + } +} + +Status GptSubgraph::UpdateFeeds( + const std::vector& last_outputs, + std::vector& next_inputs, + int current_length, + gsl::span& next_positions, + gsl::span beam_next_tokens, + gsl::span beam_indices, + int num_beams) { + // last_outputs: logits, present_0, present_1, ... + // next_inputs: input_ids, position_id, attention_mask, past_0, past_1 + + // The following updates inputs for subgraph + // TODO: Reuse buffer for input_ids and position_ids to reduce memory allocation. + + // 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); + auto element_type = DataTypeImpl::GetType(); + OrtValue input_ids; + Tensor::InitOrtValue(element_type, input_ids_shape, allocator_, input_ids); + int64_t* input_ids_data = input_ids.GetMutable()->MutableData(); + for (int i = 0; i < batch_beam_size; i++) { + input_ids_data[i] = beam_next_tokens[i]; + } + next_inputs[0] = input_ids; + + // Update position IDs + OrtValue position_ids; + Tensor::InitOrtValue(element_type, input_ids_shape, allocator_, position_ids); + int64_t* position_data = position_ids.GetMutable()->MutableData(); + for (int i = 0; i < batch_beam_size; i++) { + position_data[i] = next_positions[i]; + next_positions[i]++; + } + next_inputs[1] = position_ids; + + // Update attention mask + const OrtValue& old_mask = next_inputs[2]; + const float* old_mask_data = old_mask.Get().Data(); + int64_t mask_dims[] = {batch_beam_size, current_length}; + TensorShape mask_shape(&mask_dims[0], 2); + OrtValue attention_mask; + auto mask_type = DataTypeImpl::GetType(); + Tensor::InitOrtValue(mask_type, mask_shape, allocator_, attention_mask); + float* mask_data = attention_mask.GetMutable()->MutableData(); + for (int i = 0; i < batch_beam_size; i++) { + for (int j = 0; j < current_length - 1; j++) { + mask_data[i * current_length + j] = old_mask_data[i * (current_length - 1) + j]; + } + mask_data[i * current_length + current_length - 1] = 1.0f; + } + next_inputs[2] = attention_mask; + +#ifdef DEBUG_BEAM_SEARCH + DumpOrtValue("input_ids", input_ids); + DumpOrtValue("position_ids", position_ids); + DumpOrtValue("attention_mask", attention_mask); +#endif + + // Update past state + if (num_beams == 1) { + // feed present_* output to past_* inputs one by one + for (int i = 3; i < num_subgraph_inputs; ++i) { + next_inputs[i] = last_outputs[i - 2]; + } + } else { + PickPastState(last_outputs, next_inputs, beam_indices); + } + + return Status::OK(); +} + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/transformers/gpt_subgraph.h b/onnxruntime/contrib_ops/cpu/transformers/gpt_subgraph.h new file mode 100644 index 0000000000..75d879c7fe --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/gpt_subgraph.h @@ -0,0 +1,82 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +//#include +#include "gsl/gsl" +#include "core/framework/allocator.h" +#include "core/framework/session_state.h" +#include "core/framework/feeds_fetches_manager.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +// A class for GPT-2 subgraph inputs and outputs preparation. +struct GptSubgraph { + GptSubgraph( + const onnxruntime::Node& node_in, + const std::string& attribute_name, + const GraphViewer& subgraph_in); + + const onnxruntime::Node& node; // node that contains the subgraph + const std::string& attribute; // attribute of th node that contains the subgraph. Not used yet. + const GraphViewer& subgraph; // the subgraph + + int num_implicit_inputs; + + int num_subgraph_inputs; // same as subgraph_input_names.size(), keep it for convenience. + int num_subgraph_outputs; // same as subgraph_output_names.size() + + std::vector subgraph_input_names; + std::vector subgraph_output_names; + + // Parameters deduced from the subgraph + int num_heads; + int head_size; + int vocab_size; + int num_layers; + + // Setup exectuion + Status Setup(const SessionState& session_state, + const SessionState& subgraph_session_state); + + // Create inputs for first inference of subgraph. + void CreateInitialFeeds( + const Tensor& input_ids, + const std::vector& implicit_inputs, + int num_beams, + int pad_token_id, + gsl::span& next_positions, + std::vector& feeds); + + Status UpdateFeeds( + const std::vector& last_outputs, + std::vector& next_inputs, + int current_length, + gsl::span& next_positions, + gsl::span beam_next_tokens, + gsl::span beam_indices, + int num_beams); + + FeedsFetchesManager* GetFeedsFetchesManager() const { return feeds_fetches_manager_.get(); } + + protected: + Status Validate(const std::vector& subgraph_inputs, + const std::vector& subgraph_outputs); + + OrtValue ExpandInputs(const OrtValue& input, int num_beams) const; + + void PickPastState(const std::vector& last_outputs, + std::vector& next_inputs, + gsl::span& beam_indices); + + AllocatorPtr allocator_; + const SessionState* session_state_; + const SessionState* subgraph_session_state_; + std::unique_ptr feeds_fetches_manager_; +}; + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/transformers/logits_processor.cc b/onnxruntime/contrib_ops/cpu/transformers/logits_processor.cc new file mode 100644 index 0000000000..10ef8b6f69 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/logits_processor.cc @@ -0,0 +1,192 @@ +#include +#include "logits_processor.h" +#include "dump_tensor.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +template +gsl::span NextTokenScores::GetScores(int batch_beam_index) { + assert(batch_beam_index >= 0 && batch_beam_index < batch_beam_size); + return scores.subspan(batch_beam_index * vocab_size, vocab_size); +} + +template +void NextTokenScores::SetScore(int token_id, T score) { + assert(token_id >= 0 && token_id < vocab_size); + for (int i = 0; i < batch_beam_size; i++) { + scores[i * vocab_size + token_id] = score; + } +} + +#ifdef DEBUG_BEAM_SEARCH +template +void DumpScores(const char* name, gsl::span& scores) { + DumpString(name, 0, true); + ORT_UNUSED_PARAMETER(scores); +} +#endif + +// Interface for all scorers for beam search or beam sample. +template +MinLengthLogitsProcessor::MinLengthLogitsProcessor(int min_length, int eos_token_id) + : min_length_(min_length), eos_token_id_(eos_token_id) {} + +template +void MinLengthLogitsProcessor::Process(const ISequences* sequences, + NextTokenScores& next_token_scores) { + if (sequences->GetSequenceLength() < min_length_) { + next_token_scores.SetScore(eos_token_id_, std::numeric_limits::lowest()); + } + +#ifdef DEBUG_BEAM_SEARCH + DumpScores("MinLengthLogitsProcessor", next_token_scores.scores); +#endif +} + +template +RepetitionPenaltyLogitsProcessor::RepetitionPenaltyLogitsProcessor(float penalty) : penalty_(penalty) { +} + +template +void RepetitionPenaltyLogitsProcessor::Process(const ISequences* sequences, + NextTokenScores& next_token_scores) { + const int batch_beam_size = next_token_scores.batch_beam_size; + for (int i = 0; i < batch_beam_size; i++) { + gsl::span beam_token_scores = next_token_scores.GetScores(i); + gsl::span sequence = sequences->GetSequence(i); + + // Find unique word IDs in sequence. + std::unordered_set unique_word_ids; + for (const auto& word_id : sequence) { + unique_word_ids.insert(word_id); + } + + for (const int64_t word_id : unique_word_ids) { + T score = beam_token_scores[word_id]; + + // If score < 0, then repetition penalty > 1.0 has to multiplied to reduce the previous token probability, + // This assumes that scores are either positive (like ctrl) or negative (like GPT-2), but not a mixture. + beam_token_scores[word_id] = (score < 0 ? score * penalty_ : score / penalty_); + } + } + +#ifdef DEBUG_BEAM_SEARCH + DumpScores("RepetitionPenaltyLogitsProcessor", next_token_scores.scores); +#endif +} + +template +NoRepeatNGramLogitsProcessor::NoRepeatNGramLogitsProcessor(int ngram_size) : ngram_size_(ngram_size) { +} + +template +void NoRepeatNGramLogitsProcessor::Process(const ISequences* sequences, + NextTokenScores& next_token_scores) { + if (ngram_size_ == 0 || ngram_size_ > sequences->GetSequenceLength()) { + return; + } + + const gsl::index prefix_length = static_cast(ngram_size_ - 1); + int batch_beam_size = next_token_scores.batch_beam_size; + + for (int i = 0; i < batch_beam_size; i++) { + gsl::span beam_token_scores = next_token_scores.GetScores(i); + gsl::span sequence = sequences->GetSequence(i); + + gsl::span prefix = sequence.subspan(sequence.length() - prefix_length); + ORT_ENFORCE(prefix.length() == prefix_length); + + std::unordered_set blocked_word_ids; + for (int j = 0; j <= static_cast(sequence.length()) - ngram_size_; j++) { + // Here we use naive algorithm for matching. The complexity is O(batch_beam_size * ngram_size * sequence_length) + // TODO: build N-Gram index (hash table with prefix of length NGram - 1 as key, and list of last word of NGram as value) for fast matching. + if (ngram_size_ == 1 || prefix == sequence.subspan(j, prefix_length)) { + blocked_word_ids.insert(sequence[j + prefix_length]); + } + } + + for (const int64_t word_id : blocked_word_ids) { + beam_token_scores[word_id] = std::numeric_limits::lowest(); + } + } + +#ifdef DEBUG_BEAM_SEARCH + DumpScores("NoRepeatNGramLogitsProcessor", next_token_scores.scores); +#endif +} + +template +VocabMaskLogitsProcessor::VocabMaskLogitsProcessor(const gsl::span& vocab_mask) : vocab_mask_(vocab_mask) { +} + +template +void VocabMaskLogitsProcessor::Process(const ISequences* /*sequences*/, + NextTokenScores& next_token_scores) { + assert(!vocab_mask_.empty()); + + // Process vocabulary mask and set tokens with mask value 0 to -inf. + T* p = next_token_scores.scores.data(); + // next_token_scores shape (batch_size * num_beams, vocab_size) + // vocab_mask shape (vocab_size). TODO: support shape (batch_size, vocab_size) + for (int i = 0; i < next_token_scores.batch_beam_size; i++) { + for (int j = 0; j < next_token_scores.vocab_size; j++, p++) { + if (vocab_mask_[j] == 0) { + *p = std::numeric_limits::lowest(); + } + } + } + +#ifdef DEBUG_BEAM_SEARCH + DumpScores("VocabMaskLogitsProcessor", next_token_scores.scores); +#endif +} + +template +void LogitsProcessorList::Init(const BeamSearchParameters& parameters) { + processor_list_.clear(); + + if (parameters.repetition_penalty != 1.0f) { // 1.0 means no penalty + repetition_penalty_processor_ = std::make_unique>(parameters.repetition_penalty); + processor_list_.push_back(repetition_penalty_processor_.get()); + } + + if (parameters.no_repeat_ngram_size > 0) { + no_repeat_ngram_processor_ = std::make_unique>(parameters.no_repeat_ngram_size); + processor_list_.push_back(no_repeat_ngram_processor_.get()); + } + + if (!parameters.vocab_mask.empty()) { + vocab_mask_processor_ = std::make_unique>(parameters.vocab_mask); + processor_list_.push_back(vocab_mask_processor_.get()); + } + + if (parameters.min_length > 0) { + min_length_processor_ = std::make_unique>(parameters.min_length, parameters.eos_token_id); + processor_list_.push_back(min_length_processor_.get()); + } + + batch_beam_size_ = parameters.BatchBeamSize(); + vocab_size_ = parameters.vocab_size; +} + +template +void LogitsProcessorList::Process(const ISequences* sequences, + gsl::span& next_token_scores) { + NextTokenScores input_scores = {next_token_scores, batch_beam_size_, vocab_size_}; + for (size_t i = 0; i < processor_list_.size(); i++) { + processor_list_[i]->Process(sequences, input_scores); + } +} + +// Instantiation +template class MinLengthLogitsProcessor; +template class RepetitionPenaltyLogitsProcessor; +template class NoRepeatNGramLogitsProcessor; +template class VocabMaskLogitsProcessor; +template class LogitsProcessorList; + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h b/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h new file mode 100644 index 0000000000..78fe9acf63 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h @@ -0,0 +1,99 @@ +#pragma once +#include "sequences.h" +#include "beam_search_parameters.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +template +struct NextTokenScores { + gsl::span& scores; + int batch_beam_size; + int vocab_size; + + gsl::span GetScores(int batch_beam_index); + + void SetScore(int token_id, T score); +}; + +// Interface for all scorers for beam search or beam sample. +template +class ILogitsProcessor { + public: + virtual ~ILogitsProcessor() {} + + virtual void Process(const ISequences* sequences, + NextTokenScores& next_token_scores) = 0; +}; + +template +class MinLengthLogitsProcessor : public ILogitsProcessor { + public: + MinLengthLogitsProcessor(int min_length, int eos_token_id); + + void Process(const ISequences* sequences, + NextTokenScores& next_token_scores) override; + + private: + int min_length_; + int eos_token_id_; +}; + +template +class RepetitionPenaltyLogitsProcessor : public ILogitsProcessor { + public: + RepetitionPenaltyLogitsProcessor(float penalty); + + void Process(const ISequences* sequences, + NextTokenScores& next_token_scores) override; + + private: + float penalty_; +}; + +template +class NoRepeatNGramLogitsProcessor : public ILogitsProcessor { + public: + NoRepeatNGramLogitsProcessor(int ngram_size); + + void Process(const ISequences* sequences, + NextTokenScores& next_token_scores) override; + + private: + int ngram_size_; +}; + +template +class VocabMaskLogitsProcessor : public ILogitsProcessor { + public: + VocabMaskLogitsProcessor(const gsl::span& vocab_mask); + + void Process(const ISequences* sequences, + NextTokenScores& next_token_scores) override; + + private: + gsl::span vocab_mask_; +}; + +template +class LogitsProcessorList { +public: + LogitsProcessorList() = default ; + void Init(const BeamSearchParameters& parameters); + void Process(const ISequences* sequences, gsl::span& next_token_scores); + +private: + int batch_beam_size_; + int vocab_size_; + std::vector*> processor_list_; + + std::unique_ptr> repetition_penalty_processor_; + std::unique_ptr> no_repeat_ngram_processor_; + std::unique_ptr> vocab_mask_processor_; + std::unique_ptr> min_length_processor_; +}; + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cpu/transformers/sequences.cc b/onnxruntime/contrib_ops/cpu/transformers/sequences.cc new file mode 100644 index 0000000000..a9c70ef410 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/sequences.cc @@ -0,0 +1,76 @@ +#include "sequences.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +void Sequences::Init(AllocatorPtr allocator, const OrtValue& input_ids, int batch_beam_size, int sequence_length, int max_length) { + size_t sequences_size = SafeInt(batch_beam_size) * max_length; + size_t buffer_size = sequences_size + sequences_size; + gsl::span buffer = AllocateBuffer(allocator, sequences_space_buffer_, buffer_size, true, static_cast(0)); + + sequences[0] = buffer.subspan(0, sequences_size); + sequences[1] = buffer.subspan(sequences_size); + + // Copy input_ids to sequences[0]. + gsl::span input = input_ids.Get().DataAsSpan(); + gsl::span output = sequences[0]; + for (int i = 0; i < batch_beam_size; i++) { + gsl::span source = input.subspan(i * sequence_length, sequence_length); + gsl::span target = output.subspan(i * max_length, sequence_length); + gsl::copy(source, target); + } + current_sequences_buffer = 0; + + batch_beam_size_ = batch_beam_size; + max_length_ = max_length; + current_length_ = sequence_length; +} + +gsl::span Sequences::GetSequence(int beam_index) const { + gsl::span buffer(sequences[current_sequences_buffer].data(), sequences[current_sequences_buffer].size()); + gsl::span sequence = buffer.subspan(beam_index * max_length_, current_length_); + return sequence; +} + +int Sequences::GetSequenceLength() const { + return current_length_; +} + +void Sequences::PrintSequences() { +#ifdef DEBUG_BEAM_SEARCH + for (int i = 0; i < batch_beam_size_; i++) { + gsl::span sequence = GetSequence(i); + DumpString("sequences", i, false); + DumpTensor(nullptr, sequence.data(), 1, current_length_); + } +#endif +} + +void Sequences::AppendNextTokenToSequences( + gsl::span& beam_indices, + gsl::span& beam_next_tokens) { + gsl::span input(sequences[current_sequences_buffer].data(), sequences[current_sequences_buffer].size()); + gsl::span output = sequences[1 - current_sequences_buffer]; + + for (int i = 0; i < batch_beam_size_; i++) { + int beam_index = static_cast(beam_indices[i]); + gsl::span source = input.subspan(beam_index * max_length_, current_length_); + gsl::span target = output.subspan(i * max_length_, current_length_); + gsl::copy(source, target); + } + + // Append next token to each beam. + for (int i = 0; i < batch_beam_size_; i++) { + output[i * max_length_ + current_length_] = beam_next_tokens[i]; + } + + ++current_length_; + + // Rotate buffer for next round. + current_sequences_buffer = 1 - current_sequences_buffer; +} + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cpu/transformers/sequences.h b/onnxruntime/contrib_ops/cpu/transformers/sequences.h new file mode 100644 index 0000000000..9950330dae --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/sequences.h @@ -0,0 +1,80 @@ +#pragma once + +#include "gsl/gsl" +#include "core/common/safeint.h" +#include "core/framework/allocator.h" +#include "core/framework/ort_value.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +class ISequences { + public: + virtual ~ISequences() {} + virtual gsl::span GetSequence(int beam_index) const = 0; + virtual int GetSequenceLength() const = 0; +}; + +// This class keeps track of sequences generated. +class Sequences : public ISequences { + public: + Sequences() {} + + // Initialize the sequence with initial input_ids and related parameters. + void Init(AllocatorPtr allocator, const OrtValue& input_ids, int batch_beam_size, int sequence_length, int max_length); + + // Returns a sequence of word IDs for a given beam index ( beam_index < batch_beam_size). + gsl::span GetSequence(int beam_index) const override; + + // Returns current sequence length. + int GetSequenceLength() const override; + + // Print the sequences to StdOut in debug mode + void PrintSequences(); + + // Select sequences based on beam indices, then append next token to selected sequences. + void AppendNextTokenToSequences( + gsl::span& beam_indices, + gsl::span& beam_next_tokens); + + private: + gsl::span sequences_space; // shape (2, batch_size, num_beams, max_seq_length) + BufferUniquePtr sequences_space_buffer_; + + // Two buffers of shape (batch_size, num_beams, max_seq_length) to store sequences. + // At each time, there is only one buffer is active. The other one will be active in next token. + // Each AppendNextTokenToSequences call will trigger a rotation of active buffer. + gsl::span sequences[2]; + + // Index (either 0 or 1) of two buffers that is currently is active. + int current_sequences_buffer; + + int batch_beam_size_; + int max_length_; + int current_length_; +}; + +template +gsl::span AllocateBuffer(AllocatorPtr allocator, + BufferUniquePtr& buffer, + size_t elements, + bool fill = false, + T fill_value = T{}) { + size_t bytes = SafeInt(sizeof(T)) * elements; + void* data = allocator->Alloc(bytes); + BufferUniquePtr temp_buffer(data, BufferDeleter(allocator)); + buffer = std::move(temp_buffer); + T* first = reinterpret_cast(buffer.get()); + auto span = gsl::make_span(first, elements); + + if (fill) { + std::fill_n(first, elements, fill_value); + } + + return span; +} + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 74d3b1d74d..6eb90e969c 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -568,6 +568,139 @@ void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx } } +bool ParseScalar(const TensorProto* initializer, int& value) { + std::vector parsed_data; + if (initializer->data_type() == TensorProto::INT32) { + const auto& data = ParseData(initializer); + parsed_data.insert(parsed_data.end(), data.begin(), data.end()); + + if (parsed_data.size() == 1) { + value = parsed_data[0]; + return true; + } + } + + return false; +} + +void BeamSearchShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { + // Type inference + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); + + if (ctx.getNumOutputs() > 1) { + // Here we assume that the third output exist only if second output exists. + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 5, 1); + if (ctx.getNumOutputs() > 2) { + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 5, 2); + } + } + + // Shape inference + // input 0 (input_ids) shape: (batch_size, sequence_length) + // output 0 (sequences) shape: (batch_size, num_return_sequences, max_length) + // output 1 (sequences_scores) shape: (batch_size, num_return_sequences) + // output 2 (scores) shape: (max_length - sequence_length, batch_size, num_beams, vocab_size) + if (!hasInputShape(ctx, 0)) { + return; + } + auto& input_ids_shape = getInputShape(ctx, 0); + auto& input_ids_dims = input_ids_shape.dim(); + if (input_ids_dims.size() != 2) { + fail_shape_inference("Inputs 0 shall be 2 dimensions"); + } + if (!(input_ids_dims[0].has_dim_value() && input_ids_dims[1].has_dim_value())) { + return; + } + + int64_t batch_size = input_ids_dims[0].dim_value(); + int64_t sequence_length = input_ids_dims[1].dim_value(); + + const auto max_length = ctx.getInputData(1); + const auto num_beams = ctx.getInputData(3); + const auto num_return_sequences = ctx.getInputData(4); + if (num_beams == nullptr || max_length == nullptr || num_return_sequences == nullptr) { // not initializer + return; + } + + int max_length_value = 0; + if (!ParseScalar(max_length, max_length_value) || max_length_value <= 0) { + fail_shape_inference("Failed to parse max_length or it is not positive integer scalar"); + } + + int num_beams_value = 0; + if (!ParseScalar(num_beams, num_beams_value) || num_beams_value <= 0) { + fail_shape_inference("Failed to parse num_beams or it is not positive integer scalar"); + } + + int num_return_sequences_value = 0; + if (!ParseScalar(num_return_sequences, num_return_sequences_value) || num_return_sequences_value <= 0) { + fail_shape_inference("Failed to parse num_return_sequences or it is not positive integer scalar"); + } + + ONNX_NAMESPACE::TensorShapeProto sequences_shape; + sequences_shape.add_dim()->set_dim_value(batch_size); + sequences_shape.add_dim()->set_dim_value(num_return_sequences_value); + sequences_shape.add_dim()->set_dim_value(max_length_value); + updateOutputShape(ctx, 0, sequences_shape); + + if (ctx.getNumOutputs() > 1) { + ONNX_NAMESPACE::TensorShapeProto sequences_scores_shape; + sequences_shape.add_dim()->set_dim_value(batch_size); + sequences_shape.add_dim()->set_dim_value(num_return_sequences_value); + updateOutputShape(ctx, 1, sequences_scores_shape); + + if (ctx.getNumOutputs() > 2) { + ONNX_NAMESPACE::TensorShapeProto scores_shape; + scores_shape.add_dim()->set_dim_value(max_length_value - sequence_length); + scores_shape.add_dim()->set_dim_value(batch_size); + scores_shape.add_dim()->set_dim_value(num_beams_value); + scores_shape.add_dim(); // vocab_size is unknown + updateOutputShape(ctx, 2, scores_shape); + } + } +} + +void RegisterTextGenerationSchemas() { + ONNX_CONTRIB_OPERATOR_SCHEMA(BeamSearch) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetDoc("Beam Search for text generation. Supports GPT-2 decoder.") + .Attr("eos_token_id", "The id of the end-of-sequence token", AttributeProto::INT) + .Attr("pad_token_id", "The id of the padding token", AttributeProto::INT) + .Attr("no_repeat_ngram_size", "no repeat ngrams size", AttributeProto::INT, static_cast(0)) + .Attr("early_stopping", "early stop or not", AttributeProto::INT, static_cast(0)) + .Attr( + "body", + "The GPT-2 subgraph with input_ids, position_ids, attention_mask, past_0, past_1, ... as inputs, and logits, present_0, present_1, ... as output", + AttributeProto::GRAPH) + .Input(0, "input_ids", "The sequence used as a prompt for the generation. Shape is (batch_size, sequence_length)", "I") + .Input(1, "max_length", "The maximum length of the sequence to be generated. Shape is (1)", "I") + .Input(2, "min_length", "The minimum length below which the score of eos_token_id is set to -Inf. Shape is (1)", "I", OpSchema::Optional) + .Input(3, "num_beams", "Number of beams for beam search. 1 means no beam search. Shape is (1)", "I") + .Input(4, "num_return_sequences", "The number of returned sequences in the batch. Shape is (1)", "I") + .Input(5, "temperature", "The value used to module the next token probabilities. Accepts value > 0.0. Shape is (1)", "T") + .Input(6, "length_penalty", + "Exponential penalty to the length. Default value 1.0 means no penalty." + "Value > 1.0 encourages longer sequences, while values < 1.0 produces shorter sequences." + "Shape is (1,)", + "T", OpSchema::Optional) + .Input(7, "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(8, "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) + .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", + "Processed beam scores for each vocabulary token at each generation step." + "Beam scores consisting of log softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this beam." + "Shape is (max_length - sequence_length, batch_size, num_beams, vocab_size)", + "T", OpSchema::Optional) + .TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float tensors.") + .TypeConstraint("I", {"tensor(int32)"}, "Constrain to integer types") + .TypeConstraint("M", {"tensor(int32)"}, "Constrain mask to integer types") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + BeamSearchShapeInference(ctx); + }); +} + void RegisterBertSchemas() { static const char* Attention_ver1_doc = R"DOC( Multi-Head Self Attention that can be either unidirectional (like GPT-2) or bidirectional (like BERT). @@ -750,8 +883,7 @@ Some boolean parameters are passed by runtime input for generic purpose .TypeConstraint("B", {"tensor(bool)"}, "Constrain key_padding_mask to bool tensors.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { DecoderAttentionTypeAndShapeInference(ctx); - }); - + }); static const char* EmbedLayerNormalization_ver1_doc = R"DOC( EmbedLayerNormalization is the fusion of embedding layer in BERT model, with optional mask processing. @@ -3146,6 +3278,7 @@ It's an extension of Gelu. It takes the sum of input A and bias input B as the i RegisterNhwcSchemas(); RegisterBertSchemas(); + RegisterTextGenerationSchemas(); #ifdef BUILD_MS_EXPERIMENTAL_OPS onnxruntime::signal::RegisterSignalSchemas(); diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 03acfec47f..3e81f43850 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -2587,6 +2587,15 @@ Status Graph::VerifyNodeAndOpMatch(const ResolveOptions& options) { } } + // verify subgraphs + for (auto node_index : nodes_in_topological_order_) { + auto& node = *GetNode(node_index); + for (auto& entry : node.GetAttributeNameToMutableSubgraphMap()) { + Graph* subgraph = entry.second; + ORT_RETURN_IF_ERROR(subgraph->VerifyNodeAndOpMatch(options)); + } + } + return Status::OK(); } diff --git a/onnxruntime/core/providers/cpu/math/top_k.cc b/onnxruntime/core/providers/cpu/math/top_k.cc index db0b5086e6..f9b2c7d736 100644 --- a/onnxruntime/core/providers/cpu/math/top_k.cc +++ b/onnxruntime/core/providers/cpu/math/top_k.cc @@ -370,6 +370,56 @@ static Status TopKImpl(OpKernelContext* p_op_kernel_context, const Tensor* input return Status::OK(); } +// Wrapper over core TopK implementation +template +Status GetTopK(const Tensor* input, const int axis, const unsigned k, bool largest, bool sorted, + AllocatorPtr allocator, + onnxruntime::concurrency::ThreadPool* threadpool, + std::unique_ptr& output_values, + std::unique_ptr& output_indices) { + const TensorShape& input_shape = input->Shape(); + + // Will return axis_ as is if positive or fixes it in case it is negative + const auto axis_parsed = HandleNegativeAxis(axis, static_cast(input_shape.NumDimensions())); + + // Check to ensure k is within the bounds of what is available in that specific axis + if (input_shape[axis_parsed] < k) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "k argument [", k, + "] should not be greater than specified axis dim value [", input_shape[axis_parsed], "]"); + } + + // Resize output tensors to be the same shape as the input except + // for the specified dimension ((i.e.) axis_parsed), which will be of size k. E.x. for an input tensor + // of shape [3, 4, 5] and k=2 with axis_parsed=1, both of the outputs will be shape [3, 2, 5] + TensorShape output_shape = input_shape; + output_shape[axis_parsed] = k; + + output_values = Tensor::Create(input->DataType(), output_shape, allocator); + output_indices = Tensor::Create(DataTypeImpl::GetType(), output_shape, allocator); + + // no-op - no output buffers to fill - return silently + if (k == 0) { + return Status::OK(); + } + + if (largest) { + FindTopKElements>(input, input_shape, output_values.get(), output_indices.get(), output_shape, k, sorted, + gsl::narrow_cast(axis_parsed), threadpool); + } else { + FindTopKElements>(input, input_shape, output_values.get(), output_indices.get(), output_shape, k, sorted, + gsl::narrow_cast(axis_parsed), threadpool); + } + + return Status::OK(); +} + +// explicit instantiation +template Status GetTopK(const Tensor* input, const int axis, const unsigned k, bool largest, bool sorted, + AllocatorPtr allocator, + onnxruntime::concurrency::ThreadPool* threadpool, + std::unique_ptr& output_values, + std::unique_ptr& output_indices); + // Opset ver - 1 to 9 static void TopkOpset9ConstructorCommon(const OpKernelInfo& op_kernel_info, int& axis, unsigned int& k) { diff --git a/onnxruntime/core/providers/cpu/math/top_k.h b/onnxruntime/core/providers/cpu/math/top_k.h index 6cea1c7d6e..596f10a8e6 100644 --- a/onnxruntime/core/providers/cpu/math/top_k.h +++ b/onnxruntime/core/providers/cpu/math/top_k.h @@ -19,4 +19,11 @@ class TopK final : public OpKernel { bool largest_; // opset-11 only bool sorted_; // opset-11 only }; + +template +Status GetTopK(const Tensor* input, const int axis, const unsigned k, bool largest, bool sorted, + AllocatorPtr allocator, + onnxruntime::concurrency::ThreadPool* threadpool, + std::unique_ptr& output_values, + std::unique_ptr& output_indices); } // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/python/tools/transformers/convert_beam_search.py b/onnxruntime/python/tools/transformers/convert_beam_search.py new file mode 100644 index 0000000000..3f251fb9ad --- /dev/null +++ b/onnxruntime/python/tools/transformers/convert_beam_search.py @@ -0,0 +1,459 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import os +import time +import onnx +import logging +import argparse +from pathlib import Path +from onnx import helper +import numpy as np +from typing import List +import torch +from transformers import GPT2Config +from gpt2_helper import PRETRAINED_GPT2_MODELS +from convert_to_onnx import main as convert_gpt2_to_onnx +from benchmark_helper import Precision +""" +This converts GPT2 model to onnx with beam search operator. + +Examples: + python convert_beam_search.py -m gpt2 --gpt2_onnx .\onnx_models\gpt2_past_fp32.onnx --output .\onnx_models\gpt2_beam_search.onnx --output_sequences_scores +""" + +config: GPT2Config = None + +logger = logging.getLogger('') + + +def parse_arguments(argv=None): + parser = argparse.ArgumentParser() + + parser.add_argument('-m', + '--model_name_or_path', + required=True, + type=str, + help='Model path, or pretrained model name in the list: ' + ', '.join(PRETRAINED_GPT2_MODELS)) + + parser.add_argument('--cache_dir', + required=False, + type=str, + default=os.path.join('.', 'cache_models'), + help='Directory to cache pre-trained models') + + parser.add_argument('--gpt2_onnx', + required=True, + type=str, + help='Output directory for GPT-2 onnx model, or model path ends with .onnx') + + parser.add_argument('--output', + required=False, + type=str, + help='Output directory for beam search model, or model path ends with .onnx') + + parser.add_argument("-p", + "--precision", + required=False, + type=Precision, + default=Precision.FLOAT32, + choices=[Precision.FLOAT32, Precision.FLOAT16], + help="Precision of model to run. fp32 for full precision, fp16 for half or mixed precision") + + parser.add_argument('--use_gpu', required=False, action='store_true', help="use GPU for inference") + parser.set_defaults(use_gpu=False) + + parser.add_argument('-e', '--use_external_data_format', required=False, action='store_true') + parser.set_defaults(use_external_data_format=False) + + parser.add_argument('--disable_parity', required=False, action='store_true', help="do not run parity test") + parser.set_defaults(disable_parity=False) + + parser.add_argument('--total_runs', + required=False, + type=int, + default=1, + help='Number of times of inference for latency measurement') + + beam_search_group = parser.add_argument_group("beam search options") + + beam_search_group.add_argument('--output_sequences_scores', + required=False, + action='store_true', + help="output sequences scores") + beam_search_group.set_defaults(output_sequences_scores=False) + + beam_search_group.add_argument('--output_token_scores', + required=False, + action='store_true', + help="output token scores") + beam_search_group.set_defaults(output_token_scores=False) + + beam_search_group.add_argument('--early_stopping', required=False, action='store_true') + beam_search_group.set_defaults(early_stopping=False) + + beam_search_group.add_argument('--min_length', type=int, required=False, default=1, help='Min sequence length') + + beam_search_group.add_argument('--max_length', type=int, required=False, default=50, help='Max sequence length') + + beam_search_group.add_argument('--no_repeat_ngram_size', + type=int, + required=False, + default=0, + help='No repeat ngram size') + + beam_search_group.add_argument('--num_beams', type=int, required=False, default=4, help='Beam size') + + beam_search_group.add_argument('--num_return_sequences', + type=int, + required=False, + default=1, + help='Number of return sequence <= num_beams') + + beam_search_group.add_argument('--temperature', + type=float, + required=False, + default=1, + help='Softmax temperature for output logits.') + + beam_search_group.add_argument('--length_penalty', + type=float, + required=False, + default=1, + help='Positive. >1 to penalize and <1 to encorage short sentence.') + + beam_search_group.add_argument('--repetition_penalty', + type=float, + required=False, + default=1, + help='Positive. >1 to penalize and <1 to encorage.') + + mixed_precision_option_group = parser.add_argument_group( + "mixed precision conversion parameters that works when \"--precision fp16\" is specified") + + mixed_precision_option_group.add_argument('--io_block_list', + nargs='+', + required=False, + default=[], + help='List of inputs or outputs in float32') + + mixed_precision_option_group.add_argument( + '--op_block_list', + nargs='+', + required=False, + default=[], + help='List of operators (like Add LayerNormalization FastGelu) to compute in float32.') + + mixed_precision_option_group.add_argument('--node_block_list', + nargs='+', + required=False, + default=[], + help='List of node names to compute in float32.') + + mixed_precision_option_group.add_argument('--force_fp16_initializers', + required=False, + action='store_true', + help='Convert all float initializers to float16.') + mixed_precision_option_group.set_defaults(force_fp16_initializers=False) + + args = parser.parse_args(argv) + + return args + + +def gpt2_to_onnx(args): + model_name = args.model_name_or_path + + print(f"use convert_to_onnx.py to convert model {model_name} to onnx {args.gpt2_onnx} ...") + arguments = [ + '--model_name_or_path', model_name, '--output', args.gpt2_onnx, '--optimize_onnx', '--precision', + 'fp32' if args.precision == Precision.FLOAT32 else 'fp16', '--test_runs', '1', '--test_cases', '10' + ] + if args.use_gpu: + arguments.append('--use_gpu') + if args.use_external_data_format: + arguments.append('--use_external_data_format') + + # mixed precision conversion options + if args.precision == Precision.FLOAT16: + assert args.use_gpu, "fp16 or mixed precision model cannot run in CPU. Please add --use_gpu" + if args.io_block_list: + arguments.append('--io_block_list') + arguments.extend(args.io_block_list) + if args.op_block_list: + arguments.append('--op_block_list') + arguments.extend(args.op_block_list) + if args.node_block_list: + arguments.append('--node_block_list') + arguments.extend(args.node_block_list) + if args.force_fp16_initializers: + arguments.append('--force_fp16_initializers') + + convert_gpt2_to_onnx(arguments) + + +def shape_inference(gpt2_onnx_path): + # Run symbolic shape inference to walk around ORT shape inference issue for subgraph. + from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference + out = SymbolicShapeInference.infer_shapes(onnx.load(gpt2_onnx_path), auto_merge=True, guess_output_rank=False) + if out: + # TODO: Use external format if input has extra data. + onnx.save(out, gpt2_onnx_path) + else: + print("Failed to run symbolic shape inference on the model.") + + +def create_ort_session(model_path, use_gpu): + from onnxruntime import SessionOptions, InferenceSession, __version__ as ort_version, GraphOptimizationLevel + sess_options = SessionOptions() + sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL + execution_providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if use_gpu else ['CPUExecutionProvider'] + + ort_session = InferenceSession(model_path, sess_options, providers=execution_providers) + return ort_session + + +def convert_model(args): + if os.path.exists(args.gpt2_onnx): + print(f"skip convert_to_onnx since path existed: {args.gpt2_onnx}") + else: + gpt2_to_onnx(args) + + print(f"Run symbolic shape inference on {args.gpt2_onnx}. The file will be overwritten.") + shape_inference(args.gpt2_onnx) + + global config + config = GPT2Config.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir) + print(config) + + eos_token_id = config.eos_token_id + pad_token_id = config.eos_token_id + vocab_size = config.vocab_size + + model = onnx.load(args.gpt2_onnx) + model.graph.name = "gpt2 subgraph" + inputs = [ + "input_ids", "max_length", "min_length", "num_beams", "num_return_sequences", "temperature", "length_penalty", + "repetition_penalty", "vocab_mask" + ] + + outputs = ["sequences"] + if args.output_sequences_scores: + outputs.append("sequences_scores") + + if args.output_token_scores: + assert args.output_sequences_scores, "--output_token_scores requires --output_sequences_scores" + outputs.append("scores") + + node = helper.make_node('BeamSearch', inputs=inputs, outputs=outputs, name='BeamSearch_GPT2') + node.domain = "com.microsoft" + node.attribute.extend([ + helper.make_attribute("eos_token_id", eos_token_id), + helper.make_attribute("pad_token_id", pad_token_id), + helper.make_attribute("no_repeat_ngram_size", args.no_repeat_ngram_size), + helper.make_attribute("early_stopping", 1 if args.early_stopping else 0), + helper.make_attribute("body", model.graph), + ]) + + from onnx import TensorProto + + # graph inputs + input_ids = helper.make_tensor_value_info('input_ids', TensorProto.INT32, ['batch_size', 'sequence_length']) + max_length = helper.make_tensor_value_info('max_length', TensorProto.INT32, [1]) + min_length = helper.make_tensor_value_info('min_length', TensorProto.INT32, [1]) + num_beams = helper.make_tensor_value_info('num_beams', TensorProto.INT32, [1]) + num_return_sequences = helper.make_tensor_value_info('num_return_sequences', TensorProto.INT32, [1]) + temperature = helper.make_tensor_value_info('temperature', TensorProto.FLOAT, [1]) + length_penalty = helper.make_tensor_value_info('length_penalty', TensorProto.FLOAT, [1]) + repetition_penalty = helper.make_tensor_value_info('repetition_penalty', TensorProto.FLOAT, [1]) + vocab_mask = helper.make_tensor_value_info('vocab_mask', TensorProto.INT32, [vocab_size]) + + graph_inputs = [ + input_ids, max_length, min_length, num_beams, num_return_sequences, temperature, length_penalty, + repetition_penalty, vocab_mask + ] + + # graph outputs + sequences = helper.make_tensor_value_info('sequences', TensorProto.INT32, + ['batch_size', 'num_return_sequences', 'max_length']) + + sequences_scores = helper.make_tensor_value_info('sequences_scores', TensorProto.FLOAT, + ['batch_size', 'num_return_sequences']) + + scores = helper.make_tensor_value_info('scores', TensorProto.FLOAT, + ['max_length - sequence_length', 'batch_size', 'num_beams', vocab_size]) + + initializers = [] + + graph_outputs = [sequences] + + if args.output_sequences_scores: + graph_outputs.append(sequences_scores) + + if args.output_token_scores: + graph_outputs.append(scores) + + new_graph = helper.make_graph([node], 'gpt2-beam-search', graph_inputs, graph_outputs, initializers) + + # Create the model + new_model = helper.make_model(new_graph, producer_name='onnxruntime.transformers', opset_imports=model.opset_import) + onnx.save(new_model, args.output) + + +def test_model(args, use_vocab_mask: bool = False, sentences: List[str] = None): + from transformers import GPT2Tokenizer, GPT2LMHeadModel + + tokenizer = GPT2Tokenizer.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir) + tokenizer.padding_side = "left" + tokenizer.pad_token = tokenizer.eos_token + + model = GPT2LMHeadModel.from_pretrained(args.model_name_or_path, + cache_dir=args.cache_dir, + pad_token_id=tokenizer.eos_token_id) + + # Use different length sentences to test batching + if sentences is None: + sentences = ["The product is released", "I enjoy walking in the park", "Test best way to invest"] + + inputs = tokenizer(sentences, return_tensors='pt', padding=True) + input_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + + bad_words = "walk in park" + bad_words_ids = tokenizer.encode(bad_words, add_prefix_space=True) + bad_words_ids = [[word_id] for word_id in bad_words_ids] # Convert to list of list + if use_vocab_mask: + print("bad_words_ids", bad_words_ids) + else: + bad_words_ids = None + + global config + config = model.config + eos_token_id = config.eos_token_id + pad_token_id = config.eos_token_id + vocab_size = config.vocab_size + + torch_decoded_sequences = [] + if not args.disable_parity: + print('-' * 50) + print("Test PyTorch model and beam search with huggingface transformers...") + beam_outputs = model.generate(input_ids=input_ids, + attention_mask=attention_mask, + max_length=args.max_length, + min_length=args.min_length, + num_beams=args.num_beams, + early_stopping=args.early_stopping, + no_repeat_ngram_size=args.no_repeat_ngram_size, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + num_return_sequences=args.num_return_sequences, + temperature=args.temperature, + length_penalty=args.length_penalty, + repetition_penalty=args.repetition_penalty, + bad_words_ids=bad_words_ids, + return_dict_in_generate=True, + output_scores=True) + print("input_ids", input_ids) + print("huggingface transformers outputs:") + print("sequences", beam_outputs.sequences) + if args.output_sequences_scores: + print("sequences_scores", beam_outputs.sequences_scores) + if args.output_token_scores: + print("scores", beam_outputs.scores) + for i, sequence in enumerate(beam_outputs.sequences): + decoded_sequence = tokenizer.decode(sequence, skip_special_tokens=True) + torch_decoded_sequences.append(decoded_sequence) + print("{}: {}".format(i, decoded_sequence)) + + print('-' * 50) + print("Test ONNX model and bream search with onnxruntime...") + + ort_session = create_ort_session(args.output, args.use_gpu) + + vocab_mask = np.ones((vocab_size), dtype=np.int32) + if use_vocab_mask: + for bad_word_id in bad_words_ids: + vocab_mask[bad_word_id] = 0 + + inputs = { + "input_ids": input_ids.cpu().numpy().astype(np.int32), + "max_length": np.array([args.max_length], dtype=np.int32), + "min_length": np.array([args.min_length], dtype=np.int32), + "num_beams": np.array([args.num_beams], dtype=np.int32), + "num_return_sequences": np.array([args.num_return_sequences], dtype=np.int32), + "temperature": np.array([args.temperature], dtype=np.float32), + "length_penalty": np.array([args.length_penalty], dtype=np.float32), + "repetition_penalty": np.array([args.repetition_penalty], dtype=np.float32), + "vocab_mask": vocab_mask + } + + test_data_dir = Path(args.output).parent.as_posix() + print("test_data_dir", test_data_dir) + from bert_test_data import output_test_data + all_inputs = [inputs] + for i, inputs in enumerate(all_inputs): + dir = os.path.join(test_data_dir, 'test_data_set_' + str(i)) + output_test_data(dir, inputs) + + print("inputs", inputs) + + # Test performance + latency = [] + for _ in range(args.total_runs): + start = time.time() + result = ort_session.run(None, inputs) + latency.append(time.time() - start) + batch_size = input_ids.shape[0] + from benchmark_helper import get_latency_result + output = get_latency_result(latency, batch_size) + + print("ORT outputs:") + sequences = result[0] + print("sequences", sequences) + if args.output_sequences_scores: + print("sequences_scores", result[1]) + if args.output_token_scores: + print("scores", result[2]) + + (batch_size, num_sequences, max_length) = sequences.shape + ort_decoded_sequences = [] + for i in range(batch_size): + for j in range(num_sequences): + decoded_sequence = tokenizer.decode(sequences[i][j], skip_special_tokens=True) + ort_decoded_sequences.append(decoded_sequence) + print(f"batch {i} sequence {j}: {decoded_sequence}") + + if not args.disable_parity: + torch_sequences = beam_outputs.sequences.reshape(batch_size, args.num_return_sequences, -1) + ort_sequences = torch.LongTensor(sequences) + print("-" * 50) + print("Torch Sequences:") + print(torch_sequences) + print(torch_decoded_sequences) + print("-" * 50) + print("ORT Sequences:") + print(ort_sequences) + print(ort_decoded_sequences) + print("-" * 50) + # Compare the generated text instead of word IDs since ORT pads to max sequence length but Torch not. + is_same = (torch_decoded_sequences == ort_decoded_sequences) + print("Torch and ORT result is ", "same" if is_same else "different") + output["parity"] = is_same + + print(output) + return output + + +def main(argv=None, sentences=None): + args = parse_arguments(argv) + + if os.path.exists(args.output): + print(f"skip conversion since path existed: {args.output}") + else: + convert_model(args) + + return test_model(args, use_vocab_mask=True, sentences=sentences) + + +if __name__ == '__main__': + main() diff --git a/onnxruntime/python/tools/transformers/convert_to_onnx.py b/onnxruntime/python/tools/transformers/convert_to_onnx.py index 500a3f3b90..28120d5e80 100644 --- a/onnxruntime/python/tools/transformers/convert_to_onnx.py +++ b/onnxruntime/python/tools/transformers/convert_to_onnx.py @@ -17,7 +17,6 @@ This converts GPT2 model to onnx. Examples: import os import argparse -import coloredlogs import logging import torch import numpy diff --git a/onnxruntime/test/python/transformers/test_beam_search.py b/onnxruntime/test/python/transformers/test_beam_search.py new file mode 100644 index 0000000000..f469cdaaaa --- /dev/null +++ b/onnxruntime/test/python/transformers/test_beam_search.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python +# coding: utf-8 +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import unittest +import os +import pytest + +from parity_utilities import find_transformers_source +if find_transformers_source(): + from convert_beam_search import main as run +else: + from onnxruntime.transformers.convert_beam_search import main as run + + +class TestBeamSearch(unittest.TestCase): + def setUp(self): + #TODO: use a smaller model and enable tests in CI pipeline + self.model_name = "gpt2" + self.gpt2_onnx_path = os.path.join('.', 'onnx_models', 'gpt2_past_fp32_shape.onnx') + self.beam_search_onnx_path = os.path.join('.', 'onnx_models', 'gpt2_beam_search.onnx') + self.cpu_params = f'-m {self.model_name} --gpt2_onnx {self.gpt2_onnx_path} --output {self.beam_search_onnx_path} --output_sequences_score --repetition_penalty 2.0' + + def run_beam_search(self, arguments: str, sentences=None): + return run(arguments.split(), sentences=sentences) + + @pytest.mark.slow + def test_cpu(self): + result = self.run_beam_search(self.cpu_params + " --num_return_sequences 2", + sentences=["The product is released"]) + os.remove(self.gpt2_onnx_path) + os.remove(self.beam_search_onnx_path) + self.assertTrue(result["parity"], "ORT and PyTorch result is different") + + @pytest.mark.slow + def test_early_stopping(self): + result = self.run_beam_search(self.cpu_params + " --early_stopping") + os.remove(self.gpt2_onnx_path) + os.remove(self.beam_search_onnx_path) + self.assertTrue(result["parity"], "ORT and PyTorch result is different") + + @pytest.mark.slow + def test_temperature(self): + result = self.run_beam_search(self.cpu_params + " --temperature 0.5") + os.remove(self.gpt2_onnx_path) + os.remove(self.beam_search_onnx_path) + self.assertTrue(result["parity"], "ORT and PyTorch result is different") + + @pytest.mark.slow + def test_length_penalty(self): + result = self.run_beam_search(self.cpu_params + " --length_penalty 0.5") + os.remove(self.gpt2_onnx_path) + os.remove(self.beam_search_onnx_path) + self.assertTrue(result["parity"], "ORT and PyTorch result is different") + + @pytest.mark.slow + def test_no_repeat_ngram(self): + for ngram_size in [1, 2]: + result = self.run_beam_search(self.cpu_params + f' --no_repeat_ngram_size {ngram_size}') + os.remove(self.gpt2_onnx_path) + os.remove(self.beam_search_onnx_path) + self.assertTrue(result["parity"], "ORT and PyTorch result is different") + + +if __name__ == '__main__': + unittest.main() diff --git a/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json b/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json index 410625083c..120e6de7c0 100644 --- a/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json +++ b/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json @@ -3,6 +3,10 @@ "Affine ai.onnx CPUExecutionProvider", 7811918192248490408 ], + [ + "BeamSearch com.microsoft CPUExecutionProvider", + 6968087233460196528 + ], [ "Crop ai.onnx CPUExecutionProvider", 6914973556202621376