[AutoParallel]:fix grad_merge bug (#68727)

blacksheep-Aristotle · web-flow · commit 5c6c295d9d20 · 2024-10-16T14:19:22.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -995,7 +995,7 @@ def _parallel_pir(self, mode):
                 for job_type in self._job_plan.job_types():
                     ir_program = self._job_plan.ir_program(job_type)
                     pm.run(ir_program)
-
+        remove_unuseful_comm_op_pass(dense_program)
         self._pir_dense_main_progs[mode] = dense_program
         self._pir_dist_main_progs[mode] = dist_program
         self._pir_dist_startup_progs[mode] = startup_program
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -541,6 +541,12 @@ def remove_unuseful_comm_op_pass(program):
             if op.operand_source(0).has_one_use():
                 op.result(0).replace_all_uses_with(op.operand_source(0))
                 op.erase()
+        if (
+            op.name() == "pd_op.cast"
+            and op.result(0).dtype == op.operand_source(0).dtype
+        ):
+            op.result(0).replace_all_uses_with(op.operand_source(0))
+            op.erase()
 
 
 # In sequence_parallel, we need to transpose hidden_states
diff --git a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py
@@ -334,16 +334,15 @@ def _pir_append_gradient_merge_backward_op(
         )
         # NOTE(zhangweilong): grad may in different device in auto_parallel, so need consider all_gather op
         for used_grad_op in grad.all_used_ops():
-            if used_grad_op.name() != "pd_op.all_gather":
-                continue
-            move_to_opt_block_flag = True
-            for all_gather_result in used_grad_op.results():
-                for used_op in all_gather_result.all_used_ops():
-                    if used_op.op_role != int(OpRole.Optimize):
-                        move_to_opt_block_flag = False
-                        break
-            if move_to_opt_block_flag:
-                used_grad_op.op_role = int(OpRole.Optimize)
+            if used_grad_op.num_operands() == 1:
+                move_to_opt_block_flag = True
+                for used_op_result in used_grad_op.results():
+                    for used_op in used_op_result.all_used_ops():
+                        if used_op.op_role != int(OpRole.Optimize):
+                            move_to_opt_block_flag = False
+                            break
+                if move_to_opt_block_flag:
+                    used_grad_op.op_role = int(OpRole.Optimize)
 
         opt_ops_use_grad = [
             op
@@ -753,7 +752,7 @@ def _pir_parse_program(
     if not gradient_sync_after_accumulate:
         _pir_move_reduce_to_backward_stage(main_program, params_grads)
 
-    _pir_remove_cast_for_master_grad(main_program, params_grads)
+    # _pir_remove_cast_for_master_grad(main_program, params_grads)
 
     # step3: append scale op
     if avg:
