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,