Use a plain parameter for high-precision grouped linear weights

tests/pytorch/test_grouped_tensor.py

@@ -9,7 +9,12 @@
 import pytest
 import torch
 import transformer_engine.pytorch as te
+import transformer_engine.pytorch.ops as te_ops
+from transformer_engine.common.recipe import Float8CurrentScaling, MXFP8BlockScaling
 from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor
+from transformer_engine.pytorch.tensor.storage.grouped_tensor_storage import (
+    grouped_param_members,
+)
 from transformer_engine.pytorch import (
     Quantizer,
     Float8Quantizer,
@@ -1441,3 +1446,144 @@ def test_grouped_linear_load_state_dict_single_to_multi_param(self, tmp_path) ->
             assert torch.equal(getattr(dst, f"weight{i}"), expected_weight)
         for i, expected_bias in enumerate(expected_biases):
             assert torch.equal(getattr(dst, f"bias{i}"), expected_bias.reshape(-1))
+
+
+def _skip_without_single_param_env() -> None:
+    """Skip when the experimental single grouped parameter is not enabled."""
+    if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0":
+        pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1")
+
+
+def _make_grouped_linear_op(
+    *,
+    num_groups: int,
+    in_features: int,
+    out_features: int,
+    dtype: torch.dtype,
+    single_grouped: bool,
+) -> te_ops.GroupedLinear:
+    """Build a fusible GroupedLinear op on CUDA."""
+    return te_ops.GroupedLinear(
+        num_groups,
+        in_features,
+        out_features,
+        bias=True,
+        device="cuda",
+        dtype=dtype,
+        single_grouped_weight=single_grouped,
+        single_grouped_bias=single_grouped,
+    )
+
+
+class TestGroupedLinearHighPrecisionParam:
+    """High-precision grouped parameters are plain tensors, not GroupedTensor.
+
+    A GroupedTensor is only needed when the grouped scale/amax buffers have to
+    travel with the data. Without a quantizer there are none, so the parameter
+    stays a plain stacked tensor and keeps torch.save/DCP and generic tensor ops
+    working.
+    """
+
+    num_groups = 3
+    in_features = 64
+    out_features = 32
+
+    def _op(self, dtype: torch.dtype, single_grouped: bool = True) -> te_ops.GroupedLinear:
+        return _make_grouped_linear_op(
+            num_groups=self.num_groups,
+            in_features=self.in_features,
+            out_features=self.out_features,
+            dtype=dtype,
+            single_grouped=single_grouped,
+        )
+
+    @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+    def test_param_is_plain_tensor(self, dtype: torch.dtype) -> None:
+        """High-precision grouped weight/bias are plain stacked parameters."""
+        _skip_without_single_param_env()
+        op = self._op(dtype)
+
+        assert not isinstance(op.weight, GroupedTensor)
+        assert op.weight.shape == (self.num_groups, self.out_features, self.in_features)
+        assert op.weight.dtype == dtype
+        assert not isinstance(op.bias, GroupedTensor)
+        assert op.bias.shape == (self.num_groups, self.out_features)
+
+        # The per-group members must still be views into the one buffer.
+        members = grouped_param_members(op.weight)
+        assert len(members) == self.num_groups
+        for member in members:
+            assert member.shape == (self.out_features, self.in_features)
+            assert member.data_ptr() >= op.weight.data_ptr()
+
+    def test_matches_per_group_params(self) -> None:
+        """Forward and backward match the per-group parameter layout."""
+        _skip_without_single_param_env()
+        dtype = torch.bfloat16
+        split_sizes = torch.tensor([8, 16, 8], dtype=torch.int64)
+        total_tokens = int(split_sizes.sum())
+
+        grouped = self._op(dtype, single_grouped=True)
+        discrete = self._op(dtype, single_grouped=False)
+
+        # Give both layouts identical parameter values.
+        with torch.no_grad():
+            for idx in range(self.num_groups):
+                getattr(discrete, f"weight{idx}").copy_(grouped.weight[idx])
+                getattr(discrete, f"bias{idx}").copy_(grouped.bias[idx])
+
+        x = torch.randn(
+            total_tokens, self.in_features, dtype=dtype, device="cuda", requires_grad=True
+        )
+        x_ref = x.detach().clone().requires_grad_(True)
+        dy = torch.randn(total_tokens, self.out_features, dtype=dtype, device="cuda")
+
+        y = grouped(x, split_sizes)
+        y.backward(dy)
+        y_ref = discrete(x_ref, split_sizes)
+        y_ref.backward(dy)
+
+        assert_close(y, y_ref)
+        assert_close(x.grad, x_ref.grad)
+        for idx in range(self.num_groups):
+            assert_close(grouped.weight.grad[idx], getattr(discrete, f"weight{idx}").grad)
+            assert_close(grouped.bias.grad[idx], getattr(discrete, f"bias{idx}").grad)
+
+    def test_state_dict_round_trip(self, tmp_path) -> None:
+        """A plain grouped parameter survives torch.save/torch.load.
+
+        GroupedTensor has no ``__reduce_ex__``/``untyped_storage``, which is the
+        capability this representation restores.
+        """
+        _skip_without_single_param_env()
+        dtype = torch.bfloat16
+        src = self._op(dtype)
+        expected_weight = src.weight.detach().clone()
+        expected_bias = src.bias.detach().clone()
+
+        ckpt_path = tmp_path / "grouped_linear_op.pt"
+        torch.save(src.state_dict(), ckpt_path)
+        state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+
+        assert not isinstance(state_dict["weight"], GroupedTensor)
+        assert state_dict["weight"].shape == expected_weight.shape
+
+        dst = self._op(dtype)
+        dst.load_state_dict(state_dict)
+        assert torch.equal(dst.weight, expected_weight)
+        assert torch.equal(dst.bias, expected_bias)
+
+    def test_quantized_param_stays_grouped(self) -> None:
+        """Quantized weights keep GroupedTensor, which owns their scale buffers."""
+        _skip_without_single_param_env()
+        if mxfp8_available:
+            recipe = MXFP8BlockScaling()
+        elif fp8_available:
+            recipe = Float8CurrentScaling()
+        else:
+            pytest.skip(reason_for_no_fp8)
+        with te.quantized_model_init(enabled=True, recipe=recipe):
+            op = self._op(torch.bfloat16)
+
+        assert isinstance(op.weight, GroupedTensor)
+        assert op.weight.quantizer is not None

transformer_engine/pytorch/ops/basic/grouped_linear.py

@@ -9,7 +9,7 @@
 import contextlib
 import functools
 import math
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import torch
 
@@ -59,7 +59,13 @@
     materialize_weight_for_backward,
     materialize_weight_for_forward,
 )
-from ...tensor import GroupedTensor, GroupedTensorStorage
+from ...tensor import (
+    GroupedTensor,
+    GroupedTensorStorage,
+    grouped_param_members,
+    grouped_storage_for_gemm,
+    has_grouped_storage,
+)
 from ...triton.grouped_dbias_dscales import (
     compute_grouped_dbias,
     compute_grouped_dbias_dscales,
@@ -319,9 +325,7 @@ def backward_dw(self) -> None:
     def _get_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]:
         """Retrieve per-group bias tensors in the given dtype."""
         if self.single_grouped_bias:
-            bias_parts = self.bias.quantized_tensors
-            if bias_parts is None:
-                bias_parts = self.bias.split_into_quantized_tensors()
+            bias_parts = grouped_param_members(self.bias)
             return [maybe_dequantize(p.reshape(-1), dtype) for p in bias_parts]
         return [
             maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(self.num_groups)
@@ -451,7 +455,7 @@ def reset_parameters(self) -> None:
 
     def make_grouped_weights(self) -> None:
         """
-        Convert parameters into a GroupedTensor and re-register them as parameters.
+        Pack the per-group weights into one parameter and re-register it.
         """
 
         weights = [getattr(self, f"weight{idx}") for idx in range(self.num_groups)]
@@ -461,6 +465,20 @@ def make_grouped_weights(self) -> None:
         if recipe is not None and recipe.delayed():
             raise RuntimeError("Delayed scaling is not supported with single_grouped_weight=True")
 
+        if quantizer is None:
+            # High-precision weights need none of GroupedTensor's machinery: with
+            # uniform member shapes and no scales, the grouped GEMM only reads a
+            # data pointer, num_tensors and the logical shape, all of which a plain
+            # stacked tensor already carries. Keeping it plain preserves
+            # serialization (torch.save/DCP) and generic tensor ops.
+            with torch.no_grad():
+                grouped_weights = torch.stack([w.detach() for w in weights], dim=0).contiguous()
+            self.register_parameter("weight", torch.nn.Parameter(grouped_weights))
+            for group_idx in range(self.num_groups):
+                self.register_parameter(f"weight{group_idx}", None)
+            self._apply_delay_wgrad_param_hooks()
+            return
+
         grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes(
             num_tensors=self.num_groups,
             shapes=[(self.out_features, self.in_features)] * self.num_groups,
@@ -489,16 +507,13 @@ def make_grouped_weights(self) -> None:
         self._apply_delay_wgrad_param_hooks()
 
     def _make_grouped_biases_from_packed(self, packed_biases: torch.Tensor) -> None:
-        """Replace per-group bias parameters with one ``GroupedTensor`` (``single_grouped_bias``)."""
+        """Replace per-group bias parameters with one packed parameter (``single_grouped_bias``).
+
+        Biases are never quantized, so this is always a plain
+        ``(num_groups, out_features)`` tensor.
+        """
         bias_data = packed_biases.detach().clone().contiguous()
-        grouped_bias = GroupedTensor.make_grouped_tensor_from_rowwise_data(
-            num_tensors=self.num_groups,
-            tensor_shape=(self.out_features,),
-            rowwise_data=bias_data,
-            dtype=bias_data.dtype,
-        )
-        grouped_bias.requires_grad_(True)
-        self.register_parameter("bias", torch.nn.Parameter(grouped_bias))
+        self.register_parameter("bias", torch.nn.Parameter(bias_data))
         for group_idx in range(self.num_groups):
             self.register_parameter(f"bias{group_idx}", None)
 
@@ -723,7 +738,10 @@ def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None:
             weight_is_quantized = False
             if getattr(self, "single_grouped_weight", False):
                 weight = getattr(self, "weight", None)
-                weight_is_quantized = weight is not None and weight.quantizer is not None
+                # A high-precision grouped weight is a plain tensor with no quantizer.
+                weight_is_quantized = (
+                    weight is not None and getattr(weight, "quantizer", None) is not None
+                )
             else:
                 weight = getattr(self, f"weight{group_idx}", None)
                 weight_is_quantized = is_quantized_tensor(weight)
@@ -855,7 +873,7 @@ def _is_graph_safe_path_supported(
 
     def _get_grouped_weight_for_gemm(
         self,
-        weight_param: GroupedTensor,
+        weight_param: Union[GroupedTensor, torch.Tensor],
         weight_quantizers: list[Optional[Quantizer]],
         columnwise_usage: bool,
         with_quantized_compute: bool,
@@ -865,6 +883,25 @@ def _get_grouped_weight_for_gemm(
         Supports MXFP8/BF16/FP16 compute paths.
         """
         num_groups = self.num_groups
+        weight_shapes = [(self.out_features, self.in_features)] * num_groups
+        if not has_grouped_storage(weight_param):
+            # High-precision weights are a plain stacked parameter; describe them
+            # for the GEMM without copying.
+            if not with_quantized_compute:
+                return grouped_storage_for_gemm(
+                    weight_param,
+                    num_tensors=num_groups,
+                    shapes=weight_shapes,
+                    dtype=dtype,
+                )
+            weight_quantizer = weight_quantizers[0]
+            weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
+            return tex.group_quantize(
+                weight_param.reshape(num_groups * self.out_features, self.in_features),
+                weight_quantizer,
+                num_groups,
+                None,
+            )
         is_weight_quantized = weight_param.quantizer is not None
         if is_weight_quantized and with_quantized_compute:
             # GGEMM can use it as it is
@@ -881,7 +918,7 @@ def _get_grouped_weight_for_gemm(
                 shape=(num_groups * self.out_features, self.in_features),
                 dtype=dtype,
                 num_tensors=num_groups,
-                shapes=[(self.out_features, self.in_features)] * num_groups,
+                shapes=weight_shapes,
                 quantizer=None,
                 data=weight_data.reshape(-1),
             )
@@ -895,7 +932,7 @@ def _get_grouped_weight_for_gemm(
                 shape=(num_groups * self.out_features, self.in_features),
                 dtype=dtype,
                 num_tensors=num_groups,
-                shapes=[(self.out_features, self.in_features)] * num_groups,
+                shapes=weight_shapes,
                 quantizer=None,
                 data=weight_data.reshape(-1),
             )
@@ -976,7 +1013,7 @@ def _get_grouped_bias_for_gemm(
 
         if self.single_grouped_bias:
             # Already a contiguous (num_groups * out_features) buffer.
-            bias_data = self.bias.rowwise_data
+            bias_data = self.bias.rowwise_data if has_grouped_storage(self.bias) else self.bias
             if bias_data.dtype != dtype:
                 bias_data = bias_data.to(dtype=dtype)
         else:
@@ -1217,9 +1254,7 @@ def _fuser_forward_split_quantize(
 
         # Extract params
         if self.single_grouped_weight:
-            weights = self.weight.quantized_tensors
-            if weights is None:
-                weights = self.weight.split_into_quantized_tensors()
+            weights = grouped_param_members(self.weight)
         else:
             weights = self._forward_weight_list()  # materialized when distributed
         bs = None

transformer_engine/pytorch/tensor/__init__.py

@@ -18,7 +18,12 @@
 from .storage.mxfp8_tensor_storage import MXFP8TensorStorage
 from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage
 from .storage.nvfp4_tensor_storage import NVFP4TensorStorage
-from .storage.grouped_tensor_storage import GroupedTensorStorage
+from .storage.grouped_tensor_storage import (
+    GroupedTensorStorage,
+    grouped_param_members,
+    grouped_storage_for_gemm,
+    has_grouped_storage,
+)
 from .storage.hybrid_tensor_storage import HybridQuantizedTensorStorage
 from .float8_tensor import Float8Tensor, Float8Quantizer, Float8CurrentScalingQuantizer
 from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer
@@ -45,6 +50,9 @@
     "Float8BlockwiseQTensorStorage",
     "NVFP4TensorStorage",
     "GroupedTensorStorage",
+    "grouped_param_members",
+    "grouped_storage_for_gemm",
+    "has_grouped_storage",
     "HybridQuantizedTensorStorage",
     "IdentityTensorStorage",
     "QuantizedTensor",

transformer_engine/pytorch/tensor/storage/__init__.py

@@ -8,3 +8,6 @@
 from .float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage  # noqa: F401
 from .nvfp4_tensor_storage import NVFP4TensorStorage  # noqa: F401
 from .grouped_tensor_storage import GroupedTensorStorage  # noqa: F401
+from .grouped_tensor_storage import has_grouped_storage  # noqa: F401
+from .grouped_tensor_storage import grouped_param_members  # noqa: F401
+from .grouped_tensor_storage import grouped_storage_for_gemm  # noqa: F401