[Distributed] Support forward backward overlap schdule for VPP (#71995)

* [Distributed] Support forward_backward_overlap mode for VPP

* add

* fix name

Author: zhangyuqin1998

paddle/fluid/framework/distributed_strategy.proto

@@ -92,6 +92,7 @@ message PpConfig {
     optional bool enable_offload_queue = 11 [ default = false ];
     optional bool enable_dynamic_shape = 12 [ default = false ];
     optional bool use_dualpipev = 13 [ default = false ];
+    optional bool forward_backward_overlap_scheduler = 14 [ default = false ];
 }
 
 message DygraphShardingConfig {

python/paddle/distributed/fleet/meta_parallel/dualpipev.py

@@ -150,9 +150,7 @@ def _forward_compute(self, phase: int, micro_datasets=None) -> None:
         inputs = self._get_forward_inputs(micro_datasets, phase, acc_id)
 
         if self.overlapped_forward_backward:
-            schedule_chunk = self._layers.forward(
-                inputs, chunk_id=phase, overlap_schedule_mode=True
-            )
+            schedule_chunk = self._layers.get_schedule_chunk(chunk_id=phase)
             outputs = schedule_chunk.forward(inputs)
         else:
             schedule_chunk = None
@@ -300,9 +298,7 @@ def _forward_backward_compute(
             )
 
         # forward & backward
-        forward_chunk = self._layers.forward(
-            None, chunk_id=forward_phase, overlap_schedule_mode=True
-        )
+        forward_chunk = self._layers.get_schedule_chunk(chunk_id=forward_phase)
         backward_chunk = self.schedule_chunks[backward_phase][backward_acc_id]
         forward_outputs, forward_loss, backward_input_grads = (
             self._layers.overlapped_forward_backward(

python/paddle/distributed/fleet/meta_parallel/parallel_layers/pp_layers.py

@@ -1017,7 +1017,7 @@ def execute_func(*x):
 
         return execute_func
 
-    def forward(self, input, chunk_id=None, overlap_schedule_mode=False):
+    def update_run_function(self, chunk_id):
         if chunk_id is not None:
             assert isinstance(chunk_id, int), "chunk_id should be an int"
             assert (
@@ -1035,9 +1035,15 @@ def forward(self, input, chunk_id=None, overlap_schedule_mode=False):
             # But for interleave, self.run_function will keep updating to the target functions at every run.
             self.run_function = model_chunk.get_run_function()
 
-        if overlap_schedule_mode:
-            assert self._recompute_interval == 0
-            return self.build_schedule_nodes(0, len(self.run_function))
+    def get_schedule_chunk(self, chunk_id):
+        self.update_run_function(chunk_id)
+
+        assert self._recompute_interval == 0
+        return self.build_schedule_nodes(0, len(self.run_function))
+
+    def forward(self, input, chunk_id=None):
+        self.update_run_function(chunk_id)
+
         if self._recompute_interval == 0:
             input = self.forward_function(0, len(self.run_function))(input)
         else: