sharding stage1 V1 support Broadcast overlap Forward (#63945)

iosmers · web-flow · commit 3789d83744a7 · 2024-05-11T16:05:53.000+08:00
* sharding v1 overlap

* delete pybind

* add txt

* add b.txt

* delete test file

* add pybind

* add test case for stage1 v1 overlap

* add test case for stage1 v1 overlap

* update test case

* delete print optimizer

* update

* update models to layers

* update mlp1
diff --git a/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py
@@ -116,6 +116,20 @@ def __init__(self, optimizer, hcg):
         self._rank2params = self._partition_parameters()
         self._param2rank = self._map_param_to_rank()
 
+        self._broadcast_overlap = False
+        self._forward_pre_hook_remove_helper = []
+
+        try:
+            # The fp32 params such as layer_norm_0.w_0 will be at the end of param_list.
+            # Have to sort the params to make sure all params are in the forward using order.
+            self._broadcast_order_params = sorted(
+                self._parameter_list,
+                key=lambda x: int(x.name.split('.')[0].split('_')[-1]),
+            )
+
+        except ValueError:
+            self._broadcast_order_params = None
+
         if not self.tensor_fusion and not self.comm_overlap:
             local_params = self._rank2params[self._sharding_rank]
             self._set_inner_opt_attr('_parameter_list', local_params)
@@ -318,6 +332,13 @@ def reduce_gradients(self, parameter_list, hcg):
                         sync_op=True,
                     )
 
+    def _forward_pre_hook_function(self, tasks):
+        def __impl__(x, y):
+            for task in tasks:
+                task.wait()
+
+        return __impl__
+
     def _sharding_sync_parameters(self):
         """
         Synchronize parameter across sharding group efficiently.
@@ -334,27 +355,57 @@ def _sharding_sync_parameters(self):
             sharding_group_ranks = self._hcg.get_sharding_parallel_group().ranks
 
             broadcast_tasks = []
-            for rank, params in valid_rank_to_params.items():
-                # Compute the global source rank only once per each rank's set of parameters
-                src_rank = sharding_group_ranks[rank]
-
-                for param in params:
-                    # NOTE: We should check if the parameter is trainable, because some parameters
-                    # (e.g., freeze the parameters for training) are not trainable and should
-                    # not be broadcasted.
-                    g_var = self._get_param_grad(param)
-                    if g_var is not None:
+            if self._broadcast_overlap:
+                param2task = {}
+
+                group = self._hcg.get_sharding_parallel_group()
+                for param in self._broadcast_order_params:
+                    if param.trainable:
                         task = paddle.distributed.broadcast(
-                            param,
-                            src=src_rank,
-                            group=self._hcg.get_sharding_parallel_group(),
+                            tensor=param,
+                            src=group.ranks[self._param2rank[param.name]],
+                            group=group,
                             sync_op=False,
                         )
-                        broadcast_tasks.append(task)
+                        assert param.name not in param2task
+                        param2task[param.name] = task
+
+                for layer in self._layers.sublayers():
+                    if len(layer.sublayers()) == 0:
+                        # Register forward pre hood for leaf layers. This will get the best performance.
+                        tasks = []
+                        for param in layer.parameters():
+                            if param.trainable:
+                                if param.name in param2task:
+                                    tasks.append(param2task[param.name])
+                        self._forward_pre_hook_remove_helper.append(
+                            layer.register_forward_pre_hook(
+                                self._forward_pre_hook_function(tasks)
+                            )
+                        )
 
-            # Wait for all async broadcast tasks to complete
-            for task in broadcast_tasks:
-                task.wait()
+            else:
+                for rank, params in valid_rank_to_params.items():
+                    # Compute the global source rank only once per each rank's set of parameters
+                    src_rank = sharding_group_ranks[rank]
+
+                    for param in params:
+                        # NOTE: We should check if the parameter is trainable, because some parameters
+                        # (e.g., freeze the parameters for training) are not trainable and should
+                        # not be broadcasted.
+                        g_var = self._get_param_grad(param)
+                        if g_var is not None:
+                            task = paddle.distributed.broadcast(
+                                param,
+                                src=src_rank,
+                                group=self._hcg.get_sharding_parallel_group(),
+                                sync_op=False,
+                            )
+                            broadcast_tasks.append(task)
+
+                # Wait for all async broadcast tasks to complete
+                for task in broadcast_tasks:
+                    task.wait()
 
     def _update_trainable(self):
         """
@@ -384,10 +435,33 @@ def minimize(
 
         return result
 
+    def _set_broadcast_overlap(self, broadcast_overlap, layers=None):
+        self._broadcast_overlap = broadcast_overlap
+        if self._broadcast_overlap:
+            assert (
+                layers is not None
+            ), "To Enable Stage1 Optimizer Broadcast Overlap Forward, layers cannot be None"
+            self._layers = layers
+            warnings.warn(
+                r"Setting overlap broadcast implies that `paddle.device.cuda.synchronize()` must be manually invoked before calling `paddle.save()` and prior to inference"
+            )
+
+            if self._broadcast_order_params is None:
+                warnings.warn(
+                    r"The param name passed to the optimizer doesn't follow .+_[0-9]+\..+ patter, "
+                    "overlap broadcast may harm the performance."
+                )
+                self._broadcast_order_params = self._parameter_list
+
     @imperative_base.no_grad
     @framework.dygraph_only
     def step(self):
         # TODO Check whether the model trainable param changed and update state accordingly
+        if self._broadcast_overlap:
+            # Clear the pre forward hook in the optimizer step.
+            for hook_remove in self._forward_pre_hook_remove_helper:
+                hook_remove.remove()
+            self._forward_pre_hook_remove_helper = []
 
         target_param_list = (
             self._origin_parameter_list
diff --git a/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/hybrid_parallel_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/hybrid_parallel_optimizer.py
@@ -331,6 +331,14 @@ def __init__(self, optimizer, hcg, strategy):
                                 inner_opt._grad_clip, hcg
                             )
 
+    def _set_broadcast_overlap(self, broadcast_overlap, layers=None):
+        self._broadcast_overlap = broadcast_overlap
+        if self._broadcast_overlap:
+            self._layers = layers
+            self._inner_opt._set_broadcast_overlap(
+                self._broadcast_overlap, self._layers
+            )
+
     def _insert_sync(self, sync_var, src, mp_group, sync_mode):
         if sync_mode == "broadcast":
             paddle.distributed.broadcast(
diff --git a/test/collective/fleet/dygraph_group_sharded_stage1_bf16.py b/test/collective/fleet/dygraph_group_sharded_stage1_bf16.py
@@ -37,6 +37,7 @@ def train_mlp(
     acc_steps=1,
     use_main_grad=False,
     test_scaler=False,
+    broadcast_overlap=False,
 ):
     logging.info(
         f"-- Train Info: use_pure_bf16={use_pure_bf16}, use_main_grad={use_main_grad}, acc_steps={acc_steps}"
@@ -86,6 +87,10 @@ def train_mlp(
 
     if sharding_stage == 1:
         optimizer = fleet.distributed_optimizer(optimizer)
+        if broadcast_overlap:
+            optimizer._set_broadcast_overlap(
+                broadcast_overlap=broadcast_overlap, layers=model
+            )
 
     if sharding_stage == 1:
         model.to(device="gpu")
@@ -191,6 +196,7 @@ def _compare_bf16_o1_vs_o2(acc_steps=1):
             train_loader=train_loader,
             use_pure_bf16=False,
             acc_steps=acc_steps,
+            broadcast_overlap=False,
         )
         o2_losses, model_param_dict_o2, optimizer_state_dict_o2 = train_mlp(
             mlp2,
@@ -199,17 +205,57 @@ def _compare_bf16_o1_vs_o2(acc_steps=1):
             use_pure_bf16=True,
             use_main_grad=True,
             acc_steps=acc_steps,
+            broadcast_overlap=False,
         )
         np.testing.assert_array_equal(o2_losses, o1_losses)
         compare_state_dict(
             model_param_dict_o1, model_param_dict_o2, optimizer_state_dict_o2
         )
 
+    def _compare_bf16_broadcast_overlap(acc_steps=1):
+        mlp1 = MLP()
+        mlp2 = MLP()
+        mlp1.set_state_dict(state_dict)
+        mlp2.set_state_dict(state_dict)
+        (
+            o1_losses_overlap,
+            model_param_dict_o1_overlap,
+            optimizer_state_dict_o1_overlap,
+        ) = train_mlp(
+            mlp1,
+            sharding_stage=1,
+            train_loader=train_loader,
+            use_pure_bf16=False,
+            acc_steps=acc_steps,
+            broadcast_overlap=True,
+        )
+        mlp1.set_state_dict(state_dict)
+        (
+            o1_losses_no_overlap,
+            model_param_dict_o1_no_overlap,
+            optimizer_state_dict_o1_no_overlap,
+        ) = train_mlp(
+            mlp1,
+            sharding_stage=1,
+            train_loader=train_loader,
+            use_pure_bf16=False,
+            acc_steps=acc_steps,
+            broadcast_overlap=False,
+        )
+
+        np.testing.assert_array_equal(o1_losses_overlap, o1_losses_no_overlap)
+        np.testing.assert_array_equal(
+            model_param_dict_o1_overlap, model_param_dict_o1_no_overlap
+        )
+
     # no gradient accumulation
     _compare_bf16_o1_vs_o2(acc_steps=1)
     # gradient accumulation
     _compare_bf16_o1_vs_o2(acc_steps=2)
 
+    _compare_bf16_broadcast_overlap(acc_steps=1)
+    _compare_bf16_broadcast_overlap(acc_steps=2)
+
     # stage1 scaler test with main_grad
     mlp3 = MLP()
     mlp3.set_state_dict(state_dict)
