Fix NF4 dequantize when the scales span more than one int8 group - #2421
Open
rootkiller6788 wants to merge 2 commits into
Open
rootkiller6788 wants to merge 2 commits into
rootkiller6788 wants to merge 2 commits into
Conversation
added 2 commits
September 10, 2026 17:22
dequantize() cut the flattened int8 scale codes down to the block count before
dequantizing them. Every group of scales carries its own double_scale, so
dividing the cut 1-D vector by double_scale.unsqueeze(-1) broadcasts over the
group axis instead of pairing each scale with its own group. Harmless with a
single group, but with more than one it builds a (num_groups, num_blocks)
tensor and dequantize dies:
RuntimeError: The size of tensor a (512) must match the size of tensor b (65536)
Dequantize the scales group by group while they still have their
(num_groups, scale_block_size) shape, then drop the tail padding. The cuda and
cpu branches now share those already-aligned scales.
Signed-off-by: rootkiller6788 <17553215+rootkiller6788@user.noreply.gitee.com>
The existing qtensor test used block_size 2 with scale_block_sizes 4 on 8 elements, which works out to 8 blocks and exactly one scale group, so the group axis was never exercised. These cases span 2 groups (16 and 15 elements) and 128 groups, with the 15-element one also covering the input padding path. Signed-off-by: rootkiller6788 <17553215+rootkiller6788@user.noreply.gitee.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe NF4 dequantization path now processes grouped scale metadata before flattening and padding removal. New tests cover multiple scale groups, padded tails, exact values, shape preservation, and wide-input round trips. ChangesNF4 dequantization
Priority: ⬇️ Low Estimated code review effort: 2 (Simple) | ~15 minutes Change: Bug fix Suggested reviewers: Merge Risk: ⚪ Minimal · up to The corrected dequantization path preserves grouped scale metadata and is covered for multi-group and padded inputs. No actionable merge risk remains. 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
1 participant
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.
Ran into this trying NF4 with scale_bits=8 and a small scale_block_sizes.
dequantize() assumes the scales are a flat per-block vector, but with double quantization they are stored as int8 groups with one double_scale per group. It trimmed the flattened codes down to the block count first and dequantized after, so double_scale.unsqueeze(-1) broadcast over the group axis instead of pairing each scale with its own group:
A single scale group hides it completely, which is why it went unnoticed - that is what you get whenever scale_block_sizes is at least as large as the block count.
The fix moves the group-wise dequantization before the trim, so it happens on the (num_groups, scale_block_size) shape where the broadcast lines up, and the padding is dropped afterwards. The cuda branch needs the same thing, so both now share the aligned scales.
Testing: added tests/unit/torch/quantization/test_nf4_tensor.py. The existing qtensor test uses block_size 2 with scale_block_sizes 4 on 8 elements, which is 8 blocks and exactly one group, so nothing covered this. The new cases span 2 groups (16 and 15 elements, the latter also going through input padding) and 128 groups; all of them fail on main with the error above. I also checked the numbers against a separate from-scratch reconstruction (own block amax, own NF4 lookup, explicit per-group loop for the scales) - exact agreement for 1, 2 and 128 groups, and the recovered per-block scales land within ~1e-3 of the true block amax, i.e. just the int8 scale quantization error.
Full tests/unit/torch/quantization is 878 passed. The single failure there (test_dbrx) is pre-existing and fails without this change too.
Did not touch CHANGELOG.rst - happy to add an entry if you want one for this.
Summary by CodeRabbit
Bug Fixes
Tests