Skip to content

[PyTorch] Optionally derive FC1 dy NVFP4 amaxes from the dGLU epilogue - #3755

Open
sraman-rgb wants to merge 3 commits into
NVIDIA:mainfrom
sraman-rgb:sraman/nvfp4-fc1-dy-amax-from-dglu
Open

sraman-rgb wants to merge 3 commits into
NVIDIA:mainfrom
sraman-rgb:sraman/nvfp4-fc1-dy-amax-from-dglu

Conversation

@sraman-rgb

Copy link
Copy Markdown
Collaborator

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

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

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>
@greptile-apps

greptile-apps Bot commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerView in GreptileConfidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable defect was established.

Summary

Adds an opt-in path that reuses the dGLU epilogue’s maxima when quantizing FC1 gradients.

  • FC1 gradient quantization can reuse the FC2 dGLU amax values.

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
Loading

Reviews (3) · Last reviewed commit: "Merge remote-tracking branch 'upstream/m..." · Reviewed by Greptile

@ksivaman ksivaman self-assigned this Nov 2, 2026
sraman-rgb and others added 2 commits November 10, 2026 14:33
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

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants