Skip to content

fix(trainer): support selective logits with sequence parallel - #10039

Open
zjn20030811 wants to merge 1 commit into
modelscope:mainfrom
zjn20030811:feat/sp-logits-to-keep
Open

fix(trainer): support selective logits with sequence parallel#10039
zjn20030811 wants to merge 1 commit into
modelscope:mainfrom
zjn20030811:feat/sp-logits-to-keep

Conversation

@zjn20030811

@zjn20030811 zjn20030811 commented Sep 4, 2026

Copy link
Copy Markdown

Summary

Fixes #9765 by enabling use_logits_to_keep with sequence parallel SFT.

  • derive a local boolean selection after SP causal label shifting;
  • run the language-model head only for supervised (or shared suffix) positions;
  • scatter per-token losses back into the full local frame before GatherLoss, preserving SP/ring ordering and normalization;
  • keep loss scaling, packed boundaries, token accuracy, and all-ignored shards aligned;
  • retain the existing full-logit path for custom loss functions, label smoothing, and channel loss until those paths expose SP-aware contracts;
  • keep RLHF trainers on their established full-logit path through an explicit capability guard.

Validation

  • py -3.12 -m pytest -q tests/test_align/test_sp_logits_to_keep.py tests/test_align/test_rlhf_loss.py
  • py -3.12 -m pre_commit run --files swift/trainers/mixin.py swift/trainers/seq2seq_trainer.py swift/trainers/utils.py tests/test_align/test_sp_logits_to_keep.py
  • py -3.12 -m compileall -q swift/trainers tests/test_align/test_sp_logits_to_keep.py

The focused regression suite covers single-row masks, batched suffix selection, loss/gradient equivalence, all-ignored local shards, packed sequence boundaries, accuracy reconstruction, and the non-SFT capability guard.

@zjn20030811
zjn20030811 force-pushed the feat/sp-logits-to-keep branch from 5dd2d2d to 94c6759 Compare September 4, 2026 06:45
@zjn20030811

Copy link
Copy Markdown
Author

Validation update: the focused sequence-parallel suite covers masked rows, packed boundaries, gradient/loss equivalence, all-ignored shards, accuracy reconstruction, and the non-SFT capability guard. The branch is ready for review.

Comment thread swift/trainers/mixin.py
# 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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When will this code be used?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support use_logits_to_keep with sequence parallelism for long-context SFT

2 participants