[OMNIML-6370] [5/5] Add Qwen3.6-35B-A3B sub-4-bit AutoQuantize recipes by meenchen · Pull Request #3155 · NVIDIA/Model-Optimizer

meenchen · 2026-10-30T19:16:07Z

What does this PR do?

Type of change: New example (recipes)

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

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

Qwen3.6-35B-A3B AutoQuantize recipes at 2.75 / 3.0 / 3.25 effective bits (all linears except embeddings, lm_head and vision): routed experts choose IQ2_XXS / IQ2_XS / IQ2_S (local GPTQ + imatrix), weight-only NVFP4 (W4A16), weight-only FP8 (W8A16) or BF16; other linears stay BF16. Context-aware KL scoring at the format nearest each budget; sequential GPTQ finalizes the selection.

Usage

python examples/hf_ptq/hf_ptq.py --pyt_ckpt_path Qwen/Qwen3.6-35B-A3B \
  --recipe models/Qwen/Qwen3.6-35B-A3B/auto_quantize/iq2_gptq_nvfp4_fp8_experts_at_3p0bits \
  --auto_quantize_checkpoint ck/search.pth --export_path out   # search; rerun to finalize + export

Testing

Held-out KL vs BF16 (128 documents, 63.5k tokens; experts only, GPTQ-finalized):

Recipe	Bits	KL (AutoQuant)	KL (uniform IQ2-GPTQ curve at same bits)
`_at_2p75bits`	2.75	0.1003	≈0.106
`_at_3p0bits`	2.99	0.0766	≈0.078
`_at_3p25bits`	3.25	0.0626	0.0638 (uniform IQ2_S at 3.13 bits)

NVFP4 / FP8 are weight-only because the bit cost counts weights only and the IQ2 formats are W2A16; with W4A4 / W8A8 candidates the same search gives 0.1012 / 0.0761 / 0.0646. The weight-only KL was measured with context: lowest; in the W4A4 runs the recipes' near-budget context (1 at 3.0, 2 at 3.25) changed KL by at most 1.7%.

AA-LCR / GPQA-Diamond / SciCode subtask (reduced repeats 4 / 8 / 4; BF16-folded checkpoints on vLLM; lowest-bit-context W4A4/W8A8-search checkpoints; the 2.75 / 3.0 assignments are identical to these recipes'):

Arm	Bits	AA-LCR	GPQA	SciCode
BF16	16	69.0	83.9	39.7
Uniform IQ2_XXS-GPTQ	2.65	62.0	70.3	36.8
Uniform IQ2_XS-GPTQ	2.89	63.5	76.4	35.7
Uniform IQ2_S-GPTQ	3.13	65.5	79.2	38.2
AutoQuant @2.75	2.75	61.75	74.1	35.1
AutoQuant @3.0	2.99	68.0	78.2	38.8
AutoQuant @3.25	3.25	68.5	80.2	37.4

The recipes load in the existing recipe unit tests.

Before your PR is "Ready for review"

Is this change backward compatible?: ✅
Did you write any new necessary tests?: N/A (recipes load in the existing recipe tests)
Did you update Changelog?: N/A

🤖 Generated with Claude Code

copy-pr-bot · 2026-10-30T19:16:12Z

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 · 2026-10-30T19:16:19Z

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

…from the source weights (#3154) ### What does this PR do? Type of change: Bug fix / new feature [1/5] of the sub-4-bit AutoQuantize series (merge in order). Series (merge in order): #3154 → #3151 → #3152 → #3153 → #3155 GPTQ and imatrix overwrite the weights they quantize. `auto_quantize` rejected both, and accepting them as-is would let one candidate's calibration change the weights every later candidate — and the unquantized scoring reference — sees. AutoQuantize now snapshots (to CPU) only the weights a GPTQ/imatrix candidate's calibration rewrites: the parameters its enabled weight quantizers cover, and for imatrix alone only those imatrix encodes (not Q8_0; a fused projection when any of its experts is). After such a candidate calibrates, its GGML quantizers replay the codes it pinned while the source weights are restored for the next candidate; the snapshot is freed before sensitivity scoring. After the solve, the selected candidates' weights are left as their codes decode, as a plain GPTQ run leaves them, and no candidate is left replaying. GGML fake quant now returns exactly the reconstruction export decodes, so scoring and the exported model see the same weights (BF16 included). Refused before the model is converted: any candidate with an enabled layerwise `export_dir`, empty included, whether its options are a dict or a `LayerwiseConfig`, whose export would leave the model in export form before scoring (`hf_ptq.py` already refused it for AutoQuantize recipes); and, in a search with a weight-rewriting candidate, any candidate whose layerwise `checkpoint_dir` resume restores the layers' weights into new Parameters the source-weight restore cannot reach (`calib_mutates_weights=False` resumes are accepted). Refused before any calibration: a weight-rewriting candidate whose weight quantizer is not a single GGML quantizer or rotates the weight, or whose algorithm list runs anything besides GPTQ, imatrix and max (AWQ or SmoothQuant would transform the weight GPTQ encodes); and meta or DTensor weights it would rewrite, weights stored packed or scaled (a native FP8 checkpoint's), or ones whose storage another parameter shares (tied weights, or views of them). These checks need the converted quantizers to know which weights a candidate rewrites, so `auto_quantize` hands back the unconverted modules, and the configs that conversion changed (including nested sub-configs), when the search fails before calibrating any candidate, and the same model can be searched again. Only a search with a weight-rewriting candidate saves the modules for this, and it drops them once these checks pass, so calibration and scoring never hold the weights conversion replaces (HF Conv1D's transposed copies); other searches save nothing. A failure during calibration restores the source weights and drops the modes that calibration added. `auto_quantize` hands back the pre-quant-scale folding policy on every exit, including a failure during scoring. Weights a calibration leaves alone (imatrix on Q8_0) simply quantize the source weight. NVFP4/FP8-only searches take no snapshot and never slice fused-expert weights after the solve. `auto_quantize` now accepts `gptq`, `imatrix` and algorithm lists such as `[gptq, imatrix]`. #### Why snapshot and replay Every candidate after a GPTQ/imatrix calibration, and the unquantized scoring reference, must see the source weights. The alternatives cost more: - **Copy the model, or its state dict:** this holds every weight, about 70 GB for Qwen3.6-35B-A3B in BF16. The snapshot holds only the weights a rewriting candidate quantizes, on the host, and is freed before scoring. - **Make GPTQ/imatrix leave weights alone:** this would change calibrators that `mtq.quantize` shares, and the chosen codes would still have to be kept somewhere. The GGML quantizer already pins those codes for export, so replaying the pin makes the scored, selected and exported weights the same codes. #### Reading order 1. `_AutoQuantizeBaseSearcher.before_search`: snapshot (`_snapshot_source_weights`; its checks are `_rewritten_weight` and `_check_restorable`), then for each candidate calibrate, replay its pins and restore the source weights. 2. `_write_selected_codes_to_weights` (after the solve), and `GGMLFormat._pinned_rows` / `full_pin` (the replay). 3. `auto_quantize`: the window in which a failure hands back the unconverted modules, and the folding policy it hands back on every exit. ### Usage ```python iq2_gptq = { "quant_cfg": [{"quantizer_name": "*", "enable": False}, {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": "iq2_s", "block_sizes": {-1: 256}, "backend": "ggml", "effective_bits": 2.5625}}], "algorithm": [{"method": "gptq", "block_size": 256, "perc_damp": 0.3, "importance_weighted": True}, {"method": "imatrix"}], } model, state = mtq.auto_quantize(model, constraints={"effective_bits": 3.0}, quantization_formats=[(iq2_gptq, "IQ2_S_GPTQ"), mtq.NVFP4_DEFAULT_CFG], ...) ``` ### Testing - New tests in `tests/unit/torch/quantization/test_ggml_gptq.py`, `test_autoquant.py`, `plugins/test_fused_experts.py` and `tests/gpu/torch/quantization/test_ggml_gptq_cuda.py`. I reverted each fix in place and checked that its test fails. They cover: - other candidates' scores are unaffected by GPTQ; - mixed selection: a GPTQ layer next to a BF16 layer, a GGML entry replacing an earlier INT4 entry, and imatrix on Q8_0; - the snapshot holds only rewritten weights (a fused `down_proj` that is disabled or Q8_0), takes a whole fused `gate_up_proj` with one Q8_0 expert, and a BF16 selection hands back the source weights; - refusals (non-GGML, tied, an aliased view, a rotated GPTQ or imatrix target, AWQ then GPTQ, a layerwise `checkpoint_dir`, native FP8 weights (GPU), and fused MoE whose root, nested and experts' configs conversion changes) hand back the unconverted model and its configs, which searches again; - calibration failures (GPTQ rewrite then failure, `[gptq, layerwise imatrix]` failing midway) restore the weights, the folding policy and the mode list; - the folding policy comes back after a scoring failure; - exact BF16 replay; - no fused-expert indexing after a search without pins; - an `algorithm: []` candidate keeps the saved state restorable. - Layerwise directories (`test_huggingface.py`), on a tiny Llama with discoverable decoder layers: an FP8 candidate with an `export_dir`, with an empty `export_dir`, and with a `checkpoint_dir` beside a Q8_0 GPTQ candidate, each given as a dict and as a `LayerwiseConfig`, is refused before conversion, leaving the modules, the weights and the working directory untouched, and the same model searches again. Without the refusals, the layers are exported and scoring fails on the packed weights, or the max candidate's checkpoint is written. - Conv1D weight lifetime (`test_huggingface.py`): an INT8 search saves no module state, and a Q8_0 GPTQ search drops it before calibration, so the weights conversion replaces are dead in every calibration and scoring forward. - Regression check against the base commit (RTX 6000 Ada, BF16, FP8 + INT4 searches): a 24-layer GPT-2 (Conv1D), a 16-layer Llama and a fused-MoE Qwen3, with gradient scoring and with KL for Llama, give the same scores, selections and output logits, the same search time (within 4%) and the same memory. On GPT-2 the weights held after conversion are 0 MiB, against 576 MiB before 5acdcc4. - Unit suites `tests/unit/torch/quantization` + `tests/unit/recipe` on CPU (16 xdist workers): 2151 passed. The rest fail the same way on `main` at this branch's base on the test host: - diffusers lacks `q_bmm_quantizer` (4 failures); - `test_diffusers_wan_conv3d` fails to collect (16 errors); - one to four `test_nemotron_3_5_super_vl_autoquant_recipe` cases fail under xdist and pass serially. - `pre-commit run mypy --all-files` passes; ruff check and ruff format pass on the touched files. - Rollback cost, paid only by searches with a GPTQ/imatrix candidate: saving the module attributes takes 2 MiB and 0.16 s on Qwen3.6-35B-A3B (837 modules), and 69 MiB on a 24.6k-module non-fused MoE. - GPU (RTX 6000 Ada): `tests/gpu/torch/quantization -k "auto_quant or autoquant or ggml or gptq or imatrix"` 57 passed, including `test_fsdp2_autoquant.py` and the native FP8 refusal. The deepspeed and onnx packages on the host stop 3 files from collecting. - Stack restacked on this; the same unit suites at [5/5]: 2148 passed, with the same failures. - Qwen3.6-35B-A3B on GB300: searches over IQ2_{XXS,XS,S} (imatrix and local GPTQ), NVFP4, FP8, BF16; see PR 5 for results. ### Before your PR is "*Ready for review*" - 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?: ✅ - Did you update Changelog?: ✅ - Did you get Claude approval on this PR?: review comments addressed (007a287, 83df05c, 848bd95, 357fb38, 59ea14d, 609ed1c, 9f76d8d, 7a21227, 7e028ac, 5acdcc4, f75da89, a9f070f, 1808e13) ### Additional Information Supersedes #3082 … #3057 and #3113 … #3120 with a smaller design (~720 source lines across the series). Synced with main by a merge (99f82c8), which keeps the commits the review threads cite. #3197 made the IQ formats re-encode a weight under its pinned scales and importance ahead of the pinned-rows lookup that honours `PIN_OVERRIDES_WEIGHT`, so during a search a GPTQ or imatrix candidate would have quantized the source weight under its scales instead of replaying its codes; b7be410 puts the override first there too (caught by `test_a_replayed_pin_fake_quantizes_to_exactly_its_decoding_in_bf16`). 🤖 Generated with [Claude Code](https://claude.com/claude-code)   ## Summary by CodeRabbit * **New Features** * AutoQuantize supports GPTQ and iMatrix candidates with GGML formats. * Selected candidates retain their calibrated quantized results in the returned model. * **Bug Fixes** * Candidate scores are no longer affected by weight changes made while evaluating other candidates. * Source weights and quantization settings are restored when candidate evaluation fails; unsupported weight-rewrite configurations are rejected. * Failed conversions before calibration begins restore the model’s prior state. * Fake quantization better preserves reconstructed values in narrow data types.  --------- Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Routed experts search IQ2_XXS/IQ2_XS/IQ2_S (local GPTQ + imatrix), NVFP4, FP8 and BF16 at 2.75, 3.0 and 3.25 effective bits over all linears except embeddings, lm_head and vision; other linears stay BF16. Sequential GPTQ finalizes the selected formats. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

The bit cost counts weights only and the IQ2 candidates are W2A16, so W4A4 NVFP4 and W8A8 FP8 were scored with activation error the budget does not pay for. Use W4A16 NVFP4 and W8A16 FP8 instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from 8617993 to 1562724 Compare November 2, 2026 16:44

meenchen mentioned this pull request Nov 2, 2026

[OMNIML-6370] Fix AutoQuantize resume from a calibration-only checkpoint #3169

Draft

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from 1562724 to cdb5f9f Compare November 2, 2026 23:24

meenchen force-pushed the weimingc/autoquant-kl-context branch 2 times, most recently from 1fd2215 to 1658a6f Compare November 3, 2026 20:10

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from cdb5f9f to 7fe074e Compare November 3, 2026 20:10

meenchen force-pushed the weimingc/autoquant-kl-context branch from 1658a6f to 0986299 Compare November 3, 2026 22:33

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from 7fe074e to 11dbfb0 Compare November 3, 2026 22:33

meenchen force-pushed the weimingc/autoquant-kl-context branch from 0986299 to bfecf53 Compare November 4, 2026 05:22

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch 2 times, most recently from 06113c9 to 5491bd8 Compare November 4, 2026 07:06

meenchen force-pushed the weimingc/autoquant-kl-context branch from bfecf53 to 117544b Compare November 4, 2026 07:06

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from 5491bd8 to ba7ef3d Compare November 4, 2026 16:32

meenchen force-pushed the weimingc/autoquant-kl-context branch 2 times, most recently from 24168a1 to 4b2ee1b Compare November 4, 2026 18:52

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from ba7ef3d to b4ea7d4 Compare November 4, 2026 18:52

meenchen force-pushed the weimingc/autoquant-kl-context branch from 4b2ee1b to 1cdfb9e Compare November 4, 2026 20:20

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch 2 times, most recently from 0e6b4f0 to 8f12518 Compare November 4, 2026 21:50

meenchen force-pushed the weimingc/autoquant-kl-context branch 2 times, most recently from 52ad620 to 92d35d5 Compare November 5, 2026 17:49

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch 2 times, most recently from 5c6c5c8 to d43f209 Compare November 5, 2026 22:36

meenchen force-pushed the weimingc/autoquant-kl-context branch from 92d35d5 to 4097d58 Compare November 5, 2026 22:36

meenchen force-pushed the weimingc/autoquant-kl-context branch from 4097d58 to fbf446b Compare November 5, 2026 23:57

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from d43f209 to e6a2eff Compare November 5, 2026 23:57

meenchen force-pushed the weimingc/autoquant-kl-context branch from fbf446b to 177fad8 Compare November 6, 2026 00:09

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch 2 times, most recently from e365f56 to c1d1fc7 Compare November 6, 2026 00:24

meenchen force-pushed the weimingc/autoquant-kl-context branch 2 times, most recently from 2c167dc to e78d3cc Compare November 6, 2026 00:58

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from c1d1fc7 to 6d44deb Compare November 6, 2026 00:58

meenchen force-pushed the weimingc/autoquant-kl-context branch from e78d3cc to 5e21a27 Compare November 6, 2026 06:13

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from 6d44deb to 3d850cb Compare November 6, 2026 06:13

meenchen force-pushed the weimingc/autoquant-kl-context branch from 5e21a27 to a8d97a1 Compare November 6, 2026 21:37

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch 2 times, most recently from 2c512de to 4ca71c1 Compare November 6, 2026 23:01

meenchen force-pushed the weimingc/autoquant-kl-context branch 2 times, most recently from 15b44f5 to 92431b3 Compare November 7, 2026 02:15

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from 4ca71c1 to 9425349 Compare November 7, 2026 02:15

meenchen and others added 2 commits November 6, 2026 19:54

meenchen force-pushed the weimingc/autoquant-kl-context branch from 92431b3 to c543ea0 Compare November 7, 2026 05:32

meenchen force-pushed the weimingc/qwen36-sub4bit-recipes branch from 9425349 to 3798a0f Compare November 7, 2026 05:32

Navigation Menu

Repository navigation

[OMNIML-6370] [5/5] Add Qwen3.6-35B-A3B sub-4-bit AutoQuantize recipes - #3155

[OMNIML-6370] [5/5] Add Qwen3.6-35B-A3B sub-4-bit AutoQuantize recipes#3155
meenchen wants to merge 2 commits into
weimingc/autoquant-kl-contextfrom
weimingc/qwen36-sub4bit-recipes

meenchen commented Oct 30, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Oct 30, 2026

Uh oh!

coderabbitai Bot commented Oct 30, 2026 •

edited

Loading

Draft PR not reviewed

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

meenchen commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Usage

Testing

Before your PR is "Ready for review"

Uh oh!

copy-pr-bot Bot commented Oct 30, 2026

Uh oh!

coderabbitai Bot commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Draft PR not reviewed

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

meenchen commented Oct 30, 2026 •

edited

Loading

coderabbitai Bot commented Oct 30, 2026 •

edited

Loading