fix eval batch & non-compute_loss in pipeline

Author: SylarTiaNII

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

@@ -165,7 +165,7 @@ def _check_data_valid(self, data):
         batch_size = data.shape[0]
         assert self._micro_batch_size * self._acc_steps == batch_size, (
             "batch_size needs to be divisible by micro_batch_size. Currently, "
-            f"batch_size = {batch_size}, micro_batch_size = {self._micro_batch_size}, accumulate_steps = {self._acc_steps}."
+            f"batch_size = {batch_size}, micro_batch_size = {self._micro_batch_size}, accumulate_steps = {self._acc_steps} data_shape= {data.shape}."
         )
 
 
@@ -413,6 +413,7 @@ def __init__(self, layers, hcg, strategy):
         self.loss_fn_idx = 0
 
         self._compute_loss = True
+        self._return_host_tensor = False
         self.callbacks = pipeline_parallel_callbacks_
 
         logger.info(
@@ -991,13 +992,16 @@ def train_batch(
 
         return train_loss
 
-    def eval_batch(self, data, compute_loss=False, loss_fn_idx=0):
+    def eval_batch(self, data, compute_loss=False, loss_fn_idx=0, return_host_tensor=False):
         self.user_hooks_enabled = False
         # reset the virtual pp rank for each run
         self.set_virtual_pipeline_rank(0)
 
         self._layers.eval()
+        origin_compute_loss = self._compute_loss
         self._compute_loss = compute_loss
+        origin_return_host_tensor = self._return_host_tensor
+        self._return_host_tensor = return_host_tensor
 
         # store data id for micro_batch
         self.micro_batch_id = 0
@@ -1070,11 +1074,13 @@ def eval_batch(self, data, compute_loss=False, loss_fn_idx=0):
                 )
 
         if self._compute_loss:
-            self.train_loss = self._broadcast_final_loss()
+            train_loss = self._broadcast_final_loss()
         else:
-            self.train_loss = output_buffers
+            train_loss = output_buffers
 
-        return self.train_loss
+        self._compute_loss = origin_compute_loss
+        self._return_host_tensor = origin_return_host_tensor
+        return train_loss
 
     def _maybe_loss_compute(
         self, output_tensor, micro_dataset, overlap_schedule_mode=False
@@ -1384,6 +1390,19 @@ def _optimizer_step(self):
         if self.lr_scheduler:
             self.lr_scheduler.step()
 
+    def mark_release_tensors(self, output_tensor, can_release=True):
+        if isinstance(output_tensor, (tuple, list)):
+            for t in output:
+                setattr(t, "can_release", can_release)
+                if self._return_host_tensor:
+                    host_tensor = t.pin_memory() if hasattr(t, "pin_memory") else t.cpu()
+                    host_tensor._share_buffer_to(t)
+        else:
+            setattr(output_tensor, "can_release", can_release)
+            if self._return_host_tensor:
+                host_tensor = output_tensor.pin_memory() if hasattr(output_tensor, "pin_memory") else output_tensor.cpu()
+                host_tensor._share_buffer_to(output_tensor)
+
     def _release_output(self, output):
         def can_free(t):
             return (
@@ -1655,10 +1674,12 @@ def _get_forward_input(self, virtual_pp_rank):
         assert hasattr(self, 'output_tensors')
         if not self._forward_only:
             assert hasattr(self, 'output_tensor_grads')
-        assert len(self.input_tensors[virtual_pp_rank]) == (
-            len(self.output_tensors[virtual_pp_rank]) + 1
-        )
-        input_tensor = self.input_tensors[virtual_pp_rank][-1]
+            assert len(self.input_tensors[virtual_pp_rank]) == (
+                len(self.output_tensors[virtual_pp_rank]) + 1
+            )
+            input_tensor = self.input_tensors[virtual_pp_rank][-1]
+        else:
+            input_tensor = self.input_tensors[virtual_pp_rank].pop()
         return input_tensor
 
     def _store_forward_outputs(
@@ -1673,11 +1694,9 @@ def _store_forward_outputs(
         self.schedule_chunks[virtual_pp_rank].append(schedule_chunk)
         if self.is_pipeline_last_stage():
             self.loss_fn_chunks.append(loss_fn_node)
-
-        if self._forward_only:
-            # no need to store tensor for backward
-            self.input_tensors[virtual_pp_rank].pop()
-            self.output_tensors[virtual_pp_rank].pop()
+            # save output_tensors for return value of eval batch
+            if not self._compute_loss:
+                self.mark_release_tensors(output_tensor, False)
 
     def _forward_step_helper(
         self,
@@ -1981,7 +2000,7 @@ def forward_backward_pipeline(
         # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/schedules.py
         if not compute_loss:
             assert (
-                not forward_only
+                forward_only
             ), "compute_loss can only be set to False when forward_only is set to True"
 
         if static_scheduler:
@@ -2758,12 +2777,12 @@ def backward_async_comm(
             if self._enable_timer:
                 self.timers("broadcast_final_loss").start()
             with paddle.amp.auto_cast(enable=False):
-                train_loss = self._broadcast_final_loss(return_micro_batch_loss)
+                train_loss_or_logits = self._broadcast_final_loss(return_micro_batch_loss)
             if self._enable_timer:
                 self.timers("broadcast_final_loss").stop()
         else:
-            # else just return all intermediate output tensor for all micro steps
-            train_loss = self.output_tensors
+            # else just return logits without loss func calc
+            train_loss_or_logits = self.output_tensors.pop()
 
         if self._clear_every_step_cache:
             self._p2p_helper.clear_meta_cache()
@@ -2781,7 +2800,7 @@ def backward_async_comm(
             ), "p2p dynamic_cnt should equal to send_recv_meta_list"
             self._p2p_helper._dynamic_cnt = 0
 
-        return train_loss
+        return train_loss_or_logits
 
     def train_batch(
         self,
@@ -2812,13 +2831,16 @@ def train_batch(
 
         return train_loss
 
-    def eval_batch(self, data, compute_loss=False, loss_fn_idx=0):
+    def eval_batch(self, data, compute_loss=False, loss_fn_idx=0, return_host_tensor=False):
         self.user_hooks_enabled = False
         # reset the virtual pp rank for each run
         self.set_virtual_pipeline_rank(0)
 
         self._layers.eval()
+        origin_compute_loss = self._compute_loss
         self._compute_loss = compute_loss
+        origin_return_host_tensor = self._return_host_tensor
+        self._return_host_tensor = return_host_tensor
 
         # check loss_fn_idx is valid and loss_fn exists
         assert (
@@ -2827,7 +2849,11 @@ def eval_batch(self, data, compute_loss=False, loss_fn_idx=0):
         ), f"loss function {loss_fn_idx} should exist to compute loss"
         self.loss_fn_idx = loss_fn_idx
 
-        return self.forward_backward_pipeline(data, None, forward_only=True)
+        train_loss_or_logits = self.forward_backward_pipeline(data, None, forward_only=True, compute_loss=compute_loss)
+        self._init_buffers()
+        self._compute_loss = origin_compute_loss
+        self._return_host_tensor = origin_return_host_tensor
+        return train_loss_or_logits
 
     def get_static_scheduler(self):
         return self.forward_backward_pipeline(
@@ -2918,7 +2944,7 @@ def forward_backward_pipeline(
             get_sync_logger().info("start forward_backward_pipeline")
         if not compute_loss:
             assert (
-                not forward_only
+                forward_only
             ), "compute_loss can only be set to False when forward_only is set to True"
 
         # NOTE(shenliang03): Due to ring_exchange for pipeline with interleave, cache should be enabled
@@ -3065,12 +3091,12 @@ def forward_backward_pipeline(
             if self._enable_timer:
                 self.timers("broadcast_final_loss").start()
             with paddle.amp.auto_cast(enable=False):
-                train_loss = self._broadcast_final_loss(return_micro_batch_loss)
+                train_loss_or_logits = self._broadcast_final_loss(return_micro_batch_loss)
             if self._enable_timer:
                 self.timers("broadcast_final_loss").stop()
         else:
-            # else just return all intermediate output tensor for all micro steps
-            train_loss = self.output_tensors
+            # else just return logits without loss func calc
+            train_loss_or_logits = self.output_tensors.pop()
 
         if self._clear_every_step_cache:
             self._p2p_helper.clear_meta_cache()
@@ -3081,7 +3107,7 @@ def forward_backward_pipeline(
             get_sync_logger().info("end forward_backward_pipeline")
         self.processed_steps += 1
         self._check_user_hooks_status_at_step_end()
-        return train_loss
+        return train_loss_or_logits
 
 
 class OffloadQueue(queue.Queue):
@@ -3147,7 +3173,7 @@ def forward_backward_pipeline(
         self._reset_user_hooks_status()
         if not compute_loss:
             assert (
-                not forward_only
+                forward_only
             ), "compute_loss can only be set to False when forward_only is set to True"
         assert (
             self._using_cache
@@ -3406,12 +3432,12 @@ def forward_backward_pipeline(
             if self._enable_timer:
                 self.timers("broadcast_final_loss").start()
             with paddle.amp.auto_cast(enable=False):
-                train_loss = self._broadcast_final_loss(return_micro_batch_loss)
+                train_loss_or_logits = self._broadcast_final_loss(return_micro_batch_loss)
             if self._enable_timer:
                 self.timers("broadcast_final_loss").stop()
         else:
-            # else just return all intermediate output tensor for all micro steps
-            train_loss = self.output_tensors
+            # else just return logits without loss func calc
+            train_loss_or_logits = self.output_tensors.pop()
 
         if self._clear_every_step_cache:
             self._p2p_helper.clear_meta_cache()
@@ -3422,4 +3448,4 @@ def forward_backward_pipeline(
             get_sync_logger().info("end forward_backward_pipeline")
         self.processed_steps += 1
         self._check_user_hooks_status_at_step_end()
-        return train_loss
+        return train_loss_or_logits