Guard zero-gradient blocks in the default FP8 block-scaling recipe

tests/pytorch/test_float8_blockwise_scaling_exact.py

@@ -377,6 +377,71 @@ 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):
+    """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")
+
+    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)
+    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(
+        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):

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(

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)