diff --git a/onnxruntime/contrib_ops/cpu/layer_norm.cc b/onnxruntime/contrib_ops/cpu/layer_norm.cc index c3c26e4684..8d043635e9 100644 --- a/onnxruntime/contrib_ops/cpu/layer_norm.cc +++ b/onnxruntime/contrib_ops/cpu/layer_norm.cc @@ -31,7 +31,7 @@ namespace contrib { KernelDefBuilder() \ .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ LayerNorm); - + REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(double) @@ -50,7 +50,7 @@ Status LayerNorm::Compute(OpKernelContext* p_ctx) const { const Tensor* bias = p_ctx->Input(2); auto X_data = X->template Data(); auto scale_data = scale->template Data(); - auto bias_data = simplified ? nullptr : bias->template Data(); + auto bias_data = (simplified || nullptr == bias) ? nullptr : bias->template Data(); const TensorShape& x_shape = X->Shape(); const int64_t axis = HandleNegativeAxis(axis_, x_shape.NumDimensions()); @@ -124,6 +124,8 @@ Status LayerNorm::Compute(OpKernelContext* p_ctx) const { for (int64_t h = 0; h < norm_size; h++) { if (simplified) { p_output[h] = p_input[h] / mean_square * scale_data[h]; + } else if (nullptr == bias){ + p_output[h] = (p_input[h] - mean) / mean_square * scale_data[h]; } else { p_output[h] = (p_input[h] - mean) / mean_square * scale_data[h] + bias_data[h]; } diff --git a/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc b/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc index 29d244d64b..b84635d082 100644 --- a/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc +++ b/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc @@ -61,14 +61,16 @@ Status SkipLayerNorm::Compute(OpKernelContext* p_ctx) const { "Last dimension of gamma and input does not match"); } - const auto& beta_dims = beta->Shape().GetDims(); - if (beta_dims.size() != 1) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "beta is expected to have 1 dimension, got ", beta_dims.size()); - } - if (beta_dims[0] != input_dims[2]) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Last dimension of beta and input does not match"); + if (nullptr != beta) { + const auto& beta_dims = beta->Shape().GetDims(); + if (beta_dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "beta is expected to have 1 dimension, got ", beta_dims.size()); + } + if (beta_dims[0] != input_dims[2]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Last dimension of beta and input does not match"); + } } if (nullptr != bias) { @@ -91,7 +93,7 @@ Status SkipLayerNorm::Compute(OpKernelContext* p_ctx) const { const T* input_data = input->Data(); const T* skip_data = skip->Data(); const T* gamma_data = gamma->Data(); - const T* beta_data = beta->Data(); + const T* beta_data = beta == nullptr ? nullptr : beta->Data(); const T* bias_data = bias == nullptr ? nullptr : bias->Data(); T* output_data = output->MutableData(); @@ -119,7 +121,11 @@ Status SkipLayerNorm::Compute(OpKernelContext* p_ctx) const { mean_square = sqrt(mean_square / hidden_size - mean * mean + epsilon_); for (int64_t h = 0; h < hidden_size; h++) { - p_output[h] = (p_output[h] - mean) / mean_square * gamma_data[h] + beta_data[h]; + if (nullptr == beta_data) { + p_output[h] = (p_output[h] - mean) / mean_square * gamma_data[h]; + } else { + p_output[h] = (p_output[h] - mean) / mean_square * gamma_data[h] + beta_data[h]; + } } }, 0); diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index 16c1469f22..e3fef4c681 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -1,7 +1,7 @@ /* The implementation of this file is based on bert plugins 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"); @@ -79,7 +79,7 @@ struct KeyValuePairSum { template __device__ inline void LayerNorm( - const cub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, + const cub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, const T* gamma, const T epsilon, T* output) { // Assuming thread_data is already divided by ld @@ -101,7 +101,7 @@ __device__ inline void LayerNorm( const int idx = offset + i; const T val = output[idx]; const T g(gamma[i]); - const T b(beta[i]); + const T b = (nullptr == beta) ? (T)0 : beta[i]; output[idx] = g * (val - mu) * rsigma + b; } } @@ -129,7 +129,7 @@ __device__ inline void LayerNormSmall(const T val, const cub::KeyValuePair if (threadIdx.x < ld) { const T g(gamma[threadIdx.x]); - const T b(beta[threadIdx.x]); + const T b = (nullptr == beta) ? (T)0 : beta[threadIdx.x]; output[idx] = g * (val - mu) * rsigma + b; } } diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc index b8238f7690..2af85ca89f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc @@ -65,14 +65,16 @@ Status SkipLayerNorm::ComputeInternal(OpKernelContext* ctx) const { "Last dimension of gamma and input does not match"); } - const auto& beta_dims = beta->Shape().GetDims(); - if (beta_dims.size() != 1) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "beta is expected to have 1 dimension, got ", beta_dims.size()); - } - if (beta_dims[0] != input_dims[2]) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Last dimension of beta and input does not match"); + if (nullptr != beta) { + const auto& beta_dims = beta->Shape().GetDims(); + if (beta_dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "beta is expected to have 1 dimension, got ", beta_dims.size()); + } + if (beta_dims[0] != input_dims[2]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Last dimension of beta and input does not match"); + } } if (nullptr != bias) { @@ -98,7 +100,7 @@ Status SkipLayerNorm::ComputeInternal(OpKernelContext* ctx) const { input->template Data(), skip->template Data(), gamma->template Data(), - beta->template Data(), + beta != nullptr ? beta->template Data() : nullptr, bias != nullptr ? bias->template Data() : nullptr, epsilon_, hidden_size, diff --git a/onnxruntime/contrib_ops/cuda/layer_norm.cc b/onnxruntime/contrib_ops/cuda/layer_norm.cc index 12f37f36a0..bbd292ffa9 100644 --- a/onnxruntime/contrib_ops/cuda/layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/layer_norm.cc @@ -59,7 +59,7 @@ Status LayerNorm::ComputeInternal(OpKernelContext* ctx) const auto X_data = reinterpret_cast(X->template Data()); auto scale_data = reinterpret_cast(scale->template Data()); - auto bias_data = simplified ? nullptr: reinterpret_cast(bias->template Data()); + auto bias_data = (simplified || (nullptr == bias)) ? nullptr: reinterpret_cast(bias->template Data()); const TensorShape& x_shape = X->Shape(); const int64_t axis = HandleNegativeAxis(axis_, x_shape.NumDimensions()); @@ -91,7 +91,7 @@ Status LayerNorm::ComputeInternal(OpKernelContext* ctx) const mean_data = reinterpret_cast(mean->template MutableData()); } } - + Tensor* var = ctx->Output(output_index, TensorShape(mean_inv_std_var_dim)); CudaU* inv_var_data = nullptr; if (var != nullptr) { diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 0112963692..54bcbd2b09 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -546,7 +546,7 @@ GELU (Gaussian Error Linear Unit) approximation: Y=0.5*X*(1+tanh(0.797885*X+0.03 .Input(0, "input", "3D input tensor with shape (batch_size, sequence_length, hidden_size)", "T") .Input(1, "skip", "3D skip tensor with shape (batch_size, sequence_length, hidden_size)", "T") .Input(2, "gamma", "1D input tensor with shape (hidden_size)", "T") - .Input(3, "beta", "1D skip tensor with shape (hidden_size", "T") + .Input(3, "beta", "1D skip tensor with shape (hidden_size", "T", OpSchema::Optional) .Input(4, "bias", "1D bias tensor with shape (hidden_size", "T", OpSchema::Optional) .Output(0, "output", "3D output tensor with shape (batch_size, sequence_length, hidden_size)", "T") .Output(1, "mean", "Saved mean used during training to speed up gradient computation", "U", OpSchema::Optional) @@ -2085,7 +2085,7 @@ Example 4: .AllowUncheckedAttributes() .Input(0, "X", "Input data tensor from the previous layer.", "T") .Input(1, "scale", "Scale tensor.", "T") - .Input(2, "B", "Bias tensor.", "T") + .Input(2, "B", "Bias tensor.", "T", OpSchema::Optional) .Output(0, "Y", "Output data tensor.", "T") .Output(1, "mean", "Saved mean used during training to speed up gradient computation", "U", OpSchema::Optional) .Output(2, "inv_std_var", "Saved inverse standard variance used during training to speed up gradient computation.", "U", OpSchema::Optional) diff --git a/onnxruntime/test/contrib_ops/layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/layer_norm_op_test.cc index c8c62434d2..23dd469228 100644 --- a/onnxruntime/test/contrib_ops/layer_norm_op_test.cc +++ b/onnxruntime/test/contrib_ops/layer_norm_op_test.cc @@ -43,5 +43,28 @@ TEST(LayerNormTest, BERTLayerNorm) { tester.Run(); } +TEST(LayerNormTest, BERTLayerNorm_NoBias) { + OpTester tester("LayerNormalization", 1 /*opset_version*/); + tester.AddAttribute("axis", -1); + tester.AddAttribute("epsilon", 1e-12f); + + // create rand inputs + RandomValueGenerator random{}; + + std::vector X_dims{4, 128}; + std::vector X_data = random.Uniform(X_dims, 0.0f, 1.0f); + tester.AddInput("X", X_dims, X_data); + + std::vector scale_dims{128}; + std::vector scale_data = random.Uniform(scale_dims, 0.0f, 1.0f); + tester.AddInput("Scale", scale_dims, scale_data); + + tester.AddMissingOptionalInput(); + + tester.AddReferenceOutputs("testdata/layernorm_no_bias.onnx"); + + tester.Run(); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/layer_norm_test.cc b/onnxruntime/test/contrib_ops/layer_norm_test.cc index b963efa8d6..abe39ac4ff 100644 --- a/onnxruntime/test/contrib_ops/layer_norm_test.cc +++ b/onnxruntime/test/contrib_ops/layer_norm_test.cc @@ -29,7 +29,8 @@ static void TestLayerNorm(const std::vector& x_dims, const std::string& op, optional epsilon, int64_t axis = -1, - int64_t keep_dims = 1) { + int64_t keep_dims = 1, + bool no_bias = false) { const std::vector& n_x_m_dims = x_dims; std::vector n_dims, m_dims; ASSERT_TRUE(SplitDims(n_x_m_dims, axis, n_dims, m_dims).IsOK()); @@ -41,7 +42,7 @@ static void TestLayerNorm(const std::vector& x_dims, ASSERT_NE(keep_dims, 0); const std::vector& stats_dims = keep_dims ? n_and_ones_dims : n_dims; - + CompareOpTester test(op.c_str()); test.AddAttribute("axis", axis); test.AddAttribute("keep_dims", keep_dims); @@ -57,7 +58,7 @@ static void TestLayerNorm(const std::vector& x_dims, test.AddInput("X", n_x_m_dims, X_data); test.AddInput("scale", m_dims, scale_data, true); - if (op.compare(SIMPLIFIED_LAYER_NORM_OP) != 0) { + if (op.compare(SIMPLIFIED_LAYER_NORM_OP) != 0 && no_bias == false) { test.AddInput("B", m_dims, B_data, true); } @@ -99,6 +100,14 @@ TEST(CudaKernelTest, LayerNorm_LargeSizeTensor) { TestLayerNorm(X_dims, LAYER_NORM_OP, k_epsilon_default); } +TEST(CudaKernelTest, LayerNorm_MidSizeTensor_NoBias) { + std::vector X_dims{8, 80, 768}; + const int64_t axis = -1; + const int64_t keep_dims = 1; + const bool no_bias = true; + TestLayerNorm(X_dims, LAYER_NORM_OP, k_epsilon_default, axis, keep_dims, no_bias); +} + TEST(CudaKernelTest, SimplifiedLayerNorm_SmallSizeTensor) { const std::vector X_dims{4, 20, 128}; TestLayerNorm(X_dims, SIMPLIFIED_LAYER_NORM_OP, k_epsilon_default); diff --git a/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc b/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc index 4b53234216..d571f43628 100644 --- a/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc +++ b/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc @@ -21,7 +21,8 @@ static void RunTest( int batch_size, int sequence_length, int hidden_size, - bool use_float16 = false) { + bool use_float16 = false, + bool no_beta = false) { // Input and output shapes // Input 0 - input: (batch_size, sequence_length, hidden_size) // Input 1 - skip : (batch_size, sequence_length, hidden_size) @@ -40,7 +41,11 @@ static void RunTest( test.AddInput("input", input_dims, input_data); test.AddInput("skip", skip_dims, skip_data); test.AddInput("gamma", gamma_dims, gamma_data); - test.AddInput("beta", beta_dims, beta_data); + if (!no_beta) { + test.AddInput("beta", beta_dims, beta_data); + } else { + test.AddMissingOptionalInput(); + } test.AddAttribute("epsilon", epsilon); if (!bias_data.empty()) { test.AddInput("bias", bias_dims, bias_data); @@ -53,7 +58,11 @@ static void RunTest( test.AddInput("input", input_dims, ToFloat16(input_data)); test.AddInput("skip", skip_dims, ToFloat16(skip_data)); test.AddInput("gamma", gamma_dims, ToFloat16(gamma_data)); - test.AddInput("beta", beta_dims, ToFloat16(beta_data)); + if (!no_beta) { + test.AddInput("beta", beta_dims, ToFloat16(beta_data)); + } else { + test.AddMissingOptionalInput(); + } test.AddAttribute("epsilon", epsilon); if (!bias_data.empty()) { test.AddInput("bias", bias_dims, ToFloat16(bias_data)); @@ -138,6 +147,42 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch1_Float16) { true); } +TEST(SkipLayerNormTest, SkipLayerNormBatch1_NoBeta) { + int batch_size = 1; + int sequence_length = 2; + int hidden_size = 4; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector skip_data = { + 0.1f, -0.2f, 0.3f, 1.0f, + 0.5f, 0.1f, 0.4f, 1.6f}; + + std::vector gamma_data = { + 0.3f, 0.2f, 4.0f, 2.2f}; + + std::vector beta_data = {}; + + std::vector output_data = { + 0.08433859348297119f, -0.27090578377246857f, -1.32897164821624756f, 3.0924152374267578f, + 0.26111652255058289f, -0.31333980560302734f, -0.69631003737449646f, 1.9148544311523438f}; + + RunTest(input_data, + skip_data, + gamma_data, + beta_data, + std::vector(), + output_data, + epsilon_, + batch_size, + sequence_length, + hidden_size, + false, + true); +} + TEST(SkipLayerNormTest, SkipLayerNormBatch2) { int batch_size = 2; int sequence_length = 2; diff --git a/onnxruntime/test/testdata/layernorm_no_bias.onnx b/onnxruntime/test/testdata/layernorm_no_bias.onnx new file mode 100644 index 0000000000..3c8b5740c2 Binary files /dev/null and b/onnxruntime/test/testdata/layernorm_no_bias.onnx differ