[taking-lying-flat]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support
• Performance Optimization

PR information

Related issue: #10129. This preserves the separate-gradient-buffer fix from #10130; it does not reopen that correctness issue.

Custom SFT loss currently casts every retained logit row to FP32 and computes CE before ignoring unsupervised tokens. Multi-sample logits_to_keep suffix cropping can leave unsupervised holes. Separately, sequence-parallel chunked CE rebuilds a CE graph and calls autograd.grad for each backward chunk.

This change:

• Selects supervised rows after autoregressive label alignment, before the FP32 cast and CE, then scatters losses back to the existing per-token shape. It retains dense CE when the ignored-row count is no larger than the batch size, avoiding a near-full gather after existing token compaction
• Computes chunked backward directly as g * softmax(logits) - g * one_hot(labels), with zero upstream weights for ignored labels. It retains a separate gradient buffer and respects the forward CE compute dtype under autocast

DFT probabilities remain detached. External token weights, normalization, and channel-loss logging keep their existing semantics. The backbone context and output-head projection are unchanged. This does not add fused-linear-CE integration; chunked CE remains specific to CELOSS_PARALLEL_SIZE > 0.

Only the two production files are included; validation scripts are kept outside the repository.

AI assistance: OpenAI Codex was used for implementation, code review, local validation, and PR preparation.

Experiment results

NVIDIA RTX A1000 Laptop GPU, PyTorch 2.12.1+cu130. Timings cover loss forward and weighted backward on preallocated logits, with four warmups and the median of five groups of 15 steps, synchronized on CUDA.

Case	Before	After
BF16 sparse-supervision CE	1.441 ms	0.596 ms
BF16 sparse-supervision DFT	1.409 ms	0.640 ms
FP32 chunked CE, 1024 x 2048, chunk 64	2.105 ms	1.425 ms
FP32 chunked CE, 2048 x 8192, chunk 128	3.751 ms	2.398 ms

The sparse case uses two length-1024 samples with four assistant spans each. Existing suffix cropping retains 1922 rows, of which 256 are supervised, with vocabulary size 4096. CE-saved vocabulary-sized FP32 storage, deduplicated by underlying storage, decreases from 30.03125 MiB to 4 MiB. For ordinary CE, loss and output-head parameter gradients match exactly in this example; the DFT control also matches exactly.

For chunk size 64, CE and log-softmax calls each decrease from 32 to 16; NLL/log-softmax backward calls decrease from 16 to zero, replaced by 16 softmax calls.

The dense control (1920 supervised rows out of 1922) is approximately 3.1% slower, from 1.400 to 1.443 ms, due to valid-index detection, which includes CUDA synchronization. These are loss-phase measurements, not end-to-end training speed or memory claims. Analytical backward may round differently from log-softmax backward in FP16/BF16; AMP comparisons also used a same-logit FP64 gradient reference. No NPU/MPS/XPU validation was performed.

Validation commands (using the checkout):

python -m pytest tests/train/test_chunked_cross_entropy.py tests/train/test_cross_entropy_loss.py -q --tb=short
CELOSS_PARALLEL_SIZE=3 python -m pytest tests/sequence_parallel/test_custom_cross_entropy.py -k '1-2-1' -q --tb=short
uvx --offline pre-commit run --all-files
git diff --check

Existing trainer tests: 14 passed, including 8 subtests. Two-process CPU Gloo sequence-parallel test: 1 passed, 3 deselected. All pre-commit hooks and the whitespace check passed.

Additional local validation, excluded from the diff:

LOCAL_CHECKS=/home/cherry-cloud/tmp/ms-swift-ce-validation
python -m pytest "$LOCAL_CHECKS/test_ce_local.py" -q --tb=short
python "$LOCAL_CHECKS/bench_ce.py"

Local checks passed in two batches: 102 passed, 12 skipped (inapplicable AMP/FP32-FP64 combinations). Together with the existing tests: 117 passed. Coverage includes CPU/CUDA, FP32/FP64/FP16/BF16, AMP and GradScaler optimizer steps, CE/DFT, token weighting, channel logging, non-contiguous logits, all-ignored/dense masks, actual suffix cropping, incomplete chunks, repeated backward, auxiliary-loss sharing, and saved-logit immutability.

Checklist

• I have linked the related issue above, if there is one.
• I have reviewed and understand every changed line, and I can explain how the change works and why it is correct.
• I have run the relevant tests or reproductions locally, and included the commands and results above (if applicable).
• I am responsible for the content of this PR, regardless of how it was produced.