Skip to content

Guard blockwise FP8 backward (E5M2) wgrad against all-zero gradient blocks #7523

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
HollowMan6 wants to merge 2 commits into from
Closed

Guard blockwise FP8 backward (E5M2) wgrad against all-zero gradient blocks #7523

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e20aa97
Guard backward (E5M2) wgrad against all-zero gradient blocks
HollowMan6 Sep 11, 2026
2a32b86
Add fp8 blockwise determinism test
HollowMan6 Sep 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions megatron/core/fp8_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -797,8 +797,17 @@ def get_fp8_recipe(config: TransformerConfig):
fp8_format=fp8_format, fp8_dpa=config.fp8_dot_product_attention
)
elif config.fp8_recipe == Fp8Recipe.blockwise and is_te_min_version("2.3.0.dev0"):
fp8_recipe = transformer_engine.common.recipe.Float8BlockScaling(
fp8_format=fp8_format
# 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)

@zhongbozhu zhongbozhu Oct 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

bwd_qparams = block_scaling_recipe_cls.fp8_quant_bwd_grad
fp8_recipe.fp8_quant_bwd_grad = transformer_engine.common.recipe.QParams(

@jaredcasper jaredcasper Sep 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@HollowMan6 HollowMan6 Sep 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@jaredcasper jaredcasper Oct 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@HollowMan6 HollowMan6 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

power_2_scale=bwd_qparams.power_2_scale,
amax_epsilon=1e-12,
random_hadamard_transform=bwd_qparams.random_hadamard_transform,

@zhongbozhu zhongbozhu Oct 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

stochastic_rounding=bwd_qparams.stochastic_rounding,
fp4_2d_quantization=bwd_qparams.fp4_2d_quantization,
)
elif config.fp8_recipe == Fp8Recipe.mxfp8:
fp8_recipe = transformer_engine.common.recipe.MXFP8BlockScaling(
Expand Down
25 changes: 25 additions & 0 deletions tests/unit_tests/determinism/correctness/test_fp8_determinism.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
the runner's ``_reset_quantizer_state`` between runs
A and B (per-module ``fp8_meta`` carries amax across
forward passes).
* ``fp8-blockwise`` — blockwise FP8 scaling.
* ``fp8-mxfp8`` — Blackwell-only microscaling FP8; capability-skipped on Hopper.
* ``fp4-nvfp4`` — Blackwell-only NVFP4 block scaling; capability-skipped on Hopper.
"""
Expand All @@ -30,6 +31,7 @@
_QUANT_RECIPES = [
pytest.param({"fp8": "hybrid", "fp8_recipe": "tensorwise"}, id="fp8-tensorwise"),
pytest.param({"fp8": "hybrid", "fp8_recipe": "delayed"}, id="fp8-delayed"),
pytest.param({"fp8": "hybrid", "fp8_recipe": "blockwise"}, id="fp8-blockwise"),
pytest.param(
{"fp8": "hybrid", "fp8_recipe": "mxfp8"},
id="fp8-mxfp8",
Expand Down Expand Up @@ -99,3 +101,26 @@ def update(source):
backward=False,
what="grouped MXFP8 parameter update",
)

@pytest.mark.internal
def test_blockwise_recipe_guards_zero_backward_blocks(self):
"""The blockwise recipe must keep a finite scale for zero E5M2 wgrad blocks."""
from megatron.core.enums import Fp8Recipe
from megatron.core.fp8_utils import HAVE_TE, get_fp8_recipe, is_te_min_version
from megatron.core.transformer.transformer_config import TransformerConfig

if not HAVE_TE or not is_te_min_version("2.3.0.dev0"):
pytest.skip("blockwise FP8 recipe requires Transformer Engine >= 2.3")

config = TransformerConfig(
num_layers=1,
hidden_size=128,
ffn_hidden_size=256,
num_attention_heads=2,
fp8="hybrid",
fp8_recipe=Fp8Recipe.blockwise,
params_dtype=torch.bfloat16,
use_cpu_initialization=True,
)
recipe = get_fp8_recipe(config)
assert recipe.fp8_quant_bwd_grad.amax_epsilon == pytest.approx(1e-12)
Loading