Skip to content

feat(sft): support logits selection with Ulysses sequence parallelism - #10191

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

Excelius-Wang wants to merge 1 commit into
modelscope:mainfrom
Excelius-Wang:feat/sp-logits-to-keep

Conversation

@Excelius-Wang

Copy link
Copy Markdown
Contributor

Sequence-parallel SFT currently rejects explicit use_logits_to_keep=true. Add opt-in selection after SP has shifted and sharded labels, retaining a sentinel on prompt-only ranks. Restore scalar token losses and integer predictions to their local positions, then reuse existing gathering, weighting, normalization and accuracy. Evaluation keeps full logits; shared DPO/KTO/GKD preparation is unchanged.

Support covers native tensor-selection text causal LMs and text-only Qwen3.5/3.6 MoE inputs, with Ulysses, per-device batch one and no packing/padding-free. Media inputs, other multimodal architectures, Ring, custom loss, label smoothing, Unsloth and Liger remain excluded. Accept the all-zero modality IDs emitted by the Qwen text collator while rejecting nonzero media IDs.

Related #9765. This implements the reporter's model-family text SP path, but does not establish that their full 35B/128K, eight-H20, ZeRO3 workload fits in memory.

Validation:

  • 23 tests and 8 subtests passed. Real two-process CPU Gloo/Ulysses/DDP on tiny Qwen3 and Qwen3.5 MoE, including full training and LoRA/checkpointing, compares loss/all trainable gradients with independent full-sequence references across 72 combinations; also tests auxiliary router-loss parity, accuracy, channel loss, exclusions and shared/non-SP regressions.
  • The same MoE test on clean main fails at its original NotImplementedError.
  • Normal SwiftSft with Qwen3.5 MoE, real template/tokenizer, SP2, LoRA, checkpointing, weighted examples, accumulation two and per-step evaluation completes two steps. Selection on/off yields identical train/eval losses and all 91 final state tensors; adapters are updated.
  • Full pre-commit run --all-files and staged diff checks pass.
  • Separate single-H100 BF16 native model test with LoRA/checkpointing and SGD passes loss/gradient/update parity; this was not multi-GPU SP or trainer integration.

CPU integration uses torch 2.9.1 and transformers 5.12.1. No ZeRO/FSDP, full pretrained long-context, image/video or GPU performance claim; full upstream suite not run.

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.

1 participant