PaddlePaddle
diff --git a/‎python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_optimizer_stage2.py
Copy file name to clipboard
+27-2Lines changed: 27 additions & 2 deletions b/‎python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_optimizer_stage2.py
Copy file name to clipboard
+27-2Lines changed: 27 additions & 2 deletions
diff --git a/‎python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_stage2.py
Copy file name to clipboard
+52-17Lines changed: 52 additions & 17 deletions b/‎python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_stage2.py
Copy file name to clipboard
+52-17Lines changed: 52 additions & 17 deletions
diff --git a/‎python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_storage.py
Copy file name to clipboard
+5-1Lines changed: 5 additions & 1 deletion b/‎python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_storage.py
Copy file name to clipboard
+5-1Lines changed: 5 additions & 1 deletion
diff --git a/‎python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_utils.py
Copy file name to clipboard
+27-14Lines changed: 27 additions & 14 deletions b/‎python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_utils.py
Copy file name to clipboard
+27-14Lines changed: 27 additions & 14 deletions
diff --git a/‎test/collective/fleet/CMakeLists.txt
Copy file name to clipboard
+16-4Lines changed: 16 additions & 4 deletions b/‎test/collective/fleet/CMakeLists.txt
Copy file name to clipboard
+16-4Lines changed: 16 additions & 4 deletions
@@ -115,6 +115,17 @@ def __init__(
         else:
             self._local_params.extend(list(params))
 
+        self.use_main_grad = None
+        for param in self._local_params:
+            if self.use_main_grad is None and hasattr(param, "main_grad"):
+                self.use_main_grad = True
+            if self.use_main_grad:
+                assert hasattr(
+                    param, "main_grad"
+                ), "Params have different main grad attributes."
+        if self.use_main_grad:
+            assert not offload, "offload not support main_grad for now"
+
         self._default_device = device
         self._pfp16 = (
             len(
@@ -127,6 +138,17 @@ def __init__(
             )
             > 0
         )
+        self._pbf16 = (
+            len(
+                list(
+                    filter(
+                        lambda x: x.trainable and x.dtype == Type.bf16.value,
+                        self._local_params,
+                    )
+                )
+            )
+            > 0
+        )
 
         self._broadcast_overlap = False
         self._forward_pre_hook_remove_helper = []
@@ -294,7 +316,10 @@ def _generate_master_params(self, trainable_params):
                     )
         else:
             for param in trainable_params:
-                if param.dtype == Type.fp16.value:
+                if (
+                    param.dtype == Type.fp16.value
+                    or param.dtype == Type.bf16.value
+                ):
                     master_tensor = paddle.cast(param, Type.fp32.value)
                     master_tensor.name = param.name
                     self._optim._master_weights[param.name] = master_tensor
@@ -421,7 +446,7 @@ def _integration_params(self):
                     trainable_params = list(
                         filter(lambda x: x.trainable, params)
                     )
-                    if self._pfp16 and dst_rank == self._rank:
+                    if (self._pfp16 or self._pbf16) and dst_rank == self._rank:
                         self._generate_master_params(trainable_params)
                     if trainable_params:
                         param_storage = ParamStorage(
 
@@ -107,6 +107,14 @@ def __init__(
         self._all_params = []
         for optim in self._sharding_optimizers:
             self._all_params.extend(list(optim.local_params))
+        self.use_main_grad = None
+        for param in self._all_params:
+            if self.use_main_grad is None and hasattr(param, "main_grad"):
+                self.use_main_grad = True
+            if self.use_main_grad:
+                assert hasattr(
+                    param, "main_grad"
+                ), "Params have different main grad attributes."
 
         # sharing stage 2 comm overlap flag
         self._reduce_overlap = False
@@ -212,8 +220,11 @@ def _clear_gradients(self):
 
         # Release grads of params
         for param in self._trainable_params:
-            if param.name in self._param_grads and param.grad is not None:
-                param._zero_grads()
+            if param.name in self._param_grads:
+                if self.use_main_grad and param.main_grad is not None:
+                    param.main_grad.zero_()
+                elif param.grad is not None:
+                    param._zero_grads()
 
         # Release grads of master params with offload strategy
         if self._offload:
@@ -242,8 +253,11 @@ def _grad_scale(self):
         # Scale grads of params
         with paddle.no_grad():
             for param in self._trainable_params:
-                if param.name in self._param_grads and param.grad is not None:
-                    param.grad.scale_(scale=scale_factor)
+                if param.name in self._param_grads:
+                    if self.use_main_grad and param.main_grad is not None:
+                        param.main_grad.scale_(scale=scale_factor)
+                    elif param.grad is not None:
+                        param.grad.scale_(scale=scale_factor)
                 # param._reset_grad_inplace_version(True)
 
             # Scale grads of master params with offload strategy
@@ -362,10 +376,13 @@ def _set_reduce_overlap(self, reduce_overlap):
             ), "Only support comm overlap strategy for single optimizer"
         self._sharding_optimizers[0]._set_reduce_overlap(reduce_overlap)
 
-    def _get_scaled_grad_fn(self):
+    def _get_scaled_grad_fn(self, param):
         @paddle.autograd.no_grad()
         def scale(grad):
-            grad.scale_(self._world_size_scaling)
+            if hasattr(param, "main_grad"):
+                param.main_grad.scale_(self._world_size_scaling)
+            else:
+                grad.scale_(self._world_size_scaling)
 
         return scale
 
@@ -383,16 +400,20 @@ def reduce(*_):
                 # Skip gradient reduction, do not change status information
                 if self._grad_reduced[index]:
                     assert (
-                        param.grad is not None
-                    ), "Parameter gradient cannot be None"
+                        param.grad is not None or param.main_grad is not None
+                    ), "Parameter should have grad or main grad"
 
                     # Change reduce information
                     self._grad_reduced[index] = False
 
                     # Clear the gradient that does not belong to the current rank through the callback function
                     def cleanup():
                         if dst_rank != self._rank:
-                            param.clear_gradient(False)
+                            if self.use_main_grad:
+                                param.main_grad._clear_data()
+                                param.main_grad = None
+                            else:
+                                param.clear_gradient(False)
                         elif self._offload:
                             tmp_grad = param.grad.cast(
                                 dtype=Type.fp32.value
@@ -407,7 +428,9 @@ def cleanup():
                     # Synchronize the reduce parameter gradient asynchronize
                     self._sharding_optimizers[0]._update_task(
                         dist.reduce(
-                            tensor=param.grad,
+                            tensor=param.grad
+                            if not self.use_main_grad
+                            else param.main_grad,
                             dst=self._group.ranks[dst_rank],
                             group=self._group,
                             sync_op=not self._reduce_overlap,
@@ -426,8 +449,8 @@ def reduce(*_):
                 # Skip gradient reduction, do not change status information
                 if self._grad_reduced[index]:
                     assert (
-                        param.grad is not None
-                    ), "Parameter gradient cannot be None"
+                        param.grad is not None or param.main_grad is not None
+                    ), "Parameter should have grad or main grad"
 
                     # Change reduce information
                     self._grad_reduced[index] = False
@@ -441,7 +464,11 @@ def reduce(*_):
                         def cleanup():
                             if dst_rank != self._rank:
                                 for p in grad_storage._params:
-                                    p.clear_gradient(False)
+                                    if self.use_main_grad:
+                                        p.main_grad._clear_data()
+                                        p.main_grad = None
+                                    else:
+                                        p.clear_gradient(False)
 
                                 grad_storage.buffer._clear_data()
                             elif self._offload:
@@ -491,7 +518,7 @@ def _setup_backward_hooks(self):
             return
 
         for index, param in enumerate(self._trainable_params):
-            param._register_grad_hook(self._get_scaled_grad_fn())
+            param._register_grad_hook(self._get_scaled_grad_fn(param))
 
             dst_rank = self._trainable_param2rank[param.name]
 
@@ -519,7 +546,9 @@ def _setup_use_grad_storage(self):
             if dst_rank not in self._grad_storages[param.dtype].keys():
                 self._grad_storages[param.dtype][dst_rank] = GradStorage(
                     self._buffer_max_size[param.dtype],
-                    dtype=param.dtype,
+                    dtype=param.dtype
+                    if not self.use_main_grad
+                    else paddle.float32,
                     device=self._default_device,
                     destination=dst_rank,
                     parm2align=self._trainable_param2align,
@@ -629,11 +658,17 @@ def _dp_allreduce(self):
                             sync_op=True,
                         )
             for param in self._trainable_params:
-                if param.name in self._param_grads and param.grad is not None:
+                if param.name in self._param_grads:
+                    if self.use_main_grad and param.main_grad is None:
+                        continue
+                    elif param.grad is None:
+                        continue
                     dst_rank = self._trainable_param2rank[param.name]
                     if dst_rank == self._rank:
                         dist.all_reduce(
-                            tensor=param.grad,
+                            tensor=param.grad
+                            if not self.use_main_grad
+                            else param.main_grad,
                             group=self._dp_group,
                             sync_op=True,
                         )
 
@@ -300,7 +300,11 @@ def manumal_relase(self):
         """
         if not self._release:
             for p in self._params:
-                if p.grad is not None:
+                use_main_grad = hasattr(p, "main_grad")
+                if use_main_grad and p.main_grad is not None:
+                    p.main_grad._clear_data()
+                    p.main_grad = None
+                elif p.grad is not None:
                     p.clear_gradient(False)
 
             self.buffer = None
 
@@ -240,30 +240,43 @@ def unscale_method(self, optimizer):
         ):
             for group in optimizer._optim._param_groups:
                 for param in group['params']:
-                    if param.grad is not None:
-                        param_grads.append(param.grad)
-                        if param.grad.dtype in [
+                    tgt_grad = None
+                    if (
+                        hasattr(param, "main_grad")
+                        and param.main_grad is not None
+                    ):
+                        tgt_grad = param.main_grad
+                    elif param.grad is not None:
+                        tgt_grad = param.grad
+                    if tgt_grad is not None:
+                        param_grads.append(tgt_grad)
+                        if tgt_grad.dtype in [
                             core.VarDesc.VarType.FP16,
                             paddle.float16,
                         ]:
-                            param_grads_fp16.append(param.grad)
-                        elif param.grad.dtype in [paddle.bfloat16]:
-                            param_grads_bfp16.append(param.grad)
+                            param_grads_fp16.append(tgt_grad)
+                        elif tgt_grad.dtype in [paddle.bfloat16]:
+                            param_grads_bfp16.append(tgt_grad)
                         else:
-                            param_grads_fp32.append(param.grad)
+                            param_grads_fp32.append(tgt_grad)
         else:
             for param in optimizer._optim._parameter_list:
-                if param.grad is not None:
-                    param_grads.append(param.grad)
-                    if param.grad.dtype in [
+                tgt_grad = None
+                if hasattr(param, "main_grad") and param.main_grad is not None:
+                    tgt_grad = param.main_grad
+                elif param.grad is not None:
+                    tgt_grad = param.grad
+                if tgt_grad is not None:
+                    param_grads.append(tgt_grad)
+                    if tgt_grad.dtype in [
                         core.VarDesc.VarType.FP16,
                         paddle.float16,
                     ]:
-                        param_grads_fp16.append(param.grad)
-                    elif param.grad.dtype in [paddle.bfloat16]:
-                        param_grads_bfp16.append(param.grad)
+                        param_grads_fp16.append(tgt_grad)
+                    elif tgt_grad.dtype in [paddle.bfloat16]:
+                        param_grads_bfp16.append(tgt_grad)
                     else:
-                        param_grads_fp32.append(param.grad)
+                        param_grads_fp32.append(tgt_grad)
 
         temp_found_inf_fp16 = to_variable(np.array([0]).astype(np.bool_))
         temp_found_inf_bfp16 = to_variable(np.array([0]).astype(np.bool_))
 
@@ -296,6 +296,21 @@ if(LOCAL_ALL_ARCH AND LOCAL_ALL_PLAT)
   )
   set_tests_properties(test_dygraph_sharding_stage2 PROPERTIES TIMEOUT "200")
 endif()
+if(LOCAL_ALL_ARCH AND LOCAL_ALL_PLAT)
+  bash_test_modules(
+    test_dygraph_sharding_stage2_bf16
+    START_BASH
+    ../../legacy_test/dist_test.sh
+    TIMEOUT
+    "200"
+    LABELS
+    "RUN_TYPE=DIST"
+    ENVS
+    "PADDLE_DIST_UT_PORT=22024;NVIDIA_TF32_OVERRIDE=0;http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python"
+  )
+  set_tests_properties(test_dygraph_sharding_stage2_bf16 PROPERTIES TIMEOUT
+                                                                    "200")
+endif()
 if(LOCAL_ALL_ARCH AND LOCAL_ALL_PLAT)
   bash_test_modules(
     test_parallel_dygraph_control_flow
@@ -946,11 +961,8 @@ if((WITH_GPU) AND (LINUX))
     MODULES
     test_dygraph_save_for_auto_infer
     ENVS
-    "NVIDIA_TF32_OVERRIDE=0;http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python"
+    "NVIDIA_TF32_OVERRIDE=0;FLAGS_new_executor_micro_batching=False;http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python"
   )
   set_tests_properties(test_dygraph_save_for_auto_infer
                        PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=DIST")
-  set_tests_properties(
-    test_dygraph_save_for_auto_infer
-    PROPERTIES ENVIRONMENT "FLAGS_new_executor_micro_batching=False")
 endif()