[Distributed]Earse p2p cache for every step (#62277)

ForFishes · web-flow · commit 7aa70c1d182f · 2024-03-02T21:55:46.000+08:00
* earse cache

* earse cache

* earse cache
diff --git a/paddle/fluid/framework/distributed_strategy.proto b/paddle/fluid/framework/distributed_strategy.proto
@@ -66,6 +66,7 @@ message PpConfig {
     optional bool sharding_comm_overlap = 4 [ default = false ];
     optional bool release_gradients = 5 [ default = false ];
     optional bool overlap_p2p_comm = 6 [default = true];
+    optional bool clear_every_step_cache = 7 [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
@@ -201,6 +201,10 @@ def __init__(self, layers, hcg, strategy):
             "pp_configs"
         ].overlap_p2p_comm
 
+        self._clear_every_step_cache = self._strategy.hybrid_configs[
+            "pp_configs"
+        ].clear_every_step_cache
+
         self._batch_p2p_comm = not self._overlap_p2p_comm
 
         if self._dp_comm_overlap:
@@ -476,6 +480,10 @@ def forward_backward_pipeline(self, data, scaler=None):
             train_loss = self._broadcast_final_loss()
         if self._enable_timer:
             self.timers("broadcast_final_loss").stop()
+
+        if self._clear_every_step_cache:
+            self._p2p_helper.clear_meta_cache()
+
         self.timer_printer()
         return train_loss
 
@@ -1394,6 +1402,9 @@ def _process_bwd_buffer(step_id, tensor):
             # else just return all intermediate output tensor for all micro steps
             train_loss = self.output_tensors
 
+        if self._clear_every_step_cache:
+            self._p2p_helper.clear_meta_cache()
+
         self.timer_printer()
         return train_loss
 
@@ -1634,5 +1645,8 @@ def forward_backward_pipeline(
             # else just return all intermediate output tensor for all micro steps
             train_loss = self.output_tensors
 
+        if self._clear_every_step_cache:
+            self._p2p_helper.clear_meta_cache()
+
         self.timer_printer()
         return train_loss
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
@@ -50,6 +50,9 @@ class SendRecvMeta:
     """Mainly used to help p2p communication context information"""
 
     def __init__(self):
+        self.init_or_erase_meta()
+
+    def init_or_erase_meta(self):
         self.send_shape_message = None
         self.send_dtype_message = None
 
@@ -604,6 +607,9 @@ def _recv_meta(self):
             self._send_recv_meta.recv_meta(_hcg.get_pipe_parallel_group())
             self._send_recv_meta.has_recv_meta = self._use_cache
 
+    def clear_meta_cache(self):
+        self._send_recv_meta.init_or_erase_meta()
+
     def recv_forward(self, pp_first_stage, sync_recv=True, batch_p2p_comm=True):
         global _timers
         if _timers is not None:

Original file line number	Diff line number	Diff line change
`@@ -66,6 +66,7 @@ message PpConfig {`
`66`	`66`	`optional bool sharding_comm_overlap = 4 [ default = false ];`
`67`	`67`	`optional bool release_gradients = 5 [ default = false ];`
`68`	`68`	`optional bool overlap_p2p_comm = 6 [default = true];`
	`69`	`+ optional bool clear_every_step_cache = 7 [default = false];`
`69`	`70`	`}`
`70`	`71`
`71`	`72`	`message DygraphShardingConfig {`