From 7ec9a732021646d469eb5cb458d15b2651043bbd Mon Sep 17 00:00:00 2001
From: Bowen Bao <bowbao@microsoft.com>
Date: Wed, 1 Jul 2020 13:06:03 -0700
Subject: [PATCH] deprecate frontend layernorm postpass (#4372)

---
 orttraining/orttraining/python/postprocess.py | 19 ++++++++++---------
 .../python/onnxruntime_test_postprocess.py    |  4 ++--
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/orttraining/orttraining/python/postprocess.py b/orttraining/orttraining/python/postprocess.py
index 5b532e1605..1f406f471c 100644
--- a/orttraining/orttraining/python/postprocess.py
+++ b/orttraining/orttraining/python/postprocess.py
@@ -17,7 +17,6 @@ def run_postprocess(model):
     # this post pass is not required for pytorch > 1.6
     model = fuse_softmaxNLL_to_softmaxCE(model)
 
-    model = layer_norm_transform(model)
     model = fix_expand_shape(model)
     model = fix_expand_shape_pt_1_5(model)
     return model
@@ -207,21 +206,23 @@ def add_const(model, name, output, t_value = None, f_value = None):
     return const_node
 
 def layer_norm_transform(model):
+    # DEPRECATED: This pass is no longer needed as the transform is handled at the backend.
     # Converting below subgraph
     #
     # input
     #   |
     # ReduceMean
-    #  __|_____
-    # |        |
-    # Sub      Sub
-    # |        |
-    # |   (optional) Cast
-    # |        |
+    #   |
+    #  Sub                         Constant
+    #  _||_____                       |
+    # |        |                      |
+    # |        |                      |
+    # |   (optional) Cast      (optional) Cast
+    # |        |                      |
+    # |        |  ____________________|
+    # |        | |
     # |        Pow
     # |        |
-    # |   (optional) Cast
-    # |        |
     # |       ReduceMean
     # |        |
     # |        Add
diff --git a/orttraining/orttraining/test/python/onnxruntime_test_postprocess.py b/orttraining/orttraining/test/python/onnxruntime_test_postprocess.py
index 21c87ae68a..ffccf50943 100644
--- a/orttraining/orttraining/test/python/onnxruntime_test_postprocess.py
+++ b/orttraining/orttraining/test/python/onnxruntime_test_postprocess.py
@@ -102,7 +102,7 @@ class Test_PostPasses(unittest.TestCase):
         count_layer_norm = self.count_nodes(onnx_model, "LayerNormalization")
         count_nodes = self.count_all_nodes(onnx_model)
 
-        assert count_layer_norm == 1
+        assert count_layer_norm == 0
         assert count_nodes == 3
 
     def test_expand(self):
@@ -199,7 +199,7 @@ class Test_PostPasses(unittest.TestCase):
     def _bert_helper(self, onnx_model):
         # count layer_norm
         count_layer_norm = self.count_nodes(onnx_model, "LayerNormalization")
-        assert count_layer_norm == 12
+        assert count_layer_norm == 0
 
         # get expand node and check output shape
         expand_nodes = self.find_nodes(onnx_model, "Expand")