fix(trainer): support selective logits with sequence parallel

Author: zjn20030811

swift/trainers/mixin.py

@@ -67,6 +67,9 @@
 
 class SwiftMixin:
     FLASH_CKPT_WAIT_TIMEOUT = 1800
+    # Selective logits with SP requires a trainer-specific loss reducer.  SFT
+    # opts in below; RLHF trainers keep their established full-frame path.
+    SUPPORTS_SP_LOGITS_TO_KEEP = False
 
     def __init__(self,
                  model: PreTrainedModel,
@@ -220,6 +223,10 @@ def get_use_logits_to_keep(self, default_value: bool = True):
                 use_logits_to_keep = default_value
             self.args.use_logits_to_keep = use_logits_to_keep
             logger.info_once(f'use_logits_to_keep: {use_logits_to_keep}')
+        if (use_logits_to_keep and self.template.sequence_parallel_size > 1 and not self.SUPPORTS_SP_LOGITS_TO_KEEP):
+            logger.warning_once(
+                'Disabling use_logits_to_keep for this sequence-parallel trainer; its loss reducer is not SP-aware.')
+            use_logits_to_keep = False
         return use_logits_to_keep
 
     def _save_initial_model(self, output_dir):
@@ -1161,7 +1168,7 @@ def _get_listwise_reranker_preds(logits, labels):
         labels = torch.tensor([0] * (len(positive_indices) - 1))
         return preds, labels
 
-    def _compute_acc(self, outputs, labels, cu_seqlens=None) -> None:
+    def _compute_acc(self, outputs, labels, cu_seqlens=None, logits_to_keep=None) -> None:
         args = self.args
         logits = outputs.logits
         metrics = None
@@ -1186,9 +1193,28 @@ def _compute_acc(self, outputs, labels, cu_seqlens=None) -> None:
                     preds = torch.from_numpy(preds).to(get_current_device())
                 if isinstance(labels, np.ndarray):
                     labels = torch.from_numpy(labels).to(get_current_device())
+                # Selective lm-head projection returns only the positions
+                # addressed by ``logits_to_keep``. Reinsert those predictions
+                # into the full local shard before the existing gather path.
+                if isinstance(logits_to_keep, torch.Tensor) and logits_to_keep.dtype == torch.bool:
+                    if logits_to_keep.ndim == 1 and logits_to_keep.numel() == labels.shape[1]:
+                        selected_count = int(logits_to_keep.sum().item())
+                        if preds.shape[1] == selected_count:
+                            full_preds = torch.zeros((preds.shape[0], labels.shape[1]),
+                                                     dtype=preds.dtype,
+                                                     device=preds.device)
+                            full_preds[:, logits_to_keep] = preds
+                            preds = full_preds
+                elif isinstance(logits_to_keep, int) and 0 < logits_to_keep <= labels.shape[1]:
+                    if preds.shape[1] == logits_to_keep:
+                        full_preds = torch.zeros((preds.shape[0], labels.shape[1]),
+                                                 dtype=preds.dtype,
+                                                 device=preds.device)
+                        full_preds[:, -logits_to_keep:] = preds
+                        preds = full_preds
                 assert labels.shape[1] == preds.shape[1]
 
-                if sequence_parallel.rp_world_size > 1:
+                if (sequence_parallel.rp_world_size or 1) > 1:
                     position_ids = sequence_parallel.real_position_ids
                     position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
                 else:
@@ -1257,10 +1283,38 @@ def _evalscope_eval(self):
         return eval_dict
 
     def prepare_logits_to_keep(self, inputs):
+        """Prepare selective lm-head inputs for regular and SP SFT paths.
+
+        Sequence-parallel input preparation has already applied the causal
+        shift to labels.  Keep that full local frame intact and let the SP
+        loss function scatter selected logits back into it.
+        """
         labels = inputs['labels']
         loss_scale = inputs.get('loss_scale')
         if self.template.sequence_parallel_size > 1:
-            raise NotImplementedError()
+            # Transformers causal-LM heads accept a one-dimensional boolean
+            # sequence index for every batch row.  Keep arbitrary supervised
+            # positions for batch-size one; for a larger batch use one shared
+            # suffix that covers the earliest supervised target in any row.
+            if labels.shape[0] == 1 and not is_mp():
+                logits_to_keep = labels[0] != -100
+                # Keep one ignored position on an all-masked shard so model
+                # implementations that reject an empty lm_head input remain
+                # usable; it contributes zero to the loss.
+                if not logits_to_keep.any():
+                    logits_to_keep = logits_to_keep.clone()
+                    logits_to_keep[-1] = True
+            else:
+                supervised = labels != -100
+                first_supervised = supervised.int().argmax(dim=-1)
+                has_supervised = supervised.any(dim=-1)
+                first = first_supervised.masked_fill(~has_supervised, labels.shape[-1] - 1).min().item()
+                logits_to_keep = torch.zeros(labels.shape[-1], dtype=torch.bool, device=labels.device)
+                logits_to_keep[-max(labels.shape[-1] - first, 1):] = True
+            inputs['logits_to_keep'] = logits_to_keep
+            # Do not truncate labels/loss_scale: SP gathers a full local frame
+            # and therefore needs their original shard length.
+            return
         if labels.shape[0] == 1 and not is_mp():
             # device_map may encounter device mismatch issues.
             loss_mask = (labels != -100)[0]
@@ -1282,6 +1336,21 @@ def prepare_logits_to_keep(self, inputs):
     def get_cu_seqlens(self, position_ids, logits_to_keep) -> torch.Tensor:
         cu_seqlens = get_packed_seq_params(position_ids)['cu_seq_lens_q']
         if isinstance(logits_to_keep, torch.Tensor):
+            # SP keeps a local boolean mask while position_ids still contains
+            # the complete packed sequence. Gather the mask first so compact
+            # boundaries are computed in the global frame.
+            if (getattr(getattr(self, 'template', None), 'sequence_parallel_size', 1) > 1 and logits_to_keep.ndim == 1
+                    and logits_to_keep.numel() != position_ids.shape[-1] and (sequence_parallel.world_size or 1) > 1):
+                local_mask = logits_to_keep.unsqueeze(0)
+                gather_position_ids = None
+                if (sequence_parallel.rp_world_size or 1) > 1:
+                    gather_position_ids = sequence_parallel.real_position_ids
+                    gather_position_ids = sequence_parallel.pad(
+                        gather_position_ids, padding_value=-1, position_ids=gather_position_ids)
+                logits_to_keep = sequence_parallel.gather(local_mask, dim=1, position_ids=gather_position_ids)
+                if gather_position_ids is not None and gather_position_ids.min() == -1:
+                    logits_to_keep = logits_to_keep[gather_position_ids >= 0]
+                logits_to_keep = logits_to_keep.reshape(-1)
             kept_cumsum = logits_to_keep.to(cu_seqlens.dtype).cumsum(dim=0, dtype=cu_seqlens.dtype)
             kept_cumsum = torch.cat((cu_seqlens.new_zeros(1), kept_cumsum))
             res_cu_seqlens = kept_cumsum[cu_seqlens.long()]

swift/trainers/seq2seq_trainer.py

@@ -25,6 +25,7 @@
 
 class Seq2SeqTrainer(SwiftMixin, DataLoaderMixin, HfSeq2SeqTrainer):
     args: Seq2SeqTrainingArguments
+    SUPPORTS_SP_LOGITS_TO_KEEP = True
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -105,9 +106,20 @@ def _prepare_inputs(self, inputs):
             sequence_parallel.prepare_inputs(inputs)
 
         use_logits_to_keep = self.get_use_logits_to_keep(self.template.sequence_parallel_size == 1)
+        sp_selective_unsupported = self.template.sequence_parallel_size > 1 and (self.compute_loss_func is not None
+                                                                                 or self.label_smoother is not None
+                                                                                 or self.args.enable_channel_loss)
+        if use_logits_to_keep and sp_selective_unsupported:
+            # Custom losses and label smoothing may consume the original
+            # (unshifted) label frame.  Keep their established semantics until
+            # they provide an SP-aware selective-loss implementation.
+            logger.warning_once(
+                'Disabling use_logits_to_keep for sequence parallel custom loss/label smoothing/channel loss.')
+            use_logits_to_keep = False
         if use_logits_to_keep:
             self.prepare_logits_to_keep(inputs)
-            if args.tuner_backend == 'unsloth' and isinstance(inputs['logits_to_keep'], torch.Tensor):
+            if (args.tuner_backend == 'unsloth' and self.template.sequence_parallel_size == 1
+                    and isinstance(inputs['logits_to_keep'], torch.Tensor)):
                 inputs['logits_to_keep'] = int(inputs['logits_to_keep'].sum())
 
         base_model = self.template.get_base_model(self.model)
@@ -160,7 +172,11 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             if (self.args.enable_dft_loss or loss_scale is not None or self.args.enable_channel_loss
                     or self.template.sequence_parallel_size > 1):
                 if self.template.sequence_parallel_size > 1:
-                    outputs.loss = per_token_loss_func_sp(outputs, labels, enable_dft_loss=self.args.enable_dft_loss)
+                    outputs.loss = per_token_loss_func_sp(
+                        outputs,
+                        labels,
+                        enable_dft_loss=self.args.enable_dft_loss,
+                        logits_to_keep=inputs.get('logits_to_keep'))
                     if loss_scale is not None:
                         position_ids = sequence_parallel.real_position_ids
                         if position_ids is not None:
@@ -232,10 +248,15 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
         if (outputs.logits is not None and labels is not None and self.args.tuner_backend != 'unsloth'):
             cu_seqlens = None
             if self.template.padding_free and self.args.acc_strategy == 'seq':
-                cu_seqlens = self.get_cu_seqlens(text_position_ids, inputs.get('logits_to_keep'))
+                logits_to_keep = inputs.get('logits_to_keep')
+                # Outputs are full-frame after SP selective-loss scattering;
+                # retain full packed boundaries for sequence accuracy.
+                cu_seqlens = self.get_cu_seqlens(
+                    text_position_ids,
+                    None if self.template.sequence_parallel_size > 1 and logits_to_keep is not None else logits_to_keep)
             # Liger does not have logits
             # Unsloth has a bug with output logits
-            self._compute_acc(outputs, labels, cu_seqlens=cu_seqlens)
+            self._compute_acc(outputs, labels, cu_seqlens=cu_seqlens, logits_to_keep=inputs.get('logits_to_keep'))
         return (loss, outputs) if return_outputs else loss
 
     def training_step(self, model, inputs, *args, **kwargs):

swift/trainers/utils.py

@@ -164,27 +164,74 @@ def is_instance_of_ms_model(model: Module) -> bool:
     return False
 
 
-def per_token_loss_func_sp(outputs, labels, enable_dft_loss=False, **kwargs) -> torch.Tensor:
-    """Common loss function for sequence parallel training"""
+def per_token_loss_func_sp(outputs, labels, enable_dft_loss=False, logits_to_keep=None, **kwargs) -> torch.Tensor:
+    """Compute per-token SP loss while supporting selective lm-head logits.
+
+    Sequence-parallel labels are already causally shifted and sharded.  When
+    ``logits_to_keep`` selects a compact subset of the local hidden states,
+    CE is evaluated on that subset and scattered back into the full local
+    sequence frame before the existing gather.  This keeps packed/ring
+    ordering and collective tensor shapes unchanged while avoiding the
+    vocabulary projection for ignored tokens.
+    """
     if hasattr(outputs, 'logits'):
         logits = outputs.logits
     else:
         logits = outputs
     device = logits.device
+    labels = labels.to(device)
+    batch_size = labels.shape[0]
+    compact_selection = logits_to_keep is not None
+
+    if compact_selection:
+        local_seq_len = labels.shape[-1]
+        if isinstance(logits_to_keep, torch.Tensor):
+            if logits_to_keep.dtype != torch.bool or logits_to_keep.ndim != 1:
+                compact_selection = False
+            elif logits_to_keep.numel() != local_seq_len:
+                raise ValueError(
+                    f'logits_to_keep has length {logits_to_keep.numel()}, expected {local_seq_len} for SP labels')
+            else:
+                selected_labels = labels[:, logits_to_keep]
+        elif isinstance(logits_to_keep, int):
+            if logits_to_keep <= 0 or logits_to_keep > local_seq_len:
+                raise ValueError(f'logits_to_keep={logits_to_keep} must be in [1, {local_seq_len}] for SP labels')
+            selected_labels = labels[:, -logits_to_keep:]
+        else:
+            compact_selection = False
+
+    if compact_selection:
+        if logits.shape[1] != selected_labels.shape[1]:
+            raise ValueError(f'logits sequence length ({logits.shape[1]}) does not match selected labels '
+                             f'({selected_labels.shape[1]})')
+        logits = logits.reshape(-1, logits.shape[-1])
+        selected_labels = selected_labels.reshape(-1)
+    else:
+        logits = logits.reshape(-1, logits.shape[-1])
+        selected_labels = labels.reshape(-1)
 
-    batch_size = logits.shape[0]
-    logits = logits.view(-1, logits.shape[-1])
-    labels = labels.flatten().to(device)
     sploss_parallel_size = int(os.environ.get('CELOSS_PARALLEL_SIZE', '0'))
     if sploss_parallel_size > 0:
-        loss = ChunkedCrossEntropyLoss.apply(logits, labels, sploss_parallel_size)
+        loss = ChunkedCrossEntropyLoss.apply(logits, selected_labels, sploss_parallel_size)
     else:
         loss_fct = CrossEntropyLoss(reduction='none')
-        loss = loss_fct(logits, labels)
+        loss = loss_fct(logits, selected_labels)
     if enable_dft_loss:
         with torch.no_grad():
             target_probs = torch.exp(-loss)
         loss *= target_probs
+
+    if compact_selection:
+        # Reconstruct a full local frame for GatherLoss.  Unselected entries
+        # remain zero and correspond to -100 labels in the original frame.
+        selected_loss = loss.reshape(batch_size, -1)
+        full_loss = torch.zeros((batch_size, labels.shape[-1]), dtype=selected_loss.dtype, device=selected_loss.device)
+        if isinstance(logits_to_keep, torch.Tensor):
+            full_loss[:, logits_to_keep] = selected_loss
+        else:
+            full_loss[:, -logits_to_keep:] = selected_loss
+        loss = full_loss.reshape(-1)
+
     position_ids = sequence_parallel.real_position_ids
     if position_ids is not None:
         position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)