Fix HQQ MatMulNBits weight layout - #2589
Conversation
Quantize MatMul weights in output-channel-major K blocks so packed weights, scales, and zero points follow the MatMulNBits storage contract. Add numerical regression coverage for aligned and padded K dimensions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7271f78d-9d14-4ce5-a154-5a7ffa9a70c2
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
There was a problem hiding this comment.
Pull request overview
Fixes HQQ weight-only quantization for ONNX MatMulNBits by aligning packed weight/scales/zero-point layout with ORT’s expected output-channel-major (N-major) storage contract, preventing quantization parameters from being applied to the wrong blocks.
Changes:
- Transpose MatMul weights from
(K, N)to(N, K)before HQQ quantization and quantize along the transposed K axis to matchMatMulNBitsblock ordering. - Explicitly reject unsupported quantization axes (non-zero).
- Add an ONNX Runtime numeric regression test covering both block-aligned and padded-K cases.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
olive/passes/onnx/hqq_quantization.py |
Transposes weights and adjusts grouping/packing to produce MatMulNBits initializers in ORT’s expected (N, K_blocks, blob) layout; rejects unsupported axes. |
test/passes/onnx/test_hqq_quantization.py |
Adds ORT-based numerical regression coverage for the corrected MatMulNBits layout (aligned and padded K). |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
|
@DelwinKim please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
olive/passes/onnx/hqq_quantization.py:234
_quantize_internal_numpycan move tensors to CUDA (b_array_torch = b_array_torch.cuda()), but it returnspacked_torch/scales/zero_pointswithout forcing them back to CPU. ONNX initializers must be CPU-backed; returning CUDA tensors here is fragile (and can break ifir.tensor(...)doesn’t explicitly handle.cpu()conversion). Convert the outputs to CPU (ideally numpy) before returning so model creation/serialization is device-independent.
blob_size = block_size // 2
k_blocks = (rows + block_size - 1) // block_size
packed_torch = packed_torch.reshape(cols, k_blocks, blob_size)
return packed_torch, scales, zero_points
Describe your changes
Fix HQQ weight-only quantization for ONNX
MatMulNBits.MatMulstores its weight as[K, N], whileMatMulNBitsexpects packed weights, scales, and zero points in output-channel-major order with independent K-axis blocks. The previous HQQ path grouped the row-major[K, N]initializer directly and flattened[K_blocks, N]quantization parameters into storage interpreted as[N, K_blocks]. The resulting model was structurally valid but applied quantization parameters to the wrong weight blocks, resulting in heavy degradation in model's performance.This change:
[N, K]before HQQ quantization.[N, K_blocks]ordering.[K, N]dimensions for theMatMulNBitsKandNattributes.Checklist before requesting a review
lintrunner -aRelease note: Fix HQQ ONNX quantization producing incorrect
MatMulNBitsweight, scale, and zero-point layouts.(Optional) Issue link
No existing issue.