grouped linear single param fixes

tests/pytorch/test_grouped_mlp.py

@@ -228,6 +228,33 @@ def make_reference_and_test_tensors(
 class TestGroupedLinearOp:
     """Tests for advanced features with grouped linear basic op"""
 
+    def test_meta_single_grouped_weight_with_delayed_wgrad(self, monkeypatch) -> None:
+        """A deferred op shell must not access its grouped parent before it is attached."""
+        monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+
+        op = te.ops.GroupedLinear(
+            2,
+            16,
+            16,
+            device="meta",
+            dtype=torch.bfloat16,
+            single_grouped_weight=True,
+            delay_wgrad_compute=True,
+        )
+
+        assert op._parameters.get("weight") is None
+        assert op.weight0.device.type == "meta"
+
+        # Mirror an external caller attaching the grouped parent after constructing
+        # the parameterless shell, then verify delayed-wgrad metadata is applied.
+        grouped_weight = torch.nn.Parameter(torch.empty(2, 16, 16, device="meta"))
+        op.register_parameter("weight", grouped_weight)
+        for group_idx in range(op.num_groups):
+            op.register_parameter(f"weight{group_idx}", None)
+        op._apply_delay_wgrad_param_hooks()
+
+        assert grouped_weight.skip_backward_post_hook
+
     @pytest.mark.parametrize("bias", (False, True))
     @pytest.mark.parametrize("dtype", _dtypes)
     @pytest.mark.parametrize("quantization", _quantization_list)
@@ -1086,6 +1113,64 @@ def test_grouped_mlp_fp16(
             activation=activation,
         )
 
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    def test_single_grouped_weight_eval_preserves_columnwise_usage(
+        self,
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+        group_size: int = 2,
+        hidden_size: int = 128,
+    ) -> None:
+        """Eager eval must not drop storage still used by captured training dgrad."""
+
+        if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
+            pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
+
+        split_sizes = torch.full(
+            (group_size,),
+            128,
+            dtype=torch.int64,
+            device=device,
+        )
+        num_tokens = int(split_sizes.sum())
+        recipe = make_recipe("mxfp8")
+
+        with te.quantized_model_init(enabled=True, recipe=recipe):
+            fc1 = te.ops.GroupedLinear(
+                group_size,
+                hidden_size,
+                2 * hidden_size,
+                device=device,
+                dtype=dtype,
+                single_grouped_weight=True,
+            )
+            fc2 = te.ops.GroupedLinear(
+                group_size,
+                hidden_size,
+                hidden_size,
+                device=device,
+                dtype=dtype,
+                single_grouped_weight=True,
+            )
+            module = te.ops.Sequential(
+                fc1,
+                te.ops.ScaledSwiGLU(glu_interleave_size=32),
+                fc2,
+            )
+
+        x = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype, requires_grad=True)
+        probs = torch.ones(num_tokens, device=device, dtype=dtype)
+        with te.autocast(enabled=True, recipe=recipe):
+            module(x, split_sizes, probs, split_sizes)
+        assert fc1.weight.quantizer.columnwise_usage
+        assert fc2.weight.quantizer.columnwise_usage
+
+        with torch.no_grad(), te.autocast(enabled=True, recipe=recipe):
+            module(x.detach(), split_sizes, probs, split_sizes)
+        assert fc1.weight.quantizer.columnwise_usage
+        assert fc2.weight.quantizer.columnwise_usage
+
     @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list)
     @pytest.mark.parametrize("single_grouped_weight", (False, True))
     @pytest.mark.parametrize("accumulate_into_main_grad", (False, True))

tests/pytorch/test_sanity.py

@@ -3,6 +3,7 @@
 # See LICENSE for license information.
 
 from typing import Optional, List
+from unittest import mock
 
 import torch
 import pytest
@@ -36,6 +37,7 @@
 )
 from transformer_engine.common import recipe
 from transformer_engine.pytorch.cpp_extensions import general_gemm
+from transformer_engine.pytorch.module import grouped_linear as grouped_linear_module
 from transformer_engine.pytorch.tensor.utils import replace_raw_data
 from utils import ModelConfig, recipe_id, skip_unsupported_backward_override
 
@@ -650,6 +652,8 @@ def test_sanity_grouped_linear(
     loss = out.sum()
     loss.backward()
     assert out.shape == (num_tokens, ffn_hidden_size)
+    if single_param:
+        assert te_grouped_linear.weight.grad is not None
 
 
 @pytest.mark.parametrize("dtype", param_types)
@@ -1154,6 +1158,122 @@ def test_quantized_model_init_high_precision_init_val():
     ), "clear_high_precision_init_val() not work"
 
 
+@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+def test_grouped_linear_single_param_preserves_high_precision_init(monkeypatch):
+    """Grouped MXFP8 and discrete weights produce identical FP32 master initialization."""
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+    num_gemms = 3
+
+    def make_module(single_grouped_weight):
+        torch.manual_seed(1234)
+        torch.cuda.manual_seed(1234)
+        with quantized_model_init(
+            enabled=True,
+            recipe=recipe.MXFP8BlockScaling(),
+            preserve_high_precision_init_val=True,
+        ):
+            return GroupedLinear(
+                num_gemms=num_gemms,
+                in_features=32,
+                out_features=64,
+                bias=False,
+                params_dtype=torch.bfloat16,
+                single_grouped_weight=single_grouped_weight,
+            ).cuda()
+
+    discrete_module = make_module(False)
+    grouped_module = make_module(True)
+    discrete_init_vals = [
+        getattr(discrete_module, f"weight{i}").get_high_precision_init_val()
+        for i in range(num_gemms)
+    ]
+    expected_grouped_init = torch.stack(discrete_init_vals, dim=0)
+
+    grouped_weight = grouped_module.weight
+    assert hasattr(grouped_weight, "get_high_precision_init_val")
+    assert hasattr(grouped_weight, "clear_high_precision_init_val")
+    grouped_init = grouped_weight.get_high_precision_init_val()
+    assert grouped_init.device.type == "cpu"
+    assert grouped_init.shape == grouped_weight.shape
+    torch.testing.assert_close(
+        grouped_init,
+        expected_grouped_init,
+        rtol=0,
+        atol=0,
+    )
+
+    # This is the exact layout consumed when constructing the FP32 optimizer master.
+    discrete_master = expected_grouped_init.float()
+    grouped_master = grouped_init.float()
+    torch.testing.assert_close(grouped_master, discrete_master, rtol=0, atol=0)
+
+    grouped_weight.clear_high_precision_init_val()
+    assert grouped_weight.get_high_precision_init_val() is None
+
+
+@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+def test_grouped_linear_rejects_partial_high_precision_init(monkeypatch):
+    """Packing fails rather than mixing preserved and dequantized initialization."""
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+    with quantized_model_init(
+        enabled=True,
+        recipe=recipe.MXFP8BlockScaling(),
+        preserve_high_precision_init_val=True,
+    ):
+        module = GroupedLinear(
+            num_gemms=2,
+            in_features=32,
+            out_features=64,
+            bias=False,
+            params_dtype=torch.bfloat16,
+            single_grouped_weight=False,
+        ).cuda()
+
+    module.weight0.clear_high_precision_init_val()
+    with pytest.raises(
+        RuntimeError,
+        match="inconsistent high-precision initialization state",
+    ):
+        module.make_grouped_weights()
+
+
+@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+def test_grouped_linear_single_param_legacy_fused_wgrad(monkeypatch):
+    """Legacy GroupedLinear accumulates MXFP8 member wgrads into the grouped parent."""
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0")
+    dtype = torch.bfloat16
+
+    with quantized_model_init(enabled=True, recipe=recipe.MXFP8BlockScaling()):
+        module = GroupedLinear(
+            num_gemms=2,
+            in_features=32,
+            out_features=64,
+            bias=False,
+            params_dtype=dtype,
+            fuse_wgrad_accumulation=True,
+            single_grouped_weight=True,
+        ).cuda()
+
+    weight = module.weight
+    weight.main_grad = torch.zeros(weight.shape, dtype=dtype, device="cuda")
+    weight.grad_added_to_main_grad = False
+    inp = torch.randn(32, 32, dtype=dtype, device="cuda", requires_grad=True)
+
+    with mock.patch.object(
+        grouped_linear_module,
+        "get_dummy_wgrad",
+        wraps=grouped_linear_module.get_dummy_wgrad,
+    ) as dummy_wgrad:
+        with autocast(enabled=True, recipe=recipe.MXFP8BlockScaling()):
+            module(inp, [16, 16]).sum().backward()
+
+    assert weight.grad_added_to_main_grad
+    assert weight.grad is not None
+    assert torch.count_nonzero(weight.main_grad).item() > 0
+    dummy_wgrad.assert_any_call(list(weight.shape), weight.dtype, zero=False)
+
+
 def test_sanity_checkpointing_on_callables():
     """Test that TE checkpointing works correctly on callable modules."""
 

transformer_engine/pytorch/module/base.py

@@ -87,6 +87,27 @@
 layers_atomic_ring_exchange = []
 
 
+def _get_high_precision_init_val(parameter: torch.Tensor) -> Optional[torch.Tensor]:
+    """Return temporary pre-quantization initialization stored on a parameter."""
+    return getattr(parameter, "_high_precision_init_val", None)
+
+
+def _clear_high_precision_init_val(parameter: torch.Tensor) -> None:
+    """Release temporary pre-quantization initialization stored on a parameter."""
+    if hasattr(parameter, "_high_precision_init_val"):
+        del parameter._high_precision_init_val
+
+
+def _attach_high_precision_init_val(
+    parameter: torch.Tensor,
+    high_precision_init_val: torch.Tensor,
+) -> None:
+    """Attach TE's temporary high-precision initialization contract to a parameter."""
+    parameter._high_precision_init_val = high_precision_init_val
+    parameter.get_high_precision_init_val = MethodType(_get_high_precision_init_val, parameter)
+    parameter.clear_high_precision_init_val = MethodType(_clear_high_precision_init_val, parameter)
+
+
 def is_ub_initialized() -> bool:
     """Whether the Userbuffers communicators have been initialized."""
     return _ub_initialized
@@ -1787,21 +1808,10 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None:
                 #   should call `clear_high_precision_init_val` to remove it after master weight
                 #   is initialized.
 
-                def get(self):
-                    if hasattr(self, "_high_precision_init_val"):
-                        return self._high_precision_init_val
-                    return None
-
-                def clear(self):
-                    if hasattr(self, "_high_precision_init_val"):
-                        del self._high_precision_init_val
-
                 # DTensor.from_local() does not preserve object identity,
                 # so attach to the DTensor's local tensor when applicable.
                 target = dtensor_param._local_tensor if is_dtensor else param
-                target._high_precision_init_val = high_precision_init_val
-                target.get_high_precision_init_val = MethodType(get, target)
-                target.clear_high_precision_init_val = MethodType(clear, target)
+                _attach_high_precision_init_val(target, high_precision_init_val)
 
             if not is_dtensor:
                 self.module_setattr(name, param)