Repository navigation
[loss] fix softmax collapse with improved ce perf - #5723
Closed
ruisizhang123 wants to merge 1 commit into
Closed
ruisizhang123 wants to merge 1 commit into
ruisizhang123 wants to merge 1 commit into
Conversation
…d path faster
_LossParallelCrossEntropy computed the NLL as logsumexp(z) - z_t and a target
gradient of p_t - 1. Once the target margin exceeds ~16.6, fp32 rounds p_t to
exactly 1, so the loss and the target gradient both become exactly 0 while the
non-target gradients stay nonzero. The accurate target gradient needs the
non-target mass, which forward now computes, so forward also prepares the
backward instead of saving log-probabilities.
Numerics (TP > 1 only):
- NLL = softplus(logsumexp_{j != t}(z_j) - z_t). The owner rank drops the
target term before the sum, so no small terms are lost against exp(0) = 1.
- Target gradient: p_t - 1 below p_t = 1/4, else -p_non_target (one rounding
instead of the sum and the divide).
- Ignored labels act as targets on no shard, so their loss and gradient stay
finite and are zeroed by the label mask.
- Matches loss_parallel() up to fp32 rounding below the collapse margin; not
bitwise identical to it or to the previous code.
Memory/perf: the best form differs between eager and compiled, so the two run
different code with the same math, gated on torch.compiler.is_compiling() or
_is_non_strict_tracing() (the latter covers GraphTrainer's make_fx tracing).
- Eager cannot fuse, so forward turns one fresh fp32 [T, V_local] buffer in
place into exp, then softmax, then the saved gradient, and backward scales it
in place. The input logits are never modified.
- Compiled forward drops the target with a column mask instead of an in-place
write, so Inductor fuses the sum into one read of the logits. It saves the
input plus per-row stats, and backward rebuilds the gradient in one fused
kernel. An in-place target write would force the fp32 buffer to be
materialized, which made a compiled single-path version 10% slower than main.
- cross_entropy_loss passes pred to the Function without .float(); the Function
upcasts internally. Otherwise the compiled region saves an fp32 copy of the
logits instead of the bf16 input. Loss and gradients are bitwise unchanged.
Measured on H100, V=128256, TP=2, T=16384, loss forward + backward:
old TP CE (eager) new TP CE (eager) old TP CE (compile) new TP CE (compile)
memory 15.66 GiB 5.87 GiB 7.83 GiB 3.91 GiB
time 38.96 ms 27.11 ms 13.15 ms 6.29 ms
Eager and compiled gradients agree to ~1e-6. On real C4 (Qwen3-0.6B, TP=2 x
FSDP=4, 2000 steps) the new CE trains like main: last-100-step mean loss
4.5414 vs 4.5335 compiled and 4.5411 vs 4.5475 eager, a gap that changes sign
between modes and is ~20x below batch-to-batch noise.
The TP=1 F.cross_entropy path (and the batch-invariant policy-stats path) is
unchanged and still collapses at large margins, so TP=1 and TP>1 losses now
differ in that regime. Goldens for TP > 1 runs are regenerated (fake_pg/a10g,
real_pg/a10g, real_pg/mi350x); goldens without TP use F.cross_entropy and are
unchanged.
Test Plan:
Confident-target accuracy against an fp64 reference, ignored rows with a -inf
fallback logit, input logits untouched, and eager vs aot_eager compiled parity
over fp32 and bf16 logits:
```
pytest tests/unit_tests/cpu/components/test_loss.py
```
This PR was authored with an AI assistant (Claude Code).
meta-cla
Bot
added
the
CLA Signed
|
The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:
Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is a replicate and revision of #5661.
I don't have write access to the repo; thus, I ended up submitting a new PR here.......