support skip_check_meta in eval mode of Pipeline (#63001)

haohongxiang · web-flow · commit d648bc7442dd · 2024-03-26T09:35:15.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
@@ -727,7 +727,9 @@ def eval_batch(self, data, compute_loss=False):
 
             output_tensor = self._forward_step(input_tensor, micro_dataset)
             self._p2p_helper.send_forward(
-                output_tensor, self.is_pipeline_last_stage()
+                output_tensor,
+                self.is_pipeline_last_stage(),
+                skip_check_meta=True,
             )
 
             input_buffers.append(input_tensor)
@@ -743,7 +745,9 @@ def eval_batch(self, data, compute_loss=False):
 
             output_tensor = self._forward_step(input_tensor, micro_dataset)
             self._p2p_helper.send_forward(
-                output_tensor, self.is_pipeline_last_stage()
+                output_tensor,
+                self.is_pipeline_last_stage(),
+                skip_check_meta=True,
             )
 
             input_buffers.append(input_tensor)
diff --git a/python/paddle/distributed/fleet/meta_parallel/pp_utils/four_directions_p2p_communication.py b/python/paddle/distributed/fleet/meta_parallel/pp_utils/four_directions_p2p_communication.py
@@ -692,7 +692,7 @@ def __init__(self, use_cache=True):
         self._send_recv_meta = SendRecvMeta()
         self._use_cache = use_cache
 
-    def _send_meta(self, output_tensor):
+    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(
@@ -745,12 +745,12 @@ def recv_backward(self, pp_last_stage, sync_recv=True):
             _timers("recv_backward").stop()
         return output_tensor_grad
 
-    def send_forward(self, output_tensor, pp_last_stage):
+    def send_forward(self, output_tensor, pp_last_stage, skip_check_meta=False):
         global _timers
         if _timers is not None:
             _timers("send_forward").start()
         if not pp_last_stage:
-            self._send_meta(output_tensor)
+            self._send_meta(output_tensor, skip_check_meta=skip_check_meta)
 
             _p2p_helper(
                 tensor_send_next=output_tensor,
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
@@ -649,14 +649,14 @@ def __init__(self, use_cache=True):
         self._send_recv_meta = SendRecvMeta()
         self._use_cache = use_cache
 
-    def _send_meta(self, output_tensor):
+    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
-        else:
+        elif not skip_check_meta:
             self._send_recv_meta.check_send_message(output_tensor)
 
     def _recv_meta(self):
@@ -709,12 +709,18 @@ def recv_backward(self, pp_last_stage, sync_recv=True, batch_p2p_comm=True):
             _timers("recv_backward").stop()
         return output_tensor_grad
 
-    def send_forward(self, output_tensor, pp_last_stage, batch_p2p_comm=True):
+    def send_forward(
+        self,
+        output_tensor,
+        pp_last_stage,
+        batch_p2p_comm=True,
+        skip_check_meta=False,
+    ):
         global _timers
         if _timers is not None:
             _timers("send_forward").start()
         if not pp_last_stage:
-            self._send_meta(output_tensor)
+            self._send_meta(output_tensor, skip_check_meta=skip_check_meta)
 
             _p2p_helper(
                 tensor_send_next=output_tensor,
