Skip to content

Use a plain parameter for high-precision grouped linear weights - #3340

Draft
wujingyue wants to merge 1 commit into
NVIDIA:mainfrom
wujingyue:grouped-linear-plain-hp-param
Draft

Use a plain parameter for high-precision grouped linear weights#3340
wujingyue wants to merge 1 commit into
NVIDIA:mainfrom
wujingyue:grouped-linear-plain-hp-param

Conversation

@wujingyue

Copy link
Copy Markdown
Contributor

Problem

GroupedLinear(single_grouped_weight=True) always registers the packed weight as a
GroupedTensor — a storage-less _make_wrapper_subclass whose payload lives in Python
attributes (rowwise_data, scale_inv, first_dims, ...).

That is necessary for quantized weights, whose grouped scale/amax/scale_inv buffers
must travel with the data. It is not necessary for high-precision weights.

GroupedTensorFromPyTorchGroupedTensor (csrc/type_converters.cpp) reads only four things
when there is no quantizer:

What the GEMM reads Plain-tensor equivalent
num_tensors w.size(0)
logical_shape (G*M, N)
rowwise_data ptr + scalar_type() w.data_ptr(), w.dtype
_with_gemm_swizzled_scales constant False

Every other attribute is None and short-circuits. For uniform member shapes
make_grouped_tensor takes the branch where first_dims, last_dims and tensor_offsets
all stay None, so the descriptor carries no device-side metadata at all.

The wrapper is not free. It has none of the serialization plumbing 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.

Change

Register high-precision grouped weights and biases as plain stacked (G, M, N) 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 hands autograd back a plain (G, M, N) tensor. This applies the same transient-wrap
pattern to the parameter. Quantized weights are unchanged.

Scope

Only the fusible op (ops/basic/grouped_linear.py) is converted. module/grouped_linear.py
keeps GroupedTensor: it passes per-GEMM members directly into _GroupedLinear.apply and
holds them via ctx.origin_weight_refs weakrefs, so member lifetime and autograd wiring
there need separate treatment.

Verification

H100 (CC 9.0), cuBLASLt 130501, NVTE_GROUPED_LINEAR_SINGLE_PARAM=1.

Confirmed the fused path actually engages and stays zero-copy:

bf16 graph-safe path supported: True
weight type: Parameter (3, 32, 64)
gemm descriptor: GroupedTensorStorage
zero-copy (same storage as param): True
device-side metadata: {'first_dims': None, 'last_dims': None, 'tensor_offsets': None}

New tests (6, all passing): representation, forward/backward parity against the per-group
layout, torch.save/load round-trip, and a guard that quantized weights still get a
GroupedTensor.

No regressions — running the unmodified test file against baseline vs. this branch gives
identical results:

Suite Baseline This branch
test_grouped_tensor.py (orig file) 20 failed, 161 passed, 35 skipped 20 failed, 161 passed, 35 skipped
test_fusible_ops.py -k grouped 50 failed, 137 passed, 223 skipped 50 failed, 137 passed, 223 skipped

Those pre-existing failures are an artifact of my test environment, not of this branch: the
container ships a .so compiled at e7c550c5 while the Python is main, so newer C++
symbols (e.g. tex.scaled_swiglu) are absent. They are identical on both sides.

black --line-length=100 --preview and pylint (10.00/10) clean.

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 <jingyuew@nvidia.com>
@github-actions github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Aug 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-contribution PRs from external contributor outside the core maintainers, representing community-driven work.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant