[Auto Parallel] fix error op_role set (#67986)

AndSonder · web-flow · commit 5ee156f3e2cc · 2024-09-06T15:09:23.000+08:00
* fix

* fix

* fix

* change the default value of op_role

* fix

* fix

* fix

* fix
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1076,7 +1076,7 @@ void BindOperation(py::module *m) {
             if (int_attr) {
               return py::cast(int_attr.data());
             } else {
-              return py::cast<py::none>(Py_None);
+              return py::cast(-1);
             }
           },
           [](Operation &self, const int &op_role) {
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -8414,29 +8414,43 @@ def process_type_promotion(program):
 
 # complete the op_role of the new added ops
 @signature_safe_contextmanager
-def auto_complete_op_role(program, op_role, insert_point=None):
-    if paddle.framework.in_pir_mode():
-        initial_num_ops = program.num_ops()
-        origin_insert_point = insert_point
+def auto_complete_op_role(program, op_role):
+    def is_dist_block(block):
+        return any(op.dist_attr is not None for op in block.ops)
+
+    def validate_op_roles(block):
+        for op in block.ops:
+            if op.op_role == -1:
+                raise ValueError(
+                    f"All ops' op_role should be set before the completion. However, {op.name()}'s op_role is -1"
+                )
+
+    def set_op_roles(block, op_role, always_forward_ops):
+        for op in block.ops:
+            # TODO(luchang): Some ops are inserted during the optimization stage, and their op_role should be set to Forward.
+            # Ops like "pd_op.data" are inserted at the beginning of the block.
+            # Currently, we can't set the op_role of these ops during the optimization stage because the parallel graph cutting
+            # requires the op_role to be continuous. In the future, we should set the op_role of these ops during the
+            # optimization stage and eliminate the use of the whitelist.
+            set_op_role = (
+                op_role
+                if op.name() not in always_forward_ops
+                else int(core.op_proto_and_checker_maker.OpRole.Forward)
+            )
+            if op.op_role == -1:
+                op.op_role = set_op_role
+            for sub_block in op.blocks():
+                set_op_roles(sub_block, op_role, always_forward_ops)
+
+    block = program.global_block()
+
+    if paddle.framework.in_pir_mode() and is_dist_block(block):
+        assert op_role != -1, "Can't set op_role to -1 for new added ops"
+        validate_op_roles(block)
 
     try:
         yield
     finally:
-        if paddle.framework.in_pir_mode():
-            if insert_point is None:
-                paddle.pir.set_insertion_point_to_block_end(
-                    program.global_block()
-                )
-                insert_point = paddle.pir.get_current_insertion_point()
-            current_num_ops = program.num_ops()
-            if op_role is not None and current_num_ops > initial_num_ops:
-                for _ in range(current_num_ops - initial_num_ops):
-                    new_added_op = insert_point.prev()
-                    if new_added_op.op_role is not None:
-                        continue
-
-                    new_added_op.op_role = op_role
-                    paddle.pir.set_insertion_point(new_added_op)
-                    insert_point = paddle.pir.get_current_insertion_point()
-            if origin_insert_point:
-                paddle.pir.set_insertion_point(origin_insert_point)
+        if paddle.framework.in_pir_mode() and is_dist_block(block):
+            always_forward_ops = ["pd_op.data", "builtin.parameter"]
+            set_op_roles(block, op_role, always_forward_ops)
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -729,18 +729,21 @@ def _parallel_pir(self, mode):
                             use_master_grad=self._strategy.amp.use_master_grad,
                             use_promote=self._strategy.amp.use_promote,
                         )
-                        # bfloat16 needs no scaler
-                        scaler = paddle.amp.GradScaler(
-                            init_loss_scaling=self._strategy.amp.init_loss_scaling,
-                            incr_ratio=self._strategy.amp.incr_ratio,
-                            decr_ratio=self._strategy.amp.decr_ratio,
-                            incr_every_n_steps=self._strategy.amp.incr_every_n_steps,
-                            decr_every_n_nan_or_inf=self._strategy.amp.decr_every_n_nan_or_inf,
-                            use_dynamic_loss_scaling=self._strategy.amp.use_dynamic_loss_scaling,
-                            enable=self._strategy.amp.enable
-                            and self._strategy.amp.dtype != 'bfloat16',
-                        )
-                        scaled = scaler.scale(loss)
+                        with auto_complete_op_role(
+                            dist_program, OpRole.Forward
+                        ):
+                            # bfloat16 needs no scaler
+                            scaler = paddle.amp.GradScaler(
+                                init_loss_scaling=self._strategy.amp.init_loss_scaling,
+                                incr_ratio=self._strategy.amp.incr_ratio,
+                                decr_ratio=self._strategy.amp.decr_ratio,
+                                incr_every_n_steps=self._strategy.amp.incr_every_n_steps,
+                                decr_every_n_nan_or_inf=self._strategy.amp.decr_every_n_nan_or_inf,
+                                use_dynamic_loss_scaling=self._strategy.amp.use_dynamic_loss_scaling,
+                                enable=self._strategy.amp.enable
+                                and self._strategy.amp.dtype != 'bfloat16',
+                            )
+                            scaled = scaler.scale(loss)
                         optimizer_ops, params_grads = scaler.minimize(
                             self._optimizer, scaled
                         )
@@ -1316,6 +1319,8 @@ def _initialize(self, mode, init_parameters=True):
                             op.operand(0).set_source(reshard_var)
                 for del_op in del_ops:
                     del_op.erase()
+
+                set_all_ops_op_role(startup_prog, OpRole.Forward)
                 apply_reshard_pass(startup_prog)
                 for op in changed_ouput_op_list:
                     op.operand_source(0).persistable = True
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -16,7 +16,6 @@
 from paddle.autograd.backward_utils import ValueDict
 from paddle.base.framework import auto_complete_op_role
 from paddle.distributed.passes.pass_base import PassContext, new_pass
-from paddle.pir import get_current_insertion_point
 
 from .process_group import get_process_group
 from .reshard_funcs.base_reshard_func import (
@@ -44,9 +43,7 @@ def reshard_single_value(program, op, operand, attr):
     if prev_var.is_dist() and prev_var.dist_attr() != attr:
         operand_attr = attr.as_tensor_dist_attr()
         paddle.pir.set_insertion_point(op)
-        with auto_complete_op_role(
-            program, op.op_role, get_current_insertion_point()
-        ):
+        with auto_complete_op_role(program, op.op_role):
             # fold reshard
             if prev_var.get_defining_op().name() == 'dist_op.reshard':
                 prev_reshard = prev_var.get_defining_op()
@@ -82,9 +79,7 @@ def reshard_combine_value(program, op, operand, attr):
         )
 
     paddle.pir.set_insertion_point(op)
-    with auto_complete_op_role(
-        program, op.op_role, get_current_insertion_point()
-    ):
+    with auto_complete_op_role(program, op.op_role):
         combine_value = paddle._C_ops.builtin_combine(reshard_vars)
     return combine_value
 
@@ -119,9 +114,7 @@ def apply_partition_pass(program):
 
             # reshard input
             paddle.pir.set_insertion_point(op)
-            with auto_complete_op_role(
-                program, ref_op_role, get_current_insertion_point()
-            ):
+            with auto_complete_op_role(program, ref_op_role):
                 reshard_var = paddle._C_ops.reshard_v2(prev_var, operand_attr)
                 operand.set_source(reshard_var)
 
@@ -136,9 +129,7 @@ def apply_partition_pass(program):
             old_dist_attr = result.dist_attr()
             result.update_dist_attr(result_attr)
 
-            with auto_complete_op_role(
-                program, ref_op_role, get_current_insertion_point()
-            ):
+            with auto_complete_op_role(program, ref_op_role):
                 # reshard output to assign out input
                 reshard_var_1 = paddle._C_ops.reshard_v2(
                     result, prev_var.dist_attr()
@@ -148,14 +139,9 @@ def apply_partition_pass(program):
             if old_dist_attr == result.dist_attr():
                 continue
 
-            if ref_op_role is not None:
-                paddle.pir.set_insertion_point_after(op)
-
             reshard_var_2 = reshard_var_1
             if old_dist_attr != reshard_var_1.dist_attr():
-                with auto_complete_op_role(
-                    program, ref_op_role, get_current_insertion_point()
-                ):
+                with auto_complete_op_role(program, ref_op_role):
                     reshard_var_2 = paddle._C_ops.reshard_v2(
                         result, old_dist_attr
                     )
@@ -185,9 +171,7 @@ def apply_partition_pass(program):
                 var.update_dist_attr(attr.as_tensor_dist_attr())
 
                 # insert reshard
-                with auto_complete_op_role(
-                    program, op.op_role, get_current_insertion_point()
-                ):
+                with auto_complete_op_role(program, op.op_role):
                     reshard_var = paddle._C_ops.reshard_v2(var, old_dist_attr)
                     var.replace_all_uses_with(reshard_var)
                     reshard_var.get_defining_op().operand(0).set_source(var)
@@ -256,9 +240,7 @@ def apply_reshard_pass(dist_program, params_grads=[]):
             paddle.pir.set_insertion_point(op)
             ref_op_role = op.op_role
 
-            with auto_complete_op_role(
-                dist_program, ref_op_role, get_current_insertion_point()
-            ):
+            with auto_complete_op_role(dist_program, ref_op_role):
                 out_value = reshard_func.reshard(
                     src_dist_attr,
                     dst_dist_attr,
diff --git a/python/paddle/distributed/auto_parallel/static/utils.py b/python/paddle/distributed/auto_parallel/static/utils.py
@@ -2513,5 +2513,7 @@ def update_grad_var_to_var(program, strategy, grad_var_to_var):
 def set_all_ops_op_role(main_program, op_role):
     all_ops = main_program.global_block().ops
     for op in all_ops:
-        if op.op_role is None:
+        if op.op_role == -1:
             op.op_role = op_role
+        for sub_block in op.blocks():
+            set_all_ops_op_role(sub_block, op_role)
diff --git a/python/paddle/distributed/passes/pass_utils.py b/python/paddle/distributed/passes/pass_utils.py
@@ -718,31 +718,31 @@ def forward_complete_op_role(main_program):
     first_left_op_role = None
     first_right_op_role = None
     while iop < ops_len:
-        if all_ops[iop].op_role is not None:
+        if all_ops[iop].op_role != -1:
             first_left_op_role = all_ops[iop].op_role
             iop += 1
             continue
         else:
             right_idx = iop + 1
-            while right_idx < ops_len and all_ops[right_idx].op_role is None:
+            while right_idx < ops_len and all_ops[right_idx].op_role == -1:
                 right_idx += 1
             if right_idx >= ops_len:  # [first_left_op_role, xx, xx, xx, xx]
                 assert (
-                    first_left_op_role is not None
-                ), "first_left_op_role can't be None."
+                    first_left_op_role == -1
+                ), "first_left_op_role can't be -1."
                 for idx in range(iop, right_idx):
                     all_ops[idx].op_role = first_left_op_role
                 break
             else:  # [first_left_op_role, xx, xx, xx, xx, first_right_op_role]
                 first_right_op_role = all_ops[right_idx].op_role
                 assert (
-                    first_left_op_role is None
+                    first_left_op_role == -1
                     or first_left_op_role == first_right_op_role
                 ), f"The left and right operators of (idx[{iop}]) have different op_role."
                 for idx in range(iop, right_idx):
                     all_ops[idx].op_role = first_right_op_role
                     iop = right_idx + 1
-    if first_left_op_role is None and first_right_op_role is None:
+    if first_left_op_role == -1 and first_right_op_role == -1:
         raise ValueError("all the ops don't have the op_role.")
 
 
@@ -763,7 +763,7 @@ def _split_program_into_forward_backward_optimize(
 
     region = "opt"
     for op_idx in range(len(complete_ops) - 1, -1, -1):
-        if complete_ops[op_idx].op_role is not None:
+        if complete_ops[op_idx].op_role != -1:
             if complete_ops[op_idx].op_role == 1:
                 region = "bwd"
             elif complete_ops[op_idx].op_role == 0:
diff --git a/test/auto_parallel/hybrid_strategy/pir_reshard_nd_mesh.py b/test/auto_parallel/hybrid_strategy/pir_reshard_nd_mesh.py
@@ -19,6 +19,8 @@
 from paddle.distributed.auto_parallel.static.pir_pass import (
     apply_reshard_pass,
 )
+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
 
 
 class TestReshardNdMesh:
@@ -79,6 +81,7 @@ def create_program(self, input_shape, input_placements, output_placements):
                     dist_input, self._mesh, output_placements
                 )
             dist_program = main_program.clone()
+            set_all_ops_op_role(dist_program, OpRole.Forward)
             apply_reshard_pass(dist_program)
 
         return main_program, dist_program
diff --git a/test/auto_parallel/hybrid_strategy/pir_reshard_nd_mesh_cross_mesh.py b/test/auto_parallel/hybrid_strategy/pir_reshard_nd_mesh_cross_mesh.py
@@ -19,6 +19,8 @@
 from paddle.distributed.auto_parallel.static.pir_pass import (
     apply_reshard_pass,
 )
+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
 
 
 class TestReshardNdMeshCrossMesh:
@@ -80,6 +82,7 @@ def create_program(self, input_shape, input_placements, output_placements):
                     dist_input, self._mesh1, output_placements
                 )
             dist_program = main_program.clone()
+            set_all_ops_op_role(dist_program, OpRole.Forward)
             apply_reshard_pass(dist_program)
 
         return main_program, dist_program
diff --git a/test/auto_parallel/pir/mlp_demo_3d.py b/test/auto_parallel/pir/mlp_demo_3d.py
@@ -22,9 +22,12 @@
 import paddle
 import paddle.distributed as dist
 from paddle import nn
+from paddle.base.framework import auto_complete_op_role
 from paddle.distributed.auto_parallel.static.mix_to_dist_pass import (
     apply_mix2dist_pass,
 )
+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
 from paddle.framework import _current_expected_place
 
 BATCH_SIZE = 4
@@ -104,12 +107,15 @@ def test_to_static_program(self):
         dist_program = engine._fwd_main_progs["train"]
 
         apply_mix2dist_pass(dist_program)
+        set_all_ops_op_role(dist_program, OpRole.Forward)
         loss = dist_program.get_output_value_by_name(engine._loss_names[0])
         with paddle.static.program_guard(dist_program):
-            params_grads = paddle.autograd.ir_backward.append_backward(loss)
-            engine._optimizer._apply_optimize(
-                loss, startup_program=None, params_grads=params_grads
-            )
+            with auto_complete_op_role(dist_program, OpRole.Backward):
+                params_grads = paddle.autograd.ir_backward.append_backward(loss)
+            with auto_complete_op_role(dist_program, OpRole.Optimize):
+                engine._optimizer._apply_optimize(
+                    loss, startup_program=None, params_grads=params_grads
+                )
         from paddle.distributed.auto_parallel.static.pir_pass import (
             apply_partition_pass,
         )
diff --git a/test/auto_parallel/pir/pir_reshard_s_to_r.py b/test/auto_parallel/pir/pir_reshard_s_to_r.py
@@ -21,6 +21,8 @@
 from paddle.distributed.auto_parallel.static.pir_pass import (
     apply_reshard_pass,
 )
+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
 
 
 class TestReshardSToR:
@@ -66,6 +68,7 @@ def run_pir_test_case(self):
                 reshard_tensor = paddle._C_ops.reshard(
                     input_tensor, self._mesh, [dist.Replicate()]
                 )
+            set_all_ops_op_role(main_program, OpRole.Forward)
             apply_reshard_pass(main_program)
         ops = [op.name() for op in main_program.global_block().ops]
         if self._shard == 0:
diff --git a/test/auto_parallel/pir/pir_reshard_s_to_r_cross_mesh.py b/test/auto_parallel/pir/pir_reshard_s_to_r_cross_mesh.py
@@ -21,6 +21,8 @@
 from paddle.distributed.auto_parallel.static.pir_pass import (
     apply_reshard_pass,
 )
+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
 
 
 class TestReshardSToRCrossMesh:
@@ -60,6 +62,7 @@ def run_pir_test_case(self):
                 reshard_tensor = paddle._pir_ops.reshard(
                     input_tensor, self._out_mesh, [dist.Replicate()]
                 )
+            set_all_ops_op_role(main_program, OpRole.Forward)
             apply_reshard_pass(main_program)
 
         ops = [op.name() for op in main_program.global_block().ops]
diff --git a/test/auto_parallel/pir_reshard_r_to_s_cross_mesh.py b/test/auto_parallel/pir_reshard_r_to_s_cross_mesh.py
@@ -19,6 +19,8 @@
 from paddle.distributed.auto_parallel.static.pir_pass import (
     apply_reshard_pass,
 )
+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
 
 
 class TestReshardRToSCrossMesh:
@@ -54,7 +56,7 @@ def run_test_case(self):
 
             old_ops = [op.name() for op in main_program.global_block().ops]
             assert 'dist_op.reshard' in old_ops
-
+            set_all_ops_op_role(main_program, OpRole.Forward)
             apply_reshard_pass(main_program)
         # np.testing.assert_equal(dist_program.num_ops(), 6)
         new_ops = [op.name() for op in main_program.global_block().ops]
diff --git a/test/auto_parallel/reshard_p_to_r.py b/test/auto_parallel/reshard_p_to_r.py
@@ -29,6 +29,8 @@
 from paddle.distributed.auto_parallel.static.pir_pass import (
     apply_reshard_pass,
 )
+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
 from paddle.framework import core
 
 
@@ -88,6 +90,7 @@ def run_pir_static_test_case(self):
                 reshard_tensor = paddle._C_ops.reshard(
                     input_tensor, self._mesh, [dist.Replicate()]
                 )
+            set_all_ops_op_role(main_program, OpRole.Forward)
             apply_reshard_pass(main_program)
         np.testing.assert_equal(main_program.num_ops(), 4)
         ops = main_program.global_block().ops
diff --git a/test/auto_parallel/reshard_p_to_r_cross_mesh.py b/test/auto_parallel/reshard_p_to_r_cross_mesh.py
diff --git a/test/auto_parallel/reshard_r_to_s.py b/test/auto_parallel/reshard_r_to_s.py

Original file line number	Diff line number	Diff line change
`@@ -1076,7 +1076,7 @@ void BindOperation(py::module *m) {`
`1076`	`1076`	`if (int_attr) {`
`1077`	`1077`	`return py::cast(int_attr.data());`
`1078`	`1078`	`} else {`
`1079`		`- return py::cast<py::none>(Py_None);`
	`1079`	`+ return py::cast(-1);`
`1080`	`1080`	`}`
`1081`	`1081`	`},`
`1082`	`1082`	`[](Operation &self, const int &op_role) {`
Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,8 @@`
`19`	`19`	`from paddle.distributed.auto_parallel.static.pir_pass import (`
`20`	`20`	`apply_reshard_pass,`
`21`	`21`	`)`
	`22`	`+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role`
	`23`	`+from paddle.distributed.fleet.meta_optimizers.common import OpRole`
`22`	`24`
`23`	`25`
`24`	`26`	`class TestReshardNdMesh:`
`@@ -79,6 +81,7 @@ def create_program(self, input_shape, input_placements, output_placements):`
`79`	`81`	`dist_input, self._mesh, output_placements`
`80`	`82`	`)`
`81`	`83`	`dist_program = main_program.clone()`
	`84`	`+ set_all_ops_op_role(dist_program, OpRole.Forward)`
`82`	`85`	`apply_reshard_pass(dist_program)`
`83`	`86`
`84`	`87`	`return main_program, dist_program`
Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,8 @@`
`21`	`21`	`from paddle.distributed.auto_parallel.static.pir_pass import (`
`22`	`22`	`apply_reshard_pass,`
`23`	`23`	`)`
	`24`	`+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role`
	`25`	`+from paddle.distributed.fleet.meta_optimizers.common import OpRole`
`24`	`26`
`25`	`27`
`26`	`28`	`class TestReshardSToR:`
`@@ -66,6 +68,7 @@ def run_pir_test_case(self):`
`66`	`68`	`reshard_tensor = paddle._C_ops.reshard(`
`67`	`69`	`input_tensor, self._mesh, [dist.Replicate()]`
`68`	`70`	`)`
	`71`	`+ set_all_ops_op_role(main_program, OpRole.Forward)`
`69`	`72`	`apply_reshard_pass(main_program)`
`70`	`73`	`ops = [op.name() for op in main_program.global_block().ops]`
`71`	`74`	`if self._shard == 0:`