Add module-specific AutoQuant search spaces#1949
Conversation
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
📝 WalkthroughWalkthroughChangesAutoQuantize module search spaces
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant AutoQuantize
participant Searcher
participant ModuleGroup
participant QuantizerState
AutoQuantize->>Searcher: submit module_search_spaces
Searcher->>ModuleGroup: match patterns to runtime groups
Searcher->>QuantizerState: apply calibration recipe
QuantizerState-->>Searcher: restore active quantizers
Searcher-->>AutoQuantize: resolve candidate configuration
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1949 +/- ##
===========================================
- Coverage 77.79% 56.95% -20.85%
===========================================
Files 519 518 -1
Lines 57926 58413 +487
===========================================
- Hits 45064 33267 -11797
- Misses 12862 25146 +12284
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:
|
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>
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.
kevalmorabia97
requested review from
shengliangxu
and removed request for
kevalmorabia97
![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
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
2 participants
Type of change
New feature.
Details
Qwen3.6 search-space comparison
Compared the default W4A16 NVFP4/FP8 search space against the module-specific recipe that fixes routed experts to W4A16 NVFP4 while shared experts, linear attention, and self attention search W4A16 NVFP4 versus FP8.
Both lanes use local raw gradient scoring, active-MoE cost accounting, batch size 1, code32k calibration, full parser-on LiveCodeBench (8 repeats), and full SciCode. Active GiB/token includes scale storage.
BF16 references are 74.3667 LCB and 40.7914 SciCode, giving 1.0 percentage-point acceptance floors of 73.3667 and 39.7914. The W4-routed search space passes both benchmarks only at target 6.4; the default lane has no joint pass. Targets 6.1 and 6.7 pass LCB but miss the SciCode floor. Fixing the inexpensive routed-expert decision to W4 redirects the active-MoE budget toward always-active attention/shared-expert modules, improving accuracy while reducing active cost by about 0.12 GiB/token at the same nominal targets.
This is an end-to-end recipe comparison rather than an isolated one-factor solver ablation: the historical default lane used full linear/self-attention family decision grouping, while the module-specific lane used runtime-required grouping.
Testing
Summary by CodeRabbit
New Features
Documentation
Bug Fixes