From a122d8aca4666e485e6539054bad059583726387 Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Sat, 17 Oct 2026 02:04:12 -0700 Subject: [PATCH 1/2] [PyTorch] Quantize distributed BF16 weights in the fused grouped MLP A distributed (e.g. GTP) weight whose shards are not natively quantized is re-gathered as BF16. The fused grouped-MLP op assumed gathered weights were already quantized: forward quantized the shards before the gather, and backward read _columnwise_data / the NVFP4 amax straight off the gathered tensor, which fails for a module whose per-module recipe (e.g. MXFP8 experts under a global NVFP4 recipe) does not create native-quantized params. Gather first, then quantize whatever is still unquantized with the op's per-expert weight quantizers, mirroring module/linear.py's handling of a re-gathered BF16 distributed weight. Native-quantized shards gather as quantized tensors and pass through unchanged. The FC2 backward re-gather moves ahead of the NVFP4 amax / global-scale computation so that path no longer reads un-gathered shards. The per-expert weight quantizer lists are built once in forward and saved on the FC1/FC2 contexts; backward quantizes the re-gathered weights with those rather than reading the live op. A later forward on the same module under a different recipe, or with quantization disabled, resets the ops' quantizers before the earlier backward runs (pipelined microbatches, or a module shared across recipe scopes), which would otherwise hit None or quantize in the wrong format. Regression tests: a fake routed-expert DistributedWeight group whose materialize hooks return BF16 tensors (or MXFP8 tensors for the native-quantized case) drives the CuTeDSL fused grouped MLP under per-module MXFP8, optionally inside a global NVFP4 autocast. Distinct power-of-two scales in the forward/backward hooks make the gathered weights observable: output, dgrad and the main_grad wgrads must match a plain reference bit-exactly up to those scales. A unit check confirms already-quantized gathered weights pass through the helper unchanged. A further case runs a second forward under no recipe or NVFP4 between an MXFP8 forward and its backward and requires bit-exact agreement with a reference. Fails on the unpatched op with AttributeError on _columnwise_data in backward; passes with the fix. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren --- qa/L0_pytorch_unittest/test.sh | 1 + ...test_ops_grouped_mlp_distributed_weight.py | 368 ++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 107 +++-- 3 files changed, 448 insertions(+), 28 deletions(-) create mode 100644 tests/pytorch/test_ops_grouped_mlp_distributed_weight.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f45e9369d1..ca3d365d9a 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -91,6 +91,7 @@ PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_L PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_linear_distributed_weight.xml $TE_PATH/tests/pytorch/test_ops_grouped_linear_distributed_weight.py || test_fail "test_ops_grouped_linear_distributed_weight.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_module_grouped_linear_distributed_weight.xml $TE_PATH/tests/pytorch/test_module_grouped_linear_distributed_weight.py || test_fail "test_module_grouped_linear_distributed_weight.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_mlp_distributed_weight.xml $TE_PATH/tests/pytorch/test_ops_grouped_mlp_distributed_weight.py || test_fail "test_ops_grouped_mlp_distributed_weight.py" if [ "$RET" -ne 0 ]; then echo "Error in the following test cases:$FAILED_CASES" diff --git a/tests/pytorch/test_ops_grouped_mlp_distributed_weight.py b/tests/pytorch/test_ops_grouped_mlp_distributed_weight.py new file mode 100644 index 0000000000..02fd143110 --- /dev/null +++ b/tests/pytorch/test_ops_grouped_mlp_distributed_weight.py @@ -0,0 +1,368 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DistributedWeight dispatch in the fused grouped MLP (``ops.fused.GroupedMLP_CuTeGEMM*``). + +``test_ops_grouped_linear_distributed_weight.py`` covers the unfused ``ops.GroupedLinear`` path. +The fused op has its own weight plumbing, so it is covered here with the same kind of in-repo fake: + + * A DistributedWeight whose shards are plain BF16 (e.g. a Megatron GTP expert under a per-module + MXFP8 recipe, which creates no native-quantized params) gathers as BF16. The fused op must + gather *before* quantizing in forward, and in backward must quantize the re-gathered BF16 + weight before the dgrad GEMMs read quantized storage. + * A DistributedWeight whose shards are native-quantized gathers as quantized tensors, which + must pass through unchanged (no re-quantization). + +The per-module recipe is MXFP8 entered inside a global NVFP4 autocast, mirroring how Megatron +applies a ``--te-precision-config-file`` override around a fused expert MLP. + +As in the unfused test, the fakes apply distinct power-of-two scales in the materialize hooks so +the gathered weights are observably the ones the GEMMs consumed. A power-of-two scale commutes +exactly through MXFP8 quantization (E8M0 block scales shift, E4M3 data is unchanged) and through +BF16 rounding, so the scale relations below hold bit-exactly. FC1 gets no forward scale because the +activation between the two GEMMs is nonlinear. +""" + +from contextlib import nullcontext + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import MXFP8Quantizer, MXFP8Tensor +from transformer_engine.pytorch.ops.fused.grouped_mlp import GroupedMLP_CuTeGEMMGLU +from transformer_engine.pytorch.ops._common import is_quantized_tensor + +from utils import dtype_tols, make_recipe + +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) + +# Distinct powers of two: each commutes exactly through quantization and BF16 rounding. +FC1_BWD_SCALE = 2.0 +FC2_FWD_SCALE = 4.0 +FC2_BWD_SCALE = 8.0 + + +def _mxfp8_quantize(tensor: torch.Tensor) -> MXFP8Tensor: + """Native-quantized stand-in: both usages, like a shard the module's own quantizer created.""" + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=True) + return quantizer(tensor) + + +class _FakeDistWeight(torch.nn.Parameter): + """Single-process fake DistributedWeight shard of a routed-expert group. + + Every shard of the group is one of these (as with a real implementer such as GTP); only the + leader (``weight0``) carries ``weight_list``/``_group`` and is called back by the op. + ``materialize_group_for_forward`` / ``materialize_group_for_backward`` return the whole + group's "gathered" weights scaled by ``fwd_scale`` / ``bwd_scale``: plain BF16 tensors (a + BF16 shard all-gather) or MXFP8 tensors when ``gather_quantized`` (a native-quantized shard + all-gather). ``finalize_group_grads`` models the main-grad contract: accumulate the wgrads + into each shard's ``main_grad`` in-place and return dummies that the op discards. + ``grad_buffer`` hands the wgrad GEMM a fresh full-shape buffer in the compute dtype. + """ + + is_distributed_weight = True + # The fused op only takes the distributed path for a routed-expert group leader. + is_routed_expert = True + + def _gathered(self, scale: float): + gathered = [w.detach() * scale for w in self._group] + if self.gather_quantized: + gathered = [_mxfp8_quantize(w) for w in gathered] + return gathered + + def materialize_group_for_forward(self): + self.calls["fwd"] += 1 + return self._gathered(self.fwd_scale) + + def materialize_group_for_backward(self, **kwargs): + self.calls["bwd"] += 1 + return self._gathered(self.bwd_scale) + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls["finalize"] += 1 + wl = list(wgrads) if isinstance(wgrads, (list, tuple)) else [wgrads] + for w, g in zip(self._group, wl): + w.main_grad.add_(g.to(w.main_grad.dtype)) + w.grad_added_to_main_grad = True + return [torch.zeros_like(g) for g in wl] + + def grad_buffer(self): + return torch.empty_like(self.data) + + +def _make_fake_dist_leader(op, num_gemms, *, fwd_scale, bwd_scale, gather_quantized): + """Turn ``op.weight{i}`` into fake distributed shards; return the leader (``weight0``).""" + shards = [] + for i in range(num_gemms): + shard = _FakeDistWeight(getattr(op, f"weight{i}").data) + shard.main_grad = torch.zeros_like(shard, dtype=torch.float32) + shard.grad_added_to_main_grad = False # DDP initializes this on every param + setattr(op, f"weight{i}", shard) + shards.append(shard) + leader = shards[0] + leader.calls = {"fwd": 0, "bwd": 0, "finalize": 0} + leader.fwd_scale = fwd_scale + leader.bwd_scale = bwd_scale + leader.gather_quantized = gather_quantized + leader._group = shards + leader.weight_list = shards # routed-expert group membership, asserted by the op + return leader + + +def _skip_unless_fused_mxfp8_supported(device: str, dtype: torch.dtype) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if not GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + if dtype != torch.bfloat16: + pytest.skip("Fused grouped MLP requires BF16 data") + + +def _make_grouped_mlp(group_size, hidden_size, split_alignment, glu_interleave_size, dtype, device): + fc1 = te.ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=False, + device=device, + dtype=dtype, + split_alignment=split_alignment, + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=False, + device=device, + dtype=dtype, + split_alignment=split_alignment, + ) + module = te.ops.Sequential( + fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2 + ) + return module, fc1, fc2 + + +def _assert_fused(module) -> None: + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], GroupedMLP_CuTeGEMMGLU) + + +@pytest.mark.parametrize("outer_recipe", (None, "nvfp4")) +@pytest.mark.parametrize("gather_quantized", (False, True)) +def test_fused_grouped_mlp_distributed_weight( + outer_recipe, + gather_quantized, + *, + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + glu_interleave_size: int = 32, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", +) -> None: + """Fused grouped MLP with distributed weights: gather-then-quantize in fwd and bwd. + + ``gather_quantized=False`` is the BF16-shard case (per-module MXFP8 without native-quantized + params); ``gather_quantized=True`` is the native-quantized-shard case, whose gathered tensors + must pass through. ``outer_recipe="nvfp4"`` wraps the per-module MXFP8 autocast in a global + NVFP4 one, as Megatron does for an overridden expert MLP. + """ + _skip_unless_fused_mxfp8_supported(device, dtype) + if outer_recipe == "nvfp4" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + torch.manual_seed(0) + + split_sizes = torch.tensor( + [split_alignment, 2 * split_alignment, split_alignment, 2 * split_alignment], + dtype=torch.int, + device=device, + ) + assert split_sizes.numel() == group_size + total_tokens = int(split_sizes.sum().item()) + + module, fc1, fc2 = _make_grouped_mlp( + group_size, hidden_size, split_alignment, glu_interleave_size, dtype, device + ) + reference, ref_fc1, ref_fc2 = _make_grouped_mlp( + group_size, hidden_size, split_alignment, glu_interleave_size, dtype, device + ) + reference.load_state_dict(module.state_dict()) + + fc1_leader = _make_fake_dist_leader( + fc1, group_size, fwd_scale=1.0, bwd_scale=FC1_BWD_SCALE, gather_quantized=gather_quantized + ) + fc2_leader = _make_fake_dist_leader( + fc2, + group_size, + fwd_scale=FC2_FWD_SCALE, + bwd_scale=FC2_BWD_SCALE, + gather_quantized=gather_quantized, + ) + + x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + x.requires_grad_(True) + ref_x = x.detach().clone().requires_grad_(True) + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + + mxfp8_recipe = make_recipe("mxfp8") + outer = make_recipe(outer_recipe) if outer_recipe is not None else None + + def _run(mod, inp): + outer_ctx = te.autocast(enabled=True, recipe=outer) if outer is not None else nullcontext() + # Per-module MXFP8 override entered inside the global autocast (Megatron's pattern). + with outer_ctx, te.autocast(enabled=True, recipe=mxfp8_recipe): + y = mod(inp, split_sizes, probs, split_sizes) + y.backward(dy) + return y.detach() + + out = _run(module, x) + ref_out = _run(reference, ref_x) + + # The CuTeDSL fused op was formed on both (not the unfused fallback). + _assert_fused(module) + _assert_fused(reference) + + # Every dispatch hook fired on both groups. + for leader in (fc1_leader, fc2_leader): + assert leader.calls["fwd"] > 0 and leader.calls["bwd"] > 0 and leader.calls["finalize"] > 0 + + # fwd used the gathered FC2 weights (scale commutes exactly through quantization). + torch.testing.assert_close(out, FC2_FWD_SCALE * ref_out, rtol=0, atol=0) + # dgrad used the re-gathered FC2 and FC1 weights. + torch.testing.assert_close(x.grad, FC1_BWD_SCALE * FC2_BWD_SCALE * ref_x.grad, rtol=0, atol=0) + # wgrads were reduce-scattered into main_grad: FC2's does not depend on the weight scales, + # FC1's sees the FC2-backward scale through the activation gradient. + tols = dtype_tols(dtype) + for i in range(group_size): + fc2_w, ref_fc2_w = getattr(fc2, f"weight{i}"), getattr(ref_fc2, f"weight{i}") + fc1_w, ref_fc1_w = getattr(fc1, f"weight{i}"), getattr(ref_fc1, f"weight{i}") + torch.testing.assert_close(fc2_w.main_grad.to(dtype), ref_fc2_w.grad, **tols) + torch.testing.assert_close( + fc1_w.main_grad.to(dtype), FC2_BWD_SCALE * ref_fc1_w.grad, **tols + ) + for w in (fc1_w, fc2_w): + assert w.grad_added_to_main_grad is True + assert w.grad is not None, "distributed weights receive a dummy .grad" + + +def test_quantize_unquantized_weights_passes_through_quantized( + *, + group_size: int = 2, + hidden_size: int = 128, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", +) -> None: + """Already-quantized gathered weights are returned as-is; plain ones are quantized.""" + # Imported here so the end-to-end test above still collects on a tree without the helper. + from transformer_engine.pytorch.ops.fused.grouped_mlp import _quantize_unquantized_weights + + _skip_unless_fused_mxfp8_supported(device, dtype) + torch.manual_seed(0) + module, fc1, fc2 = _make_grouped_mlp(group_size, hidden_size, 256, 32, dtype, device) + split_sizes = torch.tensor([256] * group_size, dtype=torch.int, device=device) + total_tokens = int(split_sizes.sum().item()) + x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + + # A forward initializes the op's per-expert weight quantizers for the recipe. + with te.autocast(enabled=True, recipe=make_recipe("mxfp8")): + module(x, split_sizes, probs, split_sizes) + pre_quantized = _mxfp8_quantize(fc2.weight0.detach()) + plain = fc2.weight1.detach() + quantizers = [fc2.get_quantizer("forward", 2 * idx + 1) for idx in range(group_size)] + out = _quantize_unquantized_weights( + [pre_quantized, plain], quantizers, rowwise=True, columnwise=True + ) + + assert out[0] is pre_quantized, "quantized gathered weight must pass through unchanged" + assert is_quantized_tensor(out[1]), "BF16 gathered weight must be quantized" + torch.testing.assert_close( + out[1]._rowwise_data, _mxfp8_quantize(plain)._rowwise_data, rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("later_recipe", (None, "nvfp4"), ids=["no-recipe", "nvfp4"]) +def test_fused_grouped_mlp_distributed_weight_backward_after_recipe_change( + later_recipe, + *, + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + glu_interleave_size: int = 32, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", +) -> None: + """Backward must quantize re-gathered BF16 weights with the quantizers saved in forward. + + A later forward under another recipe (or none) resets the ops' quantizers before the earlier + backward runs, as with pipelined microbatches or a module shared across recipe scopes. Read + from the live op, backward would hit ``None`` (``later_recipe=None``) or quantize the MXFP8 + graph's weights in the wrong format (``later_recipe="nvfp4"``). + """ + _skip_unless_fused_mxfp8_supported(device, dtype) + if later_recipe == "nvfp4" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + torch.manual_seed(0) + + split_sizes = torch.tensor( + [split_alignment, 2 * split_alignment, split_alignment, 2 * split_alignment], + dtype=torch.int, + device=device, + ) + assert split_sizes.numel() == group_size + total_tokens = int(split_sizes.sum().item()) + + module, fc1, fc2 = _make_grouped_mlp( + group_size, hidden_size, split_alignment, glu_interleave_size, dtype, device + ) + reference, ref_fc1, ref_fc2 = _make_grouped_mlp( + group_size, hidden_size, split_alignment, glu_interleave_size, dtype, device + ) + reference.load_state_dict(module.state_dict()) + for fc in (fc1, fc2, ref_fc1, ref_fc2): + # BF16 shards: the gathered weights are what backward has to quantize. + _make_fake_dist_leader(fc, group_size, fwd_scale=1.0, bwd_scale=1.0, gather_quantized=False) + + x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + x.requires_grad_(True) + ref_x = x.detach().clone().requires_grad_(True) + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + mxfp8_recipe = make_recipe("mxfp8") + + with te.autocast(enabled=True, recipe=mxfp8_recipe): + out = module(x, split_sizes, probs, split_sizes) + ref_out = reference(ref_x, split_sizes, probs, split_sizes) + _assert_fused(module) + ref_out.backward(dy) + + # Another forward on the same module before its backward, under a different recipe. + later_ctx = ( + te.autocast(enabled=True, recipe=make_recipe(later_recipe)) + if later_recipe is not None + else nullcontext() + ) + with torch.no_grad(), later_ctx: + module(x, split_sizes, probs, split_sizes) + if later_recipe is None: + assert fc1.get_quantizer("forward", 1) is None, "later forward did not reset quantizers" + else: + assert not isinstance(fc1.get_quantizer("forward", 1), MXFP8Quantizer) + + out.backward(dy) + + torch.testing.assert_close(out, ref_out, rtol=0, atol=0) + torch.testing.assert_close(x.grad, ref_x.grad, rtol=0, atol=0) + for i in range(group_size): + for fc, ref_fc in ((fc1, ref_fc1), (fc2, ref_fc2)): + w, ref_w = getattr(fc, f"weight{i}"), getattr(ref_fc, f"weight{i}") + torch.testing.assert_close(w.main_grad, ref_w.main_grad, rtol=0, atol=0) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index eeb5f43335..24343b7148 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -382,6 +382,34 @@ def _group_quantize_with_amax_for_grouped_mlp( ) +def _quantize_unquantized_weights( + weights: Iterable[Any], + quantizers: Sequence[Quantizer], + *, + rowwise: bool, + columnwise: bool, +) -> list: + """Quantize plain weights with the given per-expert weight quantizers. + + Quantized entries (native-quantized params, incl. native GTP shards gathered as quantized + tensors) pass through. A distributed (e.g. GTP) BF16 weight gathers as BF16 and must be + quantized with the layer's own quantizer, mirroring ``module/linear.py``'s handling of a + re-gathered BF16 distributed weight. + + ``quantizers`` must be the ones saved on the forward context, not read from the live op: + a later forward under a different (or no) recipe replaces or clears the op's quantizers + before an earlier backward runs. + """ + out = [] + for weight, quantizer in zip(weights, quantizers): + if is_quantized_tensor(weight): + out.append(weight) + continue + quantizer.set_usage(rowwise=rowwise, columnwise=columnwise) + out.append(quantizer(weight)) + return out + + def _nvfp4_amax( tensors: GroupedTensor | Iterable[NVFP4TensorStorage], *, @@ -1358,9 +1386,15 @@ def fuser_forward( # Quantizers fc1_input_quantizer = fc1_op.get_quantizer("forward", 0) fc1_weight_quantizer = fc1_op.get_quantizer("forward", 1) + fc1_weight_quantizers = [ + fc1_op.get_quantizer("forward", 2 * idx + 1) for idx in range(num_groups) + ] fc1_grad_output_quantizer = fc1_op.get_quantizer("backward", 0) fc2_input_quantizer = fc2_op.get_quantizer("forward", 0) fc2_weight_quantizer = fc2_op.get_quantizer("forward", 1) + fc2_weight_quantizers = [ + fc2_op.get_quantizer("forward", 2 * idx + 1) for idx in range(num_groups) + ] fc2_grad_output_quantizer = fc2_op.get_quantizer("backward", 0) # Extract split sizes from extra input @@ -1482,17 +1516,14 @@ def fuser_forward( ) else: fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] - quantized_fc1_weights = [] - for idx, weight in enumerate(fc1_weights): - quantizer = fc1_op.get_quantizer("forward", 2 * idx + 1) - if not is_quantized_tensor(weight): - quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - quantized_fc1_weights.append(quantizer(weight)) - else: - quantized_fc1_weights.append(weight) - grouped_fc1_weight = quantized_fc1_weights - if fc1_is_dist: - grouped_fc1_weight = materialize_weight_for_forward(grouped_fc1_weight) + # Distributed weights gather first (native-quantized shards gather quantized, BF16 + # shards gather as BF16); anything still unquantized is quantized here. + grouped_fc1_weight = _quantize_unquantized_weights( + materialize_weight_for_forward(fc1_weights) if fc1_is_dist else fc1_weights, + fc1_weight_quantizers, + rowwise=True, + columnwise=input_requires_grad, + ) # Prepare FC2 grouped weight tensor for fused kernels. if fc2_op.single_grouped_weight: @@ -1520,15 +1551,15 @@ def fuser_forward( ) else: fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] - quantized_fc2_weights = [] - for idx, weight in enumerate(fc2_weights): - quantizer = fc2_op.get_quantizer("forward", 2 * idx + 1) - if not is_quantized_tensor(weight): - quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - quantized_fc2_weights.append(quantizer(weight)) - else: - quantized_fc2_weights.append(weight) - grouped_fc2_weight = quantized_fc2_weights + # Distributed weights are gathered (and quantized if needed) after the FC1 kernel so + # the gather overlaps FC1 compute. + grouped_fc2_weight = ( + fc2_weights + if fc2_is_dist + else _quantize_unquantized_weights( + fc2_weights, fc2_weight_quantizers, rowwise=True, columnwise=input_requires_grad + ) + ) # Some wrapper-copy paths may drop grouped storage metadata; enforce defaults. if isinstance(grouped_fc1_weight, GroupedTensor) and not hasattr( @@ -1857,7 +1888,12 @@ def fuser_forward( activation_in = activation_in.view(in_shape[0], fc1_weight_shape[0]) if fc2_is_dist: - grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight) + grouped_fc2_weight = _quantize_unquantized_weights( + materialize_weight_for_forward(grouped_fc2_weight), + fc2_weight_quantizers, + rowwise=True, + columnwise=input_requires_grad, + ) if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( grouped_fc2_weight, "_with_gemm_swizzled_scales" ): @@ -2245,7 +2281,7 @@ def fuser_forward( ) fc1_ctx.input_quantizers = [fc1_input_quantizer] - fc1_ctx.weight_quantizers = [fc1_weight_quantizer] + fc1_ctx.weight_quantizers = fc1_weight_quantizers fc1_ctx.grad_output_quantizers = [fc1_grad_output_quantizer] fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad @@ -2254,7 +2290,7 @@ def fuser_forward( fc2_ctx.input_quantizers = [fc2_input_quantizer] fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] - fc2_ctx.weight_quantizers = [fc2_weight_quantizer] + fc2_ctx.weight_quantizers = fc2_weight_quantizers fc2_ctx.dtype = dtype fc2_ctx.input_requires_grad = input_requires_grad fc2_ctx.weight_requires_grad = weight_requires_grad @@ -2505,6 +2541,18 @@ def fuser_backward( scales_tensor = scales_f32.reshape(-1, 1, 1) dscales_tensor = torch.zeros_like(scales_tensor) + # Distributed weight (e.g. GTP): re-gather before any use (the NVFP4 global scale below + # reads the weights). BF16 shards gather as BF16 and are quantized with the weight + # quantizers saved in forward (the live op's may have been replaced or cleared since). + fc2_leader = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + if is_distributed_weight(fc2_leader): + grouped_fc2_weight = _quantize_unquantized_weights( + materialize_weight_for_backward(fc2_leader), + fc2_ctx.weight_quantizers, + rowwise=False, + columnwise=True, + ) + fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn if use_nvfp4: fc2_dy_global_scale = _nvfp4_global_scale( @@ -2588,10 +2636,6 @@ def fuser_backward( if self._pass_srelu_tanh_params: fc2_dactivation_kwargs.update(tanh_clamp_scale=self._cudnn_tanh_clamp_scale) - fc2_leader = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 - if is_distributed_weight(fc2_leader): - grouped_fc2_weight = materialize_weight_for_backward(fc2_leader) - if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM fc2_weight_for_gemm = grouped_fc2_weight.copy() @@ -2876,7 +2920,14 @@ def fuser_backward( fc1_leader = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 if is_distributed_weight(fc1_leader): - grouped_fc1_weight = materialize_weight_for_backward(fc1_leader) + # BF16 GTP shards gather as BF16: quantize with the weight quantizers saved in + # forward (the live op's may have been replaced or cleared since). + grouped_fc1_weight = _quantize_unquantized_weights( + materialize_weight_for_backward(fc1_leader), + fc1_ctx.weight_quantizers, + rowwise=False, + columnwise=True, + ) fc1_dgrad_sf_override = _nvfp4_sf_dtype_override(fc1_grad_output_quantizer) From cb1ffc4afd19aa6f18d4ee51505ab1d41395474a Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Sat, 17 Oct 2026 01:59:23 -0700 Subject: [PATCH 2/2] [PyTorch] Drop NVFP4 rowwise input copies from the fused grouped MLP saved state Every wgrad path reads only the columnwise tiles of the saved FC1/FC2 inputs: the cuDNN grouped wgrad kernel, the single-group dense GEMM (general_gemm with the NT layout unwraps both operands columnwise) and the generic grouped-GEMM fallback (NVFP4 forces a TN layout by taking the columnwise data for a non-transposed A and a transposed B, which is the wgrad case). The save path already pruned the rowwise data and scales for MXFP8 but kept them for NVFP4, citing the generic fallback, so each quantized NVFP4 input was held (and, under activation offloading, copied to host and back) twice. Prune them for NVFP4 as well, unconditionally. The amax tensors are untouched. Per decoder expert layer of a Nemotron hybrid proxy (83k tokens, hidden 3072) this removes a 121 MiB data and a 15 MiB scale tensor from the ~1.9 GiB saved set, with unchanged numerics on both the cuDNN kernel and the forced generic fallback (NVTE_DISABLE_CUTEDSL_WGRAD_FUSED_GROUPED_MLP=1). Adds a saved_tensors_hooks test asserting exactly one FP4 data and one scale tensor is saved per quantized input on the fused NVFP4 path, for group sizes 1 and 4, with token counts chosen so the inputs cannot be confused with weights in the count, and registers it in the L0 unit-test runner. Grouped storage is saved as raw buffers; the single-group dense path (cuDNN frontend >= 1.27) saves a GroupedTensor wrapper whole, so the test inspects the wrapper's buffers directly and requires its rowwise data and scales to be None. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren --- qa/L0_pytorch_unittest/test.sh | 1 + ...est_ops_grouped_mlp_nvfp4_saved_tensors.py | 116 ++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 17 +-- 3 files changed, 127 insertions(+), 7 deletions(-) create mode 100644 tests/pytorch/test_ops_grouped_mlp_nvfp4_saved_tensors.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index ca3d365d9a..28f26cc1ba 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -92,6 +92,7 @@ PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_L PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_module_grouped_linear_distributed_weight.xml $TE_PATH/tests/pytorch/test_module_grouped_linear_distributed_weight.py || test_fail "test_module_grouped_linear_distributed_weight.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_mlp_distributed_weight.xml $TE_PATH/tests/pytorch/test_ops_grouped_mlp_distributed_weight.py || test_fail "test_ops_grouped_mlp_distributed_weight.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_mlp_nvfp4_saved_tensors.xml $TE_PATH/tests/pytorch/test_ops_grouped_mlp_nvfp4_saved_tensors.py || test_fail "test_ops_grouped_mlp_nvfp4_saved_tensors.py" if [ "$RET" -ne 0 ]; then echo "Error in the following test cases:$FAILED_CASES" diff --git a/tests/pytorch/test_ops_grouped_mlp_nvfp4_saved_tensors.py b/tests/pytorch/test_ops_grouped_mlp_nvfp4_saved_tensors.py new file mode 100644 index 0000000000..97f63a96c8 --- /dev/null +++ b/tests/pytorch/test_ops_grouped_mlp_nvfp4_saved_tensors.py @@ -0,0 +1,116 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""The fused grouped MLP must not save the NVFP4 rowwise input copies for backward. + +Every wgrad path (the cuDNN grouped kernel, the single-group dense GEMM and the generic grouped +GEMM fallback) reads only the columnwise tiles (data and scales) of the saved FC1 and FC2 inputs, +so the rowwise copies are dead after forward. Saving them doubled the activation memory held (and +offloaded) for the quantized inputs. + +Checked through ``saved_tensors_hooks``: the saved set must hold exactly one FP4 data tensor and +one scale tensor per quantized input. Grouped storage is saved as its raw buffers; the +single-group dense path (cuDNN frontend >= 1.27) instead saves a ``GroupedTensor`` wrapper whose +buffers are inspected directly. Numerics are covered by ``test_grouped_mlp.py``. + +Weights require grad here: with frozen weights the input quantizers skip the columnwise copy and +the fused op's ``mark_grouped_tensor`` call rejects the input before the pruning is reached, on +this change and before it alike. +""" + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.ops.fused.grouped_mlp import GroupedMLP_CuTeGEMMGLU +from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor + +from utils import make_recipe + +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) + + +# (group_size, tokens_per_group): token counts chosen so the quantized inputs' byte sizes never +# coincide with a weight's, which the saved-tensor counting below relies on. +@pytest.mark.parametrize("group_size,tokens_per_group", ((1, 768), (4, 256))) +def test_fused_grouped_mlp_nvfp4_saves_columnwise_inputs_only( + group_size, + tokens_per_group, + *, + hidden_size: int = 256, + glu_interleave_size: int = 32, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", +) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + if not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + if not GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + torch.manual_seed(0) + + total_tokens = group_size * tokens_per_group + split_sizes = torch.tensor([tokens_per_group] * group_size, dtype=torch.int, device=device) + # The fused grouped MLP requires a split alignment divisible by 256 (the default is 128). + fc1 = te.ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=False, + device=device, + dtype=dtype, + split_alignment=256, + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=False, + device=device, + dtype=dtype, + split_alignment=256, + ) + module = te.ops.Sequential( + fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2 + ) + x = (torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5).requires_grad_() + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + + saved = [] + + def pack(t): + if isinstance(t, GroupedTensor): + # Dense single-group path: the wrapper is saved whole, so count its buffers. + assert t.rowwise_data is None, "rowwise FP4 input data was saved" + assert t.scale_inv is None, "rowwise FP4 input scales were saved" + for buf in (t.columnwise_data, t.columnwise_scale_inv): + if buf is not None: + saved.append((buf.dtype, buf.numel())) + elif isinstance(t, torch.Tensor): + saved.append((t.dtype, t.numel())) + return t + + with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t): + with te.autocast(enabled=True, recipe=make_recipe("nvfp4_rht")): + y = module(x, split_sizes, probs, split_sizes) + forward_ops = module._module_groups[0]._forward_ops + assert isinstance(forward_ops[0][0], GroupedMLP_CuTeGEMMGLU) + + # Byte counts that identify the quantized FC1/FC2 inputs (both are [tokens, hidden_size] in + # NVFP4: 2 elements per byte, one E4M3 scale per 16 elements). Weights have other sizes. + fp4_data_numel = total_tokens * hidden_size // 2 + fp4_scale_numel = total_tokens * hidden_size // 16 + n_data = sum(1 for d, n in saved if d == torch.uint8 and n == fp4_data_numel) + n_scale = sum(1 for d, n in saved if d == torch.uint8 and n == fp4_scale_numel) + # One columnwise data + one columnwise scale tensor per quantized input (FC1 x, FC2 x). + assert n_data == 2, f"expected 2 FP4 input data tensors saved (columnwise only), got {n_data}" + assert n_scale == 2, f"expected 2 FP4 input scale tensors saved, got {n_scale}" + + # Backward still runs on what was saved. + y.backward(dy) + assert x.grad is not None and torch.isfinite(x.grad).all() + for p in module.parameters(): + assert p.grad is not None and torch.isfinite(p.grad).all() diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 24343b7148..8b4a7f4f4e 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -2236,13 +2236,16 @@ def fuser_forward( ) saved_grouped_fc2_x = None if recompute_srelu_fc2_x else grouped_fc2_x - # MXFP8 wgrad only needs columnwise tiles. NVFP4 generic GEMM fallbacks - # need the full grouped tensor state, including rowwise data and amax. - if not use_nvfp4: - for grouped_fc_x in (saved_fc1_x, saved_grouped_fc2_x): - if isinstance(grouped_fc_x, (GroupedTensor, GroupedTensorStorage)): - grouped_fc_x.rowwise_data = None - grouped_fc_x.scale_inv = None + # Every wgrad path reads only the columnwise tiles of the saved inputs: the cuDNN + # grouped wgrad kernel, the single-group dense GEMM (general_gemm with the NT layout + # unwraps both operands columnwise) and the generic grouped-GEMM fallback (NVFP4 + # forces a TN layout by taking the columnwise data for a non-transposed A and a + # transposed B, which is the wgrad case). Drop the rowwise data and scales for MXFP8 + # and NVFP4 alike; the amax tensors are kept. + for grouped_fc_x in (saved_fc1_x, saved_grouped_fc2_x): + if isinstance(grouped_fc_x, (GroupedTensor, GroupedTensorStorage)): + grouped_fc_x.rowwise_data = None + grouped_fc_x.scale_inv = None if cpu_offloading: activation_tensors = [