[PIR] Fix the logic for setting op_role (#67850)

AndSonder · web-flow · commit 5c7df1ed3b68 · 2024-09-02T11:32:54.000+08:00
* fix error op_role set in amp and grad clip

* remove print

* remove print
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -8410,3 +8410,33 @@ def process_type_promotion(program):
                         idx += 1
             idx += 1
     return 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
+
+    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)
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -28,7 +28,9 @@
 import paddle.distributed.auto_parallel.static.utils as auto_utils
 from paddle import pir, static, utils
 from paddle.base.executor import _to_name_str
+from paddle.base.framework import auto_complete_op_role
 from paddle.distributed import fleet
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
 from paddle.distributed.passes.pass_base import new_pass
 from paddle.distributed.passes.pass_utils import (
     _split_program_into_forward_backward_optimize,
@@ -66,14 +68,14 @@
 from .pir_pass import (
     apply_partition_pass,
     apply_reshard_pass,
-    complete_op_role,
     pipeline_pass,
     remove_other_rank_input_output_pass,
     remove_other_rank_op_pass,
     remove_unuseful_comm_op_pass,
 )
 from .planner_v2 import Planner
 from .process_group import get_all_process_groups, new_process_group
+from .utils import set_all_ops_op_role
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
@@ -695,9 +697,9 @@ def _parallel_pir(self, mode):
         # TODO(JZ-LIANG) regulization pass with pass management.
         dist_program = mix_fw_program.clone()
         apply_mix2dist_pass(dist_program)
+        set_all_ops_op_role(dist_program, OpRole.Forward)
 
         # Step 1.2: pir backward
-        last_forward_op = dist_program.global_block().ops[-1]
         if mode == "train" and self._loss and self._optimizer:
             loss = dist_program.get_output_value_by_name(self._loss_names[0])
             if loss.initialized():
@@ -739,60 +741,24 @@ def _parallel_pir(self, mode):
                             and self._strategy.amp.dtype != 'bfloat16',
                         )
                         scaled = scaler.scale(loss)
-                        last_forward_op = dist_program.global_block().ops[-1]
                         optimizer_ops, params_grads = scaler.minimize(
                             self._optimizer, scaled
                         )
-                        first_opt_op = optimizer_ops[0]
-                        backward_op_start_idx = (
-                            dist_program.global_block().ops.index(
-                                last_forward_op
-                            )
-                            + 1
-                        )
-                        opt_op_start_idx = (
-                            dist_program.global_block().ops.index(first_opt_op)
-                        )
-                        # print('after minimize', dist_program, flush=1)
-                        complete_op_role(
-                            dist_program,
-                            [
-                                [forward_op_start_idx, backward_op_start_idx],
-                                [backward_op_start_idx, opt_op_start_idx],
-                                [opt_op_start_idx, dist_program.num_ops()],
-                            ],
-                        )
                     else:
-                        params_grads = (
-                            paddle.autograd.ir_backward.append_backward(loss)
-                        )
-                        last_backward_op = dist_program.global_block().ops[-1]
-                        self._optimizer._apply_optimize(
-                            loss, startup_program, params_grads=params_grads
-                        )
-
-                        backward_op_start_idx = (
-                            dist_program.global_block().ops.index(
-                                last_forward_op
+                        with auto_complete_op_role(
+                            dist_program, OpRole.Backward
+                        ):
+                            params_grads = (
+                                paddle.autograd.ir_backward.append_backward(
+                                    loss
+                                )
                             )
-                            + 1
-                        )
-                        opt_op_start_idx = (
-                            dist_program.global_block().ops.index(
-                                last_backward_op
+                        with auto_complete_op_role(
+                            dist_program, OpRole.Optimize
+                        ):
+                            self._optimizer._apply_optimize(
+                                loss, startup_program, params_grads=params_grads
                             )
-                            + 1
-                        )
-                        complete_op_role(
-                            dist_program,
-                            [
-                                [forward_op_start_idx, backward_op_start_idx],
-                                [backward_op_start_idx, opt_op_start_idx],
-                                [opt_op_start_idx, dist_program.num_ops()],
-                            ],
-                        )
-                        # self._optimizer.minimize(loss, startup_program=startup_program)
-
             else:
                 self._logger.info(
                     "loss value is not found, skip append backward."
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -14,10 +14,10 @@
 
 import paddle
 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 ...passes.pass_utils import auto_complete_op_role
 from .process_group import get_process_group
 from .reshard_funcs.base_reshard_func import (
     choose_reshard_func,
@@ -364,6 +364,7 @@ def remove_other_rank_input_output_pass(dist_program):
         result = op.operand_source(0).get_defining_op().result(0)
         paddle.pir.set_insertion_point_after(combine_op)
         res = paddle._C_ops.builtin_combine(new_vars)
+        res.get_defining_op().op_role = op.op_role
         result.replace_all_uses_with(res)
         combine_op.erase()
         # since it is inplace op, set type of output as the same as input
@@ -419,6 +420,7 @@ def eliminate_transpose_by_reshape(program):
                 paddle.pir.set_insertion_point(op)
                 transpose_var = op.result(0)
                 reshape_var = paddle._C_ops.reshape(var, transpose_var.shape)
+                reshape_var.get_defining_op().op_role = op.op_role
                 transpose_var.replace_all_uses_with(reshape_var)
                 op.erase()
     return program
diff --git a/python/paddle/distributed/auto_parallel/static/utils.py b/python/paddle/distributed/auto_parallel/static/utils.py
@@ -23,7 +23,9 @@
 
 import paddle
 from paddle.base.framework import use_pir_api
-from paddle.base.wrapped_decorator import wrap_decorator
+from paddle.base.wrapped_decorator import (
+    wrap_decorator,
+)
 from paddle.framework import core
 from paddle.framework.io_utils import is_belong_to_optimizer, is_parameter
 from paddle.static import Variable
@@ -2506,3 +2508,10 @@ def update_grad_var_to_var(program, strategy, grad_var_to_var):
         scale_loss_grad_var_name = first_backward_op.desc.output("Out")[0]
         if scale_loss_grad_var_name not in grad_var_to_var.keys():
             grad_var_to_var[scale_loss_grad_var_name] = scale_loss_var_name
+
+
+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:
+            op.op_role = op_role
diff --git a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py
@@ -452,11 +452,23 @@ def _remove_cast_for_master_grad(main_program, dist_context):
     main_block._sync_with_cpp()
 
 
-def _pir_remove_cast_for_master_grad(main_program):
-    main_block = main_program.global_block()
-    for op in main_block.ops:
-        if _is_master_grad_cast_op(main_block, op):
-            main_program.remove_op(op)
+def _pir_remove_cast_for_master_grad(main_program, params_grads):
+    def is_cast_to_float32(op):
+        return (
+            op.name() == "pd_op.cast"
+            and op.results()[0].dtype == paddle.float32
+        )
+
+    for _, grad in params_grads:
+        if grad is None:
+            continue
+        if grad.dtype == paddle.float32:
+            continue
+
+        for op in grad.all_used_ops():
+            if is_cast_to_float32(op):
+                op.results()[0].replace_all_uses_with(grad)
+                op.erase()
 
 
 def _create_cond_block_and_update_optimizer(
@@ -644,7 +656,6 @@ def _pir_parse_program(
     avg,
     gradient_sync_after_accumulate,
 ):
-
     # step1: append gradient merge backward op to main_program
     new_params_to_grads = _pir_append_gradient_merge_backward_op(
         main_program, startup_program, params_grads
@@ -654,7 +665,7 @@ def _pir_parse_program(
     if gradient_sync_after_accumulate:
         _pir_move_reduce_to_optimizer_stage(main_program, params_grads)
 
-    _pir_remove_cast_for_master_grad(main_program)
+    _pir_remove_cast_for_master_grad(main_program, params_grads)
 
     # step3: append scale op
     if avg:
diff --git a/python/paddle/distributed/passes/auto_parallel_master_grad.py b/python/paddle/distributed/passes/auto_parallel_master_grad.py
@@ -56,21 +56,11 @@
 
 
 def _is_master_grad_cast_op(block, op):
-    in_pir_mode = paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
-        "FLAGS_enable_pir_api"
-    ]
-    if in_pir_mode:
-        op_name = op.name()
-        if op_name != "pd_op.cast":
-            return False
-        input_names = op.get_input_names()
-        output_names = op.get_output_names()
-    else:
-        op_name = op.type
-        if op_name != "cast":
-            return False
-        input_names = op.input_arg_names
-        output_names = op.output_arg_names
+    op_name = op.type
+    if op_name != "cast":
+        return False
+    input_names = op.input_arg_names
+    output_names = op.output_arg_names
 
     assert len(input_names) == 1
     assert len(output_names) == 1
diff --git a/python/paddle/distributed/passes/pass_utils.py b/python/paddle/distributed/passes/pass_utils.py
@@ -20,7 +20,6 @@
 import paddle
 from paddle.base import core
 from paddle.base.framework import Parameter, Program
-from paddle.base.wrapped_decorator import signature_safe_contextmanager
 from paddle.distributed.auto_parallel.static.dist_attribute import (
     OperatorDistAttr,
 )
@@ -32,7 +31,8 @@
     naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
     use_new_executor,
 )
-from paddle.distributed.fleet.meta_optimizers.common import OpRole
+
+from ..auto_parallel.static.utils import OpRole
 
 __not_shape_var_type__ = [
     core.VarDesc.VarType.READER,
@@ -746,30 +746,6 @@ def forward_complete_op_role(main_program):
         raise ValueError("all the ops don't have the op_role.")
 
 
-# complete the op_role of the new added ops
-@signature_safe_contextmanager
-def auto_complete_op_role(program, op_role, insert_point):
-    initial_num_ops = program.num_ops()
-    origin_insert_point = insert_point
-
-    try:
-        yield
-    finally:
-        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:
-                    break
-
-                new_added_op.op_role = op_role
-                paddle.pir.set_insertion_point(new_added_op)
-                insert_point = paddle.pir.get_current_insertion_point()
-
-            paddle.pir.set_insertion_point(origin_insert_point)
-
-
 def _split_program_into_forward_backward_optimize(
     main_program, enable_send_recv_overlap=False
 ):
diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py
@@ -904,6 +904,7 @@ def _create_master_grad(self, grad):
                 var = self._master_grads[grad]
             else:
                 var = paddle.cast(grad, 'float32')
+                var.get_defining_op().set_bool_attr('master_grad_cast', True)
                 self._master_grads[grad] = var
         else:
             if grad.name in self._master_grads:
diff --git a/python/paddle/static/amp/decorator.py b/python/paddle/static/amp/decorator.py
@@ -23,7 +23,7 @@
     program_guard,
     unique_name,
 )
-from paddle.base.framework import in_pir_mode
+from paddle.base.framework import auto_complete_op_role, in_pir_mode
 
 from .amp_nn import check_finite_and_unscale, update_loss_scaling
 from .fp16_lists import AutoMixedPrecisionLists, check_amp_dtype
@@ -34,6 +34,8 @@
 )
 from .function_overload import FunctionType, overload
 
+OpRole = core.op_proto_and_checker_maker.OpRole
+
 
 def _set_multi_precision(optimizer, multi_precision):
     if not isinstance(
@@ -740,16 +742,18 @@ def minimize(
                 "The decorated optimizer has its own `minimize` method, but it will not be executed."
             )
 
-        scaled_params_grads = self.backward(
-            loss,
-            startup_program=startup_program,
-            parameter_list=parameter_list,
-            no_grad_set=no_grad_set,
-        )
+        with auto_complete_op_role(loss.block.program, op_role=OpRole.Backward):
+            scaled_params_grads = self.backward(
+                loss,
+                startup_program=startup_program,
+                parameter_list=parameter_list,
+                no_grad_set=no_grad_set,
+            )
 
-        optimize_ops = self.apply_optimize(
-            loss, startup_program, scaled_params_grads
-        )
+        with auto_complete_op_role(loss.block.program, op_role=OpRole.Optimize):
+            optimize_ops = self.apply_optimize(
+                loss, startup_program, scaled_params_grads
+            )
 
         return optimize_ops, scaled_params_grads
 
diff --git a/python/paddle/static/amp/fp16_utils.py b/python/paddle/static/amp/fp16_utils.py
@@ -21,6 +21,7 @@
 
 import paddle
 from paddle.base import core, framework, global_scope
+from paddle.base.framework import in_pir_mode
 from paddle.base.log_helper import get_logger
 from paddle.base.wrapped_decorator import signature_safe_contextmanager
 
@@ -867,6 +868,8 @@ def _convert_to_float(place, org_array):
     framework._set_expected_place(place)
     org_tensor = paddle.to_tensor(org_array)
     fp32_array = paddle.cast(org_tensor, paddle.float32).numpy()
+    if in_pir_mode():
+        fp32_array.get_defining_op().set_bool_attr("master_grad_cast", True)
     paddle.enable_static()
     return fp32_array
 
diff --git a/test/auto_parallel/pir/mlp_demo.py b/test/auto_parallel/pir/mlp_demo.py
