feat(quantization): a write-mask so calibration can be confined to a scope [2/4]

[Fridah-nv]

Brief

What changes: Scoped calibration needs algorithms to write only their targets yet read the whole model; this enables that. Nothing changes for existing callers.

How: Adds should_process to every calibration algorithm's public API — 8 algorithms and 3 helpers — so a stage writes only to its targets. It is ANDed into each algorithm's existing is_enabled filter, defaults to None meaning the whole model, and never touches enable-state, so reads are unaffected.

---

What does this PR do?

Type of change: new feature (dormant) + bug fix

Stack 2 of 4, splitting #2292. Base: #2534.

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

A write-mask so calibration can be confined to a scope

Calibration algorithms take an optional should_process predicate and AND it into their existing is_enabled filter, so an algorithm can be restricted to part of the model. None means the whole model, which is today's behaviour, so nothing changes for any existing caller.

Three properties matter, each with a test:

• It gates writes only, never enable-state. Reads are untouched, so the activations a search-based algorithm sees are identical either way.
• It keys on module identity, not name. That is what lets it survive layerwise_calibrate handing an algorithm a reparented subtree whose module names are relative.
• It has to reach helper calls inside an algorithm, not just its top-level module loop. awq_lite calls enable_stats_collection(model) directly and would otherwise calibrate the whole model regardless of scope.

Reviewing this

One file of library change, and the question is narrow: is every write site gated? That is exactly the question the awq_lite bug above failed, which is why this is its own PR rather than buried in the plan machinery.

No caller passes a mask yet — the plan that does arrives in PR 4. This is deliberate: the threading is mechanical and error-prone, the plan is neither, and mixing them makes both harder to check.

A caveat the tests cannot cover

The guards in the max_calibrate path are mutually redundant. Removing any one of enable_stats_collection, weight_only_quantize or finish_stats_collection leaves the others covering it, so no single-guard mutation is observable. The test detects a real leak — 5 of 7 cases fail when all three are removed — but redundancy is why it cannot pin any one of them individually.

Found while mutation-testing this PR's own test, which also turned up a genuine gap: the first version only snapshotted amax, so removing GPTQ's guard changed nothing detectable, because GPTQ writes weights. The test now snapshots amax, weights and pre_quant_scale.

Testing

New tests/unit/torch/quantization/test_calib_write_mask.py — scope confinement for max/mse/awq_lite/gptq, the weight-only (no forward loop) path, and enable-state invariance.

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

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ — new optional kwarg defaulting to today's behaviour.
• 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.}

[github-actions]

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

[codecov]

Codecov Report

❌ Patch coverage is 97.05882% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 76.30%. Comparing base (3987c76) to head (7324291).

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/model_calib.py	97.05%	1 Missing ⚠️

Additional details and impacted files

@@                        Coverage Diff                         @@
##           fridah/algo-cfg-1-capabilities    #2535      +/-   ##
==================================================================
+ Coverage                           76.26%   76.30%   +0.03%     
==================================================================
  Files                                 609      609              
  Lines                               69311    69390      +79     
==================================================================
+ Hits                                52862    52945      +83     
+ Misses                              16449    16445       -4     

Flag	Coverage Δ
unit	59.13% <97.05%> (+0.19%)	⬆️

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.