[meenchen]

What does this PR do?

Type of change: New feature

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

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

At sub-4-bit budgets, scoring each (group, format) with every other group unquantized misjudges which layers matter. On Qwen3.6-35B-A3B both the gradient and KL scorers rate the first 10 layers ~30× more sensitive than layers 20–39, starve the late layers with IQ2_XXS, and lose to uniform IQ2 at every budget.

method={"method": "kl_div", "context": "lowest" | k} scores every candidate, BF16 included, with every other searched group at its lowest-bit (k-th lowest) format. All scores of a group then share the context's base damage, so the linear program over their sum ranks assignments by marginal gain at that operating point (the max-score threshold search is kept for the no-context default). Format shards all calibrate the context format; shard 0 alone scores BF16 so the merge stays exact. The recipe schema accepts method_options: {context: ...} for kl_div.

The search state records the scoring context (scoring_config). Resume and shard merge refuse a checkpoint scored with a different context, and get_auto_quantize_config(state, constraints) re-solves with it. Before this, a re-solve rebuilt the searcher with default options, so a context-scored state went through the threshold solver instead of the LP.

Testing

• Shard/merge test parametrized over context="lowest" and context=1; asserts BF16 is measured in context and that the budget is spent (fails with the threshold solver).
• The same test re-solves the merged state with get_auto_quantize_config(state, {"effective_bits": 6.0}) and asserts it matches the unsharded search (fails without the scoring_config fix). AutoQuant + Shapley unit suites pass.
• Qwen3.6-35B-A3B held-out KL (IQ2 + NVFP4 + FP8 + BF16 experts, sequential-GPTQ finalize):

bits	default gradient	context="lowest"	uniform-GPTQ curve
2.75	0.105	0.101	≈0.106
3.0	0.094	0.076	≈0.078
3.25	0.093	0.065	IQ2_S @3.13: 0.064

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ (default unchanged)
• 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

[shengliangxu]

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Changes requested: checkpoint-only finalization drops the KL context and uses the wrong solver.

Needs action:

• Persist and restore KL context for offline re-solving in algorithms.py; extend get_auto_quantize_config and HF finalization coverage (see inline).
• Reject checkpoint reuse and shard merges with incompatible scoring contexts in algorithms.py, with mismatch tests (see inline).
• Explain in the PR body why context-KL is preferred over the existing Aumann–Shapley path scorer in _auto_quantize_shapley.py.

No action needed:

• Existing test edits are justified: default cases remain covered, context calibration is shared, and KL now accepts context options.

[shengliangxu]

Bot comment.

[P1] Preserve the context for checkpoint-only re-solves

This dispatch depends on self.config['context'], but context is absent from default_state_dict, so neither the returned search state nor saved checkpoints record it. _resolve_best_recipe() constructs a KL searcher with default_search_config, making mtq.get_auto_quantize_config(state, constraints=...) use the threshold solver instead of this LP branch.

This also breaks the advertised finalization workflow: examples/hf_ptq/hf_ptq.py::_run_auto_quantize_recipe() calls that re-solve whenever finalize_algorithm resumes, even for the original budget. Rerunning the same command can therefore change the context-aware assignment before sequential GPTQ runs.

Please persist the scoring context and restore it into the offline searcher's config, retaining None semantics for legacy KL states. Extend the existing re-solve/finalization coverage to compare against the original LP assignment for both contexts, including merged checkpoints.

[shengliangxu]

Bot comment.

[P2] Validate the scoring context before reusing scores

Context changes what every raw score means, but _AutoQuantizeBaseSearcher.before_search() only checks the method, costs, formats and resolved search setup before skipping scoring for a populated checkpoint. Reusing a default-KL checkpoint with context='lowest', or changing 'lowest' to 1, silently applies the new context's solver to old-context measurements.

Likewise, _SHARD_INVARIANT_STATE does not include context. Because searched-format scores are disjoint and only shard 0 measures BF16 in context, merging shards scored under different contexts can succeed without detecting incompatible measurements.

Please record and validate a context/scoring signature on resume and in shard invariants; the existing Aumann–Shapley scoring-signature mechanism provides a precedent. Add changed-context resume and mixed-context merge cases.