【Auto-Parallel】support sharding stage1 overlap in PIR (#69390)

liym27 · web-flow · commit d7746893fa90 · 2024-11-18T19:27:05.000+08:00
* reduce_scatter overlap
* all_gather overlap
* sync param before allgather
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -155,6 +155,8 @@
     'qkv_unpack_mha',
     'hash',
     'beam_search_decode',
+    'nop',
+    'nop_',
 ]
 
 NO_NEED_GEN_STATIC_ONLY_APIS = [
@@ -241,8 +243,6 @@
     'partial_recv',
     'partial_allgather',
     'partial_allgather_',
-    'nop',
-    'nop_',
     'gemm_epilogue',
     'push_dense',
     'legacy_matmul',
diff --git a/paddle/phi/kernels/gpu/nop_kernel.cu b/paddle/phi/kernels/gpu/nop_kernel.cu
@@ -17,5 +17,11 @@
 #include "paddle/phi/core/kernel_registry.h"
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-PD_REGISTER_KERNEL(nop, GPU, ALL_LAYOUT, phi::NopKernel, float) {}
+PD_REGISTER_KERNEL(nop,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::NopKernel,
+                   float,
+                   phi::dtype::bfloat16,
+                   phi::dtype::float16) {}
 #endif
diff --git a/python/paddle/distributed/auto_parallel/sharding.py b/python/paddle/distributed/auto_parallel/sharding.py
@@ -27,10 +27,14 @@
 from paddle.distributed.auto_parallel.static.reshard_funcs.nd_mesh_reshard_func import (
     get_1D_sub_process_mesh,
 )
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
 from paddle.distributed.fleet.utils.tensor_fusion_helper import (
     align,
     get_current_device_type,
 )
+from paddle.distributed.passes.pass_utils import (
+    AutoParallelStreamType,
+)
 from paddle.framework import (
     _current_expected_place_ as _get_device,
 )
@@ -224,6 +228,9 @@ def apply_gradients(self, params_grads):
                     group_param_list.append(parameters[index])
                     group_grad_list.append(grads[index])
 
+                if self._strategy.sharding.enable_overlap:
+                    self._reduce_scatter_overlap(group_grad_list, target_block)
+
                 slice_param_dict, main_shard_fused_param, main_fused_param = (
                     self._fuse_group_param(group_param_list)
                 )
@@ -308,7 +315,10 @@ def apply_gradients(self, params_grads):
                             * align_size
                             // core.size_of_dtype(dtype)
                         )
-                pir.reset_insertion_point_to_end()
+
+                if not self._strategy.sharding.enable_overlap:
+                    pir.reset_insertion_point_to_end()
+
                 shard_size = fused_grad._local_shape[0] // self._sharding_degree
                 rank = self._sharding_group.ranks.index(dist.get_rank())
                 rank_begin = rank * shard_size
@@ -321,11 +331,23 @@ def apply_gradients(self, params_grads):
                     fused_grad, self._sharding_group.id, self._sharding_degree
                 )
 
+                if self._strategy.sharding.enable_overlap:
+                    shard_fused_grad.get_defining_op().set_execution_stream(
+                        AutoParallelStreamType.SHARDING_STREAM.value
+                    )
+                pir.reset_insertion_point_to_end()
+
                 paddle._C_ops.share_var(
                     [view_shard_fused_grad, shard_fused_grad]
                 )
+
+                slice_param_list = []
+                for slice_param, param_info in slice_param_dict.items():
+                    slice_param_list.append(slice_param)
+
                 all_gather_param_info_list.append(
                     (
+                        slice_param_list,
                         main_shard_fused_param,
                         main_fused_param,
                     )
@@ -357,17 +379,89 @@ def apply_gradients(self, params_grads):
             self._inner_opt._grad_clip.has_dist_param = has_dist_param
             self._inner_opt._grad_clip.has_not_dist_param = has_not_dist_param
         self._inner_opt.apply_gradients(new_params_grads)
+
         for (
+            slice_param_list,
             shard_param,
             fused_param,
         ) in all_gather_param_info_list:
-            allgather_value = paddle._C_ops.all_gather(
-                shard_param, self._sharding_group.id, self._sharding_degree
-            )
+            if self._strategy.sharding.enable_overlap:
+                last_idx = None
+                last_op = None
+                for op in slice_param_list[-1].all_used_ops():
+                    idx = target_block.ops.index(op)
+                    if last_idx is None or idx > last_idx:
+                        last_idx = idx
+                        last_op = op
+
+                # NOTE: add dependency between opt op and allgather_value for correctness
+                tmp = paddle._C_ops.nop(last_op.results()[0])
+                tmp.get_defining_op().set_execution_stream(
+                    AutoParallelStreamType.SHARDING_STREAM.value
+                )
+
+                allgather_value = paddle._C_ops.all_gather(
+                    shard_param, self._sharding_group.id, self._sharding_degree
+                )
+                allgather_value.get_defining_op().set_execution_stream(
+                    AutoParallelStreamType.SHARDING_STREAM.value
+                )
+            else:
+                allgather_value = paddle._C_ops.all_gather(
+                    shard_param, self._sharding_group.id, self._sharding_degree
+                )
             paddle._C_ops.share_var([fused_param, allgather_value])
+
         start_index = target_block.ops.index(last_op) + 1
         return target_block.ops[start_index:]
 
+    def _reduce_scatter_overlap(self, group_grad_list, target_block):
+        '''
+        In order to overlap computation and reduce_scatter communication, we need to:
+          a. place reduce_scatter in communication stream
+          b. place reduce_scatter op and its producer ops after the last grad define op
+        This function will complete the item b.
+        '''
+        insertion_info = {"idx": None, "op": None}
+        # 1. move ops after the grad op
+        for grad in group_grad_list:
+            stack = [grad.get_defining_op()]
+            grad_op = None
+            advance_ops = []
+            # 1.1 get the grad define op
+            while len(stack) > 0:
+                op = stack.pop()
+                if op.op_role == int(OpRole.Backward):
+                    grad_op = op
+                    break
+                if op.num_operands() == 1:  # only one operand
+                    stack.append(op.operand_source(0).get_defining_op())
+                    if op.op_role != int(OpRole.Backward):
+                        advance_ops.append(op)
+                else:
+                    break
+            if grad_op is not None:
+                new_idx = target_block.ops.index(grad_op) + 1
+                # 1.2 move ops
+                for op in advance_ops:
+                    old_idx = target_block.ops.index(op)
+                    if new_idx != old_idx:
+                        target_block.move_op(op, new_idx)
+                # 2.1 get insertion point
+                if (
+                    insertion_info["idx"] is None
+                    or new_idx > insertion_info["idx"]
+                ):
+                    insertion_info["idx"] = new_idx
+                    if len(advance_ops) > 0:
+                        insertion_info["op"] = advance_ops[-1]
+                    else:
+                        insertion_info["op"] = grad_op
+
+        # 2.2 set insertion point
+        if insertion_info["op"] is not None:
+            pir.set_insertion_point_after(insertion_info["op"])
+
     def _fuse_group_param(self, group_param_list):
         startup_program = paddle.static.default_startup_program()
         main_program = paddle.static.default_main_program()
diff --git a/test/auto_parallel/semi_auto_parallel_sharding_stage_1.py b/test/auto_parallel/semi_auto_parallel_sharding_stage_1.py
@@ -106,6 +106,31 @@ def test_sharding_stage_1_to_static(self):
             for batch_id, (image, label) in enumerate(dist_loader()):
                 loss = dist_model(image, label)
 
+    def test_sharding_stage_1_overlap_to_static(self):
+        paddle.distributed.auto_parallel.set_mesh(self._mesh)
+        data_loader = create_data_loader()
+        layer = DemoNet(self._mesh, "sharding_demonet")
+        opt = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=layer.parameters()
+        )
+        opt = dist.shard_optimizer(opt, dist.ShardingStage1(self._mesh))
+        loss_fn = nn.MSELoss()
+
+        dist_loader = dist.shard_dataloader(
+            dataloader=data_loader,
+            meshes=[self._mesh],
+            shard_dims=0,
+        )
+        strategy = dist.Strategy()
+        strategy.sharding.enable = True
+        strategy.sharding.enable_overlap = True
+        dist_model = dist.to_static(layer, dist_loader, loss_fn, opt, strategy)
+
+        dist_model.train()
+        for epoch in range(2):
+            for batch_id, (image, label) in enumerate(dist_loader()):
+                loss = dist_model(image, label)
+
     def run_test_case(self):
         if self._backend == "cpu":
             paddle.set_device("cpu")
@@ -118,6 +143,7 @@ def run_test_case(self):
         self.test_pure_sharding_stage_1()
         self.test_sharding_stage_1_to_static()
         self.test_pure_sharding_multi_mesh_stage_1()
+        self.test_sharding_stage_1_overlap_to_static()
 
 
 if __name__ == '__main__':
