Fix chunk_id error for pd op.add_n in autograd (#69654)

From00 · web-flow · commit de20e32f6038 · 2024-11-25T11:36:34.000+08:00
* Fix vpp chunk_id error for pd_op.add_n in autograd

* Format
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -1042,19 +1042,8 @@ def infer_dist_skip_op_chunk_id(op):
             op_chunk_id = -1
         return op_chunk_id
 
-    is_dist_program = False
-    for op in block.ops:
-        if op.dist_attr is not None:
-            is_dist_program = True
-            break
-
-    if not is_dist_program:
-        return
-
-    for op in block.ops:
-        if op not in state.op_to_opgrad:
-            continue
-
+    # TODO(Ruibiao): Reorganize these unclear codes about chunk_id
+    def get_op_chunk_id(op):
         if op.dist_attr is None:
             op_chunk_id = -1
             if op.name() in dist_skip_op_list:
@@ -1067,16 +1056,50 @@ def infer_dist_skip_op_chunk_id(op):
                     op_chunk_id = infer_dist_skip_op_chunk_id(prev_op)
                 else:
                     op_chunk_id = prev_op.dist_attr.chunk_id
+        return op_chunk_id
+
+    is_dist_program = False
+    for op in block.ops:
+        if op.dist_attr is not None:
+            is_dist_program = True
+            break
+
+    if not is_dist_program:
+        return
+
+    for op in reversed(block.ops):
+        if op not in state.op_to_opgrad:
+            continue
+
+        fwd_op_chunk_id = get_op_chunk_id(op)
 
         for bwd_op in state.op_to_opgrad[op]:
             if bwd_op.dist_attr is None:
                 continue
+
+            if bwd_op.name() in ["pd_op.add_", "pd_op.add_n_"]:
+                bwd_op_chunk_id = -1
+                for operand_idx in range(bwd_op.num_operands()):
+                    prev_op_chunk_id = get_op_chunk_id(
+                        bwd_op.operand_source(operand_idx).get_defining_op()
+                    )
+                    if bwd_op_chunk_id == -1:
+                        bwd_op_chunk_id = prev_op_chunk_id
+                    else:
+                        assert bwd_op_chunk_id == prev_op_chunk_id, (
+                            f"Inconsistent prev_op chunk id with {bwd_op_chunk_id} != {prev_op_chunk_id}\n"
+                            "{bwd_op.operand_source(operand_idx-1).get_defining_op()}\n"
+                            "{bwd_op.operand_source(operand_idx).get_defining_op()}"
+                        )
+            else:
+                bwd_op_chunk_id = fwd_op_chunk_id
+
             bwd_op.dist_attr = (
                 paddle.base.libpaddle.pir.create_op_dist_attribute(
                     bwd_op.dist_attr.process_mesh,
                     bwd_op.dist_attr.operands(),
                     bwd_op.dist_attr.results(),
-                    op_chunk_id,
+                    bwd_op_chunk_id,
                 )
             )
 
