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

[Fridah-nv]

What does this PR do?

Type of change: new feature

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": ["mse", {"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) or quantizer_name (when the role must be chosen explicitly).

What changed

algo_cfg.py (new) — the compile half.

• compile_algo_cfg(config, model) lowers algo_cfg + algorithm into an ordered list of AlgoStages. Pure: it reads the quantized model's structure 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 declared table per algorithm (granularity, role, requires, produces, requires_absent, supports_scoping, forward traits).
• Validation, reporting every problem in one pass: unknown algorithm, empty scope, role mismatch, a module-level algorithm given a scope that is not closed over its modules, fusible siblings split across pipelines, a stage whose every write is overwritten before being read, repeating an algorithm whose own output violates its precondition, and scoping an algorithm that cannot honour the write-mask. 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; restore is the generic quantizer-state snapshot, unchanged.

config.py — AlgoCfgEntry, CalibrationPlanConfig, QuantizeConfig.algo_cfg / .strict, MseCalibConfig.skip_max_init, GPTQCalibConfig.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.

GPTQ can now consume a preceding range search

gptq() previously re-derived amax from max unconditionally, so a range search in front of it was discarded and the compiler correctly reported the search as a dead stage. But rounding error is only compensated consistently if GPTQ works against the grid the model actually uses, so the search belongs before GPTQ — which is the order DeepCompressor's QoQ recipes ship (qoq-gchn.yaml: enable_calib_range then enable_kernel_gptq).

gptq now declares weight_amax as an input and takes skip_max_init; the executor derives the flag. mse → gptq keeps the searched amax bit-identically while differing from plain GPTQ. Declaring that one token also moved 5 of 121 ordered algorithm pairs and 105 of 1331 triples from "rejected" to "composing".

Chains verified to work

Three chains are exercised end to end and pinned by tests. All numbers are from a seeded CPU toy
model — these establish that sequencing is correct, not that any chain improves accuracy.

chain	what is verified	evidence
awq_lite → awq_clip	bit-identical to the bundled awq_full on every weight quantizer	test_awq_full_is_exactly_its_two_stage_pipeline
awq_lite → mse	MSE refines the amax AWQ left behind — moves it by up to 15.8% of scale	test_awq_then_mse_refines_the_smoothed_weights
mse → gptq	GPTQ keeps the searched amax bit-identically and the result differs from plain GPTQ	test_gptq_preserves_a_preceding_range_search

The first is the strongest signal in the PR: awq() already runs awq_lite then awq_clip
internally when asked for awq_full, and the plan surface reproduces that composite exactly as an
ordinary two-stage pipeline. An existing bundled algorithm is a pipeline; this just makes it one
the user can write. (The test also compares against awq_lite alone, so it cannot pass by
awq_clip doing nothing.)

The second is the cheap-refinement case: MSE re-searches the clipping range on AWQ's smoothed
weights with no forward pass, where awq_clip needs a full search pass. Whether it matches
awq_clip in accuracy is exactly the first experiment to run on a real model.

The third is the ordering the prior art ships (DeepCompressor's QoQ runs its range search before
the GPTQ kernel) and is what the skip_max_init work above enables.

Two more compile clean but were not executed here: gptq → mse (the reverse ordering, worth
running as the A/B control) and smoothquant → gptq (the most common chain in llm-compressor;
smoothquant finds nothing to smooth on the toy model, so it proves nothing at this scale).

Upstream bugs fixed along the way

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

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.
2. awq_lite calibrated the whole model regardless of scope, because it calls enable_stats_collection(model) directly. The write-mask has to reach helper calls inside an algorithm, not just its top-level module loop.

Selector rules, and what enforces them

Exactly one selector per entry is enforced at config-construction time by
AlgoCfgEntry._normalize_entry (a pydantic before-validator), so a malformed entry fails before
the model is ever consulted.

The semantic half — module_name for module-level algorithms, quantizer_name when the role
must be chosen explicitly — is enforced by two compile-time rules:

• a role-fixed algorithm whose quantizer_name scope matches only the wrong role is rejected
  (mse on *input_quantizer writes nothing);
• a module-level algorithm's scope must be closed under module ownership — every quantizer of
  every module it touches has to be in scope.

The second rule closes a hole worth calling out, since it is the kind of thing this whole design
is supposed to prevent. {"quantizer_name": "*weight_quantizer", "cfg": ["awq_lite"]} used to
compile clean while role_quantizers reported zero writable input quantizers — and the run then
wrote pre_quant_scale to all 14 of them. The write-mask cannot stop that: a quantizer_name
scope resolves to its parent modules, awq_lite gates on the module name, and then writes
module.input_quantizer directly. The worse half is that effective_produces and
_token_overlap are derived from the declared role sets, so the compiler understated what the
stage writes and would have missed a real conflict with a later input-quantizer stage.

Closure is the right formulation rather than a blunt "module granularity forbids quantizer_name",
because algorithm="awq_lite" lowers to a quantizer_name="*" stage — whole-model and
module_name scopes are closed by construction, so the legacy path and every shipped example are
unaffected.

Known limitation: four algorithms cannot be scoped yet

local_hessian and nvfp4_act_headroom take no should_process; svdquant and lsq swallow it via **kwargs and would run over the whole model, clobbering other stages. AlgoCapabilities.supports_scoping records this and compile rejects a scoped stage for them with a clear message rather than mis-calibrating; whole-model use is unchanged. Adding real should_process support to those four is follow-up work.

Testing

• tests/unit/torch/quantization/test_algo_cfg.py — 38 tests: lowering, plan-hash equivalence, every validation rule, derived handoff, write-mask, enable-state untouched, single recorded mode, mse → gptq preserving the searched amax, awq_lite → mse, and three fixture shapes (single-level quantizer, SequentialQuantizer, weight-only/no-forward).
• pytest tests/unit/torch/quantization/ --ignore=.../plugins → 865 passed, 1 skipped. (plugins/ does not collect in my env: test_diffusers_wan_conv3d.py fails to import diffusers — pre-existing, unrelated.)
• Save/restore round-trip on a 5-stage scoped plan reproduces the calibrated state bit-identically.

Backward compatibility. algorithm="max" and the equivalent algo_cfg compile to the same plan hash and produce bit-identical amax. Note this was initially verified too narrowly — only for a single-level weight quantizer with a forward loop. A code review found two cases where equivalence broke: SequentialQuantizer configs (W4A8 / INT4-AWQ) under a module_name scope, where sub-quantizers are grandchildren of the linear and were reachable by no module scope; and the weight-only/no-forward path, where a quantizer_name entry stripped the fallback stage of its modules and weight_only_quantize iterated nothing. Both are fixed and both now have regression tests with the fixture shapes that were missing.

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, GPU-only algorithms (gptq / svdquant compile and validate; only gptq was executed, on CPU), and any accuracy measurement — nothing here shows a chain improves a model, only that chains are expressible and correct.

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 (38 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.

Open items from an internal review, not yet addressed: derive_handoff uses any-overlap where it needs superset coverage (a narrow producer can wrongly set skip_max_init on a wider consumer); the fused-sibling check skips groups where only some members are matched by an entry; a derived handoff overrides an explicitly written kwarg; algo_cfg combined with layerwise.enable=True builds the predicate from full-model names but is applied to subtree-relative ones; and compile_algo_cfg recomputes the model index per call, which will not scale to per-layer configs on large models.

🤖 Generated with Claude Code

[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-08 22:05 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 95.11002% with 20 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.92%. Comparing base (8810eb5) to head (3cc5221).
⚠️ Report is 24 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/algo_cfg.py	94.35%	17 Missing ⚠️
modelopt/torch/quantization/mode.py	93.93%	2 Missing ⚠️
modelopt/torch/quantization/config.py	96.77%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2292      +/-   ##
==========================================
- Coverage   79.05%   77.92%   -1.14%     
==========================================
  Files         525      528       +3     
  Lines       61106    64054    +2948     
==========================================
+ Hits        48308    49912    +1604     
- Misses      12798    14142    +1344     

Flag	Coverage Δ
unit	56.31% <95.11%> (+0.49%)	⬆️

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.}