feat(quantization): run a calibration plan, one mode per stage [4/4]

[Fridah-nv]

Brief

What changes: Different algorithms can now run on different parts of a model, and chain — a range search can feed GPTQ, and an NVFP4 grid mismatch is rejected rather than silently mis-calibrating.

How: Implements plan execution by treating each algorithm as a mode: every stage is applied through the ordinary apply_mode, so the sequence is recorded in modelopt state rather than hidden inside one opaque entry. Values derived from the plan — the write-mask, the skip_max_init hand-off — travel via mode_kwargs, which reaches the convert entrypoint and is never saved.

---

What does this PR do?

Type of change: new feature

Stack 4 of 4, splitting #2292. Base: #2536. This is where the feature becomes usable.

#	PR
1	capabilities	#2534
2	write-mask	#2535
3	compile algo_cfg	#2536
4	execute the plan	← you are here

One mode per stage

Each stage is applied as its own calibration mode through the ordinary apply_mode, so it is recorded in the modelopt state under its own name — ["quantize", "max_calibrate", "mse_calibrate"] — and the mode graph sees the sequence rather than one opaque entry.

Values derived from the plan — the should_process write-mask from #2535, the hand-off kwargs — travel through mode_kwargs, which reaches the convert entrypoint and is deliberately never saved: they describe how this run was produced, not what the model is.

The hand-off

skip_max_init returns to mse, local_hessian and gptq as a function parameter the executor supplies from derive_handoff. It is not user-facing config: it is a fact about what the previous stage produced, and a user setting it by hand would be asserting something the compiler already knows.

This is what lets a range search feed GPTQ. GPTQ compensates rounding error against a specific quantization grid, so when mse searched a better one the compensation must be computed against that grid, not a fresh max one — which is the order DeepCompressor's QoQ recipes ship (qoq-gchn.yaml: enable_calib_range then enable_kernel_gptq). Previously gptq() re-derived amax unconditionally and the search in front of it was discarded.

NVFP4 grid type becomes part of the contract

The sharpest case for capabilities being per stage rather than per model. awq_lite folds a smoothing scale into the weight and needs block scales the kernel derives at run time (type: dynamic); mse/local_hessian with fp8_scale_sweep search stored per-block scales and need type: static. Running either against the wrong layout is not a tuning difference — it fails, or silently searches nothing.

requires_weight_scales declares which. A prepare hook upgrades dynamic to static at the start of the stage that needs it; static never downgrades, since that would discard a completed search. The requirement belongs to the consumer, which is why the upgrade happens at the start of the stage that needs it rather than the end of the one before — a standalone run must not be dragged into a state it never asked for.

So the intended flow — awq_lite on a dynamic grid, then mse with the sweep on a static one — is one plan, and the reverse is rejected at compile with the reason named.

Three defects the compile sweeps found, fixed here

Exhaustive sweeps (13,134 compiles) over this stack found four defects; three belong with the grid contract:

1. The rule validated a stale snapshot. prepare mutates grid type between stages, but the rule read the model as it is at compile time. [mse(fp8_scale_sweep), awq_clip] compiled clean and was rejected by the identical config once stage 0's prepare had run — so at runtime awq_clip got the static grid the rule exists to keep it away from. It now walks the plan in order. 24 of 112 static-then-dynamic chains were wrongly accepted; now 0.
2. prepare upgraded activations too. No role filter, so on the shape the shipped static presets use — nvfp4_static for *weight_quantizer, dynamic nvfp4 for *input_quantizer — all input quantizers were silently converted to static and lost their amax.
3. The same missing filter caused a false rejection, faulting awq_clip on a model whose weight grid was dynamic because its activations were static.

Testing

test_algo_cfg.py grows to 60 tests: per-stage mode recording, mse → gptq preserving the searched amax bit-identically, awq_lite → mse, awq_full reproduced exactly as its two-stage pipeline, the grid contract (upgrade, activations left alone, upgrade visible to later stages, static activations not blocking a weight-side algorithm), and legacy-path numerical equivalence.

Each defect fix is mutation-verified: the fix is reverted, the specific test confirmed to fail, and the fix restored.

pytest tests/unit/torch/quantization/ → 1195 passed, 8 skipped, plus the same two pre-existing test_huggingface.py failures and one pre-existing collection error described in #2534.

Cost of the execution model, and what it leaves open

Each stage is a separate apply_mode handed the same forward_loop, so N stages that need
activations cost N forward passes over the calibration set. Invisible on a toy model, not
invisible on a 9B with 512 sequences. Worth stating explicitly since it is a consequence of the
one-mode-per-stage choice, not an accident.

What it does not cost: chained stages could never share a forward anyway — mse → gptq means
GPTQ reads the grid MSE produced. The avoidable case is stages on disjoint scopes, which today
still pay for a pass each.

That optimisation stays reachable, and deliberately so. The write-mask predicate is a lowering of
structured information that survives on AlgoStage — selector, glob, exclusions — so an optimiser
can decide two stages are disjoint from the plan itself rather than introspecting an opaque
callable. Batching them would mean enabling collection for both stages' targets, running one pass,
and finishing each separately; that works while their quantizer sets do not overlap, which disjoint
scopes guarantee. It is a contained change in _apply_calibration_plan, needing nothing from the
mask.

Reads are whole-model by necessity, not by choice: activations arriving at a scoped module depend on
the entire network upstream, so narrowing reads would make a scoped search optimise against a
distribution the deployed model never sees. The one legitimate narrowing — layerwise calibration —
already exists and is orthogonal, which is why the mask is keyed on module identity and survives
that reparenting.

The asymmetry a future change will have to confront: the mask is per-module boolean, while the
capability model reasons per (token, role). The compiler knows a stage writes weight_amax but
not input_amax; the runtime mask can only say "this module, yes or no". _reject_partial_module_scope
exists to cover that gap. Two stages writing different roles of the same module concurrently would
require dropping that rule and making the mask per-(module, token).

Known limitations

• Three algorithms cannot be scoped yet. nvfp4_act_headroom takes no should_process; svdquant and lsq swallow it via **kwargs and would run over the whole model. scopable records this and compile rejects a scoped stage for them rather than mis-calibrating; whole-model use is unchanged.
• No NVFP4 chain test in CI. The unit suite is CPU/INT4 on a toy model. The grid-type contract is unit-tested at compile level, but no test executes an NVFP4 chain; all NVFP4 verification so far has been ad-hoc GPU scripts.
• Not covered: distributed (the rank-identical-plan property is structural but untested — a 2-GPU TP/EP test is still needed), auto_quantize, and svdquant (compiles and validates, never executed).
• Accuracy measurements from feat(quantization): scoped calibration pipelines via algo_cfg [prototype] #2292 predate the grid contract and were not re-measured; see that PR for the caveat.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ — without algo_cfg the old path, its numerics and its saved state are unchanged; algorithm="max" and the equivalent algo_cfg produce bit-identical amax.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ❌ — deferred while the stack is under review.
• Did you get Claude approval on this PR?: ❌ — draft.

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

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