[meenchen]

What does this PR do?

Type of change: New feature

[3/5] of the sub-4-bit AutoQuantize series (after #3151).

Series (merge in order): #3154 (merged) → #3151 → #3152 → #3153 → #3155

Sequential (layerwise) GPTQ makes each layer's result depend on the formats chosen for the layers before it, which a per-layer search cannot know. The search therefore scores candidates with their own local algorithms, and the frozen per-layer formats are calibrated afterwards on the source weights.

• mtq.get_auto_quantize_config(state, algorithm=...) returns the frozen config with the given calibration algorithm instead of always max.
• AutoQuantize recipes take auto_quantize.finalize_algorithm. hf_ptq.py then stops after saving the search checkpoint; rerunning the same command finds the scores, re-solves for the recipe's budget, and quantizes the freshly loaded source model with the frozen formats, the finalize algorithm and the recipe's KV-cache config before exporting.
• hf_ptq.py --auto_quantize_format_shard I/N and --auto_quantize_merge_shards run the format shards and merge them.

Usage

for i in 0 1 2 3 4; do  # one GPU each
  CUDA_VISIBLE_DEVICES=$i python hf_ptq.py --recipe <aq-recipe> ... \
    --auto_quantize_checkpoint ck/shard$i.pth --auto_quantize_format_shard $i/5 &
done; wait
python hf_ptq.py --recipe <aq-recipe> ... --auto_quantize_checkpoint ck/merged.pth \
  --auto_quantize_merge_shards ck/shard0.pth,...,ck/shard4.pth --export_path out

Testing

• End to end on GB300: Qwen3.6-35B-A3B, 3 shards → merge → sequential-GPTQ finalize (40 layers) → export.
• Unit: KL forward_step handles the early-stopped forward of layerwise candidate calibration (fails without the fix); recipe and hf_ptq helper suites pass (test_hf_ptq_args.py, test_autoquant_utils.py, test_example_utils.py: 191 tests).
• _run_auto_quantize_recipe reads the shard flags with defaults: the KV-cache AutoQuantize stage tests build args without them and failed with AttributeError (5 tests) before this.
• Finalizing memory-maps the merged checkpoint on the CPU and reads only its scores, instead of loading every searched format's calibration onto the GPU it was saved from.
• Measured effect of the finalize stage (held-out KL vs BF16, 128 docs): uniform IQ2_XXS / XS / S 0.150→0.118, 0.111→0.089, 0.081→0.064; AutoQuant 2.75 / 3.0 / 3.25 bits 0.130→0.101, 0.095→0.076, 0.078→0.065.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅
• New PIP dependency: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅

🤖 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

Tip

🎯 Customize your risk assessments

Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.

✅ Automatic approvals, with your rules

Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.