PaddlePaddle
diff --git a/‎paddle/fluid/framework/distributed_strategy.proto
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/fluid/framework/distributed_strategy.proto
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py
Copy file name to clipboard
+37-13Lines changed: 37 additions & 13 deletions b/‎python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py
Copy file name to clipboard
+37-13Lines changed: 37 additions & 13 deletions
diff --git a/‎python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py
Copy file name to clipboard
+7-5Lines changed: 7 additions & 5 deletions b/‎python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py
Copy file name to clipboard
+7-5Lines changed: 7 additions & 5 deletions
diff --git a/‎python/paddle/distributed/fleet/meta_parallel/pp_utils/utils.py
Copy file name to clipboard
-124Lines changed: 0 additions & 124 deletions b/‎python/paddle/distributed/fleet/meta_parallel/pp_utils/utils.py
Copy file name to clipboard
-124Lines changed: 0 additions & 124 deletions
@@ -68,6 +68,8 @@ message PpConfig {
 
 message DygraphShardingConfig {
   optional bool tensor_fusion = 1 [ default = false ];
+  optional int32 accumulate_steps = 2 [ default = 1 ];
+  optional bool comm_overlap = 3 [ default = false ];
 }
 
 message HybridConfig {
 
@@ -78,12 +78,23 @@ def __init__(self, optimizer, hcg):
         self.tensor_fusion = strategy.hybrid_configs[
             'sharding_configs'
         ].tensor_fusion
+        self.accumulate_steps = strategy.hybrid_configs[
+            'sharding_configs'
+        ].accumulate_steps
+        self.comm_overlap = strategy.hybrid_configs[
+            'sharding_configs'
+        ].comm_overlap
         pp_overlap = strategy.hybrid_configs['pp_configs'].sharding_comm_overlap
-        if self.tensor_fusion:
+        if self.tensor_fusion or self.comm_overlap:
             assert (
                 not pp_overlap
             ), "Can not enable pp's sharding_comm_overlap and sharding's tensor_fusion at the same time."
 
+        self._use_main_grad = hasattr(self._parameter_list[0], "main_grad")
+        self._rank2decay = {}
+        self._rank2fused = {}
+        self._comm_buffers = []
+
         self._rank2params = self._partition_parameters()
         self._param2rank = self._map_param_to_rank()
 
@@ -95,25 +106,22 @@ def __init__(self, optimizer, hcg):
                 '_param_groups', self._rank2params[self._sharding_rank]
             )
         else:
-            self._use_main_grad = hasattr(self._parameter_list[0], "main_grad")
-            self._rank2decay = {}
-            self._rank2fused = {}
             self._tensor_fusion()
 
             decay_params = [
                 p.name for p in self._rank2decay[self._sharding_rank]
             ]
-            all_params = self._rank2fused[self._sharding_rank]
+            fused_params = self._rank2fused[self._sharding_rank]
             apply_decay_param_fun = lambda x: x in decay_params
 
-            params = []
+            all_fused_params = []
             for v in self._rank2fused.values():
-                params += v
-            self._parameter_list = params
-            self._param_groups = params
+                all_fused_params += v
+            self._parameter_list = all_fused_params
+            self._param_groups = all_fused_params
 
-            self._set_inner_opt_attr('_parameter_list', all_params)
-            self._set_inner_opt_attr('_param_groups', all_params)
+            self._set_inner_opt_attr('_parameter_list', fused_params)
+            self._set_inner_opt_attr('_param_groups', fused_params)
             origin_decay_param_fun = getattr(
                 self._inner_opt, '_apply_decay_param_fun', None
             )
@@ -145,11 +153,23 @@ def clear_grad(self, set_to_zero=True):
                     p.clear_gradient(set_to_zero)
 
     def _tensor_fusion(self):
+        comm_group = self._hcg.get_sharding_parallel_group()
         for i in range(self._sharding_world_size):
             params = self._rank2params[i]
-            decay_fused, all_fused = fused_parameters(
-                params, self._use_main_grad
+            dst = comm_group.ranks[i]
+            # TODO(sharding dev): make scale_after_comm a field to be configured by user
+            decay_fused, all_fused, all_buffer = fused_parameters(
+                params,
+                use_main_grad=self._use_main_grad,
+                fuse_param=True,
+                comm_overlap=self.comm_overlap,
+                comm_group=comm_group,
+                dst=dst,
+                acc_step=self.accumulate_steps,
+                scale_after_comm=False,
             )
+            if self.comm_overlap:
+                self._comm_buffers += all_buffer
             self._rank2decay[i] = decay_fused
             self._rank2fused[i] = all_fused
             for p in all_fused:
@@ -199,6 +219,10 @@ def _map_param_to_rank(self):
     def reduce_gradients(self, parameter_list, hcg):
         # TODO merge grad / nrank with dp
         logger.debug("sharding start gradients sync")
+        if self.comm_overlap:
+            for buffer in self._comm_buffers:
+                buffer.scale_grads()
+            return
         with framework.no_grad():
             sharding_nrank = hcg.get_sharding_parallel_group().nranks
             for param in parameter_list:
 
@@ -37,11 +37,11 @@
     from .pp_utils import p2p_communication as p2p
 
 from paddle.distributed.fleet.utils.tensor_fusion_helper import (
+    HOOK_ACTION,
+    FusedCommBuffer,
     assign_group_by_size,
 )
 
-from .pp_utils.utils import HOOK_ACTION, FusedCommBuffer
-
 __all__ = []
 
 g_shard_use_reduce = int(os.environ.get("FLAGS_shard_use_reduce", 0))
@@ -334,9 +334,11 @@ def register_allreduce_overlap_hook(self, model, comm_group, acc_steps, dp):
 
             for dst in fused_parameter_group:
                 parameter_list = fused_parameter_group[dst]
-                if not dp:
+                if act != HOOK_ACTION.ALL_REDUCE:
                     # parse the relative dst rank to absolute dst rank for sharding
                     dst = comm_group.ranks[dst]
+                else:
+                    dst = -1
                 var_groups = assign_group_by_size(parameter_list)
                 for group_idx, parameters in var_groups.items():
                     buffer = FusedCommBuffer(
@@ -515,7 +517,7 @@ def forward_backward_pipeline(
         if self._comm_overlap:
             assert len(self._comm_buffers) > 0
             for buffer in self._comm_buffers:
-                buffer.scale_and_split_grads()
+                buffer.scale_grads()
 
         if self._enable_timer:
             self.timers("allreduce_shared_weight_gradients").start()
@@ -1256,7 +1258,7 @@ def forward_backward_pipeline(
             if self._comm_overlap:
                 assert len(self._comm_buffers) > 0
                 for buffer in self._comm_buffers:
-                    buffer.scale_and_split_grads()
+                    buffer.scale_grads()
 
             if static_scheduler:
                 self._reset_counter()
 
@@ -15,19 +15,10 @@
 
 import paddle
 from paddle import _legacy_C_ops
-from paddle.distributed.fleet.utils.tensor_fusion_helper import (
-    flatten_dense_tensors,
-)
-from paddle.framework import base as imperative_base
 
 __all__ = []
 
 
-class HOOK_ACTION:
-    ALL_REDUCE = 0
-    REDUCE = 1
-
-
 FLOAT_TYPE_DICT = {
     paddle.float16: "float16",
     paddle.float32: "float32",
@@ -116,118 +107,3 @@ def _all_gather(tensor, group=None, use_calc_stream=True):
         'nranks',
         nranks,
     )
-
-
-class FusedCommBuffer:
-    def __init__(self, id, params, comm_group, acc_steps=1, act=None, dst=-1):
-        self._id = id
-        self._params = params
-        self._acc_steps = acc_steps
-        self._comm_group = comm_group
-
-        self.use_main_grad = hasattr(self._params[0], "main_grad")
-
-        self._task = None
-        self._params_step_dict = {}
-        self._params_checked_in = 0
-        self._params_to_addr = {}
-
-        self._act = act
-        if self._act == HOOK_ACTION.ALL_REDUCE:
-            assert dst == -1
-        elif self._act == HOOK_ACTION.REDUCE:
-            assert dst != -1
-        else:
-            raise ValueError(
-                "The act should be allreudce for dp or reduce for sharding."
-            )
-        self._dst = dst
-
-        self._init_step_dict()
-
-        self.grad_storage = flatten_dense_tensors(
-            self._params,
-            use_main_grad=self.use_main_grad,
-            fuse_param=False,
-            warp_buffer=False,
-        ).buffer
-
-        self._record_addr()
-
-    def _record_addr(self):
-        for param in self._params:
-            addr = (
-                param.main_grad.data_ptr()
-                if self.use_main_grad
-                else param.grad.data_ptr()
-            )
-            self._params_to_addr[param.name] = addr
-
-    def _init_step_dict(self):
-        for p in self._params:
-            self._params_step_dict[p.name] = 0
-
-    def _reset_params_checked_in(self):
-        self._task = None
-        self._init_step_dict()
-        self._params_checked_in = 0
-
-    @property
-    def _all_params_checked_in(self):
-        return (
-            len(self._params) == self._params_checked_in
-            and len(self._params_step_dict) == 0
-        )
-
-    def add_grad(self, param):
-        assert param.name in self._params_step_dict
-        current_ptr = (
-            param.main_grad.data_ptr()
-            if self.use_main_grad
-            else param.grad.data_ptr()
-        )
-        if self._params_to_addr[param.name] != current_ptr:
-            raise ValueError(
-                "The address of the grad/main_grad of the param has been changed during training, "
-                "which is not allowed for dp/sharding overlap with pp. "
-                "This may be caused by some non-inplace operations on the grad/main_grad. "
-                "Please use the inplace version of the operations or disable the overlapping."
-            )
-
-        self._params_step_dict[param.name] += 1
-
-        if self._params_step_dict[param.name] == self._acc_steps:
-            self._params_checked_in += 1
-            self._params_step_dict.pop(param.name)
-
-        if self._all_params_checked_in:
-            self._comm_grads()
-
-    @imperative_base.no_grad
-    def _comm_grads(self):
-        assert self._all_params_checked_in
-
-        if self._act == HOOK_ACTION.ALL_REDUCE:
-            task = paddle.distributed.all_reduce(
-                self.grad_storage, group=self._comm_group, sync_op=False
-            )
-
-        elif self._act == HOOK_ACTION.REDUCE:
-            task = paddle.distributed.reduce(
-                self.grad_storage,
-                dst=self._dst,
-                group=self._comm_group,
-                sync_op=False,
-            )
-
-        self._task = task
-
-    @imperative_base.no_grad
-    def scale_and_split_grads(self):
-        assert self._task is not None
-        self._task.wait()
-
-        scale_factor = 1.0 / self._comm_group.nranks
-        self.grad_storage.scale_(scale_factor)
-
-        self._reset_params_checked_in()
Original file line number	Diff line number	Diff line change
`@@ -68,6 +68,8 @@ message PpConfig {`
`68`	`68`
`69`	`69`	`message DygraphShardingConfig {`
`70`	`70`	`optional bool tensor_fusion = 1 [ default = false ];`
	`71`	`+ optional int32 accumulate_steps = 2 [ default = 1 ];`
	`72`	`+ optional bool comm_overlap = 3 [ default = false ];`
`71`	`73`	`}`
`72`	`74`
`73`	`75`	`message HybridConfig {`