fix sft/ulysses eval (#4494)

Co-authored-by: tastelikefeet <[email protected]>

Author: tastelikefeet

swift/trainers/sequence_parallel/ulysses.py

@@ -118,8 +118,9 @@ def loss_scale_sp_func(outputs, labels, loss_scale=None, num_items_in_batch=None
     else:
         logits = outputs
     device = logits.device
+    if labels.shape[1] > logits.shape[1]:
+        _, _, labels, _, _, loss_scale = ulysses.pad_and_split_inputs(None, None, labels, None, None, loss_scale)
     logits = logits.view(-1, logits.shape[-1])
-    _, _, labels, _, _, loss_scale = ulysses.pad_and_split_inputs(None, None, labels, None, None, loss_scale)
 
     labels = labels.flatten().to(device)
     sploss_parallel_size = int(os.environ.get('CELOSS_PARALLEL_SIZE', '0'))
@@ -142,7 +143,7 @@ def loss_scale_sp_func(outputs, labels, loss_scale=None, num_items_in_batch=None
 
 
 @profiling_decorator
-def _prepare_inputs(self, generation_batch):
+def _prepare_inputs_grpo(self, generation_batch):
     ulysses = self.ulysses
     mode = 'train' if self.model.training else 'eval'
     if mode == 'train':
@@ -159,6 +160,14 @@ def _prepare_inputs(self, generation_batch):
     return inputs
 
 
+def _prepare_inputs(self, inputs, ulysses):
+    if 'labels' in inputs:
+        labels = inputs['labels']
+        _, _, labels, _, _, _ = ulysses.pad_and_split_inputs(None, None, labels, None, None, None)
+        inputs['labels'] = labels
+    return self._origin_prepare_inputs(inputs)
+
+
 def old_policy(self):
     ulysses = self.ulysses
     # changes: `* ulysses.sp_world_size`
@@ -171,7 +180,8 @@ def get_per_token_logps(self,
                         logits: torch.FloatTensor,
                         labels: torch.LongTensor,
                         ulysses=None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    _, _, labels, _, _, _ = ulysses.pad_and_split_inputs(None, None, labels, None, None, None)
+    if labels.shape[1] > logits.shape[1]:
+        _, _, labels, _, _, _ = ulysses.pad_and_split_inputs(None, None, labels, None, None, None)
     loss_mask = labels != self.label_pad_token_id
     labels = labels.clone()  # No need to shift, pad and split has shifted the inputs.
     labels[~loss_mask] = 0
@@ -823,9 +833,13 @@ def prepare_trainer(self, trainer):
 
         trainer.ulysses = self
         if trainer.__class__.__name__ == 'Seq2SeqTrainer':
+            trainer._origin_prepare_inputs = trainer._prepare_inputs
+            trainer._prepare_inputs = MethodType(partial(_prepare_inputs, ulysses=self), trainer)
             trainer.compute_loss_func = partial(loss_scale_sp_func, ulysses=self)
 
         elif trainer.__class__.__name__ == 'DPOTrainer':
+            trainer._origin_prepare_inputs = trainer._prepare_inputs
+            trainer._prepare_inputs = MethodType(partial(_prepare_inputs, ulysses=self), trainer)
             trainer.get_per_token_logps = MethodType(partial(get_per_token_logps, ulysses=self), trainer)
 
             def rlhf_loss_scale_sp_func(_, *args, **kwargs):
@@ -838,7 +852,7 @@ def rlhf_loss_scale_sp_func(_, *args, **kwargs):
             trainer.ulysses = self
             trainer.args.gradient_accumulation_steps = trainer.args.gradient_accumulation_steps * self.sp_world_size
             trainer.old_policy = MethodType(old_policy, trainer)
-            trainer._prepare_inputs = MethodType(_prepare_inputs, trainer)
+            trainer._prepare_inputs = MethodType(_prepare_inputs_grpo, trainer)
             trainer._get_per_token_logps = MethodType(_get_per_token_logps, trainer)
             trainer.split_by_mini_batches = MethodType(split_by_mini_batches, trainer)
 
@@ -852,7 +866,8 @@ def compute_acc(preds, labels, *args, **kwargs) -> Dict[str, List[float]]:
                 preds = torch.from_numpy(preds).to(get_current_device())
             if isinstance(labels, np.ndarray):
                 labels = torch.from_numpy(labels).to(get_current_device())
-            _, _, labels, _, _, _ = self.pad_and_split_inputs(None, None, labels, None, None, None)
+            if labels.shape[1] > preds.shape[1]:
+                _, _, labels, _, _, _ = self.pad_and_split_inputs(None, None, labels, None, None, None)
             shape0 = preds.shape[0]
             preds_output = torch.empty((shape0 * self.sp_world_size, preds.shape[1]),
                                        dtype=preds.dtype,

swift/trainers/trainers.py

@@ -229,7 +229,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
         if getattr(self.args, 'average_tokens_across_devices', False) and self.model_accepts_loss_kwargs:
             loss *= self.accelerator.num_processes
 
-        if outputs.logits is not None and labels is not None:
+        if outputs.logits is not None and labels is not None and not return_outputs:
             # Liger does not have logits
             self._compute_acc(outputs, labels)
         return (loss, outputs) if return_outputs else loss