Repository navigation
[PyTorch] Optionally derive FC1 dy NVFP4 amaxes from the dGLU epilogue - #3755
Open
sraman-rgb wants to merge 3 commits into
Open
sraman-rgb wants to merge 3 commits into
sraman-rgb wants to merge 3 commits into
Conversation
Opt-in with NVTE_NVFP4_FC1_DY_AMAX_FROM_DGLU=1. The rowwise amax is the exact dGLU epilogue amax rounded to BF16; the post-RHT columnwise amax is bounded by 4x the rowwise amax. Skips the separate grouped amax pass. 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; no actionable defect was established. SummaryAdds an opt-in path that reuses the dGLU epilogue’s maxima when quantizing FC1 gradients.
Diagram%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[FC2 backward and dGLU] --> B[BF16 FC1 gradient and maxima]
B --> C{Flag enabled and quantizer supported?}
C -->|Yes| D[Round maximum to BF16]
D --> E[Use rowwise maximum and columnwise bound]
C -->|No| F[Compute maxima from gradient]
E --> G[Quantize FC1 gradient]
F --> G
Reviews (3) · Last reviewed commit: "Merge remote-tracking branch 'upstream/m..." · Reviewed by Greptile |
Run the fused NVFP4 grouped SwiGLU MLP forward and backward with NVTE_NVFP4_FC1_DY_AMAX_FROM_DGLU off and on, over several experts including empty ones, clearing the cached flag around the test. Check that the new branch runs, that its per-expert rowwise amaxes equal the amaxes of the BF16 FC1 grad output, that the rowwise operand and the gradients that do not use the columnwise operand are unchanged, and that the FC1 wgrad error stays within 5% of the default path. Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…-amax-from-dglu Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> # Conflicts: # transformer_engine/pytorch/ops/fused/grouped_mlp.py
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.
Opt-in with NVTE_NVFP4_FC1_DY_AMAX_FROM_DGLU=1. The rowwise amax is the exact dGLU epilogue amax rounded to BF16; the post-RHT columnwise amax is bounded by 4x the rowwise amax. Skips the separate grouped amax pass.
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: