[OMNIML-5477, OMNIML-5119] Add module-specific AutoQuant search spaces#1949
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesAutoQuantize adds module-specific search spaces, optional fixed PTQ baselines, per-group no-quant controls, isolated calibration, stable cost accounting, and checkpoint compatibility validation. Recipe schemas, Hugging Face conversion, documentation, a Qwen recipe, and tests are updated. AutoQuantize module search spaces
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Recipe
participant AutoQuantize
participant Searcher
participant QuantRecipeHparam
Recipe->>AutoQuantize: provide module_search_spaces and fixed baseline
AutoQuantize->>Searcher: submit normalized search configuration
Searcher->>QuantRecipeHparam: resolve runtime-group choices
Searcher->>QuantRecipeHparam: set_calibration_recipe(recipe)
QuantRecipeHparam-->>Searcher: restore_active_quantizers()
Searcher-->>AutoQuantize: resolve configuration and checkpoint state
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1949 +/- ##
==========================================
+ Coverage 75.76% 77.79% +2.03%
==========================================
Files 518 518
Lines 58269 58498 +229
==========================================
+ Hits 44148 45510 +1362
+ Misses 14121 12988 -1133
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
The implementation looks cohesive and well tested, and I did not find a concrete correctness regression in the module-specific candidate selection, fixed-group cost accounting, calibration isolation, fused-group validation, or checkpoint signature handling. The added YAML header follows the short SPDX style, but it does not match the complete canonical text in root LICENSE_HEADER, so this GitHub PR still needs human licensing sign-off under the review policy.
Design review: the problem is to let existing AutoQuantize runtime decision groups use family-specific candidate sets, including truly fixed quantized groups, without losing sensitivity baselines or active-MoE accounting. Plausible alternatives are (1) extending the repository's existing NAS SearchSpace/wildcard-rule machinery (modelopt/torch/nas/search_space.py), (2) expressing family overrides through the existing ordered QuantizeConfig.quant_cfg/last-match-wins mechanism plus pre-partitioned AutoQuant runs, or (3) keeping one global AutoQuant search and applying a fixed PTQ config to routed experts before/after it. Already-imported Pydantic/OmegaConf can model the declarative schema, while PuLP remains the existing solver; no dependency offers the AutoQuant-specific grouped scoring directly. The PR correctly extends the existing AutoQuant/Pydantic/PuLP path rather than creating a separate solver, and its Qwen3.6 evaluation provides strong motivation, but the PR body does not explicitly compare or reject the obvious in-repo NAS/wildcard and staged-PTQ alternatives. Because design review was required, please have the owner sign off or add that rationale to the PR body.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/quantization/algorithms.py`:
- Around line 911-943: Extend the checkpoint compatibility logic around
module_search_space_signature to also record and compare a signature for
config["quantization_formats"] whenever quantizer_states or candidate_stats are
restored. Reject mismatched current formats before reusing checkpoint state, and
explicitly define the behavior for legacy checkpoints that lack the stored
format signature, consistent with the existing module-search-space handling.
In `@modelopt/torch/quantization/model_quant.py`:
- Around line 522-550: Update the public validation in the quantization entry
point before calls to _process_quantization_formats: require global and
per-module quantization_formats inputs to be valid non-empty candidate
collections, rejecting strings, mappings, and empty lists explicitly. Replace
the optimization-removable assert on processed global formats with an explicit
validation error, and apply equivalent validation to each module_search_spaces
entry before normalization.
In `@tests/unit/torch/quantization/test_autoquant.py`:
- Around line 369-370: Move the model_quant import from inside
test_auto_quantize_fixed_module_isolated_from_unrelated_calibration to the
module-level imports at the top of the test file. Remove the function-local
import while preserving the test’s existing module reference and behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4a55fb6d-15fa-4af8-b9d1-86163a14b061
📒 Files selected for processing (10)
CHANGELOG.rstexamples/hf_ptq/README.mdexamples/hf_ptq/hf_ptq.pymodelopt/recipe/config.pymodelopt/torch/quantization/algorithms.pymodelopt/torch/quantization/model_quant.pymodelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_module_spaces_at_6p0bits-active_moe.yamltests/examples/hf_ptq/test_hf_ptq_args.pytests/unit/recipe/test_loader.pytests/unit/torch/quantization/test_autoquant.py
meenchen
changed the title
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
Re-review: the implementation remains cohesive and well tested, and the three actionable CodeRabbit findings are addressed: the critical checkpoint-integrity issue now validates both global and module-specific format signatures (including legacy behavior), while candidate-collection validation and the test-local import (minor) are fixed with focused tests. I found no new correctness regression.
Design review: the problem is to let AutoQuant runtime decision groups use module-family-specific candidate sets, including fixed quantized groups, while preserving sensitivity baselines and active-MoE cost accounting. Plausible alternatives are the existing NAS SearchSpace wildcard/rule machinery in modelopt/torch/nas/search_space.py, existing ordered QuantizeConfig.quant_cfg last-match-wins overrides plus staged AutoQuant/PTQ, or a global AutoQuant search followed by fixed PTQ for routed experts. The manifest already provides Pydantic/OmegaConf for schema/composition and PuLP for solving; this PR reasonably extends that existing AutoQuant/Pydantic/PuLP path rather than adding another solver. However, the PR body still gives evaluation motivation without explaining why the obvious in-repo NAS/wildcard or staged-PTQ alternatives were rejected, so the required architectural choice needs owner sign-off or rationale added to the PR body.
Licensing also requires human sign-off: the new YAML uses only the two-line SPDX header, which does not match the complete canonical root LICENSE_HEADER. Under the GitHub licensing policy this cannot be auto-approved.
shengliangxu
left a comment
shengliangxu
left a comment
There was a problem hiding this comment.
In general looking good to me.
|
/ok to test 058601e |
|
hi @meenchen, thanks for this, will take a detailed look today/ tomo |
|
thanks Wei-Ming! overall the idea looks good to me. This is basically helping in two ways right if you can confirm
|
|
Thanks, @juhi10071998. Two nuances on the behavior:
I also added a Design rationale section to the PR body: NAS SearchSpace is not wired to AutoQuant calibration/sensitivity/cost/export, static quant_cfg wildcards cannot express candidate sets in one solve, and staged PTQ would move fixed groups outside the common sensitivity baseline and active-MoE budget. |
|
/ok to test 5627f98 |
| # One selectable format makes routed experts fixed W4A16. allow_no_quant | ||
| # controls solver choices only; their BF16-equivalent denominator cost and | ||
| # active-MoE ratio remain part of effective-bit accounting. | ||
| - module_name_patterns: |
There was a problem hiding this comment.
does it mean we don't search and only apply w4a16?
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
The config design LGTM
| - $import: w4a16_nvfp4 | ||
| - $import: fp8 | ||
| allow_no_quant: true | ||
| allow_no_quant: false |
There was a problem hiding this comment.
LGTM, I am in general curious in what cases do we want to have the unquantized bf16 for certain modules. Router is sensitive to quantization so that could be one. Have we observed other bf16 in final recipe before? Asking so if we should have any recommendation for the user generally
There was a problem hiding this comment.
That depends on the sensitivity of the module and also the candidate formats. In this case, since fp8 is able to maintain accuracy well, we don't need to explore bf16.
juhi10071998
left a comment
juhi10071998
left a comment
There was a problem hiding this comment.
Thanks Wei-Ming, left a minor comment.
So basically the cost accounting still remains same as before right for every module. We are restricting the search space for the modules and will only use the corresponding cost associated to that candidate for the given module.
|
/ok to test c2bdd81 |
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
c2bdd81 to
a60306b
Compare
|
/ok to test a60306b |
4 participants
What does this PR do?
Type of change: New feature.
This PR adds ordered, module-name-specific search spaces to AutoQuant so different runtime decision groups can use different candidate formats.
The user-facing design now separates fixed PTQ configuration from the AutoQuant search space:
Additional safeguards:
Design rationale
The fixed configuration should use the same PTQ recipe system users already use for uniform NVFP4, W4A16, FP8, and other baseline configurations. AutoQuant then only needs to describe what is genuinely searched.
For example, the Qwen3.6 recipe reuses the model-specific quantization configuration from
huggingface/qwen3_5_moe/ptq/w4a16_nvfp4-fp8_attn-kv_fp8_cast. Unmatched modules retain that PTQ baseline, while shared experts, self-attention, linear-attention, and lm_head are explicitly searched. The original PTQ recipe is unchanged, and a loader test asserts that both recipes resolve to the same fixedquantizeconfiguration.This remains one AutoQuant operation rather than staged PTQ plus AutoQuant:
The implementation therefore composes the existing PTQ recipe schema with the existing AutoQuant/PuLP path instead of introducing another search framework or solver.
Usage
The legacy global-search form remains supported by omitting quantize and supplying top-level auto_quantize.candidate_formats.
Qwen3.6 search-space comparison
Compared the default W4A16 NVFP4/FP8 search against the historical two-format module-specific recipe that keeps routed experts on the W4A16 PTQ baseline while shared experts and attention search W4A16 versus FP8. The final PR recipe keeps every matched quantizable group in W4A16 or FP8 (
allow_no_quant: false). The table below is retained as evidence for the fixed-baseline design; the final recipe validation follows it.Both lanes used local raw-gradient scoring, active-MoE accounting, batch size 1, code32k calibration, full parser-on LiveCodeBench with 8 repeats, and full SciCode. Active GiB/token includes scale storage.
BF16 references are 74.3667 LCB and 40.7914 SciCode. With a 1.5 percentage-point tolerance, the fixed-routed-expert lane jointly passes at targets 6.1, 6.4, and 6.7; the default lane has no joint pass. This comparison is end-to-end rather than an isolated solver ablation because the historical default lane used full attention-family grouping while the module-specific lane used runtime-required grouping.
Final no-BF16 recipe validation used the model-specific PTQ baseline, codeblend 16k calibration, local gradient scoring, active-MoE accounting, batch size 1, and full parser-on LiveCodeBench with 8 repeats:
Both targets pass the BF16-minus-1.5pp threshold. The model-specific baseline intentionally leaves linear-attention A/B and non-linear helper modules unquantized; all searched quantizable projections are W4A16 or FP8, and the artifact audit found no missing quant tensors.
Testing
allow_no_quant: falseupdate, the focused loader/HF mapping tests and full pre-commit recipe validation passed. The original PTQ recipe remains unchanged, and the resolved fixed baselines compare equal.Before this PR is ready for review
Additional Information
Related work: OMNIML-5477, OMNIML-5119.