Qualify dense INT2 PyTorch checkpoints

[Ti-Tai Wang (titaiwangms)]

Describe your changes

Add offline qualification coverage for the Olive side of the initial dense INT2 model-production workflow discussed in microsoft/onnxruntime#32657.

• Verify native PyTorch RTN, GPTQ, and KQuant produce real uniformly packed INT2 weights and preserve them across checkpoint save/reload.
• Verify SelectiveMixedPrecision drives each quantizer with an INT2 default, an SMP-selected INT4 projection, and an explicit INT8 projection override.
• Inspect reloaded QuantTensor settings and packed safetensors buffers instead of relying on checkpoint metadata alone.
• Check RTN weights against the original floating-point weights within the quantization-scale bound.
• Check GPTQ checkpoint reconstruction and module execution using deterministic, varied, fully local calibration sequences.
• Check KQuant checkpoint materialization against a directly recomputed KQuant result.
• Keep all tests offline with a locally constructed tiny dense Llama model and tokenizer.

No production code changes were required: the existing native RTN, GPTQ, KQuant, SMP, serialization, and reload paths already materialize dense INT2 and mixed INT2/INT4/INT8 checkpoints correctly. Mobius export qualification, ONNX-native INT2 RTN, and MoE/QMoE remain separate follow-up work.

Validation:

python -m pytest -q \
  test/passes/pytorch/test_dense_int2.py \
  test/passes/pytorch/test_rtn.py::test_rtn_int2_dense_checkpoint_packing_and_roundtrip \
  test/passes/pytorch/test_rtn.py::test_selective_mixed_precision_rtn_int2_int4_int8_checkpoint

6 passed

python -m ruff format --check \
  test/passes/pytorch/test_dense_int2.py \
  test/passes/pytorch/quantization_test_utils.py \
  test/passes/pytorch/test_rtn.py

python -m ruff check --ignore CPY001 \
  test/passes/pytorch/test_dense_int2.py \
  test/passes/pytorch/quantization_test_utils.py \
  test/passes/pytorch/test_rtn.py

git diff --check

Checklist before requesting a review

• Add unit tests for this change.
• Make sure all targeted tests can pass.
• Update documents if necessary. (No user-facing behavior changed.)
• Lint and apply fixes to your code by running lintrunner -a. (Targeted Ruff checks pass; the existing test_rtn.py has a pre-existing CPY001 finding.)
• Is this a user-facing change? If yes, give a description of this change to be included in the release notes. (Tests only; no release note required.)

(Optional) Issue link

Supports the INT2 model-production scope in microsoft/onnxruntime#32657.

[Copilot]

🟡 Changes recommended

An unresolved critical Ruff F401 finding remains.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds offline qualification tests for dense INT2 PyTorch RTN, GPTQ, KQuant, and mixed-precision checkpoints.

Changes:

• Adds local model, calibration, packing, and reload helpers.
• Verifies INT2 checkpoint packing, serialization, reconstruction, and execution.
• Adds SMP INT2/INT4/INT8 coverage.
• Critical: Remove the unused pytest import causing Ruff F401.

File summaries

File	Description
test/passes/pytorch/test_rtn.py	Adds RTN INT2 round-trip and mixed-precision coverage.
test/passes/pytorch/test_dense_int2.py	Adds GPTQ, KQuant, and SMP INT2 qualification tests.
test/passes/pytorch/quantization_test_utils.py	Provides local model, calibration, and packed-tensor helpers.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Ti-Tai Wang (titaiwangms)]

Cross-repository real-model qualification completed locally with hf-internal-testing/tiny-random-LlamaForCausalLM pinned at 9fb191250dd56d0ba7ec9785a025ed29c03d5998 (~1.03M parameters).

• Olive RTN produced a symmetric group-size-16 INT2 checkpoint.
• All 14 dense attention/MLP projections in both layers reloaded with effective bits=2.
• The checkpoint was consumed unchanged by Mobius and exported to 14 com.microsoft::MatMulNBits(bits=2) nodes.
• Source Olive packed bytes and scales matched every exported Mobius initializer exactly.
• Full-model CPU inference matched the Olive quantized PyTorch logits for a 3-token prompt: max absolute error 8.940696716308594e-08.

The generated checkpoint and ONNX artifact were kept as local qualification artifacts and were not committed.

[Ti-Tai Wang (titaiwangms)]

Real-model SMP qualification also completed with the same pinned tiny Llama revision.

Observed production behavior for this 2-layer model:

• SelectiveMixedPrecision(high_precision_mlp_down, bits=2, high_bits=4) planned model.layers.0.mlp.down_proj at INT4 and recorded an INT4 lm_head override.
• RTN does not quantize the LM head by default, so the actual checkpoint contains 14 quantized dense projections: 11 INT2, 1 INT4 (layer 0 down_proj), and 2 INT8 (explicit gate_proj overrides for both layers).
• The reloaded Olive quantized model executed successfully and produced the reference logits used for Mobius parity.

This verifies the real SMP-plan-to-RTN checkpoint path rather than assuming every layer receives the same high-precision override.