support dynamic shape in pp scheduler (#71123)

LiYuRio · web-flow · commit 70513832254e · 2025-02-14T23:05:26.000+08:00
diff --git a/paddle/fluid/framework/distributed_strategy.proto b/paddle/fluid/framework/distributed_strategy.proto
@@ -89,6 +89,7 @@ message PpConfig {
     optional bool use_batch_p2p_comm = 9 [default = true];
     optional bool best_unbalanced_scheduler = 10 [ default = false ];
     optional bool enable_offload_queue = 11 [ default = false ];
+    optional bool enable_dynamic_shape = 12 [ default = false ];
 }
 
 message DygraphShardingConfig {
diff --git a/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py b/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py
@@ -235,6 +235,14 @@ def __init__(self, layers, hcg, strategy):
         self._use_batch_p2p_comm = self._strategy.hybrid_configs[
             "pp_configs"
         ].use_batch_p2p_comm
+
+        self._dynamic_shape = self._strategy.hybrid_configs[
+            'pp_configs'
+        ].enable_dynamic_shape
+        logger.info(
+            f"Pipeline scheduler is in dynamic_shape mode={self._dynamic_shape}"
+        )
+
         if self._use_batch_p2p_comm and self._overlap_p2p_comm:
             warnings.warn(
                 "non_batch_p2p_comm should be enabled when overlap_p2p_comm is activated, setting non_batch_p2p_comm=True."
@@ -299,7 +307,9 @@ def __init__(self, layers, hcg, strategy):
         )
 
         # construct pipeline meta info
-        self._p2p_helper = p2p.P2pHelper(self._using_cache)
+        self._p2p_helper = p2p.P2pHelper(
+            self._using_cache, dynamic_shape=self._dynamic_shape
+        )
 
         self.global_rank = self._hcg.get_global_rank()
         self.micro_batch_id = 0
@@ -1739,6 +1749,7 @@ def _process_bwd_buffer(step_id, tensor):
                         batch_p2p_comm=self._use_batch_p2p_comm,
                         overlap_p2p_comm=True,
                     )
+
                     self.output_tensor_grads[self.num_model_chunks - 1].append(
                         output_tensor_grad
                     )
@@ -2211,6 +2222,10 @@ def _process_bwd_buffer(step_id, tensor):
         self.processed_steps += 1
         self._check_user_hooks_status_at_step_end()
 
+        # reset dynamic meta counter
+        if self._dynamic_shape:
+            self._p2p_helper._dynamic_cnt = 0
+
         return train_loss
 
     def train_batch(
diff --git a/python/paddle/distributed/fleet/meta_parallel/pp_utils/p2p_communication.py b/python/paddle/distributed/fleet/meta_parallel/pp_utils/p2p_communication.py
@@ -68,8 +68,11 @@ def init_or_erase_meta(self):
         self.has_send_meta = False
         self.has_recv_meta = False
 
-    def recv_meta(self, group):
-        src_rank = _hcg._get_p2p_prev_rank()
+    def recv_meta(self, group, reverse=False):
+        if reverse:
+            src_rank = _hcg._get_p2p_next_rank()
+        else:
+            src_rank = _hcg._get_p2p_prev_rank()
 
         data_numel = paddle.empty([1], dtype="int64")
         paddle.distributed.recv(data_numel, src=src_rank, group=group)
@@ -115,8 +118,11 @@ def recv_meta(self, group):
             self.recv_dtype_message = tuple(dtypes)
             self.recv_stop_gradient = tuple(stop_grads)
 
-    def send_meta(self, tensor, group):
-        dst_rank = _hcg._get_p2p_next_rank()
+    def send_meta(self, tensor, group, reverse=False):
+        if reverse:
+            dst_rank = _hcg._get_p2p_prev_rank()
+        else:
+            dst_rank = _hcg._get_p2p_next_rank()
 
         if isinstance(tensor, ((paddle.Tensor, framework.core.eager.Tensor))):
             tensor_type = 0
@@ -583,6 +589,7 @@ def _p2p_helper(
     send_recv_meta=None,
     batch_p2p_comm=True,
     wait_on_reqs=True,
+    dynamic_shape=False,
 ):
     global _hcg
 
@@ -620,6 +627,9 @@ def _p2p_helper(
             tensor_recv_prev.stop_gradient = recv_stop_gradient
 
     if recv_next:
+        if dynamic_shape:
+            send_shape_msg = send_recv_meta.recv_shape_message
+            send_dtype_msg = send_recv_meta.recv_dtype_message
         if isinstance(send_shape_msg, tuple):
             tensor_recv_next = []
             for idx, shape in enumerate(send_shape_msg):
@@ -653,24 +663,67 @@ def _p2p_helper(
 
 
 class P2pHelper:
-    def __init__(self, use_cache=True):
+    def __init__(self, use_cache=True, dynamic_shape=False):
         self._send_recv_meta = SendRecvMeta()
         self._use_cache = use_cache
-
-    def _send_meta(self, output_tensor, skip_check_meta=False):
-        if not self._send_recv_meta.has_send_meta:
-            self._send_recv_meta.set_send_message(output_tensor)
-            self._send_recv_meta.send_meta(
-                output_tensor, _hcg.get_pipe_parallel_group()
-            )
-            self._send_recv_meta.has_send_meta = self._use_cache
-        elif not skip_check_meta:
-            self._send_recv_meta.check_send_message(output_tensor)
-
-    def _recv_meta(self):
-        if not self._send_recv_meta.has_recv_meta:
-            self._send_recv_meta.recv_meta(_hcg.get_pipe_parallel_group())
-            self._send_recv_meta.has_recv_meta = self._use_cache
+        self._dynamic_shape = dynamic_shape
+
+        if dynamic_shape:
+            self._send_recv_meta_list = []
+            self._dynamic_cnt = 0
+
+    def _send_meta(self, output_tensor, skip_check_meta=False, reverse=False):
+        if not self._dynamic_shape:
+            if not self._send_recv_meta.has_send_meta:
+                self._send_recv_meta.set_send_message(output_tensor)
+                self._send_recv_meta.send_meta(
+                    output_tensor,
+                    _hcg.get_pipe_parallel_group(),
+                    reverse=reverse,
+                )
+                self._send_recv_meta.has_send_meta = self._use_cache
+            elif not skip_check_meta:
+                self._send_recv_meta.check_send_message(output_tensor)
+        else:
+            if len(self._send_recv_meta_list) <= self._dynamic_cnt:
+                meta = SendRecvMeta()
+                meta.set_send_message(output_tensor)
+                meta.send_meta(
+                    output_tensor,
+                    _hcg.get_pipe_parallel_group(),
+                    reverse=reverse,
+                )
+                meta.has_send_meta = self._use_cache
+                self._send_recv_meta_list.append(meta)
+                self._send_recv_meta = meta
+            elif not skip_check_meta:
+                meta = self._send_recv_meta_list[self._dynamic_cnt]
+                meta.check_send_message(output_tensor)
+                self._send_recv_meta = meta
+
+    def _recv_meta(self, reverse=False):
+        if not self._dynamic_shape:
+            if not self._send_recv_meta.has_recv_meta:
+                self._send_recv_meta.recv_meta(
+                    _hcg.get_pipe_parallel_group(), reverse=reverse
+                )
+                self._send_recv_meta.has_recv_meta = self._use_cache
+        else:
+            if len(self._send_recv_meta_list) <= self._dynamic_cnt:
+                meta = SendRecvMeta()
+                meta.recv_meta(_hcg.get_pipe_parallel_group(), reverse=reverse)
+                meta.has_recv_meta = self._use_cache
+                self._send_recv_meta_list.append(meta)
+                self._send_recv_meta = meta
+            elif not self._send_recv_meta_list[self._dynamic_cnt].has_recv_meta:
+                meta = self._send_recv_meta_list[self._dynamic_cnt]
+                meta.recv_meta(_hcg.get_pipe_parallel_group(), reverse=reverse)
+                meta.has_recv_meta = self._use_cache
+                self._send_recv_meta = meta
+            else:
+                self._send_recv_meta = self._send_recv_meta_list[
+                    self._dynamic_cnt
+                ]
 
     def clear_meta_cache(self):
         self._send_recv_meta.init_or_erase_meta()
@@ -693,14 +746,27 @@ def recv_forward(self, pp_first_stage, sync_recv=True, batch_p2p_comm=True):
                 send_recv_meta=self._send_recv_meta,
                 batch_p2p_comm=batch_p2p_comm,
             )
+            if self._dynamic_shape:
+                self._dynamic_cnt += 1
+
         if _timers is not None:
             _timers("recv_forward").stop()
         return input_tensor
 
-    def recv_backward(self, pp_last_stage, sync_recv=True, batch_p2p_comm=True):
+    def recv_backward(
+        self,
+        pp_last_stage,
+        sync_recv=True,
+        batch_p2p_comm=True,
+    ):
         global _timers
         if _timers is not None:
             _timers("recv_backward").start()
+
+        assert (
+            not self._dynamic_shape
+        ), "p2p_helper.recv_backward function doesn't support dynamic_shape now"
+
         if pp_last_stage:
             output_tensor_grad = None
         else:
@@ -713,8 +779,10 @@ def recv_backward(self, pp_last_stage, sync_recv=True, batch_p2p_comm=True):
                 send_recv_meta=self._send_recv_meta,
                 batch_p2p_comm=batch_p2p_comm,
             )
+
         if _timers is not None:
             _timers("recv_backward").stop()
+
         return output_tensor_grad
 
     def send_forward(
@@ -727,9 +795,13 @@ def send_forward(
         global _timers
         if _timers is not None:
             _timers("send_forward").start()
+
+        assert (
+            not self._dynamic_shape
+        ), "p2p_helper.send_forward function doesn't support dynamic_shape now"
+
         if not pp_last_stage:
             self._send_meta(output_tensor, skip_check_meta=skip_check_meta)
-
             _p2p_helper(
                 tensor_send_next=output_tensor,
                 tensor_send_prev=None,
@@ -747,6 +819,11 @@ def send_backward(
         global _timers
         if _timers is not None:
             _timers("send_backward").start()
+
+        assert (
+            not self._dynamic_shape
+        ), "p2p_helper.send_backward function doesn't support dynamic_shape now"
+
         if not pp_first_stage:
             _p2p_helper(
                 tensor_send_next=None,
@@ -765,6 +842,11 @@ def send_forward_recv_backward(
         global _timers
         if _timers is not None:
             _timers("send_forward_recv_backward").start()
+
+        assert (
+            not self._dynamic_shape
+        ), "p2p_helper.send_forward_recv_backward function doesn't support dynamic_shape now"
+
         if pp_last_stage:
             output_tensor_grad = None
         else:
@@ -786,6 +868,11 @@ def send_backward_recv_forward(
         global _timers
         if _timers is not None:
             _timers("send_backward_recv_forward").start()
+
+        assert (
+            not self._dynamic_shape
+        ), "p2p_helper.send_backward_recv_forward function doesn't support dynamic_shape now"
+
         if pp_first_stage:
             input_tensor = None
         else:
@@ -815,6 +902,10 @@ def send_forward_backward_recv_forward_backward(
         if _timers is not None:
             _timers("send_forward_backward_recv_forward_backward").start()
 
+        assert (
+            not self._dynamic_shape
+        ), "p2p_helper.send_forward_backward_recv_forward_backward function doesn't support dynamic_shape now"
+
         if output_tensor is not None:
             self._send_meta(output_tensor, skip_check_meta=skip_check_meta)
         if recv_prev:
@@ -846,11 +937,18 @@ def send_forward_recv_forward(
         if _timers is not None:
             _timers("send_forward_recv_forward").start()
 
+        need_increase_cnt = False
+
         if output_tensor is not None:
-            self._send_meta(output_tensor, skip_check_meta=skip_check_meta)
+            self._send_meta(
+                output_tensor,
+                skip_check_meta=skip_check_meta,
+            )
+            need_increase_cnt = True
 
         if recv_prev:
             self._recv_meta()
+            need_increase_cnt = True
 
         input_tensor, _, wait_handles = _p2p_helper(
             tensor_send_next=output_tensor,
@@ -865,6 +963,9 @@ def send_forward_recv_forward(
         if _timers is not None:
             _timers("send_forward_recv_forward").stop()
 
+        if self._dynamic_shape and need_increase_cnt:
+            self._dynamic_cnt += 1
+
         if overlap_p2p_comm:
             return input_tensor, wait_handles
         return input_tensor
@@ -879,6 +980,17 @@ def send_backward_recv_backward(
         global _timers
         if _timers is not None:
             _timers("send_backward_recv_backward").start()
+
+        if self._dynamic_shape:
+            need_increase_cnt = False
+            if input_tensor_grad is not None:
+                self._send_meta(input_tensor_grad, reverse=True)
+                need_increase_cnt = True
+
+            if recv_next:
+                self._recv_meta(reverse=True)
+                need_increase_cnt = True
+
         _, output_tensor_grad, wait_handles = _p2p_helper(
             tensor_send_next=None,
             tensor_send_prev=input_tensor_grad,
@@ -888,10 +1000,14 @@ def send_backward_recv_backward(
             send_recv_meta=self._send_recv_meta,
             batch_p2p_comm=batch_p2p_comm,
             wait_on_reqs=(not overlap_p2p_comm),
+            dynamic_shape=self._dynamic_shape,
         )
         if _timers is not None:
             _timers("send_backward_recv_backward").stop()
 
+        if self._dynamic_shape and need_increase_cnt:
+            self._dynamic_cnt += 1
+
         if overlap_p2p_comm:
             return output_tensor_grad, wait_handles
         return output_tensor_grad

Original file line number	Diff line number	Diff line change
`@@ -89,6 +89,7 @@ message PpConfig {`
`89`	`89`	`optional bool use_batch_p2p_comm = 9 [default = true];`
`90`	`90`	`optional bool best_unbalanced_scheduler = 10 [ default = false ];`
`91`	`91`	`optional bool enable_offload_queue = 11 [ default = false ];`
	`92`	`+ optional bool enable_dynamic_shape = 12 [ default = false ];`
`92`	`93`	`}`
`93`	`94`
`94`	`95`	`message DygraphShardingConfig {`