【AutoParallel】Sharding pass support BF16 and master_grad (#60779)

heavyrain-lzy · web-flow · commit 5d46ba044d35 · 2024-01-16T10:56:20.000+08:00
* support sharding bf16

* change sharding

* chage code according to review
diff --git a/python/paddle/distributed/auto_parallel/static/parallelizer_v2.py b/python/paddle/distributed/auto_parallel/static/parallelizer_v2.py
@@ -412,6 +412,9 @@ def _apply_post_optimization(
             config["dist_context"] = self._dist_context
             config["params_grads"] = params_grads
             config["global_rank"] = rank
+            if self._strategy.amp.enable:
+                amp_config = copy.deepcopy(self._strategy.amp.to_dict())
+                config["amp_dtype"] = amp_config['dtype']
             auto_parallel_sharding_pass = new_pass(
                 "auto_parallel_sharding", config
             )
diff --git a/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py b/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py
@@ -863,6 +863,7 @@ def insert_broadcast_ops(
 
 DtypeToSize = {
     core.VarDesc.VarType.FP16: 2,
+    core.VarDesc.VarType.BF16: 2,
     core.VarDesc.VarType.FP32: 4,
     core.VarDesc.VarType.FP64: 8,
     core.VarDesc.VarType.INT16: 2,
diff --git a/python/paddle/distributed/passes/auto_parallel_master_grad.py b/python/paddle/distributed/passes/auto_parallel_master_grad.py
@@ -33,12 +33,38 @@
 from paddle.static import program_guard
 
 from ..utils.log_utils import get_logger
-from .auto_parallel_sharding import _supported_optimizer_type
 from .pass_base import PassBase, register_pass
 
+_supported_optimizer_type = [
+    "adam",
+    "adamax",
+    "adamw",
+    "decayed_adagrad",
+    "momentum",
+    "dgc_momentum",
+    "lars_momentum",
+    "merged_momentum",
+    "lamb",
+    "sgd",
+]
+
 logger = get_logger(logging.INFO, "MasterGradPass")
 
 
+def _is_master_grad_cast_op(block, op, amp_dtype="float16"):
+    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
+
+
 def get_output_in_varlist(op, var_names) -> List[str]:
     grad_names = []
     for output_name in op.output_arg_names:
@@ -113,9 +139,9 @@ def _add_cast_op(self, cur_block, grad_names: List[str], dist_context):
                 ref_dims_mapping = ref_output_dist_attr.dims_mapping
                 ref_chunk_id = producer_op_dist_attr.chunk_id
                 grad_half_precision_name = (
-                    grad_name + '@tmp_fp16'
+                    grad_name + '@master_grad_fp16'
                     if is_fp16
-                    else grad_name + '@tmp_bf16'
+                    else grad_name + '@master_grad_bf16'
                 )
                 grad_half_precision = cur_block.create_var(
                     name=grad_half_precision_name,
diff --git a/python/paddle/distributed/passes/auto_parallel_sharding.py b/python/paddle/distributed/passes/auto_parallel_sharding.py
@@ -42,6 +42,7 @@
 from paddle.static import default_main_program, default_startup_program
 from paddle.utils import unique_name
 
+from .auto_parallel_master_grad import _is_master_grad_cast_op
 from .pass_base import PassBase, register_pass
 from .pass_utils import AutoParallelStreamType
 
@@ -72,6 +73,9 @@
 
 _logger = get_logger(logging.INFO)
 
+__amp_target_dtype__ = core.VarDesc.VarType.FP16
+__amp_target_dtype_name__ = "float16"
+
 
 def _is_reshard_op(op):
     return op.desc.has_attr(
@@ -100,6 +104,7 @@ def __init__(self):
         self.set_attr("enable_hierarchical_comm", None)
         self.set_attr("params_grads", [])
         self.set_attr("global_rank", -1)
+        self.set_attr("amp_dtype", "float16")
         self.dp_groups = set()
         self.sharding_infos = []
         self.varname_to_sharding_info = {}
@@ -181,6 +186,11 @@ def _apply_single_impl(self, main_program, startup_program, context):
             startup_program.global_block(),
         )
 
+        self.amp_dtype = self.get_attr("amp_dtype")
+        if self.amp_dtype == "bfloat16":
+            __amp_target_dtype__ = core.VarDesc.VarType.BF16
+            __amp_target_dtype_name__ = "bfloat16"
+
         # NOTE Multi / Sub-Block Support
         # we assume that only parameter are present and partitioned in main_block,
         # there is NO new param in sub_block, and all params in sub_block follows the same
@@ -1508,9 +1518,11 @@ def _is_param_grad_fp32_cast_op(block, op):
     if not is_backward_op(op):
         return False
     if not _is_desired_cast_op(
-        block, op, core.VarDesc.VarType.FP16, core.VarDesc.VarType.FP32
+        block, op, __amp_target_dtype__, core.VarDesc.VarType.FP32
     ):
         return False
+    if _is_master_grad_cast_op(block, op, __amp_target_dtype_name__):
+        return False
     output_name = op.output_arg_names[0]
     base_name = output_name[: output_name.find("@")]
     if not block.has_var(base_name):
@@ -1533,7 +1545,7 @@ def _is_desired_cast_op(
     block,
     op,
     src_var_type=core.VarDesc.VarType.FP32,
-    dst_var_type=core.VarDesc.VarType.FP16,
+    dst_var_type=__amp_target_dtype__,
 ):
     if op.type != "cast":
         return False
@@ -1552,6 +1564,8 @@ def _get_base_name_from_grad_name(grad_name):
     base_name = None
     if ".cast_fp16@GRAD" in grad_name:
         base_name = grad_name[: grad_name.find(".cast_fp16@GRAD")]
+    elif ".cast_bf16@GRAD" in grad_name:
+        base_name = grad_name[: grad_name.find(".cast_bf16@GRAD")]
     elif "@GRAD" in grad_name:
         base_name = grad_name[: grad_name.find("@GRAD")]
     return base_name
