Repository navigation
Add native PyTorch INT3 weight quantization - #2716
Merged
Merged
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
4 of 5 tasks
Contributor
There was a problem hiding this comment.
Copilot review overview
🟢 Approval recommended
Focused regression coverage addresses packing, checkpoint round-trips, and export rejection, with no unresolved blocking findings.
Review effort: Balanced
Findings: None
What changed in this PR
Adds native INT3 weight checkpoints to Olive’s PyTorch quantization workflow, without enabling INT3 ONNX export or accelerated inference.
Changes:
- Adds packed INT3 storage, mixed-precision overrides, and checkpoint reload support.
- Rejects unsupported INT3 export and external-backend paths.
- Adds regression tests and documents format limitations.
| File | Description |
|---|---|
test/passes/pytorch/test_kquant.py |
Extends quantization-quality tests to INT3. |
test/passes/pytorch/test_int3.py |
Tests native passes, reloads, overrides, and backend restrictions. |
test/passes/onnx/test_model_builder.py |
Tests early INT3 export rejection. |
test/passes/onnx/test_mobius_model_builder.py |
Tests rejection before Mobius builds. |
test/common/quant/test_utils.py |
Adds independent packing-oracle and boundary tests. |
test/common/quant/test_tensor.py |
Tests INT3 dequantization and tensor operations. |
test/common/quant/test_hf_utils.py |
Tests checkpoint buffer sizing and export validation. |
test/common/hf/test_quant.py |
Tests rejection before export-module replacement. |
skills/olive/references/quantization-onboarding.md |
Documents native INT3 availability. |
olive/passes/pytorch/quant_utils.py |
Updates bit-width configuration help. |
olive/passes/pytorch/kquant.py |
Documents INT3 support in the class docstring. |
olive/passes/pytorch/gptqmodel.py |
Rejects INT3 defaults and dynamic overrides. |
olive/passes/pytorch/autogptq.py |
Rejects unsupported INT3 configuration. |
olive/passes/pytorch/autoclip.py |
Allows INT3 clipping simulation. |
olive/passes/pytorch/autoawq.py |
Rejects unsupported native INT3. |
olive/passes/onnx/model_builder.py |
Guards native INT3 export boundaries. |
olive/passes/onnx/mobius_model_builder.py |
Validates checkpoints before invoking Mobius. |
olive/passes/onnx/inc_quantization.py |
Rejects INT3 weight-only configuration. |
olive/constants.py |
Adds the INT3 bit-width enum value. |
olive/common/quant/utils.py |
Implements INT3 packing, unpacking, and shared sizing. |
olive/common/quant/hf_utils.py |
Enables INT3 checkpoint placeholders and export validation. |
olive/common/hf/quant.py |
Rejects INT3 export-compatible conversion. |
docs/source/features/quantization.md |
Explains checkpoint layout, usage, and limitations. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Xiaoyu Z (xiaoyu-work)
approved these changes
Oct 6, 2026
Ti-Tai Wang (titaiwangms)
added a commit
to onnxruntime/mobius
that referenced
this pull request
Oct 15, 2026
## Summary Add projection-specific QMoE export for Olive fused expert checkpoints. FC1 and FC2 independently support INT2, INT3, INT4, and INT8: all 16 pairs are accepted subject to format validation, without an experimental opt-in. This is **export support, not runtime qualification**. INT3 follows the proposed raw-weight direction in microsoft/onnxruntime#32657 and consumes the native Torch INT3 checkpoint format introduced by microsoft/Olive#2716. INT3 output retains draft-format/runtime-warning metadata; this PR does not provide INT3 ORT kernels or claim that released ORT can execute it. ## Changes - Resolve FC1/FC2 widths independently instead of using an explicit pair whitelist. This intentionally includes `(2,3)`, `(3,2)`, `(3,8)`, `(8,3)`, `(2,2)`, and `(2,8)`, in addition to the other valid pairs. - Remove the experimental INT3 flag and opt-in rejection. Export requires an explicitly quantized checkpoint with compatible projection configurations; ordinary unquantized models do not become INT3 automatically. - Require exact Olive expert overrides and a supported model-wide INT4 fallback for projection-specific layouts. Preserve the existing uniform INT4 path. - Preserve row-local, LSB-first packed unsigned INT3 codes with symmetric zero point 4. For pairs containing INT3, interleave FC1 gate/up payloads lazily during serialization in public HF builds; direct preprocessing remains eager by default. Otherwise retain graph-side interleaving. Preserve FC2's K-last layout without requantization. - Validate shared power-of-two group sizes, aligned dimensions, sidecar shapes, packed zero points, and finite nonnegative scales. Reject scale conversion overflow when the adapter supplies the activation dtype. INT3-containing pairs additionally require block size 32/64/128 and FP16/BF16 activation/scale agreement. Non-INT3 layouts retain the standard loader's dtype conversion. - Apply the same numeric scale checks to the legacy uniform INT4 path without imposing the projection-specific shape ABI. Rejecting previously accepted negative, NaN, infinite, or activation-cast-overflowing scales is an intentional validation tightening. - Bind deferred FC1 weights, scales, and asymmetric zero points through uncached lazy tensors, rather than retaining transformed copies for every layer. Keep transform accounting component-scoped when prefix stripping produces identical local initializer names. - Preserve INT2/INT4/INT8 zero-point semantics and validate optional explicit offset-4 INT3 zero points before omitting them from the graph. - Emit projection-width attributes where needed and retain INT3 draft-format metadata without assigning an unapproved schema version. - Cover all 16 pairs across FP16/BF16, independent packing/dequantization oracles, CPU execution of graph-side interleaving, invalid scales, local-checkpoint public builds, and external-data save/reload. - Pass the model activation dtype through all three Qwen3.5-MoE preprocessing paths: text, VL global quantization, and VL decoder component quantization. Add 57 regression cases covering all seven INT3-containing pairs across FP16/BF16, incompatible dtypes, and non-INT3 scale conversion behavior. ## Validation Latest offline CPU validation includes the merged `main`, Qwen3.5-MoE dtype propagation, uniform scale checks, deferred serialization, and component routing: ```bash PYTHONPATH="$PWD/src" HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \ TOKENIZERS_PARALLELISM=false PYTHONDONTWRITEBYTECODE=1 CUDA_VISIBLE_DEVICES='' \ timeout 300 /home/titaiwang/miniconda3/envs/dev/bin/python -m pytest \ --rootdir="$PWD" -c "$PWD/pyproject.toml" \ src/mobius/models/int3_qmoe_test.py src/mobius/models/qwen35_test.py \ src/mobius/models/moe_test.py \ src/mobius/_weight_utils_test.py src/mobius/_flags_test.py \ src/mobius/components/_moe_test.py \ src/mobius/_model_package_test.py::TestModelPackageApplyWeights \ src/mobius/weights/_adapters_test.py \ src/mobius/integrations/_weight_loading_test.py \ src/mobius/integrations/transformers/_builder_test.py \ -q -m 'not integration and not arch_validation and not golden and not generation' \ --tb=short ``` **Passed: 728 tests, 1 deselected, in 52.36 seconds.** Coverage includes distinct per-layer payloads, source-storage aliasing, release of transformed tensors between callbacks, repeated saves, asymmetric FC1 zero points, and positive/negative component-local-name collision regressions. ```bash lintrunner src/mobius/_model_package.py src/mobius/_model_package_test.py \ src/mobius/_weight_utils.py src/mobius/integrations/_weight_loading.py \ src/mobius/integrations/transformers/_builder.py src/mobius/models/granite.py \ src/mobius/models/int3_qmoe_test.py src/mobius/models/llada.py \ src/mobius/models/moe.py src/mobius/models/qwen35.py \ src/mobius/models/qwen35_test.py src/mobius/weights/_adapters.py git diff --check ``` **Both passed.** Ruff formatting and checking also passed for the same 12 Python files. Scoped mypy reported 77 errors outside changed hunks; no project-wide type-check success or fully established clean baseline is claimed. ## Draft boundaries - All 16 pairs are exporter-supported layouts, not a runtime-qualified matrix. Execution depends on the target ORT version and execution provider. - Source checkpoint loading remains eager. Uncached lazy interleaving bounds additional FC1 payload retention, not total checkpoint or serialization peak memory. No full-checkpoint RSS or concurrent-writer qualification was run. - Direct callers opting into deferred preprocessing must forward the populated transform map when binding the returned weights. Explicit adapters and older preprocessing hooks retain their existing behavior. - No INT3 ORT kernel execution, checkpoint-to-runtime logits parity, real-model quality, or performance qualification was run. The CPU interleave tests exercise standard ONNX graph operations, not QMoE kernels. - Dense INT3 `MatMulNBits`, quantized embeddings, unsupported source layouts, and schema-unapproved tails are not enabled. - Final QMoE schema/version agreement, complete quantization provenance, and runtime qualification remain follow-up gates. This does not complete the end-to-end plan's P5 acceptance gate. - INT2 runtime/CI qualification is separate in #765 and is not included in this PR. --------- Signed-off-by: titaiwang <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: fba8a408-cd63-49f0-b8b4-5daa19ed42d4
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.
Describe your changes
Add native PyTorch INT3 weight quantization for Olive's RTN, KQuant, and GPTQ passes, including explicit mixed-precision overrides and checkpoint save/reload.
This is Torch checkpoint support, not an INT3 ONNX export or accelerated-kernel implementation. No real-model accuracy or performance qualification is claimed. Quantized fused-MoE forward execution currently requires eager experts.
Validation
Validated on an isolated branch based on upstream main at
ba856b69ec21a3080319c7affce62dc1546b94e2.Result: 712 passed.
HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 timeout 180 python -m pytest test/common/quant/test_forward_parity.py test/passes/onnx/test_model_builder.py test/passes/pytorch/test_kquant.py -k 'not real_hf and (forward_parity or onnx_export_parity or export_rejected or rejects_int3 or find_qparams or saved_test_model_path or config_only_test_dir or single_file_fallback or multi_file_output or prechecks_extra_options or structured_config or cpu_embedding or active_loader)' -q --disable-warnings --maxfail=3Result: 42 passed, 43 deselected.
After removing a redundant test import:
HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 timeout 120 python -m pytest test/passes/onnx/test_model_builder.py -k 'rejects_int3' -q --disable-warningsResult: 3 passed, 24 deselected.
Result: All commands exited successfully; Ruff checked/formatted all 21 changed Python files.
Result: Failed only with 21
RUFF/CPY001copyright-notice findings. This rule also flags existing unchanged files with the repository's copyright headers. No unrelated headers or lint configuration were changed; the same finding also affects the new test file.The full test suite, hardware-specific kernels, and real-model download-dependent tests were not run.
Checklist before requesting a review
lintrunner -aRelease notes: Native PyTorch RTN, KQuant, and GPTQ now support INT3 weight checkpoints and explicit mixed-precision INT3 overrides. INT3 uses tightly packed storage and supports save/reload and eager dequantized execution; ONNX export and accelerated INT3 inference are not yet supported.
(Optional) Issue link
Related runtime design discussion: microsoft/onnxruntime#32657. This change does not depend on or implement that runtime proposal.