Int8 conv rot - #1681
Open
Calamdor wants to merge 4 commits into
Open
Conversation
Adds LinearInt8ConvRot and a hadamard rotation helper, wired up as a new DataType option alongside the existing W8A8 variants. The Hadamard matrix is precomputed into a non-persistent buffer so forward/backward run cleanly inside torch.compile'd + gradient-checkpointed blocks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit 12b031d)
…e int8 kernel
Three changes on top of the original, all in LinearInt8ConvRot.forward.
1. CONVROT_FWD selects the forward route. The bf16 route was already present but only reachable below the
int8 row floor, so real training always took the fused kernel. The two differ in a way worth choosing
between -- measured here on a 5090, 5376 -> 7168, against the dequantized weight:
tokens fused int8 bf16
2048 1.27 ms 2.74 ms
12000 10.13 ms 11.70 ms
error 9.43e-03 3.61e-03
So bf16 is ~2.6x more accurate and 1.15-2.2x slower. Both carry the rotation's own quantization error
identically; what separates them is that the fused kernel also quantizes ACTIVATIONS per row every call.
Default stays int8, which is the faster route here and preserves existing behaviour.
Note musubi measures bf16 as more accurate AND faster on an H100 (4.10 vs 6.22 s/it), attributing it to
the fused kernel rather than int8 arithmetic. That half does not transfer to consumer hardware, where
torch._int_mm is fast: the accuracy ordering holds, the speed ordering reverses.
2. The bf16 route rotates the ACTIVATIONS instead of un-rotating the weight. The Hadamard is orthogonal so
(xR)(WR)^T == x W^T either way, but the costs are not symmetric: rotating x is tokens x in_features,
un-rotating W is out_features x in_features on every call. Worth 3.33 -> 2.74 ms at 2048 tokens.
unquantized_weight() still un-rotates, because its callers (LoRA decompose, offload sizing) want the
original basis; only this forward can use the stored basis directly.
3. Decompress before the fused kernel. LinearW8A8 inherits CompressedWeightMixin, so self.weight may be a
compressed blob, and LinearInt8ConvRotFunction consumes the int8 weight directly -- enabling weight
compression with INT_W8A8_CONVROT otherwise feeds the blob to torch._int_mm.
Also documents why the row floor exists, since it reads as an arbitrary limit and is not one. torch._int_mm
requires strictly more than 16 rows (verified on torch 2.12.0+cu130: 15 and 16 both raise "self.size(0) needs
to be greater than 16", 17 works), which also makes the bf16 branch the only path that works for small
inputs, whatever CONVROT_FWD asks for.
Verified on a clean checkout of master (e1629cc), no other branches: quantizes, the rotation buffer is
orthogonal to 5.96e-08, and against real MiniMax-H3 transformer weights it beats plain INT_W8A8 by 4.83x on
attn.qkv_proj and 4.42x on mlp.fc1 (relative error 2.1e-02 vs 1.04e-01), while being neutral on
attn.out_proj at 0.93x -- consistent with those layers' outlier ratios of 8.4x, 7.1x and 4.6x.
ConvRot inherits CompressedWeightMixin through LinearW8A8 but never exposed compression: quantize() overrides the parent and skipped _compress_weight, and unquantized_weight() read the raw weight, so a compressed ConvRot layer would have failed on the DoRA/export path. Wire both, and add the enum + UI entry so the combination is selectable. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
Author
|
Updated: |
…ences Assign the quantized weight through .data instead of rebinding self.weight to a new Parameter, matching LinearW8A8/LinearFp8/LinearNf4/LinearSVD and the compression mixin. Rebinding invalidates any reference taken before quantize. Also drop comment references to a working branch and plan document that do not exist outside the development branch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.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
2 participants
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.
Summary
Adds INT_W8A8_CONVROT: INT8 W8A8 quantization with a block-Hadamard rotation applied along the contraction dimension before quantizing, so channel outliers are spread across a block instead of dominating one channel's quantization range. The GEMM is unchanged from LinearW8A8 (torch._int_mm); the rotation is the only added op.
Two commits:
a. CONVROT_FWD selects the forward route. The bf16 route already existed but was only reachable below the int8 row floor, so real training always took the fused kernel. On a 5090 (5376→7168) bf16 is ~2.6× more accurate and 1.15–2.2× slower; default stays int8, preserving existing behaviour. musubi measures bf16 as faster on an H100 — that half doesn't transfer to consumer hardware, the accuracy half does.
b. The bf16 route rotates the activations rather than un-rotating the weight. The Hadamard is orthogonal so (xR)(WR)ᵀ == xWᵀ, but rotating x is tokens × in_features against out_features × in_features every call — 3.33 → 2.74 ms at 2048 tokens.
c. Bug fix: decompress before the fused kernel. LinearW8A8 inherits CompressedWeightMixin, so self.weight may be a compressed blob and LinearInt8ConvRotFunction consumes the int8 weight directly — enabling weight compression with INT_W8A8_CONVROT otherwise hands the blob to torch._int_mm.
Also documents why the x.shape[0] > 16 floor exists, because it reads as arbitrary and isn't: torch._int_mm requires strictly more than 16 rows (15 and 16 both raise self.size(0) needs to be greater than 16; 17 works), which also makes the bf16 branch the only path that works for small inputs regardless of CONVROT_FWD.
Worth knowing: rotation and weight compression are substitutes. On the same MiniMax-H3 transformer, ConvRot compresses to 25% saved where plain int8 saves 46% — plain int8 compresses because outliers force a large tensorwise scale and leave everything else low-entropy, which is exactly the structure ConvRot removes.
Test plan
pre-commit run --all-filespassesThe benefit tracks the outlier ratio, and is neutral-to-slightly-negative where outliers are mild — consistent with what the rotation is for.
AI assistance
(Note: commit 738cc0c is dxqb's own work, included unmodified for attribution; the declaration above covers 6515b39.)