diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc index 170f2ac0c7..2ebd9484aa 100644 --- a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc @@ -25,12 +25,14 @@ namespace contrib { REGISTER_KERNEL_TYPED(float) template -EmbedLayerNorm::EmbedLayerNorm(const OpKernelInfo& info) : OpKernel(info) {} +EmbedLayerNorm::EmbedLayerNorm(const OpKernelInfo& op_kernel_info) : OpKernel(op_kernel_info) { + ORT_ENFORCE(op_kernel_info.GetAttr("epsilon", &epsilon_).IsOK()); + ORT_ENFORCE(epsilon_ >= 0); +} 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* word_embedding = context->Input(2); @@ -112,7 +114,7 @@ Status EmbedLayerNorm::Compute(OpKernelContext* context) const { y[i] = a; sum += a * a; } - T e = sqrt(sum / hidden_size + static_cast(1.0e-13)); + T e = sqrt(sum / hidden_size + static_cast(epsilon_)); for (int i = 0; i < hidden_size; i++) { y[i] = y[i] / e * gamma_data[i] + beta_data[i]; } diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.h b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.h index dcd7683c87..5dfe6c833e 100644 --- a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.h +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.h @@ -11,8 +11,10 @@ namespace contrib { template class EmbedLayerNorm : public OpKernel { public: - explicit EmbedLayerNorm(const OpKernelInfo& info); + explicit EmbedLayerNorm(const OpKernelInfo& op_kernel_info); Status Compute(OpKernelContext* context) const override; + private: + float epsilon_; }; } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc b/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc index f8a1443944..d1be62cc0b 100644 --- a/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc +++ b/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc @@ -27,6 +27,8 @@ REGISTER_KERNEL_TYPED(double) template SkipLayerNorm::SkipLayerNorm(const OpKernelInfo& op_kernel_info) : OpKernel(op_kernel_info) { + ORT_ENFORCE(op_kernel_info.GetAttr("epsilon", &epsilon_).IsOK()); + ORT_ENFORCE(epsilon_ >= 0); } template @@ -114,7 +116,7 @@ Status SkipLayerNorm::Compute(OpKernelContext* p_ctx) const { } mean = mean / hidden_size; - mean_square = sqrt(mean_square / hidden_size - mean * mean + float(1e-12)); + 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]; diff --git a/onnxruntime/contrib_ops/cpu/skip_layer_norm.h b/onnxruntime/contrib_ops/cpu/skip_layer_norm.h index 1ede7c8f23..7723541cb6 100644 --- a/onnxruntime/contrib_ops/cpu/skip_layer_norm.h +++ b/onnxruntime/contrib_ops/cpu/skip_layer_norm.h @@ -15,6 +15,9 @@ class SkipLayerNorm final : public OpKernel { public: SkipLayerNorm(const OpKernelInfo& op_kernel_info); Status Compute(OpKernelContext* p_op_kernel_context) const override; + + private: + float epsilon_; }; } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc index d873971e8c..73cc52c822 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc @@ -30,6 +30,8 @@ using namespace ONNX_NAMESPACE; template EmbedLayerNorm::EmbedLayerNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) { + ORT_ENFORCE(op_kernel_info.GetAttr("epsilon", &epsilon_).IsOK()); + ORT_ENFORCE(epsilon_ >= 0); } template @@ -76,7 +78,8 @@ Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { word_embedding->template Data(), position_embedding->template Data(), segment_embedding->template Data(), - static_cast(hidden_size), + epsilon_, + static_cast(hidden_size), batch_size, sequence_length, element_size)) { diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.h b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.h index c9ff25c0cd..7d64460efb 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.h +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.h @@ -17,6 +17,8 @@ class EmbedLayerNorm final : public CudaKernel { public: EmbedLayerNorm(const OpKernelInfo& op_kernel_info); Status ComputeInternal(OpKernelContext* ctx) const override; + private: + float epsilon_; }; } // namespace cuda 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 2b21a44961..f5f6e64755 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -22,6 +22,7 @@ limitations under the License. #include "layer_norm.cuh" #include "embed_layer_norm_impl.h" +#include using namespace onnxruntime::cuda; using namespace cub; @@ -106,7 +107,7 @@ template __global__ void EmbedLayerNormKernel( int hidden_size, const int* input_ids, const int* segment_ids, const T* beta, const T* gamma, const T* word_embedding, const T* position_embedding, const T* segment_embedding, - T* output) { + const T epsilon, T* output) { KeyValuePairSum pair_sum; // 1. lookup word and segment of the block // blockIdx.x = position in the sequence @@ -146,7 +147,7 @@ __global__ void EmbedLayerNormKernel( } // 3. layer norm on the sum - LayerNorm(thread_data, hidden_size, output_offset, beta, gamma, output); + LayerNorm(thread_data, hidden_size, output_offset, beta, gamma, epsilon, output); } template @@ -154,13 +155,13 @@ bool EmbedSkipLayerNorm( cudaStream_t stream, int hidden_size, int batch_size, int sequence_length, const int* input_ids, const int* segment_ids, const T* beta, const T* gamma, const T* word_embedding, const T* position_embedding, const T* segment_embedding, - T* output) { + const T epsilon, T* output) { constexpr int tpb = 256; const dim3 grid(sequence_length, batch_size, 1); const dim3 block(tpb, 1, 1); EmbedLayerNormKernel - <<>>(hidden_size, input_ids, segment_ids, beta, gamma, word_embedding, position_embedding, segment_embedding, output); + <<>>(hidden_size, input_ids, segment_ids, beta, gamma, word_embedding, position_embedding, segment_embedding, epsilon, output); return CUDA_CALL(cudaPeekAtLastError()); } @@ -176,6 +177,7 @@ bool LaunchEmbedLayerNormKernel( const void* word_embedding, const void* position_embedding, const void* segment_embedding, + float epsilon, const int hidden_size, int batch_size, int sequence_length, @@ -193,13 +195,15 @@ 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(segment_embedding), + reinterpret_cast(word_embedding), reinterpret_cast(position_embedding), + reinterpret_cast(segment_embedding), __float2half_rn(epsilon), reinterpret_cast(output)); } else { 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(segment_embedding), + reinterpret_cast(word_embedding), reinterpret_cast(position_embedding), + reinterpret_cast(segment_embedding), epsilon, reinterpret_cast(output)); } } diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h index e21131bb1a..18648e6799 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h @@ -16,6 +16,7 @@ bool LaunchEmbedLayerNormKernel(void* output, // output tens const void* word_embedding, // weights for word embeddings const void* position_embedding, // weights for position embeddings const void* segment_embedding, // weights for segment (like sentence) embeddings + float epsilon, // epsilon for layer normalization const int hidden_size, // hidden size (that is head_size * num_heads) int batch_size, // batch size int sequence_length, // sequence length diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index 8a639633f0..16c1469f22 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -79,7 +79,8 @@ struct KeyValuePairSum { template __device__ inline void LayerNorm( - const cub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, const T* gamma, T* output) { + 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 using BlockReduce = cub::BlockReduce, TPB>; @@ -92,7 +93,7 @@ __device__ inline void LayerNorm( if (threadIdx.x == 0) { mu = sum_kv.key; - rsigma = Rsqrt(sum_kv.value - mu * mu); + rsigma = Rsqrt(sum_kv.value - mu * mu + epsilon); } __syncthreads(); @@ -107,7 +108,7 @@ __device__ inline void LayerNorm( template __device__ inline void LayerNormSmall(const T val, const cub::KeyValuePair& thread_data, const int ld, const int idx, - const T* beta, const T* gamma, T* output) { + const T* beta, const T* gamma, const T epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register @@ -122,7 +123,7 @@ __device__ inline void LayerNormSmall(const T val, const cub::KeyValuePair if (threadIdx.x == 0) { mu = sum_kv.key; - rsigma = Rsqrt(sum_kv.value - mu * mu); + rsigma = Rsqrt(sum_kv.value - mu * mu + epsilon); } __syncthreads(); diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc index 61a1d64ca1..dc7ed4e592 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc @@ -30,6 +30,8 @@ using namespace ONNX_NAMESPACE; template SkipLayerNorm::SkipLayerNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) { + ORT_ENFORCE(op_kernel_info.GetAttr("epsilon", &epsilon_).IsOK()); + ORT_ENFORCE(epsilon_ >= 0); } template @@ -84,6 +86,7 @@ Status SkipLayerNorm::ComputeInternal(OpKernelContext* ctx) const { "Last dimension of bias and input does not match"); } } + int sequence_length = static_cast(input_dims[1]); int hidden_size = static_cast(input_dims[2]); int64_t element_count = input_dims[0] * sequence_length * hidden_size; @@ -96,6 +99,7 @@ Status SkipLayerNorm::ComputeInternal(OpKernelContext* ctx) const { gamma->template Data(), beta->template Data(), bias != nullptr ? bias->template Data() : nullptr, + epsilon_, hidden_size, static_cast(element_count), //TODO: check range element_size)) { diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.h b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.h index c8b754a947..c66143a8c2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.h +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.h @@ -17,6 +17,9 @@ class SkipLayerNorm final : public CudaKernel { public: SkipLayerNorm(const OpKernelInfo& op_kernel_info); Status ComputeInternal(OpKernelContext* context) const override; + + private: + float epsilon_; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu index f951c752cc..9c11ff85e0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu @@ -22,6 +22,7 @@ limitations under the License. #include "layer_norm.cuh" #include "skip_layer_norm_impl.h" +#include namespace onnxruntime { namespace contrib { @@ -29,7 +30,8 @@ namespace cuda { template __global__ void SkipLayerNormKernelSmall( - const int ld, const T* input, const T* skip, const T* beta, const T* gamma, const T* bias, T* output) { + const int ld, const T* input, const T* skip, const T* beta, const T* gamma, const T* bias, + const T epsilon, T* output) { const T reverse_ld = T(1.f / ld); const int offset = blockIdx.x * ld; @@ -45,12 +47,13 @@ __global__ void SkipLayerNormKernelSmall( thread_data = pair_sum(thread_data, cub::KeyValuePair(rldval, rldval * val)); } - LayerNormSmall(val, thread_data, ld, idx, beta, gamma, output); + LayerNormSmall(val, thread_data, ld, idx, beta, gamma, epsilon, output); } template __global__ void SkipLayerNormKernel( - const int ld, const T* input, const T* skip, const T* beta, const T* gamma, const T* bias, T* output) { + const int ld, const T* input, const T* skip, const T* beta, const T* gamma, const T* bias, + const T epsilon, T* output) { const T reverse_ld = T(1.f / ld); const int offset = blockIdx.x * ld; @@ -66,13 +69,13 @@ __global__ void SkipLayerNormKernel( output[idx] = val; } - LayerNorm(thread_data, ld, offset, beta, gamma, output); + LayerNorm(thread_data, ld, offset, beta, gamma, epsilon, output); } template bool ComputeSkipLayerNorm( cudaStream_t stream, const int ld, const int n, const T* input, const T* skip, - const T* beta, const T* gamma, const T* bias, T* output) { + const T* beta, const T* gamma, const T* bias, const T epsilon, T* output) { // this must be true because n is the total size of the tensor assert(n % ld == 0); const int grid_size = n / ld; @@ -80,18 +83,18 @@ bool ComputeSkipLayerNorm( if (ld <= 32) { constexpr int block_size = 32; SkipLayerNormKernelSmall - <<>>(ld, input, skip, beta, gamma, bias, output); + <<>>(ld, input, skip, beta, gamma, bias, epsilon, output); } else if (ld <= 128) { constexpr int block_size = 128; SkipLayerNormKernelSmall - <<>>(ld, input, skip, beta, gamma, bias, output); + <<>>(ld, input, skip, beta, gamma, bias, epsilon, output); } else if (ld == 384) { constexpr int block_size = 384; SkipLayerNormKernelSmall - <<>>(ld, input, skip, beta, gamma, bias, output); + <<>>(ld, input, skip, beta, gamma, bias, epsilon, output); } else { constexpr int block_size = 256; - SkipLayerNormKernel<<>>(ld, input, skip, beta, gamma, bias, output); + SkipLayerNormKernel<<>>(ld, input, skip, beta, gamma, bias, epsilon, output); } return CUDA_CALL(cudaPeekAtLastError()); } @@ -102,7 +105,8 @@ bool LaunchSkipLayerNormKernel( const void* skip, const void* gamma, const void* beta, - const void* bias, + const void* bias, + float epsilon, int hidden_size, int element_count, size_t element_size) { @@ -119,6 +123,7 @@ bool LaunchSkipLayerNormKernel( reinterpret_cast(beta), reinterpret_cast(gamma), reinterpret_cast(bias), + __float2half_rn(epsilon), reinterpret_cast(output)); } else { return ComputeSkipLayerNorm( @@ -130,6 +135,7 @@ bool LaunchSkipLayerNormKernel( reinterpret_cast(beta), reinterpret_cast(gamma), reinterpret_cast(bias), + epsilon, reinterpret_cast(output)); } } diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.h index e0fb71b3d7..308242c010 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.h @@ -14,6 +14,7 @@ bool LaunchSkipLayerNormKernel( const void* gamma, // Layer normalization gamma tensor const void* beta, // Layer normalization beta tensor const void* bias, // Layer normalization beta tensor + float epsilon, // Layer normalization epsilon int hidden_size, // hidden size, it is the leading dimension (ld) int element_count, // number of elements in input tensor size_t element_size // element size of input tensor diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 3907612782..76a1ab4d79 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -399,6 +399,7 @@ will be calculated.)DOC"; .SinceVersion(1) .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL) .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(2, "word_embedding", "2D with shape (,hidden_size)", "T") @@ -471,6 +472,7 @@ GELU (Gaussian Error Linear Unit) approximation: Y=0.5*X*(1+tanh(0.797885*X+0.03 .SinceVersion(1) .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL) .SetDoc("Skip and Layer Normalization Fusion") + .Attr("epsilon", "The epsilon value to use to avoid division by zero.", AttributeProto::FLOAT, kDefaultSkipLayerNormEpsilon) .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") diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.h b/onnxruntime/core/graph/contrib_ops/contrib_defs.h index 10da0479d7..3dee1d0812 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.h +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.h @@ -26,5 +26,8 @@ namespace contrib { schema_func(ONNX_NAMESPACE::OpSchema(#name, __FILE__, __LINE__)) void RegisterContribSchemas(); + +constexpr const float kDefaultSkipLayerNormEpsilon = 1e-12f; +constexpr const float kDefaultEmbedLayerNormEpsilon = 1e-12f; } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 3a398fa715..2d810d5fc6 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/optimizer/initializer.h" #include "core/optimizer/embed_layer_norm_fusion.h" +#include "core/graph/contrib_ops/contrib_defs.h" #include "core/graph/graph_utils.h" #include "core/optimizer/utils.h" #include "core/framework/tensorprotoutils.h" @@ -12,7 +13,6 @@ using namespace ONNX_NAMESPACE; using namespace onnxruntime::common; namespace onnxruntime { - // Add a Cast to convert Input from int64 to int32. static NodeArg* CastToInt32(Graph& graph, NodeArg* input, ProviderType provider_type) { auto data_type = input->TypeAsProto()->tensor_type().elem_type(); @@ -699,6 +699,16 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l {layer_norm_node.MutableOutputDefs()[0], reduce_sum_node.MutableOutputDefs()[0]}, {}, kMSDomain); + // Get attribute "epsilon" from "LayerNormalization" node if available. Else, default value + // will be used. + NodeAttributes ln_attrs = layer_norm_node.GetAttributes(); + NodeAttributes::const_iterator epsilon = ln_attrs.find("epsilon"); + if (epsilon != ln_attrs.end()) { + embed_layer_norm_node.AddAttribute("epsilon", epsilon->second); + } else { + embed_layer_norm_node.AddAttribute("epsilon", contrib::kDefaultEmbedLayerNormEpsilon); + } + // Assign provider to this new node. Provider should be same as the provider for old node. embed_layer_norm_node.SetExecutionProviderType(layer_norm_node.GetExecutionProviderType()); diff --git a/onnxruntime/core/optimizer/layer_norm_fusion.cc b/onnxruntime/core/optimizer/layer_norm_fusion.cc index dddb62bf80..39a1907452 100644 --- a/onnxruntime/core/optimizer/layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/layer_norm_fusion.cc @@ -12,6 +12,8 @@ namespace onnxruntime { // LayerNorm supports limited data types. static std::vector supported_data_types{"tensor(float16)", "tensor(float)", "tensor(double)"}; +// Default epsilon +static const float DEFAULT_LAYERNORM_EPSILON = 1e-5f; static bool IsSupportedDataType(const Node& node) { for (const auto& input_arg : node.InputDefs()) { @@ -256,11 +258,12 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, // Get constant "epsilon" from "Add2" node if available. Else, default value will be used. const ONNX_NAMESPACE::TensorProto* tensor_proto = graph_utils::GetConstantInitializer(graph, add2_node.MutableInputDefs()[1]->Name()); - if (tensor_proto != nullptr) { - if (tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { + if (tensor_proto != nullptr && + tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { Initializer initializer{*tensor_proto, graph.ModelPath()}; layer_norm_node.AddAttribute("epsilon", initializer.data()[0]); - } + } else { + layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON); } // Assign provider to this new node. Provider should be same as the provider for old node. diff --git a/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc b/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc index ccf9ea8aeb..d4c93b8553 100644 --- a/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/optimizer/initializer.h" #include "core/optimizer/skip_layer_norm_fusion.h" +#include "core/graph/contrib_ops/contrib_defs.h" #include "core/graph/graph_utils.h" #include "float.h" #include @@ -236,6 +237,15 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le skip_layer_norm_input_defs, ln_node.MutableOutputDefs(), {}, kMSDomain); + // Get attribute "epsilon" from "LayerNormalization" node if available. Else, default value + // will be used. + NodeAttributes ln_attrs = ln_node.GetAttributes(); + NodeAttributes::const_iterator epsilon = ln_attrs.find("epsilon"); + if (epsilon != ln_attrs.end()) { + skip_layer_norm_node.AddAttribute("epsilon", epsilon->second); + } else { + skip_layer_norm_node.AddAttribute("epsilon", contrib::kDefaultSkipLayerNormEpsilon); + } // Assign provider to this new node. Provider should be same as the provider for old node. skip_layer_norm_node.SetExecutionProviderType(ln_node.GetExecutionProviderType()); } diff --git a/onnxruntime/python/tools/bert/BertOnnxModel.py b/onnxruntime/python/tools/bert/BertOnnxModel.py index 28ac282f14..3048b610a5 100644 --- a/onnxruntime/python/tools/bert/BertOnnxModel.py +++ b/onnxruntime/python/tools/bert/BertOnnxModel.py @@ -883,6 +883,13 @@ class BertOnnxModel(OnnxModel): embed_node.domain = "com.microsoft" + # Pass attribute "epsilon" from normalize node to EmbedLayerNormalization. + for att in normalize_node.attribute: + if att.name == 'epsilon': + embed_node.attribute.extend([att]) + # Set default value to 1e-12 if no attribute is found. + if len(embed_node.attribute) == 0: + embed_node.attribute.extend([onnx.helper.make_attribute("epsilon", 1.0E-12)]) self.replace_input_of_all_nodes(normalize_node.output[0], 'embed_output') self.remove_nodes(nodes_to_remove) @@ -1136,6 +1143,13 @@ class BertOnnxModel(OnnxModel): outputs=[node.output[0]], name=self.create_node_name("SkipLayerNormalization", name_prefix="SkipLayerNorm")) normalize_node.domain = "com.microsoft" + # Pass attribute "epsilon" from layernorm node to SkipLayerNormalization + for att in node.attribute: + if att.name == 'epsilon': + normalize_node.attribute.extend([att]) + # Set default epsilon if no epsilon exists from layernorm + if len(normalize_node.attribute) == 0: + normalize_node.attribute.extend([onnx.helper.make_attribute("epsilon", 1.0E-12)]) skip_layernorm_nodes.extend([normalize_node]) self.remove_nodes(nodes_to_remove) diff --git a/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc b/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc index 152599eefe..1275a997ac 100644 --- a/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc +++ b/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc @@ -8,6 +8,8 @@ namespace onnxruntime { namespace test { +constexpr float epsilon_ = 1e-12f; + static void RunTest( const std::vector& input_ids_data, const std::vector& segment_ids_data, @@ -19,6 +21,7 @@ static void RunTest( const std::vector& beta_data, const std::vector& output_data, const std::vector& mask_index_data, + float epsilon, int batch_size, int sequence_length, int hidden_size, @@ -69,6 +72,7 @@ static void RunTest( 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); if (has_mask) { tester.AddInput("mask", mask_dims, mask_data); } @@ -79,6 +83,7 @@ static void RunTest( 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); if (has_mask) { tester.AddInput("mask", mask_dims, mask_data); } @@ -143,6 +148,7 @@ TEST(EmbedLayerNormTest, EmbedLayerNormBatch1) { beta_data, output_data, mask_index_data, + epsilon_, batch_size, sequence_length, hidden_size); @@ -202,6 +208,7 @@ TEST(EmbedLayerNormTest, EmbedLayerNormBatch1_Float16) { beta_data, output_data, mask_index_data, + epsilon_, batch_size, sequence_length, hidden_size, @@ -272,6 +279,7 @@ TEST(EmbedLayerNormTest, EmbedLayerNormBatch2) { beta_data, output_data, mask_index_data, + epsilon_, batch_size, sequence_length, hidden_size); @@ -337,6 +345,7 @@ TEST(EmbedLayerNormTest, EmbedLayerNormBatch2_NoMask) { beta_data, output_data, mask_index_data, + epsilon_, batch_size, sequence_length, hidden_size, @@ -419,6 +428,7 @@ TEST(EmbedLayerNormTest, EmbedLayerNormLargeBatchSmallHiddenSize) { beta_data, output_data, mask_index_data, + epsilon_, batch_size, sequence_length, hidden_size); diff --git a/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc b/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc index 0d239ca4a1..4b53234216 100644 --- a/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc +++ b/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc @@ -8,6 +8,7 @@ namespace onnxruntime { namespace test { +constexpr float epsilon_ = 1e-12f; static void RunTest( const std::vector& input_data, @@ -16,6 +17,7 @@ static void RunTest( const std::vector& beta_data, const std::vector& bias_data, const std::vector& output_data, + float epsilon, int batch_size, int sequence_length, int hidden_size, @@ -39,7 +41,7 @@ static void RunTest( test.AddInput("skip", skip_dims, skip_data); test.AddInput("gamma", gamma_dims, gamma_data); test.AddInput("beta", beta_dims, beta_data); - + test.AddAttribute("epsilon", epsilon); if (!bias_data.empty()) { test.AddInput("bias", bias_dims, bias_data); } @@ -52,7 +54,7 @@ static void RunTest( test.AddInput("skip", skip_dims, ToFloat16(skip_data)); test.AddInput("gamma", gamma_dims, ToFloat16(gamma_data)); test.AddInput("beta", beta_dims, ToFloat16(beta_data)); - + test.AddAttribute("epsilon", epsilon); if (!bias_data.empty()) { test.AddInput("bias", bias_dims, ToFloat16(bias_data)); } @@ -94,6 +96,7 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch1) { beta_data, std::vector(), output_data, + epsilon_, batch_size, sequence_length, hidden_size); @@ -128,6 +131,7 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch1_Float16) { beta_data, std::vector(), output_data, + epsilon_, batch_size, sequence_length, hidden_size, @@ -169,6 +173,7 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch2) { beta_data, std::vector(), output_data, + epsilon_, batch_size, sequence_length, hidden_size); @@ -212,6 +217,7 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch2_Bias) { beta_data, bias_data, output_data, + epsilon_, batch_size, sequence_length, hidden_size);