[Auto Parallel] Support gradient_merge pass for PIR (#66641)

AndSonder · heavyrain-lzy · web-flow · commit 8caa7132a3b4 · 2024-08-02T11:18:44.000+08:00
* split program and polish the executor

* change unit test

* fix 1f1b test
---------

Co-authored-by: heavyrain_lzy &lt;1528794076@qq.com&gt;
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -710,6 +710,19 @@ void BindBlock(py::module *m) {
               None
 
         )DOC")
+      .def(
+          "move_op_to_block_end",
+          [](Block &self, Operation *op) { op->MoveTo(&self, self.end()); },
+          R"DOC(
+            Move an op to the end of the block.
+
+            Args:
+                op (pir.Operation): The operator to be moved.
+
+            Returns:
+                None
+
+            )DOC")
       .def("all_parameters",
            [](Block &self) -> py::list {
              py::list param_list;
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -26,6 +26,7 @@
 from paddle import static, utils
 from paddle.base.executor import _to_name_str
 from paddle.distributed import fleet
+from paddle.distributed.passes.pass_base import new_pass
 from paddle.framework import (
     IrGraph,
     _current_expected_place_ as _get_device,
@@ -807,12 +808,17 @@ def _parallel_pir(self, mode):
         #     'after remove_other_rank_input_output_pass', dist_program, flush=1
         # )
 
-        remove_other_rank_op_pass(dist_program)
+        remove_other_rank_op_pass(dist_program, params_grads)
 
         # print('after remove_other_rank_op_pass', dist_program, flush=1)
 
         # Part 4: Optimization Pass
         # NOTE Only those Optimization Pass that related to Parallelism (need dist attr) should be placed here and all the Pass should be Optional.
+        gradient_sync_after_accumulate = (
+            self._strategy.dp_optimization.gradient_sync_after_accumulate
+        )
+        if gradient_sync_after_accumulate:
+            global_params_grads = params_grads
 
         # TODO(xxxx) Step 4.1 DP Optimization Pass
         if self._strategy.dp_optimization.enable:
@@ -829,6 +835,24 @@ def _parallel_pir(self, mode):
             # dist_program = apply_sharding_optimization_pass(dist_program)
             pass
 
+        if mode == "train" and self._strategy.gradient_merge.enable:
+            config = copy.deepcopy(self._strategy.gradient_merge.to_dict())
+            config[
+                "gradient_sync_after_accumulate"
+            ] = gradient_sync_after_accumulate
+            config["params_grads"] = (
+                global_params_grads
+                if gradient_sync_after_accumulate
+                else params_grads
+            )
+
+            auto_parallel_gradient_merge_pass = new_pass(
+                "auto_parallel_gradient_merge_pass", config
+            )
+            auto_parallel_gradient_merge_pass.apply(
+                [dist_program], [startup_program]
+            )
+
         # TODO(JZ-LIANG) Step 4.4 Dist2Dense Pass
         # NOTE All optimization pass that need dist_attr info should be called before Dist2Dense Pass.
         dense_program = dist_program.clone()
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -229,9 +229,21 @@ def apply_reshard_pass(dist_program):
                 op.erase()
 
 
+def _remove_other_rank_params_grads(dist_params_grads):
+    cur_rank = paddle.distributed.get_rank()
+    need_remove_idx = []
+    for idx, (_, grad) in enumerate(dist_params_grads):
+        if cur_rank not in grad.dist_attr().process_mesh.process_ids:
+            need_remove_idx.append(idx)
+    for idx in need_remove_idx[::-1]:
+        dist_params_grads.pop(idx)
+
+
 # pruning op and value not belong to cur rank
-def remove_other_rank_op_pass(dist_program):
+def remove_other_rank_op_pass(dist_program, dist_params_grads):
     cur_rank = paddle.distributed.get_rank()
+
+    _remove_other_rank_params_grads(dist_params_grads)
     for op in dist_program.global_block().ops[::-1]:
         if op.name() in partition_skip_op_list:
             can_delete = True
diff --git a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py
@@ -264,12 +264,65 @@ def _append_gradient_merge_backward_op(
     return new_params_grads, grad_to_gradient_merge
 
 
+def _pir_append_gradient_merge_backward_op(
+    main_program,
+    startup_program,
+    params_grads,
+):
+    main_block = main_program.global_block()
+    startup_block = startup_program.global_block()
+
+    # {param: gradient_merge_var} to insert scale op and fill_constant op
+    new_params_grads = []
+
+    for param, grad in params_grads:
+        assert (
+            not param.is_selected_row_type()
+        ), "SELECTED_ROWS is not supported in GradientMergeOptimizer for now"
+
+        # step1: create gradient_merge var and init with 0
+        # Add persistable gradient variables in startup_program
+        with startup_block:
+            paddle.pir.set_insertion_point_to_block_end(startup_block)
+            gradient_merge_var = paddle.full(
+                shape=grad.shape, fill_value=float(0), dtype=grad.dtype
+            )
+            gradient_merge_var.persistable = True
+
+            paddle._C_ops.set_persistable_value(
+                gradient_merge_var, param.name + "@GRAD@MERGE"
+            )
+
+        # step2: Accumulate persistable gradient variables in main_program
+        # NOTE(zhaoyingli): inplace operation must be 'a = a + b', cannot be 'a = b + a'
+        with main_block:
+            gard_defining_op = grad.get_defining_op()
+            paddle.pir.set_insertion_point_after(gard_defining_op)
+
+            new_gradient_merge_var = main_block.add_kwarg(
+                param.name + "@GRAD@MERGE", grad.type()
+            )
+            new_gradient_merge_var.persistable = True
+
+            grad.replace_all_uses_with(new_gradient_merge_var)
+            new_gradient_merge_var_add = paddle._C_ops.add_(
+                new_gradient_merge_var, grad
+            )
+            new_gradient_merge_var_add.get_defining_op().op_role = (
+                gard_defining_op.op_role
+            )
+
+        # step3: Construct new_params_grads and grad_to_gradient_merge
+        new_params_grads.append((param, new_gradient_merge_var))
+
+    return new_params_grads
+
+
 def _move_reduce_to_optimizer_ops_block(
     main_program, optimize_ops_block, params_grads
 ):
     main_block = main_program.global_block()
     removed_op_idx = []
-    params_grads_name = [grad.name for _, grad in params_grads]
 
     for idx, op in list(enumerate(main_block.ops)):
         if is_data_parallel_reduce_op(op):
@@ -309,6 +362,38 @@ def _move_reduce_to_optimizer_ops_block(
     return optimize_ops_block
 
 
+def _pir_move_reduce_to_optimizer_stage(main_program):
+    main_block = main_program.global_block()
+
+    for idx, op in list(enumerate(main_block.ops)):
+        if op.name() in [
+            "pd_op.c_allreduce_sum",
+            "pd_op.c_allreduce_avg",
+            "pd_op.c_reduce_sum",
+            "pd_op.c_reduce_avg",
+        ]:
+            op_input_names = op.get_input_names()
+            # NOTE(sonder): When "@RENAME@" is in the input name, it means that the op has been renamed.
+            # Such types input names are caused by shared parameter policy.
+            # Gradient merge should accumulate the gradient of ops without renaming.
+            if "@RENAME" in op_input_names[0]:
+                continue
+
+            op.op_role = OpRole.Optimize
+            main_block.move_op_to_block_end(op)
+
+            if op.name() in ["pd_op.c_allreduce_sum", "pd_op.c_reduce_sum"]:
+                scale_index = idx + 1
+                while scale_index < len(main_block.ops):
+                    if is_data_parallel_scale_op(main_block.ops[scale_index]):
+                        main_block.ops[scale_index].op_role = OpRole.Optimize
+                        main_block.move_op_to_block_end(
+                            main_block.ops[scale_index]
+                        )
+                        break
+                    scale_index += 1
+
+
 def _remove_cast_for_master_grad(main_program, dist_context):
     rename_var_map = {}
     main_block = main_program.global_block()
@@ -349,6 +434,13 @@ 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 _create_cond_block_and_update_optimizer(
     main_program,
     cond_var,
@@ -526,12 +618,48 @@ def parse_program(
     return grad_to_gradient_merge
 
 
+def _pir_parse_program(
+    main_program,
+    startup_program,
+    params_grads,
+    k_steps,
+    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
+    )
+
+    # step2: move reduce op to optimizer_ops_block
+    if gradient_sync_after_accumulate:
+        _pir_move_reduce_to_optimizer_stage(main_program, params_grads)
+
+    _pir_remove_cast_for_master_grad(main_program)
+
+    # step3: append scale op
+    if avg:
+        main_block = main_program.global_block()
+        for op in reversed(main_block.ops):
+            if op.op_role == int(OpRole.Backward):
+                paddle.pir.set_insertion_point_after(op)
+                break
+        for _, new_grad in new_params_to_grads:
+            scale = paddle.full([], 1.0 / k_steps)
+            new_grad = paddle._C_ops.scale_(new_grad, scale, 0.0, False)
+            new_grad.get_defining_op().op_role = int(OpRole.Backward)
+            scale.get_defining_op().op_role = int(OpRole.Backward)
+
+
 @register_pass("auto_parallel_gradient_merge_pass")
 class GradientMergePass(PassBase):
     def __init__(self):
         super().__init__()
         self.set_attr("k_steps", -1)
         self.set_attr("avg", True)
+        self._in_pir_mode = paddle.base.framework.get_flags(
+            "FLAGS_enable_pir_api"
+        )["FLAGS_enable_pir_api"]
 
     def _check_self(self):
         if self.get_attr("k_steps") < 1:
@@ -547,23 +675,35 @@ def _type(self):
     def _apply_single_impl(self, main_program, startup_program, context):
         k_steps = self.get_attr("k_steps", -1)
         avg = self.get_attr("avg", False)
-        dist_context = self.get_attr("dist_context")
         params_grads = self.get_attr("params_grads")
         gradient_sync_after_accumulate = self.get_attr(
             "gradient_sync_after_accumulate", False
         )
-        grad_to_global_grad = self.get_attr("grad_to_global_grad", {})
-        with paddle.static.program_guard(main_program, startup_program):
-            grad_to_merge_grad = parse_program(
-                main_program,
-                startup_program,
-                params_grads,
-                k_steps,
-                avg,
-                dist_context,
-                gradient_sync_after_accumulate,
-            )
 
-        main_program._sync_with_cpp()
-        for k, v in grad_to_merge_grad.items():
-            grad_to_global_grad[k] = v
+        if self._in_pir_mode:
+            with paddle.static.program_guard(main_program, startup_program):
+                _pir_parse_program(
+                    main_program,
+                    startup_program,
+                    params_grads,
+                    k_steps,
+                    avg,
+                    gradient_sync_after_accumulate,
+                )
+        else:
+            dist_context = self.get_attr("dist_context")
+            grad_to_global_grad = self.get_attr("grad_to_global_grad", {})
+            with paddle.static.program_guard(main_program, startup_program):
+                grad_to_merge_grad = parse_program(
+                    main_program,
+                    startup_program,
+                    params_grads,
+                    k_steps,
+                    avg,
+                    dist_context,
+                    gradient_sync_after_accumulate,
+                )
+
+            main_program._sync_with_cpp()
+            for k, v in grad_to_merge_grad.items():
+                grad_to_global_grad[k] = v
diff --git a/python/paddle/distributed/passes/auto_parallel_master_grad.py b/python/paddle/distributed/passes/auto_parallel_master_grad.py
@@ -53,11 +53,27 @@
 
 
 def _is_master_grad_cast_op(block, op):
-    if op.type != "cast":
-        return False
-    assert len(op.input_arg_names) == 1
-    assert len(op.output_arg_names) == 1
-    input_var_name = op.input_arg_names[0]
+    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
+
+    assert len(input_names) == 1
+    assert len(output_names) == 1
+
+    input_var_name = input_names[0]
+
     return (
         "@master_grad_fp16" in input_var_name
         or "@master_grad_bf16" in input_var_name
diff --git a/python/paddle/distributed/passes/pass_utils.py b/python/paddle/distributed/passes/pass_utils.py
@@ -309,8 +309,10 @@ def set_pir_skip_gc_vars(num_micro_batches, job_types, sub_programs, jobs):
                 # if a value is renamed by shadow_output,
                 # it will be used by other sub_programs
                 type_to_var_names[type].add(op.attrs()["output_name"])
-    # NOTE(lizhiyu): After finishing the gradient merge, enbale this checking.
-    # assert "backward" not in type_to_var_names.keys(), f"The backward sub_program can't have skip_gc_vars."
+
+    assert (
+        len(type_to_var_names["backward"]) == 0
+    ), "The backward sub_program can't have skip_gc_vars."
 
     for job in jobs:
         job_type = job.type()
diff --git a/test/auto_parallel/pir/mlp_demo.py b/test/auto_parallel/pir/mlp_demo.py
