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