From e0f871c2ac9cd927777dfb2fbbc2f81504f75fc1 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 10 Aug 2026 22:48:28 +0000 Subject: [PATCH] Use a plain parameter for high-precision grouped linear weights GroupedLinear(single_grouped_weight=True) always registered the packed weight as a GroupedTensor, a storage-less _make_wrapper_subclass whose payload lives in Python attributes. That is necessary for quantized weights, whose grouped scale/amax buffers must travel with the data. It is not necessary for high-precision weights. GroupedTensorFromPyTorchGroupedTensor reads only num_tensors, logical_shape, rowwise_data and _with_gemm_swizzled_scales when there is no quantizer; every other attribute is None and short-circuits. For uniform member shapes make_grouped_tensor leaves first_dims, last_dims and tensor_offsets unset, so the descriptor holds no device-side metadata at all. All of it is derivable from a plain contiguous (G, M, N) tensor. The wrapper is not free: it has none of the serialization plumbing that QuantizedTensor carries (untyped_storage, the _to_copy handler, new_empty, __reduce_ex__), and it bans reshape-family ops, so torch.save, torch.distributed.checkpoint and generic tensor ops do not work on these parameters. Register high-precision grouped weights and biases as plain stacked parameters and build a transient GroupedTensorStorage at GEMM call time. This is already how the gradient side works: the wgrad path wraps a plain main_grad buffer and returns a plain (G, M, N) tensor to autograd. The GEMM sees the same pointer with no extra copy. Quantized weights are unchanged. Only the fusible op is converted; module/grouped_linear.py keeps GroupedTensor because it passes per-GEMM members into the autograd function and holds them by weakref, so the member lifetime and autograd wiring there need separate treatment. Signed-off-by: Jingyue Wu --- tests/pytorch/test_grouped_tensor.py | 146 ++++++++++++++++++ .../pytorch/ops/basic/grouped_linear.py | 81 +++++++--- transformer_engine/pytorch/tensor/__init__.py | 10 +- .../pytorch/tensor/storage/__init__.py | 3 + .../tensor/storage/grouped_tensor_storage.py | 49 ++++++ 5 files changed, 265 insertions(+), 24 deletions(-) diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 4dd52ee2bd..a0a835a0b2 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/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 diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index e1980d2943..1572c7d668 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/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 diff --git a/transformer_engine/pytorch/tensor/__init__.py b/transformer_engine/pytorch/tensor/__init__.py index c3355b6c62..cc36b19949 100644 --- a/transformer_engine/pytorch/tensor/__init__.py +++ b/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", diff --git a/transformer_engine/pytorch/tensor/storage/__init__.py b/transformer_engine/pytorch/tensor/storage/__init__.py index 44a77d975f..1b1cb5a15f 100644 --- a/transformer_engine/pytorch/tensor/storage/__init__.py +++ b/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 diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 3473024c03..c336f54f22 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -1316,3 +1316,52 @@ def quantize( for i in range(self.num_tensors): self.quantizer.update_quantized(tensors[i], quantized_tensors[i], noop_flag=noop_flag) return quantized_tensors + + +def has_grouped_storage(tensor: torch.Tensor) -> bool: + """Whether a grouped parameter carries grouped storage rather than plain storage. + + High-precision grouped parameters are plain ``(num_tensors, *tensor_shape)`` + tensors, so this is False for them. + """ + return isinstance(tensor, GroupedTensorStorage) + + +def grouped_param_members(tensor: torch.Tensor) -> List[torch.Tensor]: + """Split a grouped parameter into its per-group member tensors. + + Members are views into the shared buffer in both representations, so this + does not copy. + """ + if not isinstance(tensor, GroupedTensorStorage): + return list(tensor.unbind(0)) + members = tensor.quantized_tensors + if members is None: + members = tensor.split_into_quantized_tensors() + return members + + +def grouped_storage_for_gemm( + tensor: torch.Tensor, + *, + num_tensors: int, + shapes: List[Tuple[int, int]], + dtype: torch.dtype, +) -> GroupedTensorStorage: + """Describe a grouped parameter as a GroupedTensorStorage for the grouped GEMM. + + Returns ``tensor`` unchanged when it already carries grouped storage. + Otherwise wraps the plain buffer. For uniform member shapes this only builds + a Python object, since first_dims/last_dims/tensor_offsets all stay None. + """ + if isinstance(tensor, GroupedTensorStorage): + return tensor + data = tensor if tensor.dtype == dtype else tensor.to(dtype=dtype) + return GroupedTensorStorage( + shape=(sum(s[0] for s in shapes), shapes[0][1]), + dtype=dtype, + num_tensors=num_tensors, + shapes=list(shapes), + quantizer=None, + data=data.reshape(-1), + )