dynamic_shape support balanced vpp (#72356)

LiYuRio · web-flow · commit a896bf7849f4 · 2025-04-22T10:42:41.000+08:00
diff --git a/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py b/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py
@@ -2760,6 +2760,9 @@ def backward_async_comm(
 
         # reset dynamic meta counter
         if self._dynamic_shape:
+            assert self._p2p_helper._dynamic_cnt == len(
+                self._p2p_helper._send_recv_meta_list
+            ), "p2p dynamic_cnt should equal to send_recv_meta_list"
             self._p2p_helper._dynamic_cnt = 0
 
         return train_loss
@@ -3359,6 +3362,13 @@ def forward_backward_pipeline(
             backward_send_recv_buffer_queue.empty()
         ), "send_recv buffer should be empty"
 
+        # reset dynamic meta counter
+        if self._dynamic_shape:
+            assert self._p2p_helper._dynamic_cnt == len(
+                self._p2p_helper._send_recv_meta_list
+            ), "p2p dynamic_cnt should equal to send_recv_meta_list"
+            self._p2p_helper._dynamic_cnt = 0
+
         self._flush_records()
         self._sync_overlap_grads()
 
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
@@ -763,6 +763,7 @@ def recv_forward(self, pp_first_stage, sync_recv=True, batch_p2p_comm=True):
                 sync_recv=sync_recv,
                 send_recv_meta=self._send_recv_meta,
                 batch_p2p_comm=batch_p2p_comm,
+                dynamic_shape=self._dynamic_shape,
             )
             if self._dynamic_shape:
                 self._dynamic_cnt += 1
@@ -807,9 +808,6 @@ def recv_backward(
         if _timers is not None:
             _timers("recv_backward").stop()
 
-        if self._dynamic_shape and need_increase_cnt:
-            self._dynamic_cnt += 1
-
         return output_tensor_grad
 
     def send_forward(
@@ -823,10 +821,6 @@ def send_forward(
         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(
@@ -836,7 +830,11 @@ def send_forward(
                 recv_next=False,
                 send_recv_meta=self._send_recv_meta,
                 batch_p2p_comm=batch_p2p_comm,
+                dynamic_shape=self._dynamic_shape,
             )
+            if self._dynamic_shape:
+                self._dynamic_cnt += 1
+
         if _timers is not None:
             _timers("send_forward").stop()
 
@@ -847,19 +845,20 @@ def send_backward(
         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:
+            if self._dynamic_shape:
+                self._send_meta(input_tensor_grad, reverse=True)
             _p2p_helper(
                 tensor_send_next=None,
                 tensor_send_prev=input_tensor_grad,
                 recv_prev=False,
                 recv_next=False,
                 send_recv_meta=self._send_recv_meta,
                 batch_p2p_comm=batch_p2p_comm,
+                dynamic_shape=self._dynamic_shape,
             )
+            if self._dynamic_shape:
+                self._dynamic_cnt += 1
         if _timers is not None:
             _timers("send_backward").stop()
 
