perf(loss): defer the fp32 upcast in the distillation top-k path - #3496
perf(loss): defer the fp32 upcast in the distillation top-k path#3496tianyi-zhang-02 wants to merge 16 commits into
Conversation
get_distillation_topk_logprobs_from_logits upcast the whole [B, S, V] student logits to fp32 before selecting a branch, but the non-TP top-k gather reads only K columns out of it. gather-then-cast is equivalent to cast-then-gather, so the upcast is now applied on the paths that actually read the full vocabulary (zero_outside_topk, and the TP/CP paths whose chunked kernels allocate fp32 buffers from this tensor), and the remaining path upcasts the gathered [B, S, K] instead. Bitwise identical outputs and gradients on every branch. At B=1, V=151936, K=64 (the distillation_math.yaml values) the non-TP path drops from 4.06 GiB to 1.74 GiB peak at seq 2048, and the gather itself from 9.2 ms to 0.1 ms at seq 8192, since the previous form moved the full fp32 tensor through memory to keep 64 columns per position. Refs NVIDIA-NeMo#3495 Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
The gather backward scatters into the input, so with duplicate indices it would accumulate in bf16 rather than in the fp32 copy the previous code made. torch.topk never produces duplicates, so this is unreachable in practice, but the test should exercise the real shape rather than rely on randint happening not to collide. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
The comment claimed teacher_topk_indices are distinct per position. They are not: dtensor_policy_worker.py pads topk_indices with value=0, so padded sequence positions carry K copies of index 0 and the gather backward accumulates there rather than scattering. The bitwise equality holds anyway, but for a different reason -- DistillationLossFn masks those positions to exactly 0.0 before reduction, so the accumulated gradient is zero on both paths. That invariant is load-bearing and appeared nowhere in the PR. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
The existing test cannot fail if this change is reverted. Its reference is the old to(float32).gather(...) formulation, so asserting the two match holds either way -- it pins numerical equivalence, which is worth keeping, but proves nothing about the deferral. Add a test that spies on torch.Tensor.gather and asserts the vocab-wide tensor is still bf16 when it runs. Verified against a reverted copy of the function: bf16 with the change, float32 without, so it goes red on revert. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
| # Duplicate indices do occur -- padded sequence positions carry K | ||
| # copies of index 0, because dtensor_policy_worker.py pads | ||
| # topk_indices with value=0 -- so the gather backward accumulates | ||
| # there rather than scattering. The equality holds anyway, because | ||
| # DistillationLossFn masks those positions to exactly 0.0 before | ||
| # reduction, so the accumulated gradient is zero either way. |
There was a problem hiding this comment.
The loss also has an unmasked fallback (kl_loss = per_token_kl.mean()) when token_mask/sample_mask are absent, where padded duplicate-index positions would carry nonzero grads and the accumulation order would differ (bf16 post-rounding vs fp32 pre-rounding). Consider asserting that the masks are present in DistillationLossFn, so the invariant this comment relies on is enforced rather than assumed.
There was a problem hiding this comment.
Right, and that was the weak part of the comment.
DistillationLossFn fell back to kl_loss = per_token_kl.mean() when the masks were absent, and on that branch the padded positions get real gradient. Those are exactly the positions carrying K duplicate copies of index 0, since dtensor_policy_worker.py pads topk_indices with value=0. So the gather backward accumulates there instead of scattering, and it rounds in bf16 rather than fp32. That is the one case where deferring the upcast stops being bitwise equivalent. My comment claimed the equality held because the loss masks those positions — true only on the branch that has masks.
Enforced it the way you suggested: it raises now instead of falling back. Nothing reached the fallback anyway — distillation.py always sets token_mask, and setup_distillation_test_data sets both — so this turns a dead path into a loud one rather than changing behaviour. Added a CPU test that the raise fires for either mask missing. Put the silent mean() back and it goes red :)
| assert torch.equal(student.grad, reference.grad) | ||
|
|
||
|
|
||
| def test_distillation_topk_gathers_before_upcasting(monkeypatch): |
There was a problem hiding this comment.
Nit: An equivalent refactor to torch.gather(...) / index_select would keep the optimization but fail this test as a false alarm. Maybe note this in the docstring so a future refactorer knows to update the spy rather than suspect a real regression.
There was a problem hiding this comment.
Taken. The spy is bound to Tensor.gather specifically, so a rewrite to torch.gather or index_select would keep the deferral and still trip it. The docstring now says so, and points at updating the spy rather than reading it as a regression :)
|
/ok to review 006a4eb |
svcnvidia-nemo-ci
added
the
waiting-on-customer
…ies on Review pointed out that the new comment's invariant was assumed, not enforced. DistillationLossFn had an unmasked fallback (``kl_loss = per_token_kl.mean()``) when token_mask/sample_mask were absent. On that path the padded positions -- which carry K duplicate copies of index 0, since dtensor_policy_worker.py pads topk_indices with value=0 -- get real gradient, so the gather backward accumulates there. That is exactly the case where gathering in bf16 and then upcasting stops matching the old upcast-then-gather formulation, because the accumulation rounds in bf16 rather than fp32. Nothing reaches the fallback: distillation.py always sets token_mask, and the test fixture sets both. So it now raises instead, which makes the equivalence this PR relies on a property of the code rather than of its callers. Also notes on the spy test that it is coupled to ``Tensor.gather``, so an equivalent rewrite to ``torch.gather``/``index_select`` would fail it as a false alarm rather than a real regression. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
|
Thanks — the unmasked fallback point is right, and it was the weak part of that comment.
Enforced it as you suggested: the loss now raises instead of falling back. Nothing reached the fallback — Also took the docstring nit — the spy is coupled to |
svcnvidia-nemo-ci
removed
the
waiting-on-customer
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
…IA-NeMo#3496 NVIDIA-NeMo#3496 replaces DistillationLossFn's unmasked-mean fallback with a raise. Until it lands the fallback is live and must report the batch size; after it lands the branch is gone and the raise is the only correct behaviour. Asserting one of them unconditionally makes the two PRs fail as a pair while each passes alone. That is not hypothetical -- it is what happened when I merged the whole stack onto main to check exactly this. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
|
Head moved to Also: I re-measured on the 4090 and the description's memory table was wrong in our favour. It now reads 2.90 → 1.16 GiB at |
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
The existing equivalence test uses distinct indices everywhere and a uniform upstream gradient, so it never reaches the case the change is conditional on. gather's backward is a scatter-add, and padded positions carry K copies of index 0 -- dtensor_policy_worker.py pads the sequence axis with value=0. There the old formulation accumulates in fp32 and this one accumulates in bf16, which is a different number. The two agree only because DistillationLossFn masks those positions to exactly 0.0, so the incoming gradient is zero and there is nothing to accumulate. This checks both directions: masked agrees bitwise, unmasked does not. The second assertion is the point -- it is what pins the ValueError that replaced the unmasked-mean fallback, and it fails if the test ever stops reproducing the duplicate-index accumulation. Two ways to write this test so it proves nothing, both of which I hit first and both of which the docstring now names: randperm indices everywhere (no duplicates anywhere, so nothing accumulates), and a .sum() upstream gradient (log_softmax over K duplicated logits is uniform, and its gradient for an all-ones upstream is exactly zero, so the unmasked case agrees too). Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
…upcast The comment said those paths "feed chunked kernels that allocate their fp32 buffers from this tensor". That is not what happens: the forward upcasts each chunk itself, so it would take bf16 fine. The actual constraint is on the backward. ChunkedDistributedGatherLogprob.backward builds its gradient as torch.zeros_like(vocab_parallel_logits, dtype=torch.float32) and returns it, so a bf16 input makes autograd reject the fp32 gradient. Worth being exact about, because the wrong version makes the guard look wider than it is: narrowing it further is possible, it just means changing that autograd.Function rather than this line. Comment only. No behaviour change. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
8a57e3b to
87874cb
Compare
…IA-NeMo#3496 NVIDIA-NeMo#3496 replaces DistillationLossFn's unmasked-mean fallback with a raise. Until it lands the fallback is live and must report the batch size; after it lands the branch is gone and the raise is the only correct behaviour. Asserting one of them unconditionally makes the two PRs fail as a pair while each passes alone. That is not hypothetical -- it is what happened when I merged the whole stack onto main to check exactly this. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
The raise was an API change riding in a perf PR. DistillationLossDataDict already declares token_mask and sample_mask as required keys and every in-tree caller passes both, so the fallback it replaced was unreachable; what it actually guarded was a bf16 gradient-accumulation difference of order 1e-13 on positions the loss masks to zero. That is not worth turning a defensive branch into a crash, and it doubles the review surface of an otherwise mechanical change. The deferred fp32 upcast stands on its own and is what this PR is for. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
# Conflicts: # nemo_rl/distributed/model_utils.py
This reverts 4983328. The ValueError is @fujial-code's review request from 2026-08-18 -- "Consider asserting that the masks are present in DistillationLossFn, so the invariant this comment relies on is enforced rather than assumed" -- and 06954fe added it the next day in response. Removing it undid a reviewer's requested change. The reasoning for dropping it still holds in the abstract (the branch it replaces is unreachable given DistillationLossDataDict's required keys), but that is a conversation to have with the reviewer, not a unilateral revert of their ask. Keeps the merge's comment reconciliation: cp_sharder.shard_token_tensor( ..., fill=0) is now a second source of duplicate indices, so the comment names it alongside the padded-position case it already covered. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
|
Final H100 follow-up on SHA
This measures the isolated helper's allocator peak, not total end-to-end training HBM. cc @fujial-code — this is the final branch after your mask feedback. |
svcnvidia-nemo-ci
added
waiting-on-maintainers
4 participants
What does this PR do?
Defers the fp32 upcast in
get_distillation_topk_logprobs_from_logitsuntil after the non-TP top-k gather. This avoids materializing an fp32[B, S, V]tensor when onlyKlogits are used.The TP/CP and
zero_outside_topk=Truepaths keep the existing full upcast.DistillationLossFnnow requires both masks instead of falling back to an unmasked mean; every in-tree caller already supplies them, and the invariant is required for gradient equivalence at padded duplicate indices.Refs #3495.
logprob_chunk_sizeremains out of scope.Validation
Final SHA:
7534040209a051be3c7db0d0de8b582a065834ea, merged with currentmain(ccbcd4cc5).Environment: Runpod Secure Cloud,
nvcr.io/nvidia/nemo-rl:v0.7.0, 4× NVIDIA H100 80 GB HBM3 host (one GPU used per sample), driver580.126.09, Python3.13.14, PyTorch2.11.0+cu130.On a Runpod H100 80 GB, I ran the real helper in separate processes at
B=1,V=151,936,K=64, usingmain → candidate → candidate → mainfor each sequence length. Every repeat produced identical output and gradient hashes; outputs were fp32, gradients bf16, and all values were finite.These are isolated PyTorch allocator peak deltas (
max_memory_allocated - start_allocated), not an end-to-end training HBM measurement.115deselected).ruff checkandruff format --checkpassed on all changed files.