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]

🔵 Needs a closer look

Update the packing assertion to support valid padded, non-divisible shapes.

Pull request overview

Adds offline qualification tests for dense INT2 RTN checkpoint packing, serialization round-trips, and SelectiveMixedPrecision INT2/INT4/INT8 behavior.

Changes:

• Adds a local tiny dense Llama fixture.
• Validates packed buffers, metadata, reloads, and numerical fidelity.
• Tests mixed-precision RTN overrides.

File summaries

File	Summary
test/passes/pytorch/test_rtn.py	Adds dense INT2 RTN and mixed-precision qualification tests. One assertion should account for padding in non-divisible packed dimensions.

Review details

Suppressed comments (1)

test/passes/pytorch/test_rtn.py:475

• These assertions assume the unpacked last dimension is divisible by the packing factor, but pack_to_uint8 deliberately pads and stores ceil(shape[-1] / packing_factor) bytes for valid non-divisible shapes. The current 32/64-wide fixture hides this mismatch; use a ceiling-sized expected shape and compare the packed element count against that so this helper validates (rather than rejects) legal checkpoints with padding.

    packing_factor = 8 // bits
    assert quant_tensor.shape[-1] % packing_factor == 0
    assert quant_tensor.qweight.shape == (*quant_tensor.shape[:-1], quant_tensor.shape[-1] // packing_factor)
    assert quant_tensor.qweight.numel() * packing_factor == quant_tensor.numel()

• Files reviewed: 1/1 changed files
• Comments generated: 0
• 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)]

Replaced by #2671, using an upstream microsoft/Olive branch so trusted CI secrets, including the Hugging Face token, are available.