From ef78f78903b20b90ee5aaa49951b4a09048df63d Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Thu, 8 Oct 2026 14:27:17 +0300 Subject: [PATCH 1/2] Guard zero-gradient blocks in the default FP8 block-scaling recipe Set the default backward-gradient `amax_epsilon` in `Float8BlockScaling` to `1e-12`. This gives zero-amax E5M2 gradient blocks a nonzero floor in the recipe, instead of requiring downstream projects to override the recipe instance. Forward input and weight quantization defaults remain unchanged. - Use `amax_epsilon=1e-12` for the default blockwise FP8 backward-gradient quantizer. - Document the backward-only default. - Test zero-gradient quantization and Linear backward, and compare ordinary nonzero-block quantization bitwise with the previous epsilon-zero setting. - The focused TE GPU test passed on each of four H20 nodes using the patched installed Python recipe. - Both power-of-two and FP32 scale modes passed focused tests. - Megatron-LM's TP=2 blockwise FP8 determinism test passed after removing its recipe override. Signed-off-by: Hollow Man --- .../test_float8_blockwise_scaling_exact.py | 56 +++++++++++++++++++ tests/pytorch/test_recipe.py | 3 + transformer_engine/common/recipe/__init__.py | 6 +- 3 files changed, 64 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_float8_blockwise_scaling_exact.py b/tests/pytorch/test_float8_blockwise_scaling_exact.py index 09f3986ad06..e03dcc57964 100644 --- a/tests/pytorch/test_float8_blockwise_scaling_exact.py +++ b/tests/pytorch/test_float8_blockwise_scaling_exact.py @@ -377,6 +377,62 @@ def test_quantization_block_tiling_extrema_versus_reference( ) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("use_f32_scales", [False, True]) +def test_default_backward_quantizer_handles_zero_blocks(use_f32_scales): + """The default backward recipe must quantize zero E5M2 blocks with finite scales.""" + if recipe_emulated and use_f32_scales: + pytest.skip("Blackwell FP8 block scaling emulation requires power-of-two scales") + + from transformer_engine.common.recipe import Format + from transformer_engine.pytorch.quantization import Float8BlockScalingRecipeState + + recipe = Float8BlockScaling(fp8_format=Format.HYBRID, use_f32_scales=use_f32_scales) + quantizer = Float8BlockScalingRecipeState( + recipe, mode="backward", num_quantizers=1 + ).make_quantizers()[0] + assert quantizer.amax_epsilon == 1e-12 + assert quantizer.force_pow_2_scales == (not use_f32_scales) + + grad = torch.zeros((128, 128), dtype=torch.bfloat16, device="cuda") + quantized = quantizer.update_quantized( + grad, quantizer.make_empty(grad.shape, dtype=grad.dtype, device=grad.device) + ) + assert not torch.any(quantized._rowwise_data) + assert torch.isfinite(quantized._rowwise_scale_inv).all() + assert torch.all(quantized._rowwise_scale_inv > 0) + + # Nonzero blocks keep the pre-change quantization result. + old_quantizer = Float8BlockQuantizer( + fp8_dtype=TE_DType[torch.float8_e5m2], + rowwise=True, + columnwise=True, + amax_epsilon=0.0, + force_pow_2_scales=not use_f32_scales, + block_scaling_dim=recipe.grad_block_scaling_dim, + ) + grad.fill_(1) + quantized = quantizer.update_quantized( + grad, quantizer.make_empty(grad.shape, dtype=grad.dtype, device=grad.device) + ) + old_quantized = old_quantizer.update_quantized( + grad, old_quantizer.make_empty(grad.shape, dtype=grad.dtype, device=grad.device) + ) + assert torch.equal(quantized._rowwise_data, old_quantized._rowwise_data) + assert torch.equal(quantized._rowwise_scale_inv, old_quantized._rowwise_scale_inv) + + # A zero upstream gradient must give finite, zero weight and input gradients. + linear = te.Linear(128, 128, bias=False, params_dtype=torch.bfloat16, device="cuda") + grad_input = torch.ones_like(grad, requires_grad=True) + with te.autocast(recipe=recipe): + output = linear(grad_input) + output.backward(torch.zeros_like(output)) + assert torch.isfinite(linear.weight.grad).all() + assert torch.isfinite(grad_input.grad).all() + assert not torch.any(linear.weight.grad) + assert not torch.any(grad_input.grad) + + # FP8 per tesnor current scaling @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) class TestFP8BlockScalingRecipeLinear(TestFP8RecipeLinearBase): diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 446bfda8e59..e0551ba51f4 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -804,6 +804,9 @@ def test_float8_block_scaling_quant_params_follow_constructor(use_f32_scales): recipe.fp8_quant_bwd_grad, ): assert qparams.power_2_scale == (not use_f32_scales) + assert recipe.fp8_quant_fwd_inp.amax_epsilon == 0.0 + assert recipe.fp8_quant_fwd_weight.amax_epsilon == 0.0 + assert recipe.fp8_quant_bwd_grad.amax_epsilon == 1e-12 @pytest.mark.parametrize( diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index f9abf9de87a..79ae34a6b94 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -432,6 +432,10 @@ class Float8BlockScaling(Recipe): `high_precision` keeps original high-precision operands for backward, and `dequantized` dequantizes saved operands to the active high-precision compute dtype (e.g. BF16/FP16/FP32) for backward. + + The backward gradient quantizer uses an ``amax_epsilon`` of ``1e-12`` so + all-zero gradient blocks have a finite scale. Forward quantizers keep their + existing zero-epsilon behavior. """ use_f32_scales: bool = os.getenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "0") == "1" @@ -472,7 +476,7 @@ def __post_init__(self) -> None: ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." self.fp8_quant_fwd_inp = QParams(power_2_scale=not self.use_f32_scales, amax_epsilon=0.0) self.fp8_quant_fwd_weight = QParams(power_2_scale=not self.use_f32_scales, amax_epsilon=0.0) - self.fp8_quant_bwd_grad = QParams(power_2_scale=not self.use_f32_scales, amax_epsilon=0.0) + self.fp8_quant_bwd_grad = QParams(power_2_scale=not self.use_f32_scales, amax_epsilon=1e-12) def __setstate__(self, state) -> None: self.__dict__.update(state) From afc53c2fec34a6c172ef192e712f94abc9144036 Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Thu, 8 Oct 2026 14:43:54 +0300 Subject: [PATCH 2/2] address review Signed-off-by: Hollow Man --- .../test_float8_blockwise_scaling_exact.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_float8_blockwise_scaling_exact.py b/tests/pytorch/test_float8_blockwise_scaling_exact.py index e03dcc57964..42ecf4e508e 100644 --- a/tests/pytorch/test_float8_blockwise_scaling_exact.py +++ b/tests/pytorch/test_float8_blockwise_scaling_exact.py @@ -380,7 +380,7 @@ def test_quantization_block_tiling_extrema_versus_reference( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize("use_f32_scales", [False, True]) def test_default_backward_quantizer_handles_zero_blocks(use_f32_scales): - """The default backward recipe must quantize zero E5M2 blocks with finite scales.""" + """Zero E5M2 blocks must use the backward epsilon in both scaling modes.""" if recipe_emulated and use_f32_scales: pytest.skip("Blackwell FP8 block scaling emulation requires power-of-two scales") @@ -399,8 +399,17 @@ def test_default_backward_quantizer_handles_zero_blocks(use_f32_scales): grad, quantizer.make_empty(grad.shape, dtype=grad.dtype, device=grad.device) ) assert not torch.any(quantized._rowwise_data) - assert torch.isfinite(quantized._rowwise_scale_inv).all() - assert torch.all(quantized._rowwise_scale_inv > 0) + scale = torch.finfo(torch.float8_e5m2).max / 1e-12 + if not use_f32_scales: + scale = 2.0 ** math.floor(math.log2(scale)) + expected_scale_inv = 1.0 / scale + for scale_inv in (quantized._rowwise_scale_inv, quantized._columnwise_scale_inv): + torch.testing.assert_close( + scale_inv, + torch.full_like(scale_inv, expected_scale_inv), + rtol=1e-6, + atol=0.0, + ) # Nonzero blocks keep the pre-change quantization result. old_quantizer = Float8BlockQuantizer(