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

Modify EmbedLayerNorm to support distill-bert (#4666)

Browse source 
* modify cpu op

* modify cuda ops

* change is_distill to has_segment
This commit is contained in:
Ye Wang 2020-07-31 14:37:58 -07:00 committed by GitHub
parent 1eadec0eea
commit 8cf2c1c410
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 138 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc
View file

@ -34,10 +34,10 @@ template <typename T>
Status EmbedLayerNorm<T>::Compute(OpKernelContext* context) const {
ORT_RETURN_IF_ERROR(embed_layer_norm::CheckInputs(context));
const Tensor* input_ids = context->Input<Tensor>(0);
const Tensor* segment_ids = context->Input<Tensor>(1);
const Tensor* segment_ids = context->Input<Tensor>(1); // optional. nullptr if it's distill-bert
const Tensor* word_embedding = context->Input<Tensor>(2);
const Tensor* position_embedding = context->Input<Tensor>(3);
const Tensor* segment_embedding = context->Input<Tensor>(4);
const Tensor* segment_embedding = context->Input<Tensor>(4); // optional. nullptr if it's distill-bert
const Tensor* gamma = context->Input<Tensor>(5);
const Tensor* beta = context->Input<Tensor>(6);
const Tensor* mask = context->Input<Tensor>(7); // optional. nullptr if not provided
@ -56,16 +56,16 @@ Status EmbedLayerNorm<T>::Compute(OpKernelContext* context) const {
int word_embedding_length = static_cast<int>(word_embedding->Shape()[0]);
int position_embedding_length = static_cast<int>(position_embedding->Shape()[0]);
int segment_embedding_length = static_cast<int>(segment_embedding->Shape()[0]);
int segment_embedding_length = (nullptr == segment_embedding) ? 0 : static_cast<int>(segment_embedding->Shape()[0]);
auto input_ids_data = input_ids->template Data<int32_t>();
auto segment_ids_data = segment_ids->template Data<int32_t>();
auto word_embedding_data = word_embedding->template Data<T>();
auto position_embedding_data = position_embedding->template Data<T>();
auto segment_embedding_data = segment_embedding->template Data<T>();
auto gamma_data = gamma->template Data<T>();
auto beta_data = beta->template Data<T>();
auto output_data = output->template MutableData<T>();
const int32_t* input_ids_data = input_ids->template Data<int32_t>();
const int32_t* segment_ids_data = (nullptr == segment_ids) ? nullptr : segment_ids->template Data<int32_t>();
const T* word_embedding_data = word_embedding->template Data<T>();
const T* position_embedding_data = position_embedding->template Data<T>();
const T* segment_embedding_data = (nullptr == segment_embedding) ? nullptr : segment_embedding->template Data<T>();
const T* gamma_data = gamma->template Data<T>();
const T* beta_data = beta->template Data<T>();
T* output_data = output->template MutableData<T>();
// Calculate output
{
@ -83,20 +83,25 @@ Status EmbedLayerNorm<T>::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<T>(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;
}

25
onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.cc
View file

@ -11,15 +11,15 @@ namespace embed_layer_norm {
Status CheckInputs(const OpKernelContext* context) {
const Tensor* input_ids = context->Input<Tensor>(0);
const Tensor* segment_ids = context->Input<Tensor>(1);
const Tensor* segment_ids = context->Input<Tensor>(1); // optional. nullptr if it's distill-bert
const Tensor* word_embedding = context->Input<Tensor>(2);
const Tensor* position_embedding = context->Input<Tensor>(3);
const Tensor* segment_embedding = context->Input<Tensor>(4);
const Tensor* segment_embedding = context->Input<Tensor>(4); // optional. nullptr if it's distill-bert
const Tensor* gamma = context->Input<Tensor>(5);
const Tensor* beta = context->Input<Tensor>(6);
const Tensor* mask = context->Input<Tensor>(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,

8
onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc
View file

@ -39,10 +39,10 @@ Status EmbedLayerNorm<T>::ComputeInternal(OpKernelContext* context) const {
ORT_RETURN_IF_ERROR(embed_layer_norm::CheckInputs(context));
const Tensor* input_ids = context->Input<Tensor>(0);
const Tensor* segment_ids = context->Input<Tensor>(1);
const Tensor* segment_ids = context->Input<Tensor>(1); // optional. nullptr if it's distill-bert
const Tensor* word_embedding = context->Input<Tensor>(2);
const Tensor* position_embedding = context->Input<Tensor>(3);
const Tensor* segment_embedding = context->Input<Tensor>(4);
const Tensor* segment_embedding = context->Input<Tensor>(4); // optional. nullptr if it's distill-bert
const Tensor* gamma = context->Input<Tensor>(5);
const Tensor* beta = context->Input<Tensor>(6);
const Tensor* mask = context->Input<Tensor>(7); // optional. nullptr if not provided
@ -64,13 +64,13 @@ Status EmbedLayerNorm<T>::ComputeInternal(OpKernelContext* context) const {
output->template MutableData<T>(),
mask_index->template MutableData<int32_t>(),
input_ids->template Data<int32_t>(),
segment_ids->template Data<int32_t>(),
nullptr == segment_ids ? nullptr : segment_ids->template Data<int32_t>(),
nullptr == mask ? nullptr : mask->template Data<int32_t>(),
gamma->template Data<T>(),
beta->template Data<T>(),
word_embedding->template Data<T>(),
position_embedding->template Data<T>(),
segment_embedding->template Data<T>(),
nullptr == segment_embedding ? nullptr : segment_embedding->template Data<T>(),
epsilon_,
static_cast<int>(hidden_size),
batch_size,

14
onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu
View file

@ -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<half>(
stream, hidden_size, batch_size, sequence_length, input_ids, segment_ids,
reinterpret_cast<const half*>(beta), reinterpret_cast<const half*>(gamma),
reinterpret_cast<const half*>(word_embedding), reinterpret_cast<const half*>(position_embedding),
reinterpret_cast<const half*>(word_embedding), reinterpret_cast<const half*>(position_embedding),
reinterpret_cast<const half*>(segment_embedding), __float2half_rn(epsilon),
reinterpret_cast<half*>(output));
} else {

8
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

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

89
onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc
View file

@ -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<int32_t>("input_ids", input_ids_dims, input_ids_data);
tester.AddInput<int32_t>("segment_ids", segment_ids_dims, segment_ids_data);
if (!has_segment) {
tester.AddMissingOptionalInput<int32_t>();
} else {
tester.AddInput<int32_t>("segment_ids", segment_ids_dims, segment_ids_data);
}
if (use_float16) {
tester.AddInput<MLFloat16>("word_embedding", word_embedding_dims, ToFloat16(word_embedding_data));
tester.AddInput<MLFloat16>("position_embedding", position_embedding_dims, ToFloat16(position_embedding_data));
tester.AddInput<MLFloat16>("segment_embedding", segment_embedding_dims, ToFloat16(segment_embedding_data));
if (!has_segment) {
tester.AddMissingOptionalInput<MLFloat16>();
} else {
tester.AddInput<MLFloat16>("segment_embedding", segment_embedding_dims, ToFloat16(segment_embedding_data));
}
tester.AddInput<MLFloat16>("gamma", gamma_dims, ToFloat16(gamma_data));
tester.AddInput<MLFloat16>("beta", beta_dims, ToFloat16(beta_data));
tester.AddAttribute("epsilon", epsilon);
@ -80,7 +89,11 @@ static void RunTest(
} else {
tester.AddInput<float>("word_embedding", word_embedding_dims, word_embedding_data);
tester.AddInput<float>("position_embedding", position_embedding_dims, position_embedding_data);
tester.AddInput<float>("segment_embedding", segment_embedding_dims, segment_embedding_data);
if (!has_segment) {
tester.AddMissingOptionalInput<MLFloat16>();
} else {
tester.AddInput<float>("segment_embedding", segment_embedding_dims, segment_embedding_data);
}
tester.AddInput<float>("gamma", gamma_dims, gamma_data);
tester.AddInput<float>("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<int32_t> input_ids_data = {
1, 3,
1, 3,
2, 0};
std::vector<int32_t> segment_ids_data = {};
std::vector<int32_t> mask_data = {
1, 1,
1, 1,
1, 0};
std::vector<float> 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<float> 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<float> segment_embedding_data = {};
std::vector<float> gamma_data = {
0.25f, 0.15f, 0.45f, -0.66f};
std::vector<float> beta_data = {
0.6f, 0.2f, 0.5f, -0.6f};
std::vector<float> 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<int32_t> 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