From 78eff4e2072f562a7e080b8253afee295f07c55d Mon Sep 17 00:00:00 2001 From: William Date: Wed, 12 Aug 2026 00:19:12 -0400 Subject: [PATCH] [PyTorch] Decline fused grouped MLP when the backward format is not E4M3 The fused grouped MLP packs the incoming activation gradient by reinterpreting its storage as E4M3, conditioned only on NVFP4 and never on the FP8 format. Under MXFP8BlockScaling(fp8_format=Format.HYBRID) the backward quantizers emit E5M2, so those bytes are read as the wrong format rather than converted, and every gradient out of the fusion is wrong. The forward pass is unaffected, so this shows up as a model that trains too slowly instead of one that fails. Fall back to the unfused ops when the recipe's backward format is not E4M3, and raise instead of reinterpreting if such a gradient reaches the kernel path. Signed-off-by: William --- tests/pytorch/test_grouped_mlp.py | 30 +++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 16 ++++++++-- 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 499c22509c..c608759a02 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -899,6 +899,36 @@ def train_step( class TestGroupedMLPFusedOp: """Tests for grouped MLP fused op""" + def test_fusion_declined_for_e5m2_grad_output(self, monkeypatch) -> None: + """MXFP8 with Format.HYBRID must fall back to the unfused ops. + + The fused backward reinterprets the grad output's storage as E4M3, so an + E5M2 backward format would be misread rather than converted. + """ + from transformer_engine.common.recipe import Format, MXFP8BlockScaling + + fused_op_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMGLU + monkeypatch.setattr(fused_op_cls, "is_supported", classmethod(lambda cls: True)) + + # Never matches the fusion pattern, so the window scan leaves it alone. That keeps + # this test on the dispatch logic and off the kernels, which need SM100. + ops = [object(), object(), object()] + + # Declined before the window is scanned, so the original list comes straight back. + hybrid = MXFP8BlockScaling(fp8_format=Format.HYBRID) + assert ( + grouped_mlp_module.fuse_grouped_mlp_ops(ops, recipe=hybrid, fused_op_cls=fused_op_cls) + is ops + ) + + # E4M3 gets past the guard and is rebuilt by the scan. + e4m3 = MXFP8BlockScaling(fp8_format=Format.E4M3) + rebuilt = grouped_mlp_module.fuse_grouped_mlp_ops( + ops, recipe=e4m3, fused_op_cls=fused_op_cls + ) + assert rebuilt is not ops + assert rebuilt == ops + @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list) @pytest.mark.parametrize("single_grouped_weight", (False, True)) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 909e5a8a9b..31489d3bd9 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -16,7 +16,7 @@ from packaging.version import Version as PkgVersion import transformer_engine_torch as tex -from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType +from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed_weight import ( @@ -26,7 +26,7 @@ finalize_weight_grads, ) from ...module.base import _2X_ACC_WGRAD -from ...quantization import Recipe +from ...quantization import Recipe, get_fp8_torch_dtype from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer from ...tensor.grouped_tensor import GroupedTensor from ...tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor @@ -817,6 +817,11 @@ def fuse_grouped_mlp_ops( # NVFP4 fused grouped MLP uses graph-safe grouped quantize, which currently requires RHT. if recipe.nvfp4() and recipe.disable_rht: return ops + # The fused backward reinterprets the grad output's storage as E4M3, so a recipe with an + # E5M2 backward format would have its gradients misread rather than converted. NVFP4 pins + # fp8_format to E4M3, so in practice this declines MXFP8 with Format.HYBRID. + if get_fp8_torch_dtype(recipe, fprop_tensor=False) != torch.float8_e4m3fn: + return ops if activation_op_types is None: activation_op_types = (ScaledSwiGLU, ScaledClampedQGeGLU) @@ -1923,6 +1928,13 @@ def fuser_backward( or isinstance(fc1_weight_param, NVFP4Tensor) or isinstance(fc2_weight_param, NVFP4Tensor) ) + if not use_nvfp4 and fc2_grad_output_quantizer.dtype != DType.kFloat8E4M3: + # The pack below reinterprets the grad output's storage as E4M3 rather than + # converting it, so anything else would be read as the wrong format. + raise RuntimeError( + "Fused grouped MLP backward requires an E4M3 grad output, but the recipe " + f"produced {fc2_grad_output_quantizer.dtype}." + ) data_dtype = torch.float4_e2m1fn_x2 if use_nvfp4 else torch.float8_e4m3fn scale_view_dtype = torch.float8_e4m3fn if use_nvfp4 else torch.float8_e8m0fnu sf_vec_size = NVFP4_BLOCK_SCALING_SIZE if use_nvfp4 else MXFP8_BLOCK_SCALING_SIZE