From 94e34ace15f4c7011590d387c5dcabaa3e446d7f Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Tue, 27 Sep 2022 14:24:16 +0800 Subject: [PATCH] Bugfix for SimplifiedLayerNormalization (#12975) This PR is to fix https://github.com/microsoft/onnxruntime/issues/12930 and https://github.com/microsoft/onnxruntime/issues/12579. In detail: - For CPU EP, since current impl of SimplifiedLayerNormalization doesn't support input and scale having different data types, so if the sub-graph contains Cast Op, the sub-graph will not fused, this guarantee that both inputs and output data type will be same - For CUDA EP, add (fp16, float) support to (T,V) type constraints all combinations of fp16 and float can be supported in the impl With the fix, the original model can be run with SimplifiedLayerNormalization, which also helps to improve the perf. --- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 4 ++ onnxruntime/contrib_ops/cuda/layer_norm.cc | 1 + .../contrib_ops/rocm/rocm_contrib_kernels.cc | 4 ++ .../core/optimizer/layer_norm_fusion.cc | 58 +++++++++++-------- .../core/providers/cuda/nn/layer_norm.cc | 3 + .../core/providers/cuda/nn/layer_norm_impl.cu | 2 + .../test/optimizer/graph_transform_test.cc | 30 +++++++++- .../cuda/cuda_training_kernels.cc | 6 ++ .../training_ops/cuda/nn/layer_norm.cc | 1 + .../training_ops/cuda/nn/layer_norm_impl.cu | 6 +- .../rocm/rocm_training_kernels.cc | 6 ++ 11 files changed, 91 insertions(+), 30 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 1af113f4c5..f63490f5e7 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -80,11 +80,13 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kO class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 16, double_double_double, LayerNormalization); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 16, MLFloat16_float_MLFloat16, LayerNormalization); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 16, float_float_MLFloat16, LayerNormalization); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 16, MLFloat16_float_float, LayerNormalization); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 16, BFloat16_float_BFloat16, LayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float_float_float, SimplifiedLayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double_double_double, SimplifiedLayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16_float_MLFloat16, SimplifiedLayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float_float_MLFloat16, SimplifiedLayerNormalization); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16_float_float, SimplifiedLayerNormalization); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Inverse); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Trilu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int8_t_MLFloat16, QuantizeLinear); @@ -185,11 +187,13 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/layer_norm.cc b/onnxruntime/contrib_ops/cuda/layer_norm.cc index 7736051fd7..fad9ce7a6b 100644 --- a/onnxruntime/contrib_ops/cuda/layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/layer_norm.cc @@ -29,6 +29,7 @@ REGISTER_KERNEL_TYPED(float, float, float) REGISTER_KERNEL_TYPED(double, double, double) REGISTER_KERNEL_TYPED(MLFloat16, float, MLFloat16) REGISTER_KERNEL_TYPED(float, float, MLFloat16) +REGISTER_KERNEL_TYPED(MLFloat16, float, float) REGISTER_KERNEL_TYPED(BFloat16, float, BFloat16) } // namespace cuda diff --git a/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc b/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc index 0837866cc3..d997213a0a 100644 --- a/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc @@ -78,11 +78,13 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kO class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 16, double_double_double, LayerNormalization); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 16, MLFloat16_float_MLFloat16, LayerNormalization); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 16, float_float_MLFloat16, LayerNormalization); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 16, MLFloat16_float_float, LayerNormalization); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 16, BFloat16_float_BFloat16, LayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, float_float_float, SimplifiedLayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, double_double_double, SimplifiedLayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, MLFloat16_float_MLFloat16, SimplifiedLayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, float_float_MLFloat16, SimplifiedLayerNormalization); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, MLFloat16_float_float, SimplifiedLayerNormalization); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, Inverse); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, Trilu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, int8_t_MLFloat16, QuantizeLinear); @@ -182,11 +184,13 @@ Status RegisterRocmContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/optimizer/layer_norm_fusion.cc b/onnxruntime/core/optimizer/layer_norm_fusion.cc index 63e530cf1a..ef47f9b064 100644 --- a/onnxruntime/core/optimizer/layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/layer_norm_fusion.cc @@ -447,23 +447,22 @@ X ------> SimplifiedLayerNormalization ^ Scale --------| */ -Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const { +Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, + const logging::Logger& logger) const { GraphViewer graph_viewer(graph); const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder(); InlinedVector> nodes_to_remove; for (auto node_index : node_topology_list) { nodes_to_remove.clear(); auto* p_pow = graph.GetNode(node_index); - if (p_pow == nullptr) - continue; // we removed the node as part of an earlier fusion + if (p_pow == nullptr) continue; // we removed the node as part of an earlier fusion Node& pow_node = *p_pow; ORT_RETURN_IF_ERROR(Recurse(pow_node, modified, graph_level, logger)); if (!graph_utils::IsSupportedOptypeVersionAndDomain(pow_node, "Pow", {7, 12, 13, 15}) || !graph_utils::IsSupportedProvider(pow_node, GetCompatibleExecutionProviders()) || - !optimizer_utils::CheckOutputEdges(graph, pow_node, 1) || - graph.NodeProducesGraphOutput(pow_node) || + !optimizer_utils::CheckOutputEdges(graph, pow_node, 1) || graph.NodeProducesGraphOutput(pow_node) || !IsSupportedDataType(pow_node)) { continue; } @@ -478,8 +477,7 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr Node& reduce_mean_node = *graph.GetNode(p_reduce_mean->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(reduce_mean_node, "ReduceMean", {1, 11, 13}) || reduce_mean_node.GetExecutionProviderType() != pow_node.GetExecutionProviderType() || - !optimizer_utils::CheckOutputEdges(graph, reduce_mean_node, 1) || - !IsSupportedDataType(reduce_mean_node) || + !optimizer_utils::CheckOutputEdges(graph, reduce_mean_node, 1) || !IsSupportedDataType(reduce_mean_node) || reduce_mean_node.GetInputEdgesCount() == 0) { continue; } @@ -492,8 +490,7 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr Node& add_node = *graph.GetNode(p_add->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(add_node, "Add", {7, 13, 14}) || add_node.GetExecutionProviderType() != pow_node.GetExecutionProviderType() || - !optimizer_utils::CheckOutputEdges(graph, add_node, 1) || - !IsSupportedDataType(add_node)) { + !optimizer_utils::CheckOutputEdges(graph, add_node, 1) || !IsSupportedDataType(add_node)) { continue; } nodes_to_remove.push_back(add_node); @@ -506,8 +503,7 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr if (!graph_utils::IsSupportedOptypeVersionAndDomain(sqrt_node, "Sqrt", {6, 13}) || sqrt_node.GetExecutionProviderType() != pow_node.GetExecutionProviderType() || - !optimizer_utils::CheckOutputEdges(graph, sqrt_node, 1) || - !IsSupportedDataType(sqrt_node) || + !optimizer_utils::CheckOutputEdges(graph, sqrt_node, 1) || !IsSupportedDataType(sqrt_node) || sqrt_node.GetInputEdgesCount() == 0) { continue; } @@ -520,8 +516,7 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr Node& div_node = *graph.GetNode(p_div->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(div_node, "Div", {7, 13, 14}) || div_node.GetExecutionProviderType() != pow_node.GetExecutionProviderType() || - !optimizer_utils::CheckOutputEdges(graph, div_node, 1) || - !IsSupportedDataType(div_node)) { + !optimizer_utils::CheckOutputEdges(graph, div_node, 1) || !IsSupportedDataType(div_node)) { continue; } nodes_to_remove.push_back(div_node); @@ -533,9 +528,22 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr continue; } + // There are only 4 possible cases (x=Pow->ReduceMean->Add->Sqrt->Div and cannot on fp16 type, y=Mul): + // 1. Cast(to:float)->x->Cast(to:fp16)->y : SimplifiedLayerNorm(T:fp16,V:fp16) + // 2. Cast(to:float)->x->y : SimplifiedLayerNorm(T:fp16,V:float) + // 3. x->Cast(to:fp16)->y : SimplifiedLayerNorm(T:float,V:fp16) + // 4. x->y : SimplifiedLayerNorm(T:float,V:float) + // They all work for GPU EP. + // For CPU EP, we have only SimplifiedlayerNorm(T:float,V:float) implementation, so only #4 works. But for #1 and + // #2, if we treat the entry Cast as a normal node, meaning has_leading_cast is false, then for #2, we can still + // fuse it to "Cast(to:float)->SimplifiedlayerNorm(T:float,V:float)" (same as applying #4 to the x->y after + // Cast), so the condition for CPU EP to fuse or not is always setting has_leading_cast to false and checking if + // there is a Cast between x and y. Having Cast between means cannot fuse. const Node* p_pow_input_node = graph_utils::GetInputNode(pow_node, 0); bool has_leading_cast = false; - if (p_pow_input_node) { + bool is_gpu_ep = pow_node.GetExecutionProviderType() == kCudaExecutionProvider || + pow_node.GetExecutionProviderType() == kRocmExecutionProvider; + if (is_gpu_ep && p_pow_input_node) { Node& pow_input_node = *graph.GetNode(p_pow_input_node->Index()); // If input to Pow is a Cast, and the Cast has 2 consumers only (Pow, Div) if (graph_utils::IsSupportedOptypeVersionAndDomain(pow_input_node, "Cast", {9, 13}) && @@ -550,14 +558,14 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr Node* next_node = graph.GetNode(div_node.OutputNodesBegin()->Index()); if (graph_utils::IsSupportedOptypeVersionAndDomain(*next_node, "Cast", {9, 13}) && optimizer_utils::CheckOutputEdges(graph, *next_node, 1)) { + if (!is_gpu_ep) continue; nodes_to_remove.push_back(*next_node); next_node = graph.GetNode(next_node->OutputNodesBegin()->Index()); } Node& mul_node = *next_node; if (!graph_utils::IsSupportedOptypeVersionAndDomain(mul_node, "Mul", {7, 13, 14}) || - mul_node.GetExecutionProviderType() != pow_node.GetExecutionProviderType() || - !IsSupportedDataType(mul_node)) { + mul_node.GetExecutionProviderType() != pow_node.GetExecutionProviderType() || !IsSupportedDataType(mul_node)) { continue; } nodes_to_remove.push_back(mul_node); @@ -597,16 +605,14 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr NodeArg* x_input = has_leading_cast ? graph.GetNode(p_pow_input_node->Index())->MutableInputDefs()[0] : pow_node.MutableInputDefs()[0]; InlinedVector layer_norm_input_defs{x_input, scale}; - Node& layer_norm_node = graph.AddNode(graph.GenerateNodeName("SimplifiedLayerNormalization"), - "SimplifiedLayerNormalization", - "fused LayerNorm subgraphs ", - layer_norm_input_defs, - {}, {}, kOnnxDomain); + Node& layer_norm_node = + graph.AddNode(graph.GenerateNodeName("SimplifiedLayerNormalization"), "SimplifiedLayerNormalization", + "fused LayerNorm subgraphs ", layer_norm_input_defs, {}, {}, kOnnxDomain); // Get constant "epsilon" from "Add" node if available. Else, default value will be used. - const ONNX_NAMESPACE::TensorProto* tensor_proto = graph_utils::GetConstantInitializer(graph, add_node.MutableInputDefs()[1]->Name()); - if (tensor_proto != nullptr && - tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { + const ONNX_NAMESPACE::TensorProto* tensor_proto = + graph_utils::GetConstantInitializer(graph, add_node.MutableInputDefs()[1]->Name()); + 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 { @@ -629,11 +635,13 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr #ifdef ENABLE_TRAINING // add one extra output def, so we have 2 output defs that match what gradient builder expected - layer_norm_node.MutableOutputDefs().push_back(&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("saved_inv_std_var"), nullptr)); + layer_norm_node.MutableOutputDefs().push_back( + &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("saved_inv_std_var"), nullptr)); #endif modified = true; } return Status::OK(); } + } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/nn/layer_norm.cc b/onnxruntime/core/providers/cuda/nn/layer_norm.cc index 981dfc988d..365bf197e8 100644 --- a/onnxruntime/core/providers/cuda/nn/layer_norm.cc +++ b/onnxruntime/core/providers/cuda/nn/layer_norm.cc @@ -103,15 +103,18 @@ Status LayerNorm::ComputeInternal(OpKernelContext* ctx) con template class LayerNorm; // contrib op usage +LAYERNORM_IMPL(float, float, float, false) LAYERNORM_IMPL(double, double, double, false) LAYERNORM_IMPL(MLFloat16, float, MLFloat16, false) LAYERNORM_IMPL(float, float, MLFloat16, false) +LAYERNORM_IMPL(MLFloat16, float, float, false) LAYERNORM_IMPL(BFloat16, float, BFloat16, false) LAYERNORM_IMPL(float, float, float, true) LAYERNORM_IMPL(double, double, double, true) LAYERNORM_IMPL(MLFloat16, float, MLFloat16, true) LAYERNORM_IMPL(float, float, MLFloat16, true) +LAYERNORM_IMPL(MLFloat16, float, float, true) LAYERNORM_IMPL(BFloat16, float, BFloat16, true) #endif } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/nn/layer_norm_impl.cu b/onnxruntime/core/providers/cuda/nn/layer_norm_impl.cu index 6f338ab03e..717b994001 100644 --- a/onnxruntime/core/providers/cuda/nn/layer_norm_impl.cu +++ b/onnxruntime/core/providers/cuda/nn/layer_norm_impl.cu @@ -390,11 +390,13 @@ LAYERNORM_LINEAR_IMPL(float, float, float, true) LAYERNORM_LINEAR_IMPL(half, float, half, true) LAYERNORM_LINEAR_IMPL(double, double, double, true) LAYERNORM_LINEAR_IMPL(float, float, half, true) +LAYERNORM_LINEAR_IMPL(half, float, float, true) LAYERNORM_LINEAR_IMPL(float, float, float, false) LAYERNORM_LINEAR_IMPL(half, float, half, false) LAYERNORM_LINEAR_IMPL(double, double, double, false) LAYERNORM_LINEAR_IMPL(double, float, double, false) LAYERNORM_LINEAR_IMPL(float, float, half, false) +LAYERNORM_LINEAR_IMPL(half, float, float, false) LAYERNORM_LINEAR_IMPL(BFloat16, float, BFloat16, true) LAYERNORM_LINEAR_IMPL(BFloat16, float, BFloat16, false) diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 63d40717d9..730dcf16c2 100755 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -3968,6 +3968,8 @@ TEST_F(GraphTransformationTests, SimplifiedLayerNormFusionTest) { } } +// If EP is non-GPU EP or unknown, the sub-graph will be not fused because CPU impl for SimplifiedLayerNormalization +// doesn't support input and scale having different data types. TEST_F(GraphTransformationTests, SimplifiedLayerNormWithCastsFusionTest) { auto model_uri = MODEL_FOLDER "fusion/simplified_layer_norm_with_casts.onnx"; std::shared_ptr p_model; @@ -3976,7 +3978,27 @@ TEST_F(GraphTransformationTests, SimplifiedLayerNormWithCastsFusionTest) { InlinedHashSet compatible_eps; onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; - ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(compatible_eps), TransformerLevel::Level2)); + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(compatible_eps), + TransformerLevel::Level2)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_)); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_TRUE(op_to_count["SimplifiedLayerNormalization"] == 0); +} + +TEST_F(GraphTransformationTests, SimplifiedLayerNormWithCastsFusionTestCudaEp) { + auto model_uri = MODEL_FOLDER "fusion/simplified_layer_norm_with_casts.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + for (auto& node : graph.Nodes()) { + node.SetExecutionProviderType(kCudaExecutionProvider); + } + + InlinedHashSet compatible_eps; + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(compatible_eps), + TransformerLevel::Level2)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_)); std::map op_to_count = CountOpsInGraph(graph); @@ -3991,10 +4013,12 @@ TEST_F(GraphTransformationTests, SimplifiedLayerNormWithCastsFusionTest) { for (const Node& node : graph.Nodes()) { if (node.OpType() == "SimplifiedLayerNormalization") { // LayerNormalization should have two inputs. - EXPECT_EQ(node.InputDefs().size(), 2u) << "LayerNormalization number of inputs does not equal to 2. Got:" << node.InputDefs().size(); + EXPECT_EQ(node.InputDefs().size(), 2u) + << "LayerNormalization number of inputs does not equal to 2. Got:" << node.InputDefs().size(); // LayerNormalization input "scale" and "bias" should have the same dimension. const TensorShapeProto* scale_shape = node.InputDefs()[1]->Shape(); - EXPECT_EQ(scale_shape->dim_size(), 1) << "LayerNormalization scale should be 1D. Got: " << scale_shape->dim_size(); + EXPECT_EQ(scale_shape->dim_size(), 1) + << "LayerNormalization scale should be 1D. Got: " << scale_shape->dim_size(); } else if (node.OpType() == "Cast") { continue; } else { diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index 4652eb3d86..5265f908fe 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc @@ -133,14 +133,17 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double_double_double, LayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_float_MLFloat16, LayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_MLFloat16, LayerNormalizationGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_float_float, LayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_float, SimplifiedLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double_double_double, SimplifiedLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_float_MLFloat16, SimplifiedLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_MLFloat16, SimplifiedLayerNormalizationGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_float_float, SimplifiedLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_float, InvertibleLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double_double_double, InvertibleLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_float_MLFloat16, InvertibleLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_MLFloat16, InvertibleLayerNormalizationGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_float_float, InvertibleLayerNormalizationGrad); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, SliceGrad); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, GatherElementsGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Scale); @@ -352,14 +355,17 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc b/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc index c46b32791a..61913d3260 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc +++ b/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc @@ -33,6 +33,7 @@ REGISTER_GRADIENT_KERNEL_TYPED(float, float, float) REGISTER_GRADIENT_KERNEL_TYPED(double, double, double) REGISTER_GRADIENT_KERNEL_TYPED(MLFloat16, float, MLFloat16) REGISTER_GRADIENT_KERNEL_TYPED(float, float, MLFloat16) +REGISTER_GRADIENT_KERNEL_TYPED(MLFloat16, float, float) REGISTER_GRADIENT_KERNEL_TYPED(BFloat16, float, BFloat16) template diff --git a/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu b/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu index 11eab0ca58..2d89ed0571 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu @@ -49,7 +49,7 @@ namespace { // https://github.com/NVIDIA/apex/issues/246 template struct SharedMemory; - + template <> struct SharedMemory { __device__ float* getPointer() { @@ -57,7 +57,7 @@ namespace { return s_float; } }; - + template <> struct SharedMemory { __device__ double* getPointer() { @@ -609,10 +609,12 @@ LAYERNORMGRAD_IMPL(float, float, float, true) LAYERNORMGRAD_IMPL(double, double, double, true) LAYERNORMGRAD_IMPL(half, float, half, true) LAYERNORMGRAD_IMPL(float, float, half, true) +LAYERNORMGRAD_IMPL(half, float, float, true) LAYERNORMGRAD_IMPL(float, float, float, false) LAYERNORMGRAD_IMPL(double, double, double, false) LAYERNORMGRAD_IMPL(half, float, half, false) LAYERNORMGRAD_IMPL(float, float, half, false) +LAYERNORMGRAD_IMPL(half, float, float, false) LAYERNORMGRAD_IMPL(BFloat16, float, BFloat16, true) LAYERNORMGRAD_IMPL(BFloat16, float, BFloat16, false) diff --git a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc index 3ffdd8ecfe..f3964a3173 100644 --- a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc +++ b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc @@ -130,14 +130,17 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double_double_double, LayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_float_MLFloat16, LayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_float_MLFloat16, LayerNormalizationGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_float_float, LayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_float_float, SimplifiedLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double_double_double, SimplifiedLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_float_MLFloat16, SimplifiedLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_float_MLFloat16, SimplifiedLayerNormalizationGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_float_float, SimplifiedLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_float_float, InvertibleLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double_double_double, InvertibleLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_float_MLFloat16, InvertibleLayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_float_MLFloat16, InvertibleLayerNormalizationGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_float_float, InvertibleLayerNormalizationGrad); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, SliceGrad); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, GatherElementsGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, Scale); @@ -318,14 +321,17 @@ Status RegisterRocmTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo,