【AutoParallelism】Add 'Eager1 F1B" pipeline strategy (#57605)

heavyrain-lzy · web-flow · commit cb3c681e195a · 2023-09-27T17:26:03.000+08:00
* delete sync operation in pipeline strategy

* add Eager1F1B pipeline strategy
diff --git a/paddle/fluid/operators/collective/recv_v2_op.cu.cc b/paddle/fluid/operators/collective/recv_v2_op.cu.cc
@@ -150,7 +150,7 @@ class RecvOpV2CUDAKernel : public framework::OpKernel<T> {
       std::vector<phi::DenseTensor> out_tensor;
       auto out_shape = ctx.Attr<std::vector<int>>("out_shape");
       auto out = ctx.Output<phi::DenseTensor>("Out");
-      auto out_dims = out->dims();
+      // auto out_dims = out->dims();
 
       if (dynamic_shape) {
         VLOG(3) << "recv_v2 will use dynamic shape with send_v2 for switch";
@@ -162,9 +162,9 @@ class RecvOpV2CUDAKernel : public framework::OpKernel<T> {
                             peer,
                             pg);
         out->Resize(new_dim);
-        out->mutable_data<T>(new_dim, place);
+        ctx.cuda_device_context().Alloc<T>(out);
       } else {
-        out->mutable_data<T>(out_dims, place);
+        ctx.cuda_device_context().Alloc<T>(out);
       }
 
       out_tensor.emplace_back(*out);
@@ -230,7 +230,7 @@ class RecvOpV2CUDAKernel : public framework::OpKernel<T> {
         VLOG(3) << "LodTensorArray: idx(" << idx << ")";
         auto out = &out_array->at(idx);
         auto out_dims = out->dims();
-        out->mutable_data<T>(out_dims, place, 0);
+        ctx.cuda_device_context().Alloc<T>(out);
         auto numel = out->numel();
         if (comm_ctx) {
           comm_ctx->Recv(out, numel, peer, stream);
@@ -246,7 +246,7 @@ class RecvOpV2CUDAKernel : public framework::OpKernel<T> {
 
     auto out_shape = ctx.Attr<std::vector<int>>("out_shape");
     auto out = ctx.Output<phi::DenseTensor>("Out");
-    auto out_dims = out->dims();
+    // auto out_dims = out->dims();
     auto numel = out->numel();
 
     if (dynamic_shape) {
@@ -259,9 +259,9 @@ class RecvOpV2CUDAKernel : public framework::OpKernel<T> {
                                                 /* ProcessGroup* */ nullptr);
       out->Resize(new_dim);
       numel = out->numel();
-      out->mutable_data<T>(new_dim, place);
+      ctx.cuda_device_context().Alloc<T>(out);
     } else {
-      out->mutable_data<T>(out_dims, place);
+      ctx.cuda_device_context().Alloc<T>(out);
     }
     if (comm_ctx) {
       comm_ctx->Recv(out, numel, peer, stream);
diff --git a/python/paddle/distributed/auto_parallel/constants.py b/python/paddle/distributed/auto_parallel/constants.py
@@ -111,6 +111,7 @@ def set_field_default_config(category, field, default_value):
 set_field_default_config(PIPELINE, "micro_batch_size", 1)
 set_field_default_config(PIPELINE, "accumulate_steps", 1)
 set_field_default_config(PIPELINE, "generation_batch_size", 1)
+set_field_default_config(PIPELINE, "enable_send_recv_overlap", False)
 
 #########################################
 # quantization configuration
diff --git a/python/paddle/distributed/passes/pass_utils.py b/python/paddle/distributed/passes/pass_utils.py
@@ -444,11 +444,39 @@ def _insert_sync_for_fthenb_1f1b(program):
         block._sync_with_cpp()
 
 
-def _program_for_fthenb_and_1f1b(program):
+def _overlap_send_recv(program):
+    """
+    This function is used to replace the function '_insert_sync_for_fthenb_1f1b'.
+    The finally target of this function is as follows:
+        1. no need to insert the 'c_sync_calc' and 'c_sync_calc' operators
+        2. 'send_v2' operator uses 'dist_attr.execution_stream' to set stream of its own.
+        3. 'recv_v2' opeator uses 'dist_attr.execution_stream' to set stream of its own.
+    """
+    for block in program.blocks:
+        for op in block.ops:
+            if op.type == 'send_v2':
+                op._set_attr("dynamic_shape", False)
+                op._set_attr("use_calc_stream", True)
+                ring_id = op.attr("ring_id")
+                op.dist_attr.execution_stream = "send_stream_" + str(ring_id)
+                op.dist_attr.stream_priority = 0
+            elif op.type == 'recv_v2':
+                op._set_attr("dynamic_shape", False)
+                op._set_attr("use_calc_stream", True)
+                op.dist_attr.execution_stream = "recv_stream"
+                op.dist_attr.stream_priority = 0
+            else:
+                pass
+
+
+def _program_for_fthenb_and_1f1b(program, enable_send_recv_overlap=False):
     """
     This implementation is for fthenb and 1f1b programs and is called in partial_programs function.
     """
-    _insert_sync_for_fthenb_1f1b(program)
+    if enable_send_recv_overlap:
+        _overlap_send_recv(program)
+    else:
+        _insert_sync_for_fthenb_1f1b(program)
 
     lr_prog = Program()
     fwd_prog = Program()
diff --git a/python/paddle/distributed/passes/pipeline_scheduler_pass.py b/python/paddle/distributed/passes/pipeline_scheduler_pass.py
@@ -72,8 +72,12 @@ def _create_job_list(self):
         return job_list
 
     def _partial_programs(self, program):
+        # NOTE: The flag "enable_send_recv_overlap" may increase the reserved memory of GPUs.
+        enable_send_recv_overlap = self.get_attr("enable_send_recv_overlap")
         types = [LR, FORWARD, BACKWARD, OPT]
-        sub_program_list = _program_for_fthenb_and_1f1b(program)
+        sub_program_list = _program_for_fthenb_and_1f1b(
+            program, enable_send_recv_overlap
+        )
         return types, sub_program_list
 
 
@@ -347,8 +351,12 @@ def _op_cost(self, op):
             return 0.0
 
     def _partial_programs(self, program):
+        # NOTE: The flag "enable_send_recv_overlap" may increase the reserved memory of GPUs.
+        enable_send_recv_overlap = self.get_attr("enable_send_recv_overlap")
         types = [LR, FORWARD, BACKWARD, OPT]
-        sub_programs = _program_for_fthenb_and_1f1b(program)
+        sub_programs = _program_for_fthenb_and_1f1b(
+            program, enable_send_recv_overlap
+        )
 
         enable_backward_forward_overlap = self.get_attr(
             "enable_backward_forward_overlap"
@@ -401,11 +409,72 @@ def is_comm_op_valid_to_overlap(self, op):
         )
 
 
+@register_pass("pipeline_scheduler_Eager1F1B")
+class PipelineEager1F1BPass(PipelinePassBase):
+    def __init__(self):
+        super().__init__()
+
+    def _create_job_list(self):
+        num_micro_batches = self.get_attr("num_micro_batches")
+        pp_stage = self.get_attr("pp_stage")
+        pp_degree = self.get_attr("pp_degree")
+
+        job_list = []
+        lr_job = core.Job("lr")
+        job_list.append(lr_job)
+
+        assert (
+            2 * (pp_degree - pp_stage) - 1 <= num_micro_batches
+        ), "Num of micro batches should larger than 2 * (pp_degree - pp_stage) - 1."
+
+        micro_batch_in_warmup = 2 * (pp_degree - pp_stage) - 1
+        micro_batch_in_1f1b = num_micro_batches - micro_batch_in_warmup
+
+        forward_micro_batch_id = 0
+        for i in range(micro_batch_in_warmup):
+            forward_job = core.Job("forward")
+            forward_job.set_micro_batch_id(forward_micro_batch_id)
+            job_list.append(forward_job)
+            forward_micro_batch_id += 1
+
+        backward_micro_batch_id = 0
+        for i in range(micro_batch_in_1f1b):
+            backward_job = core.Job("backward")
+            backward_job.set_micro_batch_id(backward_micro_batch_id)
+            job_list.append(backward_job)
+            backward_micro_batch_id += 1
+            forward_job = core.Job("forward")
+            forward_job.set_micro_batch_id(forward_micro_batch_id)
+            job_list.append(forward_job)
+            forward_micro_batch_id += 1
+
+        for i in range(micro_batch_in_warmup):
+            backward_job = core.Job("backward")
+            backward_job.set_micro_batch_id(backward_micro_batch_id)
+            job_list.append(backward_job)
+            backward_micro_batch_id += 1
+
+        opt_job = core.Job("optimizer")
+        job_list.append(opt_job)
+        return job_list
+
+    def _partial_programs(self, program):
+        # NOTE: The flag "enable_send_recv_overlap" may increase the reserved memory of GPUs.
+        enable_send_recv_overlap = self.get_attr("enable_send_recv_overlap")
+        # TODO: More function will be added later. Now it uses the same logic as FTthenB and 1F1B.
+        types = ["lr", "forward", "backward", "optimizer"]
+        sub_program_list = _program_for_fthenb_and_1f1b(
+            program, enable_send_recv_overlap
+        )
+        return types, sub_program_list
+
+
 def apply_pass(main_program, startup_program, pass_name, pass_attr={}):
     assert pass_name in [
         "FThenB",
         "1F1B",
-    ], f"pipeline scheduler only support FThenB and 1F1B, but recieve {pass_name}"
+        "Eager1F1B",
+    ], f"pipeline scheduler only support FThenB, 1F1B and Eager1F1B, but recieve {pass_name}"
 
     if pass_name == "1F1B":
         # TODO(Ruibiao): Move FLAGS_1f1b_backward_forward_overlap and
diff --git a/test/auto_parallel/pipeline_scheduler_unittest.py b/test/auto_parallel/pipeline_scheduler_unittest.py
@@ -26,15 +26,16 @@
 paddle.enable_static()
 
 
-def apply_pass(schedule_mode="FThenB"):
+def apply_pass(schedule_mode="FThenB", enable_send_recv_overlap=False):
     strategy = auto.Strategy()
     strategy.auto_mode = "semi"
     strategy.reinit = True
 
     pipeline = strategy.pipeline
     pipeline.enable = True
     pipeline.schedule_mode = schedule_mode
-    pipeline.accumulate_steps = 2
+    pipeline.accumulate_steps = 4
+    pipeline.enable_send_recv_overlap = enable_send_recv_overlap
 
     return strategy
 
@@ -48,7 +49,7 @@ class Test1F1BPass(unittest.TestCase):
     def setUp(self):
         self.rtol = 1e-5
         self.atol = 1e-8
-        self.batch_size = 2
+        self.batch_size = 4
         self.batch_num = 10
         self.clip_norm = 0.2
         self.dataset = FakeDataset(self.batch_size * self.batch_num)
@@ -61,10 +62,12 @@ def init(self, engine):
         place = paddle.base.CUDAPlace(ParallelEnv().dev_id)
         engine._executor = paddle.static.Executor(place)
 
-    def get_engine(self, schedule_mode="FThenB"):
+    def get_engine(
+        self, schedule_mode="FThenB", enable_send_recv_overlap=False
+    ):
         reset_prog()
 
-        strategy = apply_pass(schedule_mode)
+        strategy = apply_pass(schedule_mode, enable_send_recv_overlap)
         clip = paddle.nn.ClipGradByGlobalNorm(self.clip_norm)
         opt = paddle.optimizer.AdamW(learning_rate=0.00001, grad_clip=clip)
         model, loss = generate_model("pp")
@@ -112,19 +115,50 @@ def test_pp_pass(self):
         assert engine_1f1b._strategy.pipeline.schedule_mode == "1F1B"
         assert os.environ.get('FLAGS_new_executor_micro_batching') == "True"
 
+        # pp2 eager1f1b training with standalone executor
+        os.environ['FLAGS_new_executor_micro_batching'] = 'True'
+        engine_eager1f1b = self.get_engine(schedule_mode="Eager1F1B")
+        history_eager1f1b = engine_eager1f1b.fit(
+            self.dataset, 3, batch_size=self.batch_size, log_freq=1
+        )
+        assert engine_eager1f1b._strategy.pipeline.schedule_mode == "Eager1F1B"
+        assert os.environ.get('FLAGS_new_executor_micro_batching') == "True"
+
+        # pp2 1f1b training with standalone executor
+        os.environ['FLAGS_new_executor_micro_batching'] = 'True'
+        engine_1f1b_overlap = self.get_engine(
+            schedule_mode="1F1B", enable_send_recv_overlap=True
+        )
+        history_1f1b_overlap = engine_1f1b_overlap.fit(
+            self.dataset, 3, batch_size=self.batch_size, log_freq=1
+        )
+        assert engine_1f1b_overlap._strategy.pipeline.schedule_mode == "1F1B"
+        assert (
+            engine_1f1b_overlap._strategy.pipeline.enable_send_recv_overlap
+            is True
+        )
+        assert os.environ.get('FLAGS_new_executor_micro_batching') == "True"
+
         # NOTE: every sample data from dataset is all the same
         if paddle.distributed.get_rank() == 1:
             losses_fleet_1f1b = np.array(history_fleet_1f1b.history["loss"])
             losses_fthenb = np.array(history_fthenb.history["loss"])
             losses_1f1b = np.array(history_1f1b.history["loss"])
-            # accumulate_steps is 2
-            assert losses_fthenb[0].shape[0] == 2
-            assert losses_1f1b[0].shape[0] == 2
+            losses_eager1f1b = np.array(history_eager1f1b.history["loss"])
+            losses_1f1b_overlap = np.array(history_1f1b_overlap.history["loss"])
+            # accumulate_steps is 4
+            assert losses_fthenb[0].shape[0] == 4
+            assert losses_1f1b[0].shape[0] == 4
+            assert losses_eager1f1b[0].shape[0] == 4
+            assert losses_1f1b_overlap[0].shape[0] == 4
             # losses_fleet_1f1b is the last loss of accumulate_steps
             # losses_fthenb is all the losses of accumulate_steps
             # losses_1f1b is alla the losses of accumulate_steps
             self.check_results(losses_fleet_1f1b[0], losses_fthenb[0][-1])
             self.check_results(losses_fleet_1f1b[0], losses_1f1b[0][-1])
+            self.check_results(losses_fleet_1f1b[0], losses_eager1f1b[0][-1])
+            # self.check_results(losses_fleet_1f1b[0], losses_1f1b_overlap[0][-1])
+            self.check_results(losses_fleet_1f1b[0], losses_1f1b_overlap[0][-1])
 
 
 if __name__ == "__main__":
