[meenchen]

What does this PR do?

Type of change: Bug fix

auto_quantize saves a checkpoint after calibration, before scoring. Resuming from it loads candidate_stats with a weights-only load, which returns a plain dict instead of the defaultdict(dict) the searcher creates. Scoring then runs to the end and initialize_candidate_stats raises KeyError on the first group, losing all scores. This hit a preempted 7-format sub-4-bit search on a 74B MoE after ~8 h of scoring.

The fix creates each group's entry explicitly (setdefault) before writing its formats and scores.

Independent of the sub-4-bit series (#3154 → #3151 → #3152 → #3153 → #3155); based on main.

Testing

• Extended test_auto_quantize_calibration_only_checkpoint_validates_module_search_spaces: it now resumes a copy of the calibration-only checkpoint and asserts every group's candidate_stats entry. Fails with KeyError without the fix; passes with it.
• tests/unit/torch/quantization/test_autoquant.py: 79 passed.
• The same one-line fix let the interrupted 74B shard resume and finish scoring.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅
• 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?: ✅ (extended an existing test)
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: ❌

Additional Information

Found while running the sub-4-bit AutoQuantize series (OMNIML-6370).

🤖 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

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-3169/
Built to branch gh-pages at 2026-11-02 16:50 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.34%. Comparing base (c92993c) to head (03d9ef9).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3169      +/-   ##
==========================================
- Coverage   80.35%   80.34%   -0.01%     
==========================================
  Files         675      676       +1     
  Lines       77281    77380      +99     
==========================================
+ Hits        62099    62173      +74     
- Misses      15182    15207      +25     

Flag	Coverage Δ
unit	62.61% <100.00%> (+0.01%)	⬆️

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.