From 7638a3c75ee7d6c8b8120bca16664248ec076a30 Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Sat, 17 Oct 2026 23:57:35 -0700 Subject: [PATCH 1/3] [PyTorch] Keep DistributedWeight objects on ctx across saved-tensor hooks _Linear and _GroupedLinear save the DistributedWeight parameter itself for backward and gate their distributed-weight path on getting that object back (is_distributed_weight(saved_weight)). That holds only while no torch.autograd.graph.saved_tensors_hooks are installed: with hooks active, SavedVariable::unpack rebuilds a fresh plain tensor from the unpack hook's result, dropping the Python subclass and its is_distributed_weight marker. The backward then takes the plain-parameter branch, whose weakrefs point at the transient all-gathered weights, and fails with "weight was removed while fuse_wgrad_accumulation=True". Seen with Megatron-Core fine-grained activation offloading (which installs such hooks around expert fc1/act) on bf16 GTP-sharded GroupedLinear experts. Quantized weights were unaffected because their storage objects already stay on ctx via prepare_for_saving/restore_from_saved, and the op-fuser path reads the module's own weight attribute instead of the saved tensor. Keep the DistributedWeight objects as ordinary LinearBwdArgs / GroupedLinearBwdArgs / GroupedLinearGraphSafeBwdArgs references and prefer them over the saved tensors in backward, on both the split-quantize and the native grouped-tensor GroupedLinear paths. Add tests/pytorch/test_module_distributed_weight_saved_tensor_hooks.py, using an in-repo fake DistributedWeight implementer, covering Linear and both GroupedLinear paths with identity saved-tensor hooks installed. The grouped-tensor case skips when the native path is unsupported on the device or cuBLASLt rather than silently duplicating the split-quantize case. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Xiaowei Ren --- qa/L0_pytorch_unittest/test.sh | 1 + ...e_distributed_weight_saved_tensor_hooks.py | 237 ++++++++++++++++++ .../pytorch/module/grouped_linear.py | 17 ++ transformer_engine/pytorch/module/linear.py | 9 + 4 files changed, 264 insertions(+) create mode 100644 tests/pytorch/test_module_distributed_weight_saved_tensor_hooks.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f45e9369d10..a032683f506 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -55,6 +55,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor.xml python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/test_fusible_ops.py || test_fail "test_fusible_ops.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_selective_activation_checkpoint.xml $TE_PATH/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py || test_fail "test_selective_activation_checkpoint.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_distributed_weight.xml $TE_PATH/tests/pytorch/test_distributed_weight.py || test_fail "test_distributed_weight.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_module_distributed_weight_saved_tensor_hooks.xml $TE_PATH/tests/pytorch/test_module_distributed_weight_saved_tensor_hooks.py || test_fail "test_module_distributed_weight_saved_tensor_hooks.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cross_entropy.xml $TE_PATH/tests/pytorch/test_cross_entropy.py || test_fail "test_cross_entropy.py" diff --git a/tests/pytorch/test_module_distributed_weight_saved_tensor_hooks.py b/tests/pytorch/test_module_distributed_weight_saved_tensor_hooks.py new file mode 100644 index 00000000000..a1273bcc9c9 --- /dev/null +++ b/tests/pytorch/test_module_distributed_weight_saved_tensor_hooks.py @@ -0,0 +1,237 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DistributedWeight dispatch in ``Linear`` / ``GroupedLinear`` under saved-tensor hooks. + +Both modules pass the DistributedWeight parameter itself through ``save_for_backward`` and +used to gate their backward on ``is_distributed_weight(saved_weight)``. That only holds while +no ``torch.autograd.graph.saved_tensors_hooks`` are installed: with hooks active, autograd +unpacks a saved leaf as a *fresh plain tensor* (whatever the unpack hook returns), dropping the +Python subclass and its ``is_distributed_weight`` marker. Backward then took the plain-parameter +branch, whose weakrefs point at the transient all-gathered weights, and failed with +"weight was removed while fuse_wgrad_accumulation=True" (or silently used the wrong weights). + +Megatron-LM's fine-grained activation offloading installs exactly such hooks around expert +layers, which is how this surfaced. The fix keeps the DistributedWeight objects on the autograd +context and prefers them in backward. These tests pin that: with identity hooks installed, the +distributed path must still be taken and produce the same results as without hooks. + +TE ships no DistributedWeight implementer (they live in the caller, e.g. Megatron-LM's GTP), so +an in-repo fake is used. It applies distinct scales in the forward / backward materialize hooks, +so taking the wrong branch fails a specific assertion rather than drifting numerically. +""" + +import pytest +import torch +from torch.autograd.graph import saved_tensors_hooks + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.module.grouped_linear import ( + is_module_grouped_tensor_path_supported, +) + +FWD_SCALE = 2.0 +BWD_SCALE = 4.0 +IN_F, OUT_F, TOKENS = 256, 512, 128 +DTYPE, DEVICE = torch.bfloat16, "cuda" + +FUSE_WGRAD = pytest.mark.parametrize("fuse_wgrad", [False, True], ids=["unfused", "fused-wgrad"]) + + +class _FakeDistWeight(torch.nn.Parameter): + """Fake distributed weight; the group leader holds every member in ``_group``.""" + + is_distributed_weight = True + + def materialize_group_for_forward(self): + self.calls["fwd"] += 1 + return [(w.detach() * FWD_SCALE).requires_grad_(True) for w in self._group] + + def materialize_group_for_backward(self, **kwargs): + self.calls["bwd"] += 1 + return [w.detach() * BWD_SCALE for w in self._group] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls["finalize"] += 1 + wl = list(wgrads) if isinstance(wgrads, (list, tuple)) else [wgrads] + if not self.fuse_wgrad_accumulation: + return [g.clone() for g in wl] + for w, g in zip(self._group, wl): + w.main_grad.add_(g.to(w.main_grad.dtype)) # stands in for the reduce-scatter + w.grad_added_to_main_grad = True + return [torch.zeros_like(g) for g in wl] # real grad is in main_grad + + def grad_buffer(self): + return self.wgrad_scratch + + +def _identity_hooks(): + """Hooks that change nothing; their mere presence makes autograd re-wrap saved leaves.""" + return saved_tensors_hooks(lambda t: t, lambda t: t) + + +def _install_fakes(module, weight_names, fuse_wgrad): + """Replace the module's weights with fakes sharing one group; return the leader.""" + fakes = [] + for name in weight_names: + fake = _FakeDistWeight(getattr(module, name).data) + fake.calls = {"fwd": 0, "bwd": 0, "finalize": 0} + fake.fuse_wgrad_accumulation = fuse_wgrad + fake.main_grad = torch.zeros((OUT_F, IN_F), dtype=torch.float32, device=DEVICE) + fake.wgrad_scratch = torch.zeros_like(fake.main_grad) + fake.grad_added_to_main_grad = False + setattr(module, name, fake) + fakes.append(fake) + for fake in fakes: + fake._group = fakes + return fakes[0] + + +def _check_dist_path_kept( + module, weight_names, leader, reference, ref_leader, out, ref_out, x, ref_x, fuse_wgrad +): + """Hooked run must have used every DistributedWeight hook and match the unhooked run.""" + assert leader.calls["fwd"] > 0, "materialize_group_for_forward never called" + assert leader.calls["bwd"] > 0, ( + "materialize_group_for_backward never called: backward lost the DistributedWeight " + "after saved_tensors_hooks unpacked the saved weight as a plain tensor" + ) + assert leader.calls["finalize"] > 0, "finalize_group_grads never called" + assert leader.calls == ref_leader.calls + # Same kernels on the same data: results are bitwise equal to the unhooked run. + 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 name in weight_names: + w, ref_w = getattr(module, name), getattr(reference, name) + if fuse_wgrad: + assert w.grad_added_to_main_grad is True + torch.testing.assert_close(w.main_grad, ref_w.main_grad, rtol=0, atol=0) + assert torch.count_nonzero(w.main_grad) > 0 + else: + assert w.grad is not None and ref_w.grad is not None + torch.testing.assert_close(w.grad, ref_w.grad, rtol=0, atol=0) + + +def _skip_without_cuda(): + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + +def test_saved_tensor_hooks_unpack_subclass_as_plain_tensor(): + """Premise: with hooks installed, a saved DistributedWeight comes back as a plain tensor.""" + _skip_without_cuda() + seen = {} + + class _Probe(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w): + ctx.save_for_backward(w) + return x * 2 + + @staticmethod + def backward(ctx, g): + (w,) = ctx.saved_tensors + seen["type"] = type(w) + seen["marked"] = bool(getattr(w, "is_distributed_weight", False)) + return g * 2, None + + w = _FakeDistWeight(torch.ones(4, device=DEVICE)) + x = torch.ones(4, device=DEVICE, requires_grad=True) + + _Probe.apply(x, w).sum().backward() + assert seen["type"] is _FakeDistWeight and seen["marked"] + + with _identity_hooks(): + _Probe.apply(x, w).sum().backward() + assert seen["type"] is torch.Tensor and not seen["marked"] + + +@FUSE_WGRAD +def test_linear_distributed_weight_under_saved_tensor_hooks(fuse_wgrad): + """``Linear`` must keep the distributed path in backward when saved-tensor hooks are on.""" + _skip_without_cuda() + torch.manual_seed(0) + module, reference = ( + te.Linear( + IN_F, + OUT_F, + bias=False, + device=DEVICE, + params_dtype=DTYPE, + fuse_wgrad_accumulation=fuse_wgrad, + ) + for _ in range(2) + ) + reference.load_state_dict(module.state_dict()) + leader = _install_fakes(module, ["weight"], fuse_wgrad) + ref_leader = _install_fakes(reference, ["weight"], fuse_wgrad) + + x = torch.randn(TOKENS, IN_F, dtype=DTYPE, device=DEVICE, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + + ref_out = reference(ref_x) + ref_out.sum().backward() + with _identity_hooks(): + out = module(x) + out.sum().backward() + + _check_dist_path_kept( + module, ["weight"], leader, reference, ref_leader, out, ref_out, x, ref_x, fuse_wgrad + ) + # The scales prove which weights were used: FWD_SCALE in forward, BWD_SCALE in dgrad. + plain = te.Linear(IN_F, OUT_F, bias=False, device=DEVICE, params_dtype=DTYPE) + plain.load_state_dict({"weight": leader.data}, strict=False) + px = x.detach().clone().requires_grad_(True) + pout = plain(px) + pout.sum().backward() + torch.testing.assert_close(out.float(), FWD_SCALE * pout.float(), rtol=1e-5, atol=1e-5) + torch.testing.assert_close(x.grad.float(), BWD_SCALE * px.grad.float(), rtol=1e-5, atol=1e-5) + + +@FUSE_WGRAD +@pytest.mark.parametrize("num_gemms", [2, 4]) +@pytest.mark.parametrize( + "use_grouped_tensor", [False, True], ids=["split-quantize", "grouped-tensor"] +) +def test_grouped_linear_distributed_weight_under_saved_tensor_hooks( + num_gemms, fuse_wgrad, use_grouped_tensor +): + """``GroupedLinear`` (both GEMM paths) must keep the distributed path under hooks.""" + _skip_without_cuda() + if use_grouped_tensor and not is_module_grouped_tensor_path_supported(None, DTYPE): + # GroupedLinear would silently fall back to split-quantize, which the other case + # already covers; skip rather than report coverage of the native backward. + pytest.skip("native grouped-tensor path unsupported on this device / cuBLASLt") + torch.manual_seed(0) + names = [f"weight{i}" for i in range(num_gemms)] + module, reference = ( + te.GroupedLinear( + num_gemms, + IN_F, + OUT_F, + bias=False, + device=DEVICE, + params_dtype=DTYPE, + fuse_wgrad_accumulation=fuse_wgrad, + use_grouped_tensor=use_grouped_tensor, + ) + for _ in range(2) + ) + reference.load_state_dict(module.state_dict()) + leader = _install_fakes(module, names, fuse_wgrad) + ref_leader = _install_fakes(reference, names, fuse_wgrad) + + m_splits = torch.full((num_gemms,), TOKENS, dtype=torch.int64, device=DEVICE) + x = torch.randn(num_gemms * TOKENS, IN_F, dtype=DTYPE, device=DEVICE, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + + ref_out = reference(ref_x, m_splits) + ref_out.sum().backward() + with _identity_hooks(): + out = module(x, m_splits) + out.sum().backward() + + _check_dist_path_kept( + module, names, leader, reference, ref_leader, out, ref_out, x, ref_x, fuse_wgrad + ) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 828ec74d9be..7279ec1d051 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -241,6 +241,11 @@ class GroupedLinearBwdArgs: origin_weight_refs: Optional[Any] = None origin_weights_overwrite_main_grad: bool = False main_grad_funcs: Optional[Any] = None + # DistributedWeight objects kept as Python references. They are also passed through + # save_for_backward (as saved_weights), but when saved-tensor hooks are active (e.g. + # activation offloading) autograd unpacks saved leaves as fresh plain tensors, dropping + # the subclass and its ``is_distributed_weight`` marker. + dist_weights: Optional[List[Any]] = None reduce_and_update_bwd_fp8_tensors: bool = False @@ -466,6 +471,7 @@ def _grouped_linear_setup_ctx( bwd_args.wgrad_store = fwd_args.wgrad_store bwd_args.cpu_offloading = fwd_args.cpu_offloading bwd_args.dgrad_out = fwd_args.dgrad_out + bwd_args.dist_weights = list(weights) if is_distributed_weight(weights[0]) else None if fwd_args.fuse_wgrad_accumulation and fwd_args.weights_requires_grad: # Keep weakrefs to weights to preserve attributes like main_grad @@ -520,6 +526,9 @@ def _grouped_linear_backward_impl( inputmats = args.inputmats weights = args.weights_fp8 saved_weights = args.saved_weights + if args.dist_weights is not None: + # Prefer the objects kept in forward; see GroupedLinearBwdArgs.dist_weights. + saved_weights = args.dist_weights biases = args.biases device = grad_output.device in_features = args.weights_shape_1 @@ -755,6 +764,9 @@ class GroupedLinearGraphSafeBwdArgs: weight_quantizers: List[Quantizer] = None grad_output_quantizers: List[Quantizer] = None is_dist_weight: bool = False + # See GroupedLinearBwdArgs.dist_weights: the saved shards may be unpacked as plain + # tensors under saved-tensor hooks, so keep the DistributedWeight objects here. + dist_weights: Optional[List[Any]] = None num_gemms: int = 0 in_features: int = 0 out_features: int = 0 @@ -964,6 +976,7 @@ def _grouped_linear_graph_safe_setup( bwd_args.input_quantizers = fwd_args.input_quantizers bwd_args.weight_quantizers = fwd_args.weight_quantizers bwd_args.is_dist_weight = is_distributed_weight(fwd_args.weights[0]) + bwd_args.dist_weights = list(fwd_args.weights) if bwd_args.is_dist_weight else None bwd_args.grad_output_quantizers = fwd_args.grad_output_quantizers bwd_args.m_splits_tensor = fwd_args.m_splits_tensor bwd_args.num_gemms = fwd_args.num_gemms @@ -1038,6 +1051,10 @@ def _grouped_linear_graph_safe_backward( main_grads = [None] * num_weight_args is_dist_weight = args.is_dist_weight if is_dist_weight: + # Forward saved the shards, not the gathered copies. Prefer the originals kept + # on args: saved-tensor hooks may have unpacked the shards as plain tensors. + if args.dist_weights is not None: + weight_tensors = args.dist_weights origin_weights = list(weight_tensors) if args.fuse_wgrad_accumulation and args.weights_requires_grad: main_grads = [main_grad_func() for main_grad_func in args.main_grad_funcs] diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 73c9503026f..bb2268a264c 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -291,6 +291,10 @@ class LinearBwdArgs: fuse_wgrad_accumulation: bool = False wgrad_store: Optional[Any] = None origin_weight_ref: Optional[Any] = None + # DistributedWeight object kept as a Python reference: saved-tensor hooks (e.g. activation + # offloading) make autograd unpack ``saved_weight`` as a fresh plain tensor, dropping the + # subclass and its ``is_distributed_weight`` marker. + dist_weight: Optional[Any] = None origin_weight_overwrites_main_grad: bool = False main_grad_func: Optional[Callable[[], torch.Tensor]] = None @@ -1060,6 +1064,7 @@ def _linear_setup_ctx( bwd_args.is_first_microbatch = fwd_args.is_first_microbatch bwd_args.fuse_wgrad_accumulation = fuse_wgrad_accumulation bwd_args.wgrad_store = fwd_args.wgrad_store + bwd_args.dist_weight = weight if is_distributed_weight(weight) else None if fuse_wgrad_accumulation and fwd_args.weight_requires_grad: bwd_args.origin_weight_ref = weakref.ref(weight) bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) @@ -1116,6 +1121,10 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None inputmat = args.inputmat weight_fp8 = args.weight_fp8 saved_weight = args.saved_weight + if args.dist_weight is not None: + # Prefer the object kept in forward; see LinearBwdArgs.dist_weight. + saved_weight = args.dist_weight + args.dist_weight = None is_dist_weight = is_distributed_weight(saved_weight) bias = args.bias input_quantizer = args.input_quantizer From a9a8420c8fc7e636d2a5bd83bd1ca4800183013c Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Sat, 17 Oct 2026 02:04:12 -0700 Subject: [PATCH 2/3] [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 a032683f506..b0265f90ebe 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_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 00000000000..02fd1431102 --- /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 eeb5f43335f..24343b71485 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 250974fde43484a9bd901f568ee9111d92ba9fdc Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Sat, 17 Oct 2026 01:59:23 -0700 Subject: [PATCH 3/3] [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 b0265f90ebe..b96ff992aa6 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -93,6 +93,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 00000000000..97f63a96c83 --- /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 24343b71485..8b4a7f4f4e5 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 = [