eliminate transpose (#63948)

zhiqiu · web-flow · commit 2a2f635dafb0 · 2024-04-29T10:26:56.000+08:00
* replace transpose by reshape in some cases

* support transpose backward

* add to engine

* remove program_guard

* add ut
diff --git a/python/paddle/base/executor.py b/python/paddle/base/executor.py
@@ -1089,7 +1089,18 @@ def _get_program_and_executor(self, cached_data):
         ):
             pm = pir.PassManager()
             for p in new_program._pass_opt['pass_list']:
-                pm.add_pass(p, {})
+                # Temporary implementation, it will be refined when auto_parallel refactored
+                if p == 'eliminate_transpose':
+                    from paddle.distributed.auto_parallel.static.pir_pass import (
+                        eliminate_transpose_by_reshape,
+                    )
+
+                    for job_type in plan.job_types():
+                        ir_program = plan.ir_program(job_type)
+                        eliminate_transpose_by_reshape(ir_program)
+                else:
+                    pm.add_pass(p, {})
+
             for job_type in plan.job_types():
                 ir_program = plan.ir_program(job_type)
                 pm.run(ir_program)
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -632,7 +632,6 @@ def _parallel_pir(self, mode):
         # Part 1: Complete program
         # Step 1.1: Mix2Dense Pass
         # TODO(JZ-LIANG) regulization pass with pass management.
-
         dist_program = paddle.base.libpaddle.pir.apply_mix2dist_pass(
             mix_fw_program
         )
diff --git a/python/paddle/distributed/auto_parallel/static/parallelizer_v2.py b/python/paddle/distributed/auto_parallel/static/parallelizer_v2.py
@@ -32,12 +32,16 @@
     use_new_executor,
 )
 
-NEW_IR_PASS = [
+PIR_PASS = [
     'fused_gemm_epilogue_pass',
     'fused_linear_param_grad_add_pass',
     'fused_dropout_add_pass',
 ]
 
+PIR_PYTHON_PASS = [
+    'eliminate_transpose',
+]
+
 
 class Parallelizer:
     def __init__(self, mode, completer, dist_context):
@@ -513,13 +517,13 @@ def _apply_post_optimization(
         ir_pass_list = []
         if self.is_train and self._strategy.fused_passes.enable:
             if len(self._strategy.fused_passes.fused_passes_list) > 0:
-                new_pass_list = []
+                program_pass_list = []
                 for p in self._strategy.fused_passes.fused_passes_list:
-                    if p in NEW_IR_PASS and enable_ir:
+                    if enable_ir and p in (PIR_PASS + PIR_PYTHON_PASS):
                         ir_pass_list.append(p)
                     else:
-                        new_pass_list.append(new_pass(p))
-                pass_manager = PassManager(new_pass_list)
+                        program_pass_list.append(new_pass(p))
+                pass_manager = PassManager(program_pass_list)
                 pass_manager.apply([main_program], [startup_program])
 
         main_program._pass_opt = {}
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -154,3 +154,34 @@ def apply_reshard_pass(program):
                 )
 
     return new_program
+
+
+# In sequence_parallel, we need to transpose hidden_states
+# from [bs, seq, hidden] to [seq, bs, hidden] to perform
+# split and allgather at dim 0.
+# The transpose may lead to about 3% performance
+# in llama-70B model (tp4pp8).
+# We found that, when bs=1, which is the common case in llm
+# training, the transpose is equal to reshape.
+# So, this pass is to haddle the specific case.
+def eliminate_transpose_by_reshape(program):
+    for op in program.global_block().ops:
+        if (
+            op.name() == 'pd_op.transpose'
+            or op.name() == 'pd_op.transpose_grad'
+        ):
+            var = op.operand(0).source()
+            rank = len(var.shape)
+            perm = op.attrs()['perm']
+            perm = [p + rank if p < 0 else p for p in perm]
+            # only support transpose dim 0 and dim 1
+            expected_perm = [1, 0] + [i + 2 for i in range(rank - 2)]
+            if perm == expected_perm and (
+                var.shape[0] == 1 or var.shape[1] == 1
+            ):
+                paddle.pir.set_insertion_point(op)
+                transpose_var = op.result(0)
+                reshape_var = paddle._C_ops.reshape(var, transpose_var.shape)
+                transpose_var.replace_all_uses_with(reshape_var)
+                program.global_block().remove_op(op)
+    return program
diff --git a/test/auto_parallel/pir/CMakeLists.txt b/test/auto_parallel/pir/CMakeLists.txt
@@ -22,4 +22,7 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
                                            60)
   set_tests_properties(test_semi_auto_parallel_dist_to_static_pir
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 30)
+  py_test_modules(
+    test_eliminate_transpose_pass MODULES test_eliminate_transpose_pass ENVS
+    FLAGS_enable_pir_in_executor=1)
 endif()
diff --git a/test/auto_parallel/pir/test_eliminate_transpose_pass.py b/test/auto_parallel/pir/test_eliminate_transpose_pass.py
@@ -0,0 +1,46 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+
+paddle.enable_static()
+
+
+class TestEliminateTransposePass(unittest.TestCase):
+    def test_to_static_program(self):
+        a = paddle.rand([1, 1024, 4096])
+        b = paddle.reshape(a, [1024, 1, 4096])
+        c = paddle.transpose(a, [1, 0, 2])
+
+        paddle.seed(2024)
+        program = paddle.static.default_main_program()
+        exe = paddle.static.Executor()
+        (
+            b1,
+            c1,
+        ) = exe.run(program, fetch_list=[b, c])
+        np.testing.assert_equal(b1, c1)
+
+        paddle.seed(2024)
+        program._pass_opt = {'pass_list': ['eliminate_transpose']}
+        (c2,) = exe.run(program, fetch_list=[c])
+        np.testing.assert_equal(c1, c2)
+
+
+if __name__ == "__main__":
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -632,7 +632,6 @@ def _parallel_pir(self, mode):`
`632`	`632`	`# Part 1: Complete program`
`633`	`633`	`# Step 1.1: Mix2Dense Pass`
`634`	`634`	`# TODO(JZ-LIANG) regulization pass with pass management.`
`635`		`-`
`636`	`635`	`dist_program = paddle.base.libpaddle.pir.apply_mix2dist_pass(`
`637`	`636`	`mix_fw_program`
`638`	`637`	`)`