Skip to content

fix(loss): reuse aligned cross entropy under sequence parallelism - #10185

Open
Excelius-Wang wants to merge 1 commit into
modelscope:mainfrom
Excelius-Wang:fix/sp-custom-cross-entropy
Open

Excelius-Wang wants to merge 1 commit into
modelscope:mainfrom
Excelius-Wang:fix/sp-custom-cross-entropy

Conversation

@Excelius-Wang

Copy link
Copy Markdown
Contributor

Explicit --loss_type cross_entropy recomputes ordinary causal cross entropy from labels that sequence parallelism has already shifted and sharded. This shifts the targets again and loses the gathered loss; with token weights, the gathered weights can also have a different length from the locally recomputed loss.

Reuse the per-token loss already aligned, gathered and weighted by Seq2SeqTrainer. When the callback infers the token count, preserve the default trainer's distributed normalization and training/evaluation scaling. The non-sequence-parallel path is unchanged.

Validation:

  • CPU Gloo regression coverage for Ulysses, Ring and hybrid layouts, plus combined data/sequence parallelism with unequal valid-token counts, aligned/tail-padded/packed sequences, absent/nonuniform/zero weights, explicit/inferred counts, token averaging on/off and training/evaluation. Compare losses and gradients against both the default trainer path and an independent PyTorch CE reference (288 combinations).
  • OMP_NUM_THREADS=1 /tmp/ms-swift-core-test-venv/bin/python -m pytest tests/sequence_parallel/test_custom_cross_entropy.py tests/train/test_cross_entropy_loss.py -q: 5 passed, 8 subtests passed.
  • /tmp/ms-swift-lint-venv/bin/pre-commit run --all-files: all hooks passed.
  • Additional local probe using Swift Qwen3 text-template encoding and tiny randomly initialized Qwen3 logits reproduces the original loss mismatch and weighted shape error; the fix matches default loss and local logits gradients.

Coverage uses CPU Gloo with one or two data-parallel replicas. It does not exercise GPU distributed attention or a full SFT run. Related #10100 added token weighting for custom CE but did not validate sequence parallelism.

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