[AutoParallel][PIR] fix vpp pass when open fused op (#68383)

AndSonder · web-flow · commit fcd9e31d40bc · 2024-09-26T11:00:36.000+08:00
* fix for vpp when open fused pass

* fix

* fix

* Update ir_backward.py

* format code
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -996,6 +996,60 @@ def create_backward_prune_set(
     return outputs_set, inputs_set, no_gradvar_set
 
 
+def _complete_grad_op_chunk_id(block, state):
+    dist_skip_op_list = ["builtin.split", "builtin.combine"]
+
+    def infer_dist_skip_op_chunk_id(op):
+        if op.name() == "builtin.split":
+            op_chunk_id = (
+                op.operand_source(0).get_defining_op().dist_attr.chunk_id
+            )
+        elif op.name() == "builtin.combine":
+            op_chunk_id = op.result(0).get_defining_op().dist_attr.chunk_id
+        else:
+            # TODO(luchang): need to support more ops such as pd_op.pylayer and so on
+            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
+
+        if op.dist_attr is None:
+            op_chunk_id = -1
+            if op.name() in dist_skip_op_list:
+                op_chunk_id = infer_dist_skip_op_chunk_id(op)
+        else:
+            op_chunk_id = op.dist_attr.chunk_id
+            if op_chunk_id == -1 and op.name() == "dist_op.reshard":
+                prev_op = op.operand_source(0).get_defining_op()
+                if prev_op.name() in dist_skip_op_list:
+                    op_chunk_id = infer_dist_skip_op_chunk_id(prev_op)
+                else:
+                    op_chunk_id = prev_op.dist_attr.chunk_id
+
+        for bwd_op in state.op_to_opgrad[op]:
+            if bwd_op.dist_attr is None:
+                continue
+            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,
+                )
+            )
+
+
 def calc_gradient_helper(
     outputs: Value | Sequence[Value],
     inputs: Value | Sequence[Value],
@@ -1057,29 +1111,8 @@ def calc_gradient_helper(
         outputs_fwd_set, inputs_fwd_set, no_grad_set, state
     )
 
-    # set struct name for grad ops
-    for op in block.ops:
-        if op in state.op_to_opgrad:
-            if op.dist_attr is None:
-                continue
-
-            op_chunk_id = op.dist_attr.chunk_id
-            if op_chunk_id == -1 and op.name() == "dist_op.reshard":
-                op_chunk_id = (
-                    op.operand_source(0).get_defining_op().dist_attr.chunk_id
-                )
-
-            for bwd_op in state.op_to_opgrad[op]:
-                if bwd_op.dist_attr is None:
-                    continue
-                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,
-                    )
-                )
+    # set chunk id for grad ops
+    _complete_grad_op_chunk_id(block, state)
 
     remove_ops = []
     if not is_inplace_net(backward_ops) and inputs:
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -30,6 +30,7 @@
     choose_reshard_func,
     copy_dist_attr_with_new_member,
     copy_op_attr_with_new_member,
+    copy_process_mesh_with_new_member,
 )
 from .reshard_funcs.reshard_func_register import register_reshard_funcs
 from .utils import (
@@ -148,11 +149,19 @@ def apply_partition_pass(program):
             result.update_dist_attr(result_attr)
 
             with auto_complete_op_role(program, ref_op_role):
+                prev_op = prev_var.get_defining_op()
+
                 # reshard output to assign out input
                 reshard_var_1 = paddle._C_ops.reshard_v2(
                     result, prev_var.dist_attr()
                 )
-                paddle.assign(reshard_var_1, prev_var)
+                assign_out = paddle._C_ops.assign_out_(reshard_var_1, prev_var)
+                assign_out.get_defining_op().dist_attr = (
+                    copy_op_attr_with_new_member(
+                        assign_out.get_defining_op().dist_attr,
+                        new_chunk_id=prev_op.dist_attr.chunk_id,
+                    )
+                )
 
             if old_dist_attr == result.dist_attr():
                 continue
@@ -687,7 +696,7 @@ def _get_seg_struct_names(ops, seg_method):
     seg_op_mesh = collections.OrderedDict()
 
     for i in range(fwd_start_op_index, fwd_end_op_index + 1):
-        if ops[i].name() == "builtin.combine":
+        if ops[i].name() in dist_skip_op_list:
             continue
 
         struct_name = _extract_seg_method(ops[i], seg_method)
@@ -960,10 +969,19 @@ def complete_chunk_id(dist_program, pipeline_strategy):
             new_dst_dist_attr = copy_dist_attr_with_new_member(
                 dst_dist_attr, new_process_mesh=dst_process_mesh
             )
+            new_process_ids = (
+                src_process_mesh.process_ids + dst_process_mesh.process_ids
+            )
+            new_process_mesh = copy_process_mesh_with_new_member(
+                op.dist_attr.process_mesh,
+                new_process_ids=new_process_ids,
+            )
+
             op.dist_attr = copy_op_attr_with_new_member(
                 op_dist_attr,
                 new_operands=[new_src_dist_attr],
                 new_results=[new_dst_dist_attr],
+                new_process_mesh=new_process_mesh,
             )
         elif reshard_func_name == "SameStatusReshardFunction":
             op.result(0).replace_all_uses_with(var)
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/r_to_p_reshard_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/r_to_p_reshard_func.py
@@ -60,10 +60,11 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
         else:
             dst_value = paddle.assign(src_value)
 
+        src_chunk_id = src_value.get_defining_op().dist_attr.chunk_id
         dst_value.set_type(dst_type)
         dst_value.get_defining_op().dist_attr = (
             paddle.base.libpaddle.pir.create_op_dist_attribute(
-                dst_mesh, [src_dist_attr], [dst_dist_attr]
+                dst_mesh, [src_dist_attr], [dst_dist_attr], src_chunk_id
             )
         )
 

Original file line number	Diff line number	Diff line change
`@@ -60,10 +60,11 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):`
`60`	`60`	`else:`
`61`	`61`	`dst_value = paddle.assign(src_value)`
`62`	`62`
	`63`	`+ src_chunk_id = src_value.get_defining_op().dist_attr.chunk_id`
`63`	`64`	`dst_value.set_type(dst_type)`
`64`	`65`	`dst_value.get_defining_op().dist_attr = (`
`65`	`66`	`paddle.base.libpaddle.pir.create_op_dist_attribute(`
`66`		`- dst_mesh, [src_dist_attr], [dst_dist_attr]`
	`67`	`+ dst_mesh, [src_dist_attr], [dst_dist_attr], src_chunk_id`
`67`	`68`	`)`
`68`	`69`	`)`
`69`	`70`