diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f45e9369d1..28f26cc1ba 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -91,6 +91,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_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/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)