feat(quantization): scoped calibration pipelines via algo_cfg [prototype]

[Fridah-nv]

What does this PR do?

Type of change: new feature (prototype — draft, not for merge as-is)

Adds an opt-in algo_cfg key that assigns an ordered calibration pipeline per scope, instead of a single model-wide algorithm. This expresses two things the current surface cannot: different algorithms for different parts of the model (parallel), and an ordered pipeline on the same targets where each stage consumes the previous one's mutated weights/scales (sequential).

Usage

config = {
    "quant_cfg": [...],                                    # unchanged
    "algo_cfg": [
        {"module_name": "*self_attn*",         "cfg": ["awq_lite", "mse"]},
        {"module_name": "*mlp*",               "cfg": ["max", {"method": "gptq", "block_size": 64}]},
        {"quantizer_name": "*input_quantizer", "cfg": ["max"]},
    ],
    "algorithm": "max",   # fallback for anything no entry matches
}
mtq.quantize(model, config, forward_loop)

An algo_cfg entry has the same {<selector>, "cfg": ...} shape as a quant_cfg entry: quant_cfg entries carry quantizer attributes, algo_cfg entries carry the ordered algorithms. Exactly one selector per entry — module_name (module/weight-level algorithms, role implied by the algorithm) or quantizer_name (when the role must be chosen explicitly, e.g. max on inputs only).

What changed

algo_cfg.py (new) — the compile half.

• compile_algo_cfg(config, model) lowers algo_cfg + algorithm into an ordered list of AlgoStages. It is pure: it reads the quantized model's structure (quantizer/linear names) to resolve globs and validate, but mutates nothing, runs no forward and touches no data. So bad configs fail before any expensive calibration, it is testable without running a model, and the plan is a pure function of (config, structure) — hence identical on every rank, which is what keeps predicate scoping from desynchronizing collectives.
• ALGO_CAPABILITIES — a small declared table per algorithm (granularity, role, requires, produces, requires_absent, forward traits).
• Validation, reporting every problem in one pass: unknown algorithm, empty scope, role mismatch, fusible siblings split across pipelines, a stage whose every write is overwritten before being read, and repeating an algorithm whose own output violates its precondition. The last two are derived from the capability table; overlap is judged per state token and per quantizer role, so two stages sharing a module do not conflict if they write different roles.
• derive_handoff() — a stage whose inputs an earlier stage already produced is told to skip its own initialization (skip_max_init), derived from the declared capabilities rather than a hard-coded algorithm pair.
• plan_hash() — excludes provenance, so equivalent plans written two ways hash the same; intended as the rank-identical-plan assert.

mode.py — the execute half. A calibration_plan mode whose convert compiles the plan, then runs each stage in order through the existing wrapped_calib_func with a should_process write-mask built from the stage's scope. Records one mode (per-stage modes would bloat state); restore is the generic quantizer-state snapshot, unchanged.

config.py — AlgoCfgEntry, CalibrationPlanConfig, QuantizeConfig.algo_cfg / .strict, MseCalibConfig.skip_max_init; need_calibration considers algo_cfg.

model_quant.py — calibrate(..., algo_cfg=, strict=); quantize passes them through.

model_calib.py — should_process write-mask threaded into the module-iteration points of max / mse / awq / awq_clip / gptq / smoothquant. Default None means "whole model", i.e. today's behaviour. The mask gates writes only and never toggles enable-state, so the activations search-based algorithms see are unchanged.

algorithm lowers through the same path as its all-"*" case, so there is no second engine. With no algo_cfg the old path, its numerics and its saved state are untouched.

Two upstream bugs fixed along the way

Both reproduce on today's un-scoped algorithm=[...] list — they are not artifacts of the scoped plan, but they block sequencing, which is the point of this change.

1. Anything sequenced after mse crashed. _mse_calibrate_weights assigned weight_quantizer._calibrator to its search calibrator and never restored it, so the next stage that collects stats re-entered a spent calibrator: algorithm=['max','mse','max'] → TypeError: unsupported operand type(s) for *: 'NoneType' and 'Tensor'. Now restored in a finally. The structural fix is still to make the per-tensor strategy config-selected rather than swapping the object.
2. awq_lite calibrated the whole model regardless of scope, because it calls enable_stats_collection(model) / finish_stats_collection(model) directly. The write-mask has to reach helper calls inside an algorithm, not just its top-level module loop — worth knowing for any similar work.

Testing

• tests/unit/torch/quantization/test_algo_cfg.py — 25 new tests: lowering, plan-hash equivalence, every validation rule, derived handoff, write-mask, enable-state untouched, single recorded mode, and "a stage can follow mse". ~1s.
• pytest tests/unit/torch/quantization/ --ignore=.../plugins → 827 passed, 1 skipped, unchanged from main. (plugins/ does not collect in my env: test_diffusers_wan_conv3d.py fails to import diffusers — pre-existing, unrelated.)
• Backward compatibility checked numerically: algorithm="max" and the equivalent algo_cfg produce bit-identical amax on every weight quantizer, and compile to the same plan hash.
• Save/restore round-trip on a 5-stage scoped plan reproduces the calibrated state bit-identically.

Not covered: distributed (no multi-GPU available — the rank-identical-plan property is structural but untested; a 2-GPU TP/EP test is still needed), shared-forward batching across independent stages, auto_quantize per-layer algorithm, and GPU-only algorithms (gptq / svdquant compile and validate but were not executed).

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ — algo_cfg is opt-in; without it the old path, numerics and saved state are unchanged.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A — no copied code, no new dependencies.
• Did you write any new necessary tests?: ✅ — test_algo_cfg.py (25 tests).
• Did you update Changelog?: ❌ — deferred while the config surface is under design review; will add before this leaves draft.
• Did you get Claude approval on this PR?: ❌ — draft.

Additional Information

Draft on purpose. The config surface (algo_cfg, per-entry cfg) and how much of the capability contract belongs in the first cut are what I would most like feedback on before polishing this for merge.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2292/
Built to branch gh-pages at 2026-09-01 00:23 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 94.27083% with 22 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.14%. Comparing base (8810eb5) to head (233bbfd).

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/algo_cfg.py	93.92%	17 Missing ⚠️
modelopt/torch/quantization/config.py	93.54%	2 Missing ⚠️
modelopt/torch/quantization/mode.py	93.33%	2 Missing ⚠️
modelopt/torch/quantization/model_calib.py	97.56%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2292      +/-   ##
==========================================
+ Coverage   79.05%   79.14%   +0.08%     
==========================================
  Files         525      526       +1     
  Lines       61106    61457     +351     
==========================================
+ Hits        48308    48640     +332     
- Misses      12798    12817      +19     

Flag	Coverage Δ
unit	56.04% <94.27%> (+0.22%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{Comment @coderabbitai help to get the list of available commands.}