From 8cf2c1c41075e455895f49e198bb3fe8e9e8231f Mon Sep 17 00:00:00 2001 From: Ye Wang <52801275+wangyems@users.noreply.github.com> Date: Fri, 31 Jul 2020 14:37:58 -0700 Subject: [PATCH] Modify EmbedLayerNorm to support distill-bert (#4666) * modify cpu op * modify cuda ops * change is_distill to has_segment --- .../contrib_ops/cpu/bert/embed_layer_norm.cc | 39 ++++---- .../cpu/bert/embed_layer_norm_helper.cc | 25 +++--- .../contrib_ops/cuda/bert/embed_layer_norm.cc | 8 +- .../cuda/bert/embed_layer_norm_impl.cu | 14 ++- .../core/graph/contrib_ops/contrib_defs.cc | 8 +- .../contrib_ops/embedlayernorm_op_test.cc | 89 ++++++++++++++++++- 6 files changed, 138 insertions(+), 45 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc index cdf484fb07..952ad156c1 100644 --- a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc @@ -34,10 +34,10 @@ template Status EmbedLayerNorm::Compute(OpKernelContext* context) const { ORT_RETURN_IF_ERROR(embed_layer_norm::CheckInputs(context)); const Tensor* input_ids = context->Input(0); - const Tensor* segment_ids = context->Input(1); + const Tensor* segment_ids = context->Input(1); // optional. nullptr if it's distill-bert const Tensor* word_embedding = context->Input(2); const Tensor* position_embedding = context->Input(3); - const Tensor* segment_embedding = context->Input(4); + const Tensor* segment_embedding = context->Input(4); // optional. nullptr if it's distill-bert const Tensor* gamma = context->Input(5); const Tensor* beta = context->Input(6); const Tensor* mask = context->Input(7); // optional. nullptr if not provided @@ -56,16 +56,16 @@ Status EmbedLayerNorm::Compute(OpKernelContext* context) const { int word_embedding_length = static_cast(word_embedding->Shape()[0]); int position_embedding_length = static_cast(position_embedding->Shape()[0]); - int segment_embedding_length = static_cast(segment_embedding->Shape()[0]); + int segment_embedding_length = (nullptr == segment_embedding) ? 0 : static_cast(segment_embedding->Shape()[0]); - auto input_ids_data = input_ids->template Data(); - auto segment_ids_data = segment_ids->template Data(); - auto word_embedding_data = word_embedding->template Data(); - auto position_embedding_data = position_embedding->template Data(); - auto segment_embedding_data = segment_embedding->template Data(); - auto gamma_data = gamma->template Data(); - auto beta_data = beta->template Data(); - auto output_data = output->template MutableData(); + const int32_t* input_ids_data = input_ids->template Data(); + const int32_t* segment_ids_data = (nullptr == segment_ids) ? nullptr : segment_ids->template Data(); + const T* word_embedding_data = word_embedding->template Data(); + const T* position_embedding_data = position_embedding->template Data(); + const T* segment_embedding_data = (nullptr == segment_embedding) ? nullptr : segment_embedding->template Data(); + const T* gamma_data = gamma->template Data(); + const T* beta_data = beta->template Data(); + T* output_data = output->template MutableData(); // Calculate output { @@ -83,20 +83,25 @@ Status EmbedLayerNorm::Compute(OpKernelContext* context) const { failed.store(true, std::memory_order_release); return; } - int segment_col_index = segment_ids_data[index]; - if (segment_col_index < 0 || segment_col_index >= segment_embedding_length) { - failed.store(true, std::memory_order_release); - return; + int segment_col_index = 0; + if (nullptr != segment_ids_data) { + segment_col_index = segment_ids_data[index]; + if (segment_col_index < 0 || segment_col_index >= segment_embedding_length) { + failed.store(true, std::memory_order_release); + return; + } } T* y = output_data + index * hidden_size; const T* input_word_embedding = word_embedding_data + word_col_index * hidden_size; const T* input_position_embedding = position_embedding_data + position_col_index * hidden_size; - const T* input_segment_embedding = segment_embedding_data + segment_col_index * hidden_size; + const T* input_segment_embedding = (nullptr == segment_embedding_data) ? nullptr : segment_embedding_data + segment_col_index * hidden_size; T sum = static_cast(0); for (int i = 0; i < hidden_size; i++) { - T subtotal = input_word_embedding[i] + input_position_embedding[i] + input_segment_embedding[i]; + T subtotal = input_word_embedding[i] + input_position_embedding[i]; + if (nullptr != segment_embedding_data) + subtotal += input_segment_embedding[i]; y[i] = subtotal; sum += subtotal; } diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.cc b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.cc index 0b5bdda71f..6dd13a056c 100644 --- a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.cc +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.cc @@ -11,15 +11,15 @@ namespace embed_layer_norm { Status CheckInputs(const OpKernelContext* context) { const Tensor* input_ids = context->Input(0); - const Tensor* segment_ids = context->Input(1); + const Tensor* segment_ids = context->Input(1); // optional. nullptr if it's distill-bert const Tensor* word_embedding = context->Input(2); const Tensor* position_embedding = context->Input(3); - const Tensor* segment_embedding = context->Input(4); + const Tensor* segment_embedding = context->Input(4); // optional. nullptr if it's distill-bert const Tensor* gamma = context->Input(5); const Tensor* beta = context->Input(6); const Tensor* mask = context->Input(7); // optional. nullptr if not provided - if (input_ids->Shape() != segment_ids->Shape()) { + if (nullptr != segment_ids && input_ids->Shape() != segment_ids->Shape()) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 0 and 1 shall have same shape"); } @@ -48,10 +48,16 @@ Status CheckInputs(const OpKernelContext* context) { "position_embedding is expected to have 2 dimensions, got ", position_embedding_dims.size()); } - const auto& segment_embedding_dims = segment_embedding->Shape().GetDims(); - if (segment_embedding_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "segment_embedding is expected to have 2 dimensions, got ", segment_embedding_dims.size()); + if (nullptr != segment_embedding) { + const auto& segment_embedding_dims = segment_embedding->Shape().GetDims(); + if (segment_embedding_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "segment_embedding is expected to have 2 dimensions, got ", segment_embedding_dims.size()); + } + if (word_embedding_dims[1] != segment_embedding_dims[1]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "word_embedding and segment_embedding shall have same dimension 1"); + } } if (word_embedding_dims[1] != position_embedding_dims[1]) { @@ -59,11 +65,6 @@ Status CheckInputs(const OpKernelContext* context) { "word_embedding and position_embedding shall have same dimension 1"); } - if (word_embedding_dims[1] != segment_embedding_dims[1]) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "word_embedding and segment_embedding shall have same dimension 1"); - } - const auto& beta_dims = beta->Shape().GetDims(); if (beta_dims.size() != 1) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc index 38e1166a58..8adffa85ed 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc @@ -39,10 +39,10 @@ Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_ERROR(embed_layer_norm::CheckInputs(context)); const Tensor* input_ids = context->Input(0); - const Tensor* segment_ids = context->Input(1); + const Tensor* segment_ids = context->Input(1); // optional. nullptr if it's distill-bert const Tensor* word_embedding = context->Input(2); const Tensor* position_embedding = context->Input(3); - const Tensor* segment_embedding = context->Input(4); + const Tensor* segment_embedding = context->Input(4); // optional. nullptr if it's distill-bert const Tensor* gamma = context->Input(5); const Tensor* beta = context->Input(6); const Tensor* mask = context->Input(7); // optional. nullptr if not provided @@ -64,13 +64,13 @@ Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { output->template MutableData(), mask_index->template MutableData(), input_ids->template Data(), - segment_ids->template Data(), + nullptr == segment_ids ? nullptr : segment_ids->template Data(), nullptr == mask ? nullptr : mask->template Data(), gamma->template Data(), beta->template Data(), word_embedding->template Data(), position_embedding->template Data(), - segment_embedding->template Data(), + nullptr == segment_embedding ? nullptr : segment_embedding->template Data(), epsilon_, static_cast(hidden_size), batch_size, diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu index f5f6e64755..9e856e2e35 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -1,7 +1,7 @@ /* The implementation of this file is based on embLayerNorm plugin in TensorRT demo: https://github.com/NVIDIA/TensorRT/tree/release/5.1/demo/BERT/ - + Copyright 2019 NVIDIA Corporation Licensed under the Apache License, Version 2.0 (the "License"); @@ -121,7 +121,11 @@ __global__ void EmbedLayerNormKernel( const int sequence_position = blockIdx.y * gridDim.x + blockIdx.x; if (threadIdx.x == 0) { word_id = input_ids[sequence_position]; - segment_id = segment_ids[sequence_position]; + if (nullptr == segment_ids) { + segment_id = 0; + } else { + segment_id = segment_ids[sequence_position];; + } } __syncthreads(); @@ -137,7 +141,9 @@ __global__ void EmbedLayerNormKernel( for (int it = threadIdx.x; it < hidden_size; it += TPB) { const T w(word_embedding[word_offset + it]); - const T t(segment_embedding[segment_offset + it]); + T t(0); + if (nullptr != segment_embedding) + t = segment_embedding[segment_offset + it]; const T p(position_embedding[position_offset + it]); const T val = w + t + p; @@ -195,7 +201,7 @@ bool LaunchEmbedLayerNormKernel( return EmbedSkipLayerNorm( stream, hidden_size, batch_size, sequence_length, input_ids, segment_ids, reinterpret_cast(beta), reinterpret_cast(gamma), - reinterpret_cast(word_embedding), reinterpret_cast(position_embedding), + reinterpret_cast(word_embedding), reinterpret_cast(position_embedding), reinterpret_cast(segment_embedding), __float2half_rn(epsilon), reinterpret_cast(output)); } else { diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index fd005926ff..4645bd71ad 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -294,8 +294,8 @@ void RegisterBertSchemas() { Multi-Head Self Attention that can be either unidirectional (like GPT-2) or bidirectional (like BERT). The mask_index input is optional. Besides raw attention mask with shape (batch_size, past_sequence_length + sequence_length), we also support other two formats: When input has right-side padding, mask_index is one dimension with shape (batch_size), -where value of each element is the end position, or valid length of actual sequence excluding padding. When input has -left-side padding, mask_index has shape (2 * batch_size), where the values are the exclusive end positions followed by +where value of each element is the end position, or valid length of actual sequence excluding padding. When input has +left-side padding, mask_index has shape (2 * batch_size), where the values are the exclusive end positions followed by the inclusive start positions. When unidirectional is 1, and each token only attend to previous tokens. For GPT-2, both past and present state are optional. Present state could appear in output even when past state is not in input. )DOC"; @@ -461,10 +461,10 @@ will be calculated.)DOC"; .SetDoc(EmbedLayerNormalization_ver1_doc) .Attr("epsilon", "The epsilon value to use to avoid division by zero.", AttributeProto::FLOAT, kDefaultEmbedLayerNormEpsilon) .Input(0, "input_ids", "2D words IDs with shape (batch_size, sequence_length)", "T1") - .Input(1, "segment_ids", "2D segment IDs with shape (batch_size, sequence_length)", "T1") + .Input(1, "segment_ids", "2D segment IDs with shape (batch_size, sequence_length)", "T1", OpSchema::Optional) .Input(2, "word_embedding", "2D with shape (,hidden_size)", "T") .Input(3, "position_embedding", "2D with shape (, hidden_size)", "T") - .Input(4, "segment_embedding", "2D with shape (, hidden_size)", "T") + .Input(4, "segment_embedding", "2D with shape (, hidden_size)", "T", OpSchema::Optional) .Input(5, "gamma", "1D gamma tensor for layer normalization with shape (hidden_size)", "T") .Input(6, "beta", "1D beta tensor for layer normalization with shape (hidden_size)", "T") .Input(7, "mask", "2D attention mask with shape (batch_size, sequence_length)", "T1", OpSchema::Optional) diff --git a/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc b/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc index 1275a997ac..177d8d745b 100644 --- a/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc +++ b/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc @@ -26,7 +26,8 @@ static void RunTest( int sequence_length, int hidden_size, bool use_float16 = false, - bool has_mask = true) { + bool has_mask = true, + bool has_segment = true) { int min_cuda_architecture = use_float16 ? 530 : 0; bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); @@ -65,11 +66,19 @@ static void RunTest( OpTester tester("EmbedLayerNormalization", 1, onnxruntime::kMSDomain); tester.AddInput("input_ids", input_ids_dims, input_ids_data); - tester.AddInput("segment_ids", segment_ids_dims, segment_ids_data); + if (!has_segment) { + tester.AddMissingOptionalInput(); + } else { + tester.AddInput("segment_ids", segment_ids_dims, segment_ids_data); + } if (use_float16) { tester.AddInput("word_embedding", word_embedding_dims, ToFloat16(word_embedding_data)); tester.AddInput("position_embedding", position_embedding_dims, ToFloat16(position_embedding_data)); - tester.AddInput("segment_embedding", segment_embedding_dims, ToFloat16(segment_embedding_data)); + if (!has_segment) { + tester.AddMissingOptionalInput(); + } else { + tester.AddInput("segment_embedding", segment_embedding_dims, ToFloat16(segment_embedding_data)); + } tester.AddInput("gamma", gamma_dims, ToFloat16(gamma_data)); tester.AddInput("beta", beta_dims, ToFloat16(beta_data)); tester.AddAttribute("epsilon", epsilon); @@ -80,7 +89,11 @@ static void RunTest( } else { tester.AddInput("word_embedding", word_embedding_dims, word_embedding_data); tester.AddInput("position_embedding", position_embedding_dims, position_embedding_data); - tester.AddInput("segment_embedding", segment_embedding_dims, segment_embedding_data); + if (!has_segment) { + tester.AddMissingOptionalInput(); + } else { + tester.AddInput("segment_embedding", segment_embedding_dims, segment_embedding_data); + } tester.AddInput("gamma", gamma_dims, gamma_data); tester.AddInput("beta", beta_dims, beta_data); tester.AddAttribute("epsilon", epsilon); @@ -433,5 +446,73 @@ TEST(EmbedLayerNormTest, EmbedLayerNormLargeBatchSmallHiddenSize) { sequence_length, hidden_size); } + +TEST(EmbedLayerNormTest, EmbedLayerNormBatch_Distill) { + int batch_size = 3; + int sequence_length = 2; + int hidden_size = 4; + + std::vector input_ids_data = { + 1, 3, + 1, 3, + 2, 0}; + + std::vector segment_ids_data = {}; + + std::vector mask_data = { + 1, 1, + 1, 1, + 1, 0}; + + std::vector word_embedding_data = { + 0.2f, 0.1f, 0.4f, -0.6f, + 0.3f, 0.2f, 0.5f, 0.6f, + 0.6f, 0.7f, 0.0f, -0.1f, + 0.8f, 0.6f, 0.9f, 1.2f, + 0.1f, 0.3f, 0.5f, 0.9f, + 1.0f, -2.0f, 1.1f, 0.8f}; + + std::vector position_embedding_data = { + 0.1f, 0.1f, 0.4f, 0.6f, + 0.6f, 0.0f, 0.8f, 0.6f, + 0.3f, 0.9f, -2.0f, 0.8f}; + + std::vector segment_embedding_data = {}; + + std::vector gamma_data = { + 0.25f, 0.15f, 0.45f, -0.66f}; + + std::vector beta_data = { + 0.6f, 0.2f, 0.5f, -0.6f}; + + std::vector output_data = { + 0.39587587118148804, 0.03670068085193634, 0.7449488639831543, -1.4981462955474854, + 0.61326867341995239, -0.046796366572380066, 0.81048583984375, -1.1954958438873291, + 0.39587587118148804, 0.03670068085193634, 0.7449488639831543, -1.4981462955474854, + 0.61326867341995239, -0.046796366572380066, 0.81048583984375, -1.1954958438873291, + 0.75811392068862915, 0.38973665237426758, -0.069209933280944824, -0.18257927894592285, + 0.73836749792098999, 0.071695566177368164, 1.111332893371582, 0.097372293472290039}; + + std::vector mask_index_data = { + 2, 2, 1}; + + RunTest(input_ids_data, + segment_ids_data, + mask_data, + word_embedding_data, + position_embedding_data, + segment_embedding_data, + gamma_data, + beta_data, + output_data, + mask_index_data, + epsilon_, + batch_size, + sequence_length, + hidden_size, + false, + true, + false); +} } // namespace test } // namespace onnxruntime