Repository navigation
Fix the Top-K KD loss running out of memory at long sequences with large vocabularies - #3080
Conversation
…s of the logits TopLogitsKLLoss kept about five float32 copies of the student logits alive at the backward peak, which runs out of memory at long sequences with large vocabularies. The normalizers now run over row chunks, and a custom autograd function saves only the logits and writes their gradient chunk by chunk. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Chunked normalization and autograd modelopt/torch/distill/plugins/megatron.py |
New operations compute full-vocabulary log-sum-exp and local Top-K values in fp32 chunks. The custom autograd operation computes student normalization and selected logits, then computes gradients in chunks during backward, including tensor-parallel gradient reduction. |
Loss integration and validation modelopt/torch/distill/plugins/megatron.py, tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py, CHANGELOG.rst |
TopLogitsKLLoss uses the chunked operations. Tests cover reference gradients, peak memory, shared-offset invariance, and empty inputs. The changelog records the memory change and accuracy with a large common logit offset. |
Priority:
Estimated code review effort: 3 (Moderate) | ~25 minutes
Change: Bug fix
Sequence Diagram(s)
sequenceDiagram
participant TopLogitsKLLoss
participant ScaledLogSumExp
participant ChunkedAutograd
participant TensorParallelGroup
TopLogitsKLLoss->>ScaledLogSumExp: Compute teacher normalizer
TopLogitsKLLoss->>ChunkedAutograd: Compute student normalizer and selected logits
ChunkedAutograd->>TensorParallelGroup: Reduce gradients when tensor parallelism is active
ChunkedAutograd->>ChunkedAutograd: Compute student gradients in chunks during backward
Possibly related PRs
- NVIDIA/Model-Optimizer#2988: Changes tensor-parallel gradient scaling for
TopLogitsKLLoss, which the chunked backward path must preserve. - NVIDIA/Model-Optimizer#3072: Includes PR
#2988’s tensor-parallelTopLogitsKLLossgradient-scaling change.
Merge Risk: 
685af
-
Chunked normalization reduces memory use while preserving the loss definition. No merge-blocking issue was identified.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title names the main change: the Top-K KD loss in TopLogitsKLLoss ran out of memory at long sequences with large vocabularies, and the change computes normalizers in row chunks to cut peak memory.… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns PASS. The pull request changes only CHANGELOG.rst,modelopt/torch/distill/plugins/megatron.py, and a test file. The added Python code contains notorch.load(..., weights_only=False), `numpy.load…Full details: Title check
-
Explanation
The title names the main change: the Top-K KD loss in TopLogitsKLLoss ran out of memory at long sequences with large vocabularies, and the change computes normalizers in row chunks to cut peak memory. The raw summary and objectives both describe this chunked computation as the primary change, and the title is short, specific, and free of file lists or vague wording.
Full details: Security Anti-Patterns
-
Explanation
PASS. The pull request changes only
CHANGELOG.rst,modelopt/torch/distill/plugins/megatron.py, and a test file. The added Python code contains notorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# nosec. No dependency manifest changed. The onlyeval()match is the unchangedself._teacher_model.eval()call.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 6 functions across 2 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches: 💡 1
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #3080 +/- ##
==========================================
+ Coverage 73.57% 79.12% +5.54%
==========================================
Files 661 662 +1
Lines 75682 77569 +1887
==========================================
+ Hits 55685 61377 +5692
+ Misses 19997 16192 -3805
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
ChenhanYu
left a comment
ChenhanYu
left a comment
There was a problem hiding this comment.
Pensieve comment review using chenhany's profile (revision 7).
Changes requested: the custom normalizer backward loses numerical stability for logits with large common offsets.
Needs action:
- Fix
_ScaledLogSumExpAndGather.backwardusing centered normalization statistics, and add an offset-invariance gradient regression; see the inline finding. - Confirm repository utility reuse and current-head GPU/TP test results; repository navigation and test execution were unavailable.
No action needed:
- Existing reference-test edits preserve assertions and move backward inside the chunk-size patch, appropriately exercising chunked backward.
- The 78 added core-logic lines are under budget; no licensing changes are visible.
- No duplicate issue appeared in supplied discussion; bounded history cannot establish that older discussions contain none.
Review context was incomplete.
[warning] modelopt/torch/distill/plugins/megatron.py:468 — Reconstructing softmax as exp(x/T - lse) loses normalization information when the rounded fp32 log-sum-exp absorbs its small correction. For example, with two equal logits of 224 and T=1, lse rounds to 224, so this computes [1, 1] instead of [0.5, 0.5]. The previous centered exp/sum computation differentiates to [0.5, 0.5]. Thus the new backward can change gradients substantially under a common logit offset. Save the per-row maximum and centered exponential sum, then recompute exp(x/T - max)/sum in backward. Add a regression comparing gradients against the previous centered reference with large shared offsets.
Review run: 81450d9861de4623af804fd5c2cc7810
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM; no correctness issues found, and the justified existing-test edit strengthens chunked tensor-parallel gradient coverage but requires human sign-off.
Needs action:
- Sign off on the numerical-reference test edit in
tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py: backward now runs inside the chunk-size patch, preserving assertions while exercising multiple chunks.
Work relative to the row max: the normalizer and the Top-K values are both shifted by it, so it cancels out of the log-probs, and backward rebuilds the softmax from the shifted logits instead of exp(x - lse) with a rounded lse. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
|
@ChenhanYu thanks, the offset finding was real. It is fixed in c6d9e53.
On your second item, I ran this at the current head:
On utility reuse: nothing in the repo computes a chunked log-sum-exp with a custom backward. |
ChenhanYu
left a comment
ChenhanYu
left a comment
There was a problem hiding this comment.
Pensieve nudge review using chenhany's profile (revision 7).
Nudge: the previous numerical-stability issue is resolved, but full-file integration remains unverified from the bounded snapshot.
Needs action:
- Confirm the unchanged
TopLogitsKLLossintegration preserves tensor-parallel gradient scaling with the new backward collective; surrounding implementation was unavailable for review.
No action needed:
- ✔️ Review 5679420659's offset concern is resolved: forward and backward use centered statistics, with a new loss-and-gradient offset regression.
- Reply 6464613891 addresses utility reuse and reports 27 current-head tests passing on two GPUs.
- The existing reference-test edit preserves assertions and moves backward inside the chunk-size patch, appropriately exercising chunked backward.
- Duplicate checking used supplied history only; omitted discussions, repository navigation, and independent test execution were unavailable.
Review context was incomplete.
Review run: 0603102d69864116b994635b11504a88
|
@ChenhanYu Confirmed: TP gradient scaling is the same as on
|
kevalmorabia97
left a comment
kevalmorabia97
left a comment
There was a problem hiding this comment.
Pensieve approve review using kmorabia's profile (revision 2).
Approve: the chunked Top-K normalizers fix the OOM with unchanged loss and gradients. The earlier offset-stability and TP-gradient concerns are both resolved.
No action needed:
- ✔️ Resolved since the last review: the offset instability (centered forward and backward statistics, plus
test_topk_logits_kl_loss_offset_invariant) and the TP gradient scaling question. - The backward all-reduce of
grad_lsein_ScaledLogSumExpAndGathermatches the backward of the autograd all-reduce it replaces. Author reply 6466281874 shows the TP=2 reference test fails if that all-reduce is dropped or the gradient doubled. I read the code to confirm the collective; I did not run the tests. - The gradient
softmax * grad_lse + grad_vals, divided by T, is correct, and the scatter indices are unique Top-K entries. - The existing reference-test edit keeps its assertions and moves backward inside the chunk-size patch, so the chunk loop and TP backward are exercised.
- 81 added core lines, under budget.
modelopt/torch/is not a setup-owned path, and there are no licensing changes. - Prior reviews at this head used bounded history, so duplicate checking is limited to what was supplied.
Review run: e0bdcec9383a41e9876cb20fdbb76ee4
ChenhanYu
left a comment
ChenhanYu
left a comment
There was a problem hiding this comment.
Pensieve nudge review using chenhany's profile (revision 7).
Nudge: prior correctness concerns are resolved, but repository-wide verification remains unavailable in this bounded review.
Needs action:
- Confirm repository review requirements and surrounding integration at
c6d9e53, and sign off on the justified reference-test edit intests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py.
No action needed:
- ✔️ Resolved: review 5679420659's offset instability through centered statistics and a regression test; review 5680900716's TP-gradient question through reply 6466281874's collective explanation and reported mutation tests.
- The existing reference-test assertions remain intact; moving backward inside the chunk-size patch exercises chunked backward.
- No new concrete defect found. Reply 6464613891 addresses utility reuse and reports 27 passing current-head tests.
- Repository navigation and independent execution were unavailable. Duplicate checking was limited to supplied history, which may omit older discussions.
Review context was incomplete.
Review run: 390bacf139ac4b888f467a3251c8a428
ChenhanYu
left a comment
ChenhanYu
left a comment
There was a problem hiding this comment.
Pensieve nudge review using chenhany's profile (revision 7).
Nudge: prior correctness concerns are resolved, but the bounded snapshot leaves repository requirements and surrounding integration unverified.
Needs action:
- Confirm repository requirements and full-file integration at
c6d9e53, continuing review 5686210730; include human sign-off on the justified reference-test edit.
No action needed:
- ✔️ Resolved: review 5679420659’s offset instability through centered statistics and regression coverage; review 5680900716’s TP-gradient concern through the matching backward collective and reply 6466281874.
- Existing reference assertions remain intact; moving backward inside the chunk-size patch exercises chunked backward.
- Reply 6464613891 addresses utility reuse and reports 27 passing current-head GPU tests. Independent execution and repository navigation were unavailable.
- No new concrete defect found. Duplicate checking was limited to supplied history, which may omit older discussions.
Review context was incomplete.
Review run: cc920c2c15174d9d8cff540c6f84f663
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Changes requested: _scaled_logsumexp now rejects empty batch or sequence inputs that the previous loss implementation accepted.
Needs action:
- Fix empty-row shape restoration in
modelopt/torch/distill/plugins/megatron.pyand add forward/backward coverage; see the inline finding. - Correct the “identical loss and gradients” claim in
CHANGELOG.rst; the offset-stability fix intentionally changes inaccurate results. - Sign off on the justified TP-reference-test edit in
tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py.
No action needed:
- ✔️ Resolved since the previous reviews: offset instability, TP-gradient scaling, and the full-file integration question.
- The reference-test assertions remain intact; moving backward inside the chunk-size patch strengthens coverage.
Reshape the normalizer results with explicit widths, since -1 cannot be inferred for zero-element tensors, and add forward/backward coverage. Reword the changelog entry: results now differ from before for logits with a large common offset. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM; prior correctness concerns are resolved, with human sign-off required only for the justified TP-reference-test edit.
Needs action:
- Sign off on the TP-reference-test edit in
tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py: assertions are unchanged, and backward now runs inside the chunk-size patch.
No action needed:
- ✔️ Resolved since the last review: empty-input forward/backward handling and the inaccurate changelog claim.
ChenhanYu
left a comment
ChenhanYu
left a comment
There was a problem hiding this comment.
Pensieve nudge review using chenhany's profile (revision 7).
Nudge: the supplied snapshot includes a later commit than the requested head, preventing an exact-commit review.
Needs action:
- Supply the diff from
580ffff1cddb9b20a38b7434e6290654ebf96a8fto requested headc6d9e538d9c8f7d6dc9cd3f6308d07bceaacf420; the supplied snapshot instead targets685afb47df7acbbf4c70e667ec1b8f4f60d337f4. - Confirm repository requirements and integration, continuing review 5686249770; obtain human sign-off on the justified TP-reference-test edit, which preserves assertions while exercising chunked backward.
No action needed:
- ✔️ The supplied later snapshot visibly addresses empty-input comment 4403422324 and changelog comment 4403422338. These fixes cannot establish resolution at the requested earlier head.
- Independent inspection found centered normalization and the matching backward gradient collective consistent with the intended calculation; no new concrete defect was identified in the supplied patch.
- Reviewer: Codex (GPT-6). Only the supplied snapshot and bounded discussion were inspected. Repository navigation, independent tests, and exact-target verification were unavailable; duplicate checking is limited to supplied history.
Review context was incomplete.
Review run: ff78be1906f64fea9edb0b2233a15377
What does this PR do?
Type of change: Bug fix
Since #2459,
TopLogitsKLLoss(--logit_kl_top_k/logit_kl_topk) normalizes the student and teacher over the full vocabulary. It computes those normalizers on float32 copies of the logits, so about five float32 tensors of shape [sequence, vocab] are alive at the backward peak. With large vocabularies this runs out of memory: for Qwen3.6-35B-A3B (248K vocab), each copy is 30 GiB at a 32K sequence and 61 GiB at 64K, so QAD at the tutorial's settings no longer fit on a 276 GiB GB300.The normalizers now run over row chunks of the logits (1 GiB of float32 per chunk), relative to each row's max:
The loss is unchanged. The tensor-parallel collectives are the same (max and sum all-reduces in forward), plus an all-reduce of the [sequence, 1] normalizer gradient in backward, which replaces the autograd all-reduce's.
Usage
No API change.
Testing
Forward and backward on [4096, 1, 248320] bf16 logits with K = 4096 on an RTX 6000 Ada:
Loss within 8e-7 relative of before; 99.99% of gradient elements bitwise identical, the rest differ by one bf16 ulp.
New
test_topk_logits_kl_loss_memory: forward and backward stay under 1.5x the size of the logits (before this fix, about 10x).test_topk_logits_kl_loss_numerics_ghost_token_reference(TP=2, against a full-vocab reference, with Top-P, reverse KL and temperature) now runs with 2-row chunks, so it covers the chunk loop and the tensor-parallel backward.New
test_topk_logits_kl_loss_offset_invariant: adding 2**20 to the logits changes neither the loss nor the gradient. Working relative to the row max makes this exact; onmain, the rounded log-normalizer moves the loss by 3% and the gradient by 0.035 at that offset.tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py: 27 passed on 2 GPUs innemo:26.08.Qwen3.6-35B-A3B NVFP4 QAD at a 64K sequence (32 nodes × 4 GB300, TP1 EP8, GBS 256, top-4096 KD, full recomputation, with Add a fake-quantized weight cache for QAT and QAD #3034's weight cache): runs at 20.3 s/iteration with a peak of 197.7 GiB allocated, where it ran out of memory before.
Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Fixes the memory regression from #2459.
LogitsKLLoss(full-vocab KL, nologit_kl_topk) is unchanged.🤖 Generated with Claude Code
Summary by CodeRabbit