Guard blockwise FP8 backward (E5M2) wgrad against all-zero gradient blocks by HollowMan6 · Pull Request #7523 · NVIDIA/Megatron-LM

HollowMan6 · 2026-09-20T00:22:59Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Guard the blockwise FP8 recipe's backward (E5M2) weight-gradient quantization against all-zero gradient blocks, where amax = 0 gives an infinite scale and NaN gradients.

Split out of #7054 as discussed there: a small fp8_utils.py fix that is independent of the GLM-5.3-Flash model work. Sibling PRs: #7521 (mHC precision controls) and #7522 (KPool DSA indexer); #7054 now carries only the KDA layers.

Changes

Float8BlockScaling is now constructed with an explicit fp8_quant_bwd_grad=QParams(power_2_scale=<recipe default>, amax_epsilon=1e-12). amax_epsilon floors the per-block amax, so a gradient block that is entirely zero gets a finite scale instead of inf; blocks with a non-zero amax are unaffected, and the forward quantization parameters are unchanged. power_2_scale is read from the recipe class default so the guard follows Transformer Engine's setting.

Motivation

With fp8_recipe=blockwise, a weight-gradient block that is entirely zero has amax = 0, so its E5M2 scale becomes inf and the quantized gradient turns into NaN, which then poisons the whole weight gradient. This came up in the GLM-5.3-Flash FP8 work in #7054 but is not specific to it.

Testing

No unit test added: the change is a recipe parameter and only affects blocks whose gradient amax is exactly zero (previously NaN).

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue: none (small fix split out of #7054)

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

copy-pr-bot · 2026-09-20T00:23:03Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

HollowMan6 · 2026-09-21T07:40:27Z

Thank you @Phlip79 ! Please let me know if there are any other comments, thanks!

Phlip79 · 2026-09-22T17:29:54Z

/ok to test f7ea550

jaredcasper

What tests have been done to ensure this doesn't impact other training runs in fp8? What, if any, is the speed impact of this?

jaredcasper · 2026-09-23T21:37:53Z

+                block_scaling_recipe_cls = transformer_engine.common.recipe.Float8BlockScaling
+                fp8_recipe = block_scaling_recipe_cls(fp8_format=fp8_format)
+                bwd_qparams = block_scaling_recipe_cls.fp8_quant_bwd_grad
+                fp8_recipe.fp8_quant_bwd_grad = transformer_engine.common.recipe.QParams(


This monkey patching TE? Shouldn't this just be fixed in TE?

Yes, the current code updates only the recipe instance to make the guard effective Bbut it's only related to the default recipe choice so I think eventually we should be flexible to customize the recipe. To avoid the monkey patch, TE ideally need to accept additional fp8_quant_bwd_grad for its transformer_engine.common.recipe.Float8BlockScaling, so that something like this can work: f7ea550

Sorry for delay. So why not fix this in TE and avoid the monkey patch? If it's just the default recipe then fix the default recipe in TE? Adding the flexibility can come later if needed.

Thank you! I opened NVIDIA/TransformerEngine#3651 on the TE, I prefer keeping this PR opened until that one is discussed/merged/released for other's reference. Thanks again.

HollowMan6 · 2026-09-23T22:17:31Z

What tests have been done to ensure this doesn't impact other training runs in fp8? What, if any, is the speed impact of this?

The change is restricted to the Float8BlockScaling recipe path. This modification only constructs the recipe and updates backward quantization parameters during initialization with Transformer Engine API, so I don't expect it will introduce impact on speed.

Signed-off-by: Hollow Man <hollowman@opensuse.org>

dimapihtar · 2026-10-06T12:39:46Z

/ok to test 2a32b86

zhongbozhu · 2026-10-08T20:19:50Z

+                fp8_recipe.fp8_quant_bwd_grad = transformer_engine.common.recipe.QParams(
+                    power_2_scale=bwd_qparams.power_2_scale,
+                    amax_epsilon=1e-12,
+                    random_hadamard_transform=bwd_qparams.random_hadamard_transform,


that's confusing code, because there is no hadamard transform to fp8 block scaling at all, so it's supposed to be False unless we are talking about nvfp4

zhongbozhu · 2026-10-08T20:20:53Z

+                # Guard backward (E5M2) wgrad against all-zero gradient blocks:
+                # amax=0 → scale=inf → NaN.
+                block_scaling_recipe_cls = transformer_engine.common.recipe.Float8BlockScaling
+                fp8_recipe = block_scaling_recipe_cls(fp8_format=fp8_format)


I don't think we should use E5M2 for backward pass whenever we are using block scaling recipe, E5M2 is not gonna be useful for this recipe, so what was the motivation for using E5M2 in the first place?

zhongbozhu · 2026-10-08T20:33:25Z

Guard the blockwise FP8 recipe's backward (E5M2) weight-gradient quantization against all-zero gradient blocks, where amax = 0 gives an infinite scale and NaN gradients.

This should be fixed TE agreed, although I don't think we should use E5M2 in the first place.

Also it's a bit surprising to me this issue would happen because there is already a guard in the kernel to take care of this: https://github.com/NVIDIA/TransformerEngine/blob/62b13895c2e6edabfae7cfe6d2ae761358503c3d/transformer_engine/common/recipe/recipe_common.cuh#L24 So I would be curious to understand why would this issue happen, thanks.

HollowMan6 · 2026-10-08T22:05:46Z

Thank you @zhongbozhu for pointing this out that there is already a guard in the kernel. For context, I was working on FP8 E2E RL support for GLM5.3 Flash in verl. I was referring to https://github.com/verl-project/verl-recipe/blob/4da0651ebfc7ddbf240e41bace809e05911a5526/low_precision/run_dapo_qwen3_moe_30b_megatron_fp8e2e.sh#L96-L99 (E5M2 backward + blockwise)

The NaN was kind of random and hard to say if there's a deterministic way to reproduce it, and right now I'm not sure if it was caused by some issues in the sibling PRs instead (#7521 / #7522 / #7054) and whether it has been fixed later in those PRs. I will try to remove the patches here/in TE and let you know if I found a deterministic way to reproduce it.

zhongbozhu · 2026-10-14T00:51:59Z

@HollowMan6 Hi, thanks for the pointer, it's an interesting choice for them to use E5M2, since my impression is that you do not need E5M2 in blockwise recipe at all, you just need E4m3, the E5M2 backward recipe should be only be used in a per-tensor based recipe, not blockwise recipe.

Also the fp8 blockwise recipe is only recommended on Hopper, and in Blackwell we should just do mxfp8, and with E4M3 for both fwd and bwd.

Faithful dev port of main PR NVIDIA#7523 at 2a32b86: construct the blockwise FP8 recipe with an explicit backward-gradient QParams whose amax_epsilon is 1e-12, so an all-zero E5M2 wgrad block gets a finite scale instead of inf and no longer turns the weight gradient into NaN. Forward quantization parameters and the other backward QParams fields keep the Transformer Engine recipe defaults. Add the blockwise case to the FP8 determinism test and a recipe check. Every line the PR adds or removes in fp8_utils.py and test_fp8_determinism.py is applied verbatim. Dev adaptations: - test_fp8_determinism.py: dev lacks main's Blackwell-only test_grouped_mxfp8_parameter_updates_replay_bit_exact (and its inspect and determinism.kernels.harness imports), so test_blockwise_recipe_guards_zero_backward_blocks is appended after dev's test_bit_exact_under_quantization instead of after that test. Not ported: none. Already on dev: none. Upstream-PR: NVIDIA#7523 Upstream-Base: daa0074 Upstream-Head: 2a32b86 Scope: faithful dev port; deviations are listed above Ported-by: Robin Zhang <robinz@nvidia.com> Signed-off-by: Hollow Man <hollowman@opensuse.org> Signed-off-by: Robin Zhang <robinz@nvidia.com>

HollowMan6 · 2026-10-25T18:32:37Z

Closing as it seems this defensive change is no longer needed.

nemo-automation-bot Bot added the community-request label Sep 20, 2026

HollowMan6 marked this pull request as ready for review September 20, 2026 00:30

HollowMan6 requested review from a team as code owners September 20, 2026 00:30

Copilot AI lite review requested due to automatic review settings September 20, 2026 00:30

Copilot AI reviewed Sep 20, 2026

svcnvidia-nemo-ci requested a review from a team September 20, 2026 00:31

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Sep 20, 2026

HollowMan6 force-pushed the fp8-e5m2-wgrad-guard branch from 5dfc4f9 to 7386024 Compare September 20, 2026 13:51

This was referenced Sep 20, 2026

[vllm] feat: Support LoRA for DeepSeek V4 & DeepSeek V4.1 & GLM5.3 Flash verl-project/verl#7483

Open

feat(model): add GLM-5.3-Flash support NVIDIA-NeMo/Megatron-Bridge#6044

Merged

Phlip79 approved these changes Sep 21, 2026

View reviewed changes

Comment thread megatron/core/fp8_utils.py Resolved

svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label Sep 21, 2026

HollowMan6 force-pushed the fp8-e5m2-wgrad-guard branch from 7386024 to 865a5cb Compare September 21, 2026 07:37

svcnvidia-nemo-ci removed the waiting-on-customer Waiting on the original author to respond label Sep 21, 2026

HollowMan6 force-pushed the fp8-e5m2-wgrad-guard branch from 865a5cb to f7ea550 Compare September 22, 2026 11:08

copy-pr-bot Bot temporarily deployed to public September 22, 2026 17:33 Inactive

HollowMan6 force-pushed the fp8-e5m2-wgrad-guard branch from f7ea550 to 2f68553 Compare September 22, 2026 21:59

Phlip79 added the 26.10 label Sep 22, 2026

jaredcasper reviewed Sep 23, 2026

View reviewed changes

kevalmorabia97 mentioned this pull request Sep 23, 2026

[3/3] Add GLM-5.3-Flash support to Megatron-Bridge PTQ and HF export NVIDIA/Model-Optimizer#2539

Draft

Phlip79 removed request for a team October 5, 2026 17:26

HollowMan6 added 2 commits October 5, 2026 21:43

Guard backward (E5M2) wgrad against all-zero gradient blocks

e20aa97

Signed-off-by: Hollow Man <hollowman@opensuse.org>

Add fp8 blockwise determinism test

2a32b86

Signed-off-by: Hollow Man <hollowman@opensuse.org>

HollowMan6 force-pushed the fp8-e5m2-wgrad-guard branch from 4d81cd0 to 2a32b86 Compare October 5, 2026 18:43

copy-pr-bot Bot temporarily deployed to test October 6, 2026 12:41 Inactive

copy-pr-bot Bot temporarily deployed to public October 6, 2026 12:45 Inactive

svcnvidia-nemo-ci added waiting-on-customer Waiting on the original author to respond and removed waiting-on-maintainers Waiting on maintainers to respond waiting-on-customer Waiting on the original author to respond labels Oct 6, 2026

zhongbozhu reviewed Oct 8, 2026

View reviewed changes

zhongbozhu requested changes Oct 8, 2026

View reviewed changes

svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label Oct 9, 2026

HollowMan6 mentioned this pull request Oct 10, 2026

Guard zero-gradient blocks in the default FP8 block-scaling recipe NVIDIA/TransformerEngine#3651

Closed

8 of 13 tasks

HollowMan6 marked this pull request as draft October 10, 2026 08:17

svcnvidia-nemo-ci removed the waiting-on-customer Waiting on the original author to respond label Oct 10, 2026

Phlip79 removed the 26.10 label Oct 13, 2026

cuichenx mentioned this pull request Oct 14, 2026

[tracking] Kimi K3: migrate Bridge model code to native Megatron-Core NVIDIA-NeMo/Megatron-Bridge#6451

Open

14 tasks

buptzyb mentioned this pull request Oct 21, 2026

[Dev] [DO NOT MERGE] GLM-5.3-Flash integration stack tracker #8400

Draft

HollowMan6 closed this Oct 25, 2026

HollowMan6 deleted the fp8-e5m2-wgrad-guard branch October 25, 2026 18:32

Conversation

HollowMan6 commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Changes

Motivation

Testing

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Sep 20, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

HollowMan6 commented Sep 21, 2026

Uh oh!

Phlip79 commented Sep 22, 2026

Uh oh!

jaredcasper left a comment

Choose a reason for hiding this comment

Uh oh!

jaredcasper Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

HollowMan6 Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

jaredcasper Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

HollowMan6 Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

HollowMan6 commented Sep 23, 2026

Uh oh!

dimapihtar commented Oct 6, 2026

Uh oh!

zhongbozhu Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

zhongbozhu Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

zhongbozhu commented Oct 8, 2026

Uh oh!

HollowMan6 commented Oct 8, 2026

Uh oh!

zhongbozhu commented Oct 14, 2026

Uh oh!

HollowMan6 commented Oct 25, 2026

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

HollowMan6 commented Sep 20, 2026 •

edited

Loading