Use a plain parameter for high-precision grouped linear weights - #3340
Draft
wujingyue wants to merge 1 commit into
Draft
Use a plain parameter for high-precision grouped linear weights#3340wujingyue wants to merge 1 commit into
wujingyue wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
GroupedLinear(single_grouped_weight=True)always registers the packed weight as aGroupedTensor— a storage-less_make_wrapper_subclasswhose payload lives in Pythonattributes (
rowwise_data,scale_inv,first_dims, ...).That is necessary for quantized weights, whose grouped
scale/amax/scale_invbuffersmust travel with the data. It is not necessary for high-precision weights.
GroupedTensorFromPyTorchGroupedTensor(csrc/type_converters.cpp) reads only four thingswhen there is no quantizer:
num_tensorsw.size(0)logical_shape(G*M, N)rowwise_dataptr +scalar_type()w.data_ptr(),w.dtype_with_gemm_swizzled_scalesFalseEvery other attribute is
Noneand short-circuits. For uniform member shapesmake_grouped_tensortakes the branch wherefirst_dims,last_dimsandtensor_offsetsall stay
None, so the descriptor carries no device-side metadata at all.The wrapper is not free. It has none of the serialization plumbing
QuantizedTensorcarries(
untyped_storage(), the_to_copyhandler,new_empty,__reduce_ex__), and it bansreshape-family ops, so
torch.save,torch.distributed.checkpointand generic tensor ops donot work on these parameters.
Change
Register high-precision grouped weights and biases as plain stacked
(G, M, N)parameters,and build a transient
GroupedTensorStorageat GEMM call time.This is already how the gradient side works — the wgrad path wraps a plain
main_gradbufferand hands autograd back a plain
(G, M, N)tensor. This applies the same transient-wrappattern to the parameter. Quantized weights are unchanged.
Scope
Only the fusible op (
ops/basic/grouped_linear.py) is converted.module/grouped_linear.pykeeps
GroupedTensor: it passes per-GEMM members directly into_GroupedLinear.applyandholds them via
ctx.origin_weight_refsweakrefs, so member lifetime and autograd wiringthere 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:
New tests (6, all passing): representation, forward/backward parity against the per-group
layout,
torch.save/loadround-trip, and a guard that quantized weights still get aGroupedTensor.No regressions — running the unmodified test file against baseline vs. this branch gives
identical results:
test_grouped_tensor.py(orig file)test_fusible_ops.py -k groupedThose pre-existing failures are an artifact of my test environment, not of this branch: the
container ships a
.socompiled ate7c550c5while the Python ismain, so newer C++symbols (e.g.
tex.scaled_swiglu) are absent. They are identical on both sides.black --line-length=100 --previewandpylint(10.00/10) clean.