【AutoParalle】Optimize cast inserted by master_grad_pass (#63228)

heavyrain-lzy · web-flow · commit 02251ccf00ad · 2024-04-08T16:10:15.000+08:00
* optimize cast op

* revert executor.py
diff --git a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py
@@ -24,6 +24,7 @@
     get_world_process_group,
 )
 from paddle.distributed.auto_parallel.static.utils import (
+    is_backward_op,
     is_forward_op,
     is_optimize_op,
     naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
@@ -37,6 +38,7 @@
 from paddle.framework import core
 from paddle.static import device_guard
 
+from .auto_parallel_master_grad import _is_master_grad_cast_op
 from .pass_base import PassBase, PassType, register_pass
 
 world_process_group = get_world_process_group()
@@ -220,7 +222,7 @@ def _append_gradient_merge_backward_op(
                     outputs={"Out": startup_gradient_merge_var},
                     attrs={
                         "shape": grad.shape,
-                        "dtype": grad.dtype,
+                        "dtype": startup_gradient_merge_var.dtype,
                         "value": float(0),
                     },
                 )
@@ -307,6 +309,46 @@ def _move_reduce_to_optimizer_ops_block(
     return optimize_ops_block
 
 
+def _remove_cast_for_master_grad(main_program, dist_context):
+    rename_var_map = {}
+    main_block = main_program.global_block()
+    for idx, op in reversed(list(enumerate(main_block.ops))):
+        if _is_master_grad_cast_op(main_block, op):
+            input_var_name = op.input_arg_names[0]
+            output_var_name = op.output_arg_names[0]
+            rename_var_map[input_var_name] = output_var_name
+            in_var = main_block.var(input_var_name)
+            out_var = main_block.var(output_var_name)
+            out_var.desc.set_dtype(in_var.dtype)
+            main_block._remove_op(idx, sync=False)
+            main_block._remove_var(input_var_name)
+
+    # rename "xxx@GRAD@master_grad_fp16" --> "xxx@GRAD"
+    if len(rename_var_map) > 0:
+        for op in reversed(main_block.ops):
+            if is_forward_op(op):
+                break
+            if is_backward_op(op):
+                output_var_names = op.output_arg_names
+                op_dist_attr = dist_context.get_op_dist_attr_for_program(op)
+                for output_var_name in output_var_names:
+                    if output_var_name in rename_var_map:
+                        out_dims_mapping = op_dist_attr.get_output_dims_mapping(
+                            output_var_name
+                        )
+                        op.desc._rename_output(
+                            output_var_name, rename_var_map[output_var_name]
+                        )
+                        op_dist_attr.set_output_dims_mapping(
+                            rename_var_map[output_var_name], out_dims_mapping
+                        )
+                        del rename_var_map[output_var_name]
+        assert (
+            len(rename_var_map) == 0
+        ), f"rename_var_map must be empty, but it is: {rename_var_map}"
+    main_block._sync_with_cpp()
+
+
 def _create_cond_block_and_update_optimizer(
     main_program,
     cond_var,
@@ -464,6 +506,8 @@ def parse_program(
             main_program, optimize_ops_block, params_grads
         )
 
+    _remove_cast_for_master_grad(main_program, dist_context)
+
     # 4 create gradient_merge_cond
     cond_var = _get_gm_cond_var(main_program, k_steps, dist_context)
 
diff --git a/python/paddle/distributed/passes/auto_parallel_master_grad.py b/python/paddle/distributed/passes/auto_parallel_master_grad.py
@@ -52,18 +52,16 @@
 logger = get_logger(logging.INFO, "MasterGradPass")
 
 
-def _is_master_grad_cast_op(block, op, amp_dtype="float16"):
+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]
-    if amp_dtype == "float16":
-        return "@master_grad_fp16" in input_var_name
-    elif amp_dtype == "bfloat16":
-        return "@master_grad_bf16" in input_var_name
-    else:
-        return False
+    return (
+        "@master_grad_fp16" in input_var_name
+        or "@master_grad_bf16" in input_var_name
+    )
 
 
 def get_output_in_varlist(op, var_names) -> List[str]:
@@ -158,8 +156,19 @@ def _add_cast_op(self, cur_block, grad_names: List[str], dist_context):
                     ref_mesh,
                     chunk_id=ref_chunk_id,
                 )
+
+                producer_op_dist_attr = (
+                    dist_context.get_op_dist_attr_for_program(producer_op)
+                )
+                origin_out_dims_mapping = (
+                    producer_op_dist_attr.get_output_dims_mapping(grad_name)
+                )
                 producer_op._rename_output(grad_name, grad_half_precision.name)
+                producer_op_dist_attr.set_output_dims_mapping(
+                    grad_half_precision.name, origin_out_dims_mapping
+                )
                 grad_var.desc.set_dtype(core.VarDesc.VarType.FP32)
+
                 cast_op = cur_block._insert_op_without_sync(
                     idx + 1,
                     type="cast",
diff --git a/python/paddle/distributed/passes/auto_parallel_sharding.py b/python/paddle/distributed/passes/auto_parallel_sharding.py
@@ -1603,7 +1603,7 @@ def _is_param_grad_fp32_cast_op(block, op):
         block, op, __amp_target_dtype__, core.VarDesc.VarType.FP32
     ):
         return False
-    if _is_master_grad_cast_op(block, op, __amp_target_dtype_name__):
+    if _is_master_grad_cast_op(block, op):
         return False
     output_name = op.output_arg_names[0]
     base_name = output_name[: output_name.find("@")]
