diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f45e9369d1..b96ff992aa 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" @@ -91,6 +92,8 @@ 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" +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_module_distributed_weight_saved_tensor_hooks.py b/tests/pytorch/test_module_distributed_weight_saved_tensor_hooks.py new file mode 100644 index 0000000000..a1273bcc9c --- /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/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/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/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 828ec74d9b..7279ec1d05 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 73c9503026..bb2268a264 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 diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index eeb5f43335..8b4a7f4f4e 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" ): @@ -2200,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 = [ @@ -2245,7 +2284,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 +2293,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 +2544,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 +2639,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 +2923,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)