Repository navigation
[PyTorch] Grouped MLP: emit FP32 scale gradient with no cast to BF16 - #3870
Open
sraman-rgb wants to merge 11 commits into
Open
sraman-rgb wants to merge 11 commits into
sraman-rgb wants to merge 11 commits into
Conversation
…2 scales The block-scaled fused grouped MLP backward rounds the FP32 scale (router probability) gradient to the compute dtype with `.to(dtype)`. When the scales input is FP32 (e.g. FP32 dispatched router probabilities), autograd then casts that BF16 gradient back to FP32: two cast kernels per layer and microbatch. Return float32(round(grad, compute dtype)) directly from one Triton kernel (`round_through_dtype`) when the scales are FP32. The values are bitwise identical to the previous `.to(bf16)` + autograd `.to(fp32)` pair; BF16/FP16 scales keep the old path. Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Contributor
 [Medium impact] The PR appears safe to merge. SummaryThe grouped MLP keeps FP32 scale gradients without rounding them to BF16 first.
Both earlier findings are addressed. No new actionable issues were found. GPU tests were not run. Reviews (9) · Last reviewed commit: "Revert properly" · Reviewed by Greptile |
vthumbe1503
reviewed
Nov 13, 2026
sraman-rgb
changed the title
Review feedback: the block-scaled fused grouped MLP rounded the FP32 scale (router probability) gradient from the cuDNN kernels to the compute dtype with .to(dtype), a leftover from when the scales were always BF16. With FP32 scales, autograd then cast it back to FP32, so the gradient lost precision for no reason and cost two cast kernels. Return the gradient in the dtype of the scales input instead, as the BF16 fused grouped MLP path already does: FP32 scales get the kernels' FP32 values unchanged (no cast), BF16 scales get one cast as before. This drops the round_through_dtype Triton kernel. The test checks the FP32 gradient is returned bitwise unchanged with no cast in the backward, and the BF16 case. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Siddhartha Raman S <sraman@nvidia.com>
…casts Review feedback: the scale gradient from the cuDNN kernels is already FP32, so return it as is (no helper; autograd casts it for lower-precision scales), and pass the scales to the cuDNN kernels in FP32 in the forward for every recipe, as the NVFP4 path already did, instead of rounding them to the compute dtype for MXFP8. Drops the cast-counting test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Siddhartha Raman S <sraman@nvidia.com>
Collaborator
|
/te-ci pytorch |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Siddhartha Raman S <sraman@nvidia.com>
sraman-rgb
dismissed stale reviews from ksivaman and vthumbe1503
via
November 14, 2026 02:57
7dc8e0c
sraman-rgb
changed the title
vthumbe1503
previously approved these changes
Nov 14, 2026
Collaborator
|
/te-ci pytorch |
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Ensure fc1_prob_tensor is converted to fp32 only when using NVFP4. Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Collaborator
|
/te-ci pytorch |
This branch has not been deployed
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.
The block-scaled fused grouped MLP backward rounds the FP32 scale (router probability) gradient to the compute dtype with
.to(dtype). When the scales input is FP32 (e.g. FP32 dispatched router probabilities), autograd then casts that BF16 gradient back to FP32: two cast kernels per layer and microbatch.Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: