fix grad merge (#68460)

AndSonder · web-flow · commit 6ec2af12085c · 2024-09-26T11:08:24.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -238,10 +238,6 @@ def fold_reshard_pass(dist_program):
 
     @staticmethod
     def reshard_op_pass(dist_program, params_grads=[]):
-        # {grad.id: grad}
-        sharded_grad = {}
-        grad_ids = [grad.id for _, grad in params_grads if grad is not None]
-
         for op in dist_program.global_block().ops:
             if op.name() == 'dist_op.reshard':
                 var = op.operand_source(0)
@@ -277,27 +273,10 @@ def reshard_op_pass(dist_program, params_grads=[]):
 
                 if out_value is not None:
                     op.result(0).replace_all_uses_with(out_value)
-                    if var.id in grad_ids:
-                        if var.get_defining_op().has_attr(
-                            "replace_all_uses_with_reshard_var"
-                        ):
-                            sharded_grad[var.id] = out_value
 
                 if op.result(0).use_empty():
                     op.erase()
 
-                if out_value is not None and var.use_empty():
-                    if var.id in grad_ids:
-                        sharded_grad[var.id] = out_value
-
-        # update params_grads with sharded grad
-        for idx, (param, grad) in enumerate(params_grads):
-            if grad is None:
-                continue
-
-            if grad.id in sharded_grad:
-                params_grads[idx] = (param, sharded_grad[grad.id])
-
     @staticmethod
     def apply_reshard_pass(dist_program, params_grads=[]):
         ReshardPasses.fold_reshard_pass(dist_program)
