fix (#67907)

AndSonder · web-flow · commit f5075bcff89f · 2024-09-03T15:09:25.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -795,11 +795,7 @@ def _parallel_pir(self, mode):
         # TODO(hitywt) Step 3.2: Reshard Pass
         #   resolute the reshard op into special collective operation.
         #   collect the communicator created during resolution.
-        gradient_sync_after_accumulate = (
-            self._strategy.dp_optimization.gradient_sync_after_accumulate
-        )
-        if gradient_sync_after_accumulate:
-            global_params_grads = params_grads
+        global_params_grads = params_grads
 
         apply_reshard_pass(dist_program, params_grads)
         # print('after reshard', dist_program, flush=1)
@@ -840,14 +836,8 @@ def _parallel_pir(self, mode):
 
         if mode == "train" and self._strategy.gradient_merge.enable:
             config = copy.deepcopy(self._strategy.gradient_merge.to_dict())
-            config["gradient_sync_after_accumulate"] = (
-                gradient_sync_after_accumulate
-            )
-            config["params_grads"] = (
-                global_params_grads
-                if gradient_sync_after_accumulate
-                else params_grads
-            )
+            config["gradient_sync_after_accumulate"] = True
+            config["params_grads"] = global_params_grads
 
             auto_parallel_gradient_merge_pass = new_pass(
                 "auto_parallel_gradient_merge_pass", config
diff --git a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py
@@ -21,7 +21,6 @@
     is_data_parallel_reduce_op,
     is_data_parallel_scale_op,
 )
-from paddle.distributed.auto_parallel.static.pir_pass import comm_ops
 from paddle.distributed.auto_parallel.static.process_group import (
     get_world_process_group,
 )
@@ -322,19 +321,37 @@ def _pir_append_gradient_merge_backward_op(
             for op in grad.all_used_ops()
             if op.op_role == int(OpRole.Optimize)
         ]
+
         grad.replace_grad_users_with(
             new_gradient_merge_var, set(opt_ops_use_grad)
         )
 
+        for opt_op in opt_ops_use_grad:
+            if opt_op.name() == "pd_op.c_allreduce_sum":
+                paddle.pir.set_insertion_point_after(opt_op)
+                allreduce_sum_out = opt_op.result(0)
+
+                scale = paddle.full([], 0.5)
+                scale_out = paddle._C_ops.scale_(
+                    allreduce_sum_out, scale, 0.0, False
+                )
+
+                scale.get_defining_op().op_role = int(OpRole.Optimize)
+                scale_out.get_defining_op().op_role = int(OpRole.Optimize)
+
         # reset gradient merge var to zero after finishing optimization
         paddle.pir.set_insertion_point_to_block_end(main_block)
-        new_gradient_merge_var_zero = paddle._C_ops.fill_(
-            new_gradient_merge_var, 0.0
+        set_value = paddle.full(
+            shape=[1], fill_value=float(0), dtype=grad.dtype
         )
-        new_gradient_merge_var_zero.get_defining_op().op_role = int(
-            OpRole.Optimize
+        new_gradient_merge_var_zero = paddle._C_ops.set_value_with_tensor_(
+            new_gradient_merge_var, set_value, [], [], [], [], [], []
         )
 
+        set_value_op = new_gradient_merge_var_zero.get_defining_op()
+        set_value_op.op_role = int(OpRole.Optimize)
+        set_value.get_defining_op().op_role = int(OpRole.Optimize)
+
         # step3: Construct new_params_grads and grad_to_gradient_merge
         new_params_grads.append((param, new_gradient_merge_var))
 
@@ -385,31 +402,8 @@ def _move_reduce_to_optimizer_ops_block(
     return optimize_ops_block
 
 
-def _pir_move_reduce_to_optimizer_stage(main_program):
-    main_block = main_program.global_block()
-
-    for idx, op in list(enumerate(main_block.ops)):
-        if op.name() in comm_ops:
-            op_input_names = op.get_input_names()
-            # NOTE(sonder): When "@RENAME@" is in the input name, it means that the op has been renamed.
-            # Such types input names are caused by shared parameter policy.
-            # Gradient merge should accumulate the gradient of ops without renaming.
-            if "@RENAME" in op_input_names[0]:
-                continue
-
-            op.op_role = OpRole.Optimize
-            main_block.move_op_to_block_end(op)
-
-            if op.name() in ["pd_op.c_allreduce_sum", "pd_op.c_reduce_sum"]:
-                scale_index = idx + 1
-                while scale_index < len(main_block.ops):
-                    if is_data_parallel_scale_op(main_block.ops[scale_index]):
-                        main_block.ops[scale_index].op_role = OpRole.Optimize
-                        main_block.move_op_to_block_end(
-                            main_block.ops[scale_index]
-                        )
-                        break
-                    scale_index += 1
+def _pir_move_reduce_to_backward_stage(main_program):
+    pass
 
 
 def _remove_cast_for_master_grad(main_program, dist_context):
@@ -661,9 +655,9 @@ def _pir_parse_program(
         main_program, startup_program, params_grads
     )
 
-    # step2: move reduce op to optimizer_ops_block
-    if gradient_sync_after_accumulate:
-        _pir_move_reduce_to_optimizer_stage(main_program, params_grads)
+    # step2: move back reduce op to backward stage
+    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)
 
