LayerNormFusion - Cast support (#4320)

* cast support for layernormfusion

* cast support for layernormfusion

* bug fix

* fix build

* bug fix

* fix test

* minor refactor

* on comments

* on comments

* on comments

* on comments

Co-authored-by: Ethan Tao <ettao@microsoft.com>

onnxruntime/core/optimizer/layer_norm_fusion.cc

@@ -45,6 +45,16 @@ X --> ReduceMean --> Sub --> Pow --> ReduceMean --> Add --> Sqrt --> Div --> Mul
 |                     ^
 |                     |
 +---------------------+
+
+In recent pytorch, Cast nodes may be inserted before Pow to ensure that both inputs 'base' and 'power' are the same type 
+due to restriction in older opsets. Therefore, Layer Normalization will also handle the case below :
++---------------------+
+|                     |
+|                     v
+X --> ReduceMean --> Sub --> Cast --> Pow --> ReduceMean --> Add --> Sqrt --> Div --> Mul --> Add
+                      |                                                        ^
+                      |                                                        |
+                      +--------------------------------------------------------+
 */
 Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const {
   GraphViewer graph_viewer(graph);
@@ -76,7 +86,6 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
       if ((*iter).OpType().compare("Sub") == 0) {
         if (subCnt == 0) {
           p_sub_node = &(*iter);
-
         } else {
           p_sub_node_dup = &(*iter);
         }
@@ -95,7 +104,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
     Node& sub_node = *graph.GetNode(p_sub_node->Index());
     if (!graph_utils::IsSupportedOptypeVersionAndDomain(sub_node, "Sub", {7}) ||
         sub_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() ||
-        !optimizer_utils::CheckOutputEdges(graph, sub_node, subCnt == 1 ? 2u: 1u) ||
+        !optimizer_utils::CheckOutputEdges(graph, sub_node, subCnt == 1 ? 2u : 1u) ||
         !IsSupportedDataType(sub_node)) {
       continue;
     }
@@ -174,7 +183,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
       continue;
     }
     nodes_to_remove.push_back(reduce_mean2_node);
-    
+
     // Traceback the reduceMean node to find pow --> reduceMean
     Node& pow_node = *graph.GetNode(reduce_mean2_node.InputNodesBegin()->Index());
     if (!graph_utils::IsSupportedOptypeVersionAndDomain(pow_node, "Pow", {7, 12}) ||
@@ -185,8 +194,21 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
     }
     nodes_to_remove.push_back(pow_node);
 
-    // Traceback the pow node to find sub --> pow
-    const Node* p_sub2_node = graph_utils::FirstParentByType(pow_node, "Sub");
+    // check if Cast node exists between sub and pow
+    const Node* p_cast_node = graph_utils::FirstParentByType(pow_node, "Cast");
+    if (p_cast_node != nullptr) {
+      Node& cast_node = *graph.GetNode(pow_node.InputNodesBegin()->Index());
+      if (!graph_utils::IsSupportedOptypeVersionAndDomain(cast_node, "Cast", {9}) ||
+          cast_node.GetExecutionProviderType() != cast_node.GetExecutionProviderType() ||
+          !optimizer_utils::CheckOutputEdges(graph, cast_node, 1) ||
+          !IsSupportedDataType(cast_node)) {
+        continue;
+      }
+      nodes_to_remove.push_back(cast_node);
+    }
+
+    // Traceback from the last node in vector to find sub --> pow  or  sub --> cast
+    const Node* p_sub2_node = graph_utils::FirstParentByType(nodes_to_remove.back(), "Sub");
     if (p_sub2_node == nullptr ||
         (p_sub2_node != p_sub_node && p_sub2_node != p_sub_node_dup)) {
       continue;
@@ -212,36 +234,66 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
     }
     nodes_to_remove.push_back(last_add_node);
 
+    // get axes attributes
+    const onnxruntime::NodeAttributes& attributes = reduce_mean_node.GetAttributes();
+    std::vector<int64_t> axes_values;
+    if (attributes.find("axes") != attributes.end()) {
+      axes_values = RetrieveValues<int64_t>(attributes.at("axes"));
+    }
+
     // Get the inputs for the new LayerNormalization node.
+    // scale and bias could be multi-dims; we only support it for training at the moment
+    // because SkipLayerNorm kernel, for example, has dependency on single dim size
     NodeArg* scale = nullptr;
     NodeArg* bias = nullptr;
     for (size_t i = 0; i < mul_node.MutableInputDefs().size(); i++) {
       if (graph_utils::NodeArgIsConstant(graph, *(mul_node.MutableInputDefs()[i])) ||
           graph_utils::IsGraphInput(graph, mul_node.MutableInputDefs()[i])) {
+#ifdef ENABLE_TRAINING
+        if (axes_values.empty() ||
+            mul_node.MutableInputDefs()[i]->Shape()->dim_size() == static_cast<int>(axes_values.size())) {
+          scale = mul_node.MutableInputDefs()[i];
+        }
+#else
         // Scale must be 1d.
         if (mul_node.MutableInputDefs()[i]->Shape()->dim_size() == 1) {
           scale = mul_node.MutableInputDefs()[i];
         }
+#endif
       }
     }
 
     for (size_t i = 0; i < last_add_node.MutableInputDefs().size(); i++) {
       if (graph_utils::NodeArgIsConstant(graph, *(last_add_node.MutableInputDefs()[i])) ||
           graph_utils::IsGraphInput(graph, last_add_node.MutableInputDefs()[i])) {
+#ifdef ENABLE_TRAINING
+        if (axes_values.empty() ||
+            last_add_node.MutableInputDefs()[i]->Shape()->dim_size() == static_cast<int>(axes_values.size())) {
+          bias = last_add_node.MutableInputDefs()[i];
+        }
+#else
         // Bias must be 1d.
         if (last_add_node.MutableInputDefs()[i]->Shape()->dim_size() == 1) {
           bias = last_add_node.MutableInputDefs()[i];
         }
+#endif
       }
     }
     if (scale == nullptr || bias == nullptr) {
       continue;
     }
 
-    // Scale and bias must have the same dimension.
-    if (scale->Shape()->dim(0).dim_value() != bias->Shape()->dim(0).dim_value()) {
-      continue;
+    // Scale and bias must have the same shape.
+    bool same_dim = true;
+    for (int i = 0; i < scale->Shape()->dim_size(); i++) {
+      if (scale->Shape()->dim(i).dim_value() != bias->Shape()->dim(i).dim_value()) {
+        same_dim = false;
+        break;
+      }
     }
+    if (!same_dim)
+      continue;
+
     const std::vector<NodeArg*> layer_norm_input_defs{reduce_mean_node.MutableInputDefs()[0], scale, bias};
     Node& layer_norm_node = graph.AddNode(graph.GenerateNodeName("LayerNormalization"),
                                           "LayerNormalization",
@@ -252,9 +304,9 @@ 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 &&
-      tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
-        Initializer initializer{*tensor_proto, graph.ModelPath()};
-        layer_norm_node.AddAttribute("epsilon", initializer.data<float>()[0]);
+        tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
+      Initializer initializer{*tensor_proto, graph.ModelPath()};
+      layer_norm_node.AddAttribute("epsilon", initializer.data<float>()[0]);
     } else {
       layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON);
     }

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -2123,6 +2123,28 @@ TEST_F(GraphTransformationTests, LayerNormFusionTest) {
   }
 }
 
+TEST_F(GraphTransformationTests, LayerNormWithCastFusionTest) {
+  auto model_uri = MODEL_FOLDER "fusion/layer_norm_with_cast.onnx";
+  std::shared_ptr<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+  graph_transformation_mgr.Register(onnxruntime::make_unique<LayerNormFusion>(), TransformerLevel::Level2);
+  auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_);
+  ASSERT_TRUE(ret.IsOK());
+
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+
+#ifdef ENABLE_TRAINING
+  ASSERT_TRUE(op_to_count["Cast"] == 0);
+  ASSERT_TRUE(op_to_count["LayerNormalization"] == 1);
+#else
+  ASSERT_TRUE(op_to_count["Cast"] == 1);
+  ASSERT_TRUE(op_to_count["LayerNormalization"] == 0);
+#endif
+}
+
 TEST_F(GraphTransformationTests, LayerNormWithSubDupFusionTest) {
   auto model_uri = MODEL_FOLDER "fusion/layer_norm_sub_dup.onnx";
   std::shared_ptr<Model> p_model;