release input of pipeline in FthenB (#73398) (#73508)

Author: ForFishes

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

@@ -1680,7 +1680,11 @@ def _store_forward_outputs(
             self.output_tensors[virtual_pp_rank].pop()
 
     def _forward_step_helper(
-        self, micro_dataset, micro_step, overlap_schedule_mode=False
+        self,
+        micro_dataset,
+        micro_step,
+        overlap_schedule_mode=False,
+        release_input=False,
     ):
         virtual_pp_rank = self._get_virtual_pp_rank(micro_step, forward=True)
         self.set_virtual_pipeline_rank(virtual_pp_rank)
@@ -1698,6 +1702,10 @@ def _forward_step_helper(
         self._store_forward_outputs(
             virtual_pp_rank, output_tensor, schedule_chunk, loss_fn_node
         )
+
+        if release_input:
+            return output_tensor, input_tensor
+
         return output_tensor
 
     def _overlap_comm_grads(self):
@@ -3202,7 +3210,9 @@ def forward_backward_pipeline(
         # to simplify the code logic of stage 1F1B.
         for micro_step in range(startup_steps):
             self._record_stamp("F", micro_step, '"B"', forward=True)
-            output_tensor = self._forward_step_helper(micro_dataset, micro_step)
+            output_tensor, input_t = self._forward_step_helper(
+                micro_dataset, micro_step, release_input=True
+            )
             self._record_stamp("F", micro_step, '"E"', forward=True)
             next_forward_virtual_pp_rank = self._get_virtual_pp_rank(
                 micro_step + 1, forward=True
@@ -3233,6 +3243,7 @@ def forward_backward_pipeline(
                 input_tensor
             )
             self._release_output(output_tensor)
+            self._release_output(input_t)
 
         if self.is_pipeline_first_stage(ignore_virtual=True):
             assert (
@@ -3247,8 +3258,8 @@ def forward_backward_pipeline(
             backward_micro_step_id = micro_step
 
             self._record_stamp("F", forward_micro_step_id, '"B"', forward=True)
-            output_tensor = self._forward_step_helper(
-                micro_dataset, forward_micro_step_id
+            output_tensor, input_t = self._forward_step_helper(
+                micro_dataset, forward_micro_step_id, release_input=True
             )
             self._record_stamp("F", forward_micro_step_id, '"E"', forward=True)
 
@@ -3259,6 +3270,10 @@ def forward_backward_pipeline(
                 self.is_pipeline_last_stage(ignore_virtual=True),
                 batch_p2p_comm=self._use_batch_p2p_comm,
             )
+
+            self._release_output(output_tensor)
+            self._release_output(input_t)
+
             output_tensor_grad = self._p2p_helper.recv_backward(
                 self.is_pipeline_last_stage(ignore_virtual=True),
                 batch_p2p_comm=self._use_batch_p2p_comm,