Repository navigation
Conversation
meta-cla
Bot
added
the
CLA Signed
|
@claude review this pr |
|
Claude finished @ruisizhang123's task in 1m 59s —— View job PR Review: fix softmax collapse with improved CE perf
SummaryThe math holds up.
There are two things to resolve before merge (CI goldens and compile behavior), plus a few smaller points. Blocking / must-address1. TP golden numerics will change. Regenerate them and add a convergence proof.
2. Check the in-place mutation of a saved tensor under
If compile does object, a non-mutating Non-blocking suggestions3. TP and non-TP CE are now numerically different by design.
4. The comment could also be sharper. By Sterbenz, 5. Optional memory win: skip the full fp32 copy of the logits. 6. Using a version-counter error as a feature. There is also a related subtlety. For fp32 inputs, 7. Tests.
Nits
Overall the direction is good and the numerics are careful. The main asks are the golden/convergence evidence and a compile check. |
|
Claude finished @ruisizhang123's task in 3m 38s —— View job Request changes -- review (answers the golden-update question: export via |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Verdict: Request changes -- the TP numerics change still needs regenerated goldens and a convergence run. Splitting off the perf half would let most of the win land without either.
Golden numerics (your question). You don't need the hardware, because CI produces the values. Only the TP>1 goldens are affected:
| suite | files |
|---|---|
fake_pg/a10g |
llama3, gpt_oss, qwen3 |
real_pg/a10g |
llama3, llama3_pp, gpt_oss, qwen3 |
real_pg/mi350x |
llama3, llama3_pp, gpt_oss, qwen3 |
These are unaffected: deepseek_v3*, gpt_oss_pp, muse_glimmer, sft, kimi_k3_b200 (all run without TP), and rl_grpo_cuda.txt (its batch-invariant path all-gathers the vocab and uses F.cross_entropy).
- real_pg/a10g: the branch is on
pytorch/torchtitan, so you can rungh workflow run integration_test.yaml --ref ce_fix -f export_results=true. Download theintegration-8gpu-{features,models}-outputsartifacts; the finished files, headers included, are underreal_pg/a10g/. - fake_pg/a10g: the fake-PG job runs only on a
ciflow/fake-pgtag push, and export is always off there. When the comparison fails,loss_compare.pyprints a "Baseline metrics in import file format" block in the job log. Copy it from the run your label already triggers, and keep the existing# parallelism:line, because the printed block leaves it out. - real_pg/mi350x:
integration_test_rocm.yamlhardcodesEXPORT_RESULTS="false". Add theciflow/rocmlabel and copy the same block from the failing logs.
@tianyu-l may prefer another flow.
Example. Take a row whose target margin z_t - logsumexp_{j!=t} z_j is 20:
- On
main(fp32):sumexp = 1 + 2e-9, which rounds to 1. The NLL is 0.0 and the target grad is 0, but the non-target grads still sum to about 2e-9, so the row's gradient no longer sums to 0. - With this PR: NLL =
softplus(-20), about 2.06e-9, and the target grad is about -2.06e-9, which balances the rest of the row.
Must fix
- Two concerns in one PR.
- Perf half: save
softmax - onehotin forward and scale it in backward. This is the same arithmetic as the old backwardexp(log_probs) + (-1 at target), so it should stay bitwise identical, with no golden churn. It also carries most of the memory and time win, which comes from fewer[T, V_local]temporaries. - Collapse fix: non-target sumexp, softplus, and the
-p_ntbranch. This changes numerics, and only for TP>1. The TP=1 path (F.cross_entropy,loss.py:81) keeps collapsing in exactly the regime this PR targets. - Please split them. The collapse fix needs a maintainer decision on whether both paths get it, plus goldens and a C4 TP>1 run. If you keep them together, the goldens above and the convergence comparison are required.
- Perf half: save
test_loss.py:723-- the inductor branch never runs in the CPU suite (see the inline comment).
Still open from the last review
loss.py:274: the in-placemul_on a saved tensor. The description still doesn't say whether the 15.66 -> 7.83 GiB and 38.96 -> 27.82 ms numbers were measured with thelossregion compiled. Unverified:aot_eagermay turn the mutation into a copy, and the new test checks values only.- The description doesn't say that TP=1 and TP>1 losses now differ.
loss.py:92:HIGH_PROB_THRESHOLDis still a public name in a shared module.test_loss.py:692: the test still asserts that a second backward raises. That error is a side effect of the version counter, not a contract; use@once_differentiableor drop the assertion.- Should consider: with bf16 RL input,
logits.float()pluslogits - local_maxcreates two fp32[T, V_local]buffers.
Not verified: I couldn't run the unit tests (pytest isn't permitted in this job; the CPU, GPU and Integration CI runs were still queued). I also didn't confirm that the integration goldens fail; that is inferred from the TP configs.
| backends = ["aot_eager"] | ||
| if self.device_type == "cuda": | ||
| backends.append("inductor") |
There was a problem hiding this comment.
This test lives in the CPU suite, where device_type is cpu, so the inductor branch never runs in CI (see .claude/rules/tests.md: GPU-only tests don't belong in CPU suites). The only compiled coverage is aot_eager, which checks values but not whether functionalization turns the in-place mul_ on the saved tensor into a copy. Inductor is the backend the loss region actually uses, and that copy is the risk to the memory win. Either drop the dead branch or move an inductor case to tests/unit_tests/gpu/.
| backends = ["aot_eager"] | |
| if self.device_type == "cuda": | |
| backends.append("inductor") | |
| backends = ["aot_eager"] |
cb2cdcf to
ce2c123
Compare
|
@claude review this again, I have updated the golden numbers and convergence run.
The aot_eager is mainly for making sure the code compiles in cpu test. It's not for performance but for test purpose only. |
|
Claude finished @ruisizhang123's task in 3m 7s —— View job Request changes: three |
There was a problem hiding this comment.
Verdict: Request changes -- three TP goldens that the full suite checks were not regenerated, and on the default compiled path the new CE is about 10% slower.
What this PR does. For TP>1 it computes the CE as softplus(logsumexp_{j!=t} z_j - z_t), so confident targets no longer collapse to a loss and target grad of exactly 0. It also forms softmax - onehot in forward, so backward is a single in-place scale.
Example. A row with target margin 20, in fp32:
main: the NLL is 0 and the target grad is 0, but the non-target grads still sum to about 2e-9.- This PR: the NLL is about 2.06e-9 and the target grad is about -2.06e-9, so the row's gradient sums to 0.
Resolved since the last review
- Goldens for fake_pg/a10g, real_pg/a10g
llama3_pp, and all of real_pg/mi350x are regenerated. - The C4 TP=2 x FSDP=4 convergence run (eager and compiled) matches
main: the last-100 mean differs by 0.1-0.2%, well inside batch noise. - On
aot_eager: agreed. The compile test only needs to prove that the region traces withfullgraph=True. The compiled memory numbers (7.83 GiB in both columns) show the in-placemul_doesn't add a copy. - The dead inductor branch and the second-backward assertion are gone,
_HIGH_PROB_THRESHOLDis private, and the bf16 path keeps a single fp32 buffer.
Must fix
tests/assets/losses/real_pg/a10g/{llama3,gpt_oss,qwen3}.txtare not regenerated.- Your green Integration run came from the
ciflow/fake-pgtag. On tag pushes the real-PG job runs only thereal_pg_requiredsubset (integration_test.yaml:154), so it compared just the*_ppgoldens. - The full suite (pushes to main and the 6-hour cron) runs
llama3_fsdp+tp+cp,qwen3, andgpt_osswith TP=2 against those three files (models.py:632,724,772). - Those configs changed on fake_pg/a10g and on mi350x, so they will fail after merge.
- Fix: run
gh workflow run integration_test.yaml --ref ce_fix -f export_results=true(no tag ref, so you get the full suite), then take the three files from theintegration-8gpu-models-outputsartifact.
- Your green Integration run came from the
torchtitan/components/loss.py:241-262-- buildingsoftmax - onehotin forward slows the compiledlossregion from 13.15 to 14.51 ms, and memory doesn't improve. That region is on by default in every model'slocal_compile.regions. See the inline comment.
Should consider
torchtitan/components/loss.py:112-- the docstring describes the change relative to the old code ("no longer bitwise"); describe the current property instead (inline suggestion).
Questions
- (maintainers) Should the TP=1 path (
F.cross_entropy,loss.py:81) and the batch-invariant RL path get the same fix? Today TP=1 and TP>1 runs give different results in exactly the overconfident regime this PR targets. - Is there a run on
mainwhere the collapse actually hurt training, e.g. a late-stage or fine-tuning loss plateau? The C4 run shows the fix is neutral, but it doesn't show the problem exists.
Not verified: that the three a10g real_pg goldens actually fail (I inferred it from the mi350x and fake_pg diffs on the same configs), and that the smaller backward design recovers the compiled time. The one CPU unit-test failure (test_every_cp_recipe_passes_the_gate: TrainingConfig has no disable_cuda_graphs) appears unrelated to this PR.
| # Save the CE gradient softmax - onehot, so backward is one scale. Every | ||
| # class gets its softmax probability; the owner rank's target slot gets | ||
| # p_t - 1, formed by the formula that rounds once: p_t - 1 cancels once | ||
| # the target dominates, so use -p_non_target from _HIGH_PROB_THRESHOLD | ||
| # on. | ||
| target_exp = torch.exp(target_shifted) | ||
| sum_exp = non_target_sumexp + torch.where(ignored.unsqueeze(-1), 0, target_exp) | ||
| exp_shifted.div_(sum_exp) | ||
| target_prob = target_exp / sum_exp | ||
| target_grad = torch.where( | ||
| target_prob < _HIGH_PROB_THRESHOLD, | ||
| target_prob - 1, | ||
| -non_target_sumexp / sum_exp, | ||
| ) | ||
| # Rows whose target is not on this rank rewrite lane 0 with its own | ||
| # softmax value, leaving it unchanged. | ||
| local_lane_prob = exp_shifted[row_idx, local_labels] | ||
| exp_shifted[row_idx, local_labels] = torch.where( | ||
| out_of_range, local_lane_prob, target_grad.squeeze(-1) | ||
| ) | ||
|
|
||
| # Save local-shard log probabilities for the fused CE backward. | ||
| ctx.save_for_backward(log_probs, labels) | ||
| ctx.save_for_backward(exp_shifted, labels) |
There was a problem hiding this comment.
Forming softmax - onehot in forward only helps eager. By your own table it makes the compiled path slower: 13.15 -> 14.51 ms (+10%), with memory unchanged at 7.83 GiB. Compiled is the default path, because "loss" is in the default local_compile.regions of llama3, qwen3, deepseek_v3, gpt_oss and the others. So the main training path regresses, and the "improved ce perf" in the title holds only for eager.
The likely cause is that forward now runs extra in-place passes over [T, V_local]: the -inf scatter, exp_, div_, and the second scatter. On main, Inductor fused the whole old backward into a single kernel.
Smallest fix (unverified): keep the softplus NLL, but save shifted (or the exp of it) plus the [T, 1] stats sum_exp, non_target_sumexp, target_prob, and build the gradient in backward as exp(shifted) / sum_exp with the target slot set to target_grad, then scale by g. Inductor can fuse that into one kernel the way it did before. Alternatively, show that the compiled numbers recover. Eager is not the path to optimize for (.claude/rules/design.md: "Debug-only and eager fallback paths don't deserve optimization").
| margin exceeds ~16.6) while non-target gradients survive. This is no longer | ||
| bitwise identical to ``loss_parallel()``. |
There was a problem hiding this comment.
"No longer bitwise identical" tells the history of the code rather than its current state, and readers won't have the old version to compare against. Describe the property instead:
| margin exceeds ~16.6) while non-target gradients survive. This is no longer | |
| bitwise identical to ``loss_parallel()``. | |
| margin exceeds ~16.6) while non-target gradients survive. It matches | |
| ``loss_parallel()`` up to fp32 rounding below that margin. |
ruisizhang123
added
the
ciflow/real-pg
28501f8 to
191a4d4
Compare
|
@claude added the golden files, and fixed the review points. recheck this PR |
|
Claude finished @ruisizhang123's task in 2m 58s —— View job Request changes: the goldens and the compiled-path slowdown are fixed, but the CE now has separate eager and compiled implementations, chosen through a private API ( |
There was a problem hiding this comment.
Verdict: Request changes -- the goldens and the compiled-path slowdown are fixed. The new commit, though, splits the CE into an eager implementation and a compiled one, chosen through a private tracing API, and the eager one runs only off the main path.
What this PR does. For TP>1 the CE is now computed as softplus(logsumexp_{j!=t} z_j - z_t), so confident targets keep a nonzero loss and a target gradient that balances the rest of the row. There are two implementations. When compiled, it saves the bf16 input plus [T, 1] row stats and recomputes the softmax in backward. In eager, it builds softmax - onehot in place during forward and scales it in backward. The caller no longer upcasts pred.
Example. A row with target margin 20, fp32:
main: NLL = 0 and target grad = 0, while the non-target grads sum to about 2e-9.- This PR: NLL is about 2.06e-9 and the target grad is about -2.06e-9, so the row's gradient sums to 0.
Resolved since the last review
real_pg/a10g/{llama3,gpt_oss,qwen3}.txtare regenerated, so every TP>1 golden has now been updated.- The compiled path is now faster: 13.15 -> 6.29 ms and 7.83 -> 3.91 GiB, according to the updated description.
- The "no longer bitwise" docstring now states the property instead.
Must fix
torchtitan/components/loss.py:227-255-- there are two CE implementations, gated onis_compiling() or _is_non_strict_tracing(). The eager branch only runs when thelossregion isn't compiled, so it is an optimized eager fallback (.claude/rules/design.md). It also brings a private upstream API with no TODO. Most integration goldens run that eager branch, becausetorchtitan_recipes/tests/__init__.py:30clearslocal_compile.regionsunder SPMD type checking, so they don't guard the production path. Fix: keep only the recompute form (see inline).
Questions
- (maintainers, still open) Should the TP=1
F.cross_entropypath (loss.py:81) and the batch-invariant RL path get the same fix? Today TP=1 and TP>1 give different results in the overconfident regime.
Details
- Recompute form in eager, rough comparison with
main. Forward: the fp32 upcast, the exp, and the maskedwhereare transients (mainhad shifted, exp andlog_probs). Saved: the bf16 input plus four[T, 1]tensors (mainsaved fp32log_probs). Backward: exp/div/where, plus an in-placemul_on a fresh tensor. So eager should end up close tomain, not worse. Timing is unverified. - With a single implementation,
test_loss_parallel_cross_entropy_compiledbecomes a plain fullgraph-traces-and-matches-eager check of one implementation, and the eager-only claim in the docstring oftest_..._ignored_rows_and_inputs("saves its gradient in a fresh buffer") should be dropped.
Not verified: eager time and memory of the recompute-only form, and which integration configs run the loss region compiled vs eager. I read the code and couldn't run the GPU tests.
| # GraphTrainer traces with make_fx, which is_compiling() does not report. | ||
| recompute = ( | ||
| torch.compiler.is_compiling() or torch.compiler._is_non_strict_tracing() | ||
| ) | ||
| if recompute: | ||
| # Under Inductor the masked sum fuses into one read of the logits, | ||
| # and backward recomputes the softmax from the saved input, so no | ||
| # [T, V_local] fp32 buffer is written in forward. An in-place | ||
| # target patch would force that buffer to be materialized. | ||
| is_target = _target_mask(target_lane, local_vocab_size) | ||
| exp_shifted = torch.exp(logits.float() - local_max) | ||
| non_target_sumexp = torch.where(is_target, 0, exp_shifted).sum( | ||
| dim=-1, keepdim=True | ||
| ) | ||
| else: | ||
| # Eager cannot fuse, so instead one fresh fp32 buffer is turned in | ||
| # place into exp, then softmax, then the saved gradient. A non-fp32 | ||
| # input is upcast straight into it. The input logits are never | ||
| # modified (callers reuse them, e.g. for entropy). Rows whose target | ||
| # is not on this rank index lane 0, a real class, and keep it. | ||
| if logits_dtype == torch.float32: | ||
| shifted = logits - local_max | ||
| else: | ||
| shifted = logits.float().sub_(local_max) | ||
| shifted[row_idx, local_labels] = torch.where( | ||
| out_of_range, local_target_shifted, float("-inf") | ||
| ) | ||
| exp_shifted = shifted.exp_() | ||
| non_target_sumexp = exp_shifted.sum(dim=-1, keepdim=True) |
There was a problem hiding this comment.
There are now two CE implementations, picked by whether we're tracing. Only the recompute branch runs on the production path: both callers sit inside @local_compile("loss"), and "loss" is in every model's default regions. The else branch, plus its save/backward halves at L283-291 and L773, runs only when the loss region is not compiled. .claude/rules/design.md covers this case: "Debug-only and eager fallback paths don't deserve optimization ... when compiled/kernel paths exist."
What it costs to keep:
- About 25 lines of in-place index/scatter logic (lane-0 rewrite,
-infpatch,div_, mutating a saved tensor in backward), each of which needs its own reasoning and tests. - A private upstream API,
torch.compiler._is_non_strict_tracing(), added only to choose between the two branches, with no TODO or upstream issue (CLAUDE.md principle 3). - The goldens mostly don't cover the code that runs in production.
torchtitan_recipes/tests/__init__.py:30clearslocal_compile.regionswhenever SPMD type checking is on, so those integration configs run the eager branch. The default compiled branch is covered only by the CPUaot_eagertest.
Smallest fix: drop the gate and the else branches, and always use the masked-sum + recompute form (forward L236-240, save L280-282, backward L769-771). That also removes ctx.recompute, row_idx in forward, and the private API call. In eager the recompute form should cost about what main costs today: a few [T, V_local] fp32 temporaries and a saved low-precision input instead of saved fp32 log_probs. Unverified: please re-measure eager, and regenerate any goldens that move.
There was a problem hiding this comment.
@ruisizhang123 let's do what claude suggested, namely, dropping the eager only optimization
| ``logsumexp(z) - z_t``, so confident targets keep a nonzero loss and target | ||
| gradient instead of both rounding to exactly 0 (fp32 rounds once the target | ||
| margin exceeds ~16.6) while non-target gradients survive. It matches | ||
| ``loss_parallel()`` up to fp32 rounding below that margin. |
There was a problem hiding this comment.
Please remove loss_parallel() reference as it's no longer used in titan
| # GraphTrainer traces with make_fx, which is_compiling() does not report. | ||
| recompute = ( | ||
| torch.compiler.is_compiling() or torch.compiler._is_non_strict_tracing() | ||
| ) | ||
| if recompute: | ||
| # Under Inductor the masked sum fuses into one read of the logits, | ||
| # and backward recomputes the softmax from the saved input, so no | ||
| # [T, V_local] fp32 buffer is written in forward. An in-place | ||
| # target patch would force that buffer to be materialized. | ||
| is_target = _target_mask(target_lane, local_vocab_size) | ||
| exp_shifted = torch.exp(logits.float() - local_max) | ||
| non_target_sumexp = torch.where(is_target, 0, exp_shifted).sum( | ||
| dim=-1, keepdim=True | ||
| ) | ||
| else: | ||
| # Eager cannot fuse, so instead one fresh fp32 buffer is turned in | ||
| # place into exp, then softmax, then the saved gradient. A non-fp32 | ||
| # input is upcast straight into it. The input logits are never | ||
| # modified (callers reuse them, e.g. for entropy). Rows whose target | ||
| # is not on this rank index lane 0, a real class, and keep it. | ||
| if logits_dtype == torch.float32: | ||
| shifted = logits - local_max | ||
| else: | ||
| shifted = logits.float().sub_(local_max) | ||
| shifted[row_idx, local_labels] = torch.where( | ||
| out_of_range, local_target_shifted, float("-inf") | ||
| ) | ||
| exp_shifted = shifted.exp_() | ||
| non_target_sumexp = exp_shifted.sum(dim=-1, keepdim=True) |
There was a problem hiding this comment.
@ruisizhang123 let's do what claude suggested, namely, dropping the eager only optimization
…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).
|
@claude review again |
There was a problem hiding this comment.
Verdict: Approve -- please address the items below before merging. The eager-only branch and the private _is_non_strict_tracing() call are gone, so there is now one implementation. What's left is stale text in the description and in one docstring. A maintainer still needs to review, since this touches core loss code.
What this PR does. For TP>1, the CE is now softplus(logsumexp_{j!=t} z_j - z_t). Forward saves the input logits (bf16 stays bf16) and the [T, 1] row stats. Backward recomputes softmax, puts the precomputed target gradient in the target slot, and scales the result. Eager and compiled run the same code.
Example. A row with target margin 20, fp32:
main: the NLL and the target grad are both 0, but the non-target grads sum to about 2e-9.- This PR: the NLL is about 2.06e-9 and the target grad is about -2.06e-9, so the row's gradient sums to 0.
Resolved since the last review
- Single recompute path, as @tianyu-l asked (
loss.py:224-231,273-276).ctx.recompute,row_idx, and the private API are removed. - @tianyu-l's request to drop the "matches
loss_parallel()" sentence from the docstring is done. - The goldens still hold. In eager, the removed branch summed and divided the same values as the recompute form, so the integration configs, which run the loss region in eager under SPMD type checking, see the same numbers. CI at HEAD agrees: Integration (one run), ROCm, and the CPU and GPU unit tests are green.
Must fix
- The PR description and commit message describe the removed design.
- Item 2 of the description says the PR "move[s] softmax-onehot calculate to fwd pass". The commit message also describes the eager/compiled gate on
is_compiling() or _is_non_strict_tracing(). - The eager column (5.87 GiB / 27.11 ms) was measured on the in-place form that this PR no longer contains.
- Please update both texts, then re-measure eager or drop that column. A squash merge keeps the commit message.
- Item 2 of the description says the PR "move[s] softmax-onehot calculate to fwd pass". The commit message also describes the eager/compiled gate on
tests/unit_tests/cpu/components/test_loss.py:660-- this docstring still says the loss "saves its gradient in a fresh buffer". That is no longer true; see the inline comment.
Should consider
torchtitan/components/loss.py:107-- the class docstring still opens with "Replacestorch.distributed.tensor.parallel.loss_parallel()". @tianyu-l saidloss_parallel()is no longer used in titan, so this reference could go too.
Questions
- (maintainers, still open) Should the TP=1
F.cross_entropypath (loss.py:83) and the batch-invariant RL path get the same fix? Today TP=1 and TP>1 give different results when a target is overconfident.
Not verified: the eager time and memory of the recompute-only form, and the result of the Integration run still in progress.
| logit at the fallback label 0 is -inf, and the input logits are not | ||
| modified (the loss saves its gradient in a fresh buffer). |
There was a problem hiding this comment.
This parenthetical describes the eager in-place form that was removed. Forward now saves the input logits and the [T, 1] row stats, and backward builds a fresh gradient. Drop it:
| logit at the fallback label 0 is -inf, and the input logits are not | |
| modified (the loss saves its gradient in a fresh buffer). | |
| logit at the fallback label 0 is -inf, and the input logits are not | |
| modified. |
This PR does two things for TP cross entropy calculation:
use
F.softplusfor ce calculation to address the softmax collapse issue when the model is overconfidentmove softmax-onehot calculate to fwd pass; as a result, we don't need to redo softmax-related ops in bwd, and buffer movement, which would improve the customized eager TP autograd perf.
It changes numerics for TP > 1 case, since we are using a different implementations of softmax + cross_entropy
The numbers are as follows: (T=16384, V=128256, TP=2)
Convergence on Qwen 0.6B with/without compile