Single grouped weight fixes

tests/pytorch/test_grouped_mlp.py

@@ -232,6 +232,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"))
+        assert not hasattr(grouped_weight, "skip_backward_post_hook")
+        op.register_parameter("weight", grouped_weight)
+        for group_idx in range(op.num_groups):
+            op.register_parameter(f"weight{group_idx}", None)
+
+        assert grouped_weight.skip_backward_post_hook
+
     @pytest.mark.parametrize("bias", (False, True))
     @pytest.mark.parametrize("dtype", _dtypes)
     @pytest.mark.parametrize("quantization", _quantization_list)
@@ -1104,6 +1131,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

@@ -1154,6 +1154,85 @@ 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()
+
+
 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)

transformer_engine/pytorch/module/grouped_linear.py

@@ -27,6 +27,9 @@
     _2X_ACC_FPROP,
     _2X_ACC_DGRAD,
     _2X_ACC_WGRAD,
+    _attach_high_precision_init_val,
+    _clear_high_precision_init_val,
+    _get_high_precision_init_val,
 )
 from ._common import WeightGradStore
 from ..quantization import FP8GlobalStateManager, QuantizerRole
@@ -1522,6 +1525,19 @@ def make_grouped_weights(self, defer_init=False) -> None:
 
         weights = [getattr(self, f"weight{i}") for i in range(self.num_gemms)]
 
+        # TE preserves the original BF16/FP16 initialization on each quantized
+        # parameter so distributed optimizers can construct lossless FP32 masters.
+        # Packing the parameters must transfer those values to the new registered
+        # grouped parameter; otherwise its master is initialized by dequantizing
+        # MXFP8 and starts from a different value than the discrete-weight layout.
+        high_precision_init_vals = [_get_high_precision_init_val(weight) for weight in weights]
+        if any(value is not None for value in high_precision_init_vals) and not all(
+            value is not None for value in high_precision_init_vals
+        ):
+            raise RuntimeError(
+                "Grouped weights have inconsistent high-precision initialization state"
+            )
+
         # Create the weight storage.
         grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes(
             num_tensors=self.num_gemms,
@@ -1545,9 +1561,18 @@ def make_grouped_weights(self, defer_init=False) -> None:
             and (weight_quantizers[0] is None or not weight_quantizers[0].internal)
         ):
             raise RuntimeError("Found internal quantizer with `single_grouped_weight=True`.")
+        grouped_parameter = torch.nn.Parameter(grouped_weights)
+        if all(value is not None for value in high_precision_init_vals):
+            _attach_high_precision_init_val(
+                grouped_parameter,
+                torch.stack(high_precision_init_vals, dim=0),
+            )
+            for weight in weights:
+                _clear_high_precision_init_val(weight)
+
         self.register_parameter(
             "weight",
-            torch.nn.Parameter(grouped_weights),
+            grouped_parameter,
             init_fn=self.init_method,
             get_rng_state_tracker=self.get_rng_state_tracker,
             fp8_meta_index=self._offsets["weight"],

transformer_engine/pytorch/ops/basic/grouped_linear.py

@@ -220,12 +220,33 @@ def __init__(
 
         self._apply_delay_wgrad_param_hooks()
 
+    def register_parameter(
+        self,
+        name: str,
+        param: Optional[torch.nn.Parameter],
+    ) -> None:
+        """Register a parameter and apply delayed-wgrad metadata when needed."""
+        super().register_parameter(name, param)
+
+        # A single-grouped-weight op may be constructed on the meta device and
+        # receive its grouped parent later. Mark that parent at attachment time,
+        # before DDP/FSDP inspects the parameter to install backward hooks.
+        if name == "weight" and param is not None and getattr(self, "single_grouped_weight", False):
+            wgrad_store = getattr(self, "wgrad_store", None)
+            if wgrad_store is not None and wgrad_store.delay_wgrad_compute():
+                param.skip_backward_post_hook = True
+
     def _apply_delay_wgrad_param_hooks(self) -> None:
         """Set ``skip_backward_post_hook`` on weights when delaying wgrad (bias uses main backward)."""
         if not self.wgrad_store.delay_wgrad_compute():
             return
         if self.single_grouped_weight:
-            self.weight.skip_backward_post_hook = True
+            # A meta-device op may be created as a parameterless shell and have its
+            # grouped parent attached after construction. In that case there is no
+            # ``weight`` parameter to mark yet.
+            weight = self._parameters.get("weight")
+            if weight is not None:
+                weight.skip_backward_post_hook = True
         else:
             for group_idx in range(self.num_groups):
                 getattr(self, f"weight{group_idx}").skip_backward_post_hook = True

transformer_engine/pytorch/ops/fused/grouped_mlp.py

@@ -973,7 +973,13 @@ def fuser_forward(
                     "FC1 expected GroupedTensor weight with single_grouped_weight=True."
                 )
             if fc1_op.weight.quantizer is not None:
-                fc1_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad)
+                # Full-iteration CUDA graph replay does not rerun this Python metadata
+                # update. Remember that a grad-enabled forward requested columnwise
+                # storage so eager eval cannot drop buffers still used by captured dgrad.
+                fc1_weight_quantizer.set_usage(
+                    rowwise=True,
+                    columnwise=input_requires_grad or fc1_weight_quantizer.columnwise_usage,
+                )
                 fc1_op.weight.quantizer = fc1_weight_quantizer
                 grouped_fc1_weight = fc1_op.weight
             else:
@@ -1005,7 +1011,10 @@ def fuser_forward(
                     "FC2 expected GroupedTensor weight with single_grouped_weight=True."
                 )
             if fc2_op.weight.quantizer is not None:
-                fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad)
+                fc2_weight_quantizer.set_usage(
+                    rowwise=True,
+                    columnwise=input_requires_grad or fc2_weight_quantizer.columnwise_usage,
+                )
                 fc2_op.weight.quantizer = fc2_weight_quantizer
                 grouped_fc2_weight = fc2_op.weight
             else: