Repository navigation
[OMNIML-3817] Move model-specific PTQ modeling into modelopt/torch/models [1/4] - #2576
Conversation
|
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. |
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (2)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 9 remain after this review. 📝 WalkthroughWalkthroughThe change adds model specs for Falcon, Llama 4, Qwen3-VL MoE, Step-3.5, and Step-3.7. The models package imports the new packages to register their specs. MoE layout tests now cover Llama 4 and Qwen3-VL MoE. ChangesModel spec registrations
Priority: ➖ Normal Estimated code review effort: 2 (Simple) | ~15 minutes Change: Refactor Merge Risk: ⚪ Minimal · up to The new model registrations remain compatible with the declared Transformers range, and the inspected Qwen quantization and export paths handle both layouts. No concrete merge-blocking behavior issue is established; merge after normal checks. 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings
🧪 Generate unit tests (beta)
🛠️ Fix failing CI checks 💡
Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2576 +/- ##
==========================================
+ Coverage 69.18% 78.46% +9.27%
==========================================
Files 620 628 +8
Lines 69622 69638 +16
==========================================
+ Hits 48171 54644 +6473
+ Misses 21451 14994 -6457
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:
|
dc72b71 to
4706a4d
Compare
|
764bf11 to
f149be0
Compare
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: the registry additions follow the existing design and preserve export policy; human sign-off is required for licensing and justified test edits.
Needs action:
- Confirm the new files’ 2024 copyright dates against the 2026
LICENSE_HEADER, and sign off on the exhaustive test-table extensions.
No action needed:
- The design rationale supports extending existing dataclass-based
ModelSpecpackages rather than retaining centralized metadata or introducing Pydantic/OmegaConf configuration. - Existing parametrized tests exercise the new MoE names and accessors, while registry-wide checks cover remote-code declarations. Test edits extend coverage without weakening assertions.
- Grouped export remains disabled for the new MoE specs; Step models intentionally declare no unsupported MoE layout.
|
@cjluo-nv re: the copyright item in your review — fixed in 132783c: the 10 new files in this PR ( |
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: the copyright concern is resolved; the registry additions are sound, with human sign-off still required for justified existing-test edits.
Needs action:
- Sign off on the exhaustive table extensions in
tests/unit/torch/models/test_model_specs.py; they cover the new MoE specs without weakening existing assertions.
No action needed:
- ✔️ Resolved since the last review: all ten new files now match the canonical 2026
LICENSE_HEADER. - The design extends the existing dataclass-based
ModelSpecregistry rather than introducing a parallel metadata system; the description explains model colocation and export constraints. - Existing parametrized tests cover new block names and expert accessors; registry-wide checks cover Step’s remote-code declarations. Grouped export remains disabled for the new MoE specs.
132783c to
187669a
Compare
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: the prior copyright concern is resolved, and the registry additions are sound; justified existing-test extensions still require human sign-off.
Needs action:
- Sign off on the exhaustive table extensions in
tests/unit/torch/models/test_model_specs.py; they add the new MoE layouts and root names without weakening existing assertions.
No action needed:
- ✔️ Resolved since the last review: all ten new source files match the canonical 2026
LICENSE_HEADER. - The design extends the existing dataclass-based registry rather than introducing a parallel system; model colocation and export constraints justify the approach.
- Parametrized tests cover the new layouts and expert-name accessors; registry-wide tests check Step’s remote-code declarations. Grouped export remains disabled for the new MoE specs.
- Tests inspected, not executed.
ed8c199 to
e4726f3
Compare
e4726f3 to
73db571
Compare
These model types have model-specific PTQ support but no ModelSpec. Register one for each, recording only facts checked against transformers 4.57 and 5.14: - llama4, qwen3_vl_moe: fused MoE layout (gate_up_proj/down_proj), no gate/up pair, grouped export off. Both blocks were already detected as MoE, so this is data only. - falcon: dense, no sections. - step3p5, step3p7: remote code, no MoESpec. Step's expert projections sit on the MoE MLP itself with no experts container; declaring the block would send AWQ export into get_experts_list, which does not support that layout. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The new llama4, qwen3_vl_moe, falcon, step3p5 and step3p7 packages copied a 2024 header from an existing spec; new files carry the year in LICENSE_HEADER (2026). Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
73db571 to
b811de2
Compare
transformers 5.16 added the step3p7 package, which also registers the step3p5 text model type, so the specs no longer describe remote-code-only models. Set min_transformers_version to 5.16 instead of modeling_source=remote_code; older Step checkpoints still load as remote code and are covered by step3p5/modeling_ptq.py. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
) ### What does this PR do? Type of change: Refactor ### Summary of the series Model-specific PTQ modeling moves out of `modelopt/torch/quantization/plugins/huggingface.py` into the per-model packages under `modelopt/torch/models/<model_type>/modeling_ptq.py`, next to each model's `specs.py`. By the end of the series `huggingface.py` shrinks from ~1880 to ~1500 lines and keeps only the code shared across models: sequential/fused MoE auto-detection, attention, `FP8Linear`, `CompressedLinear`, parallel linears, and the transposed-quantization helpers that gpt_oss and llama4 share. Model types that gain PTQ modeling but had no `ModelSpec` get one. - Class and function bodies move verbatim. The only rewrites: on-the-fly callbacks are added to `CUSTOM_MODEL_PLUGINS` by their own module, and `is_homogeneous_hf_model` imports `is_nemotron_h_model` lazily. - `huggingface.py` imports every `modeling_ptq` from an explicit list. The list sits after the generic wrappers these modules build on and before the homogeneous decoder discoverer is registered, so Nemotron-H's more specific discoverer still matches first. - Each package's `__init__.py` keeps importing only `specs`, so `import modelopt.torch.models` does not pull in quantization or transformers. - No public API or quantization behavior changes; only private (`_`-prefixed) names change module. Merge in order. #2576 targets `main`; each later PR targets the one before it. 1. #2576 — Add model specs for llama4, qwen3_vl_moe, falcon, step3p5 and step3p7 2. #2577 — Add the `modeling_ptq.py` convention; move Nemotron-H, Falcon and Llama4 ← **this PR** 3. #2578 — Move GPT-OSS 4. #2580 — Move the Step family (step3p5, shared by step3p7) The DBRX slice (#2579) was dropped: #2649 removed DBRX support, so there was nothing left to move. Likewise, the Qwen3-VL-MoE half of [3/4] was dropped: #2670 removed its pre-5.12 wrapper. ### This PR [2/4] Introduces `<model_type>/modeling_ptq.py` and moves the first three models there: - `nemotron_h`: `is_nemotron_h_model` / `get_nemotron_h_decoder_layers` and their layerwise-calibration registration. - `falcon`: the `FalconLinear` registration and `register_falcon_linears_on_the_fly`. - `llama4`: `_QuantLlama4TextExperts` and its registration. The registrations move in the form #2670 gave them (a top-level transformers import plus `if ... not in QuantModuleRegistry`), and `huggingface.py` drops the imports it no longer uses. It also adds the `importlib` loop in `huggingface.py` that imports each `modeling_ptq`, documents the convention in `modelopt/torch/models/README.md` and the package docstring, and adds `tests/unit/torch/models/test_modeling_ptq_registration.py`, which checks in a fresh interpreter per case that registration does not depend on which module is imported first. It also raises the `llama4`, `qwen3_vl_moe` and `falcon` specs from #2576 to the transformers 5.5 floor that #2670 set for every other spec, and drops the `qwen3_vl_moe` spec's pointer to the pre-5.12 wrapper #2670 removed. ### Usage No API change. To add PTQ support for a new model, put its wrapper in `modelopt/torch/models/<model_type>/modeling_ptq.py` and add `<model_type>` to the import loop at the end of `quantization/plugins/huggingface.py`. ### Testing Each branch of the stack checked out and tested on `main` at 90ba9fb, CPU, torch 2.11. With transformers 5.18: `pytest tests/unit/torch/models tests/unit/torch/quantization tests/unit/torch/export tests/unit/recipe`, 2370 passed, 1 skipped. With transformers 5.5 (the new floor): `tests/unit/torch/models` plus `test_moe_linear.py` and `test_huggingface.py`, 205 passed, 1 skipped. GPU tests are left to CI. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ Only private (`_`-prefixed) names change module; public API and quantization behavior are unchanged. - 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?: N/A (verbatim move; covered by existing tests) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (internal refactor) - Did you get Claude approval on this PR?: ❌ ### Additional Information Part of a 5-PR stack; see the series list above. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added post-training quantization support for Llama 4 expert layers. * Expanded Falcon quantization coverage to include linear layers discovered at runtime and older remote-code checkpoints. * Added Nemotron-H model detection and decoder-layer discovery for quantization. * **Bug Fixes** * Improved Qwen3-VL mixture-of-experts compatibility by applying the legacy quantization wrapper only when needed. * **Documentation** * Clarified how model-specific quantization support is loaded and where model support information belongs. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…dels [3/4] (#2578) ### What does this PR do? Type of change: Refactor ### Summary of the series Model-specific PTQ modeling moves out of `modelopt/torch/quantization/plugins/huggingface.py` into the per-model packages under `modelopt/torch/models/<model_type>/modeling_ptq.py`, next to each model's `specs.py`. By the end of the series `huggingface.py` shrinks from ~1880 to ~1500 lines and keeps only the code shared across models: sequential/fused MoE auto-detection, attention, `FP8Linear`, `CompressedLinear`, parallel linears, and the transposed-quantization helpers that gpt_oss and llama4 share. Model types that gain PTQ modeling but had no `ModelSpec` get one. - Class and function bodies move verbatim. The only rewrites: on-the-fly callbacks are added to `CUSTOM_MODEL_PLUGINS` by their own module, and `is_homogeneous_hf_model` imports `is_nemotron_h_model` lazily. - `huggingface.py` imports every `modeling_ptq` from an explicit list. The list sits after the generic wrappers these modules build on and before the homogeneous decoder discoverer is registered, so Nemotron-H's more specific discoverer still matches first. - Each package's `__init__.py` keeps importing only `specs`, so `import modelopt.torch.models` does not pull in quantization or transformers. - No public API or quantization behavior changes; only private (`_`-prefixed) names change module. Merge in order. #2576 and #2577 have merged; #2578 now targets `main`, and #2580 targets #2578's branch. 1. #2576 — Add model specs for llama4, qwen3_vl_moe, falcon, step3p5 and step3p7 (merged) 2. #2577 — Add the `modeling_ptq.py` convention; move Nemotron-H, Falcon and Llama4 (merged) 3. #2578 — Move GPT-OSS ← **this PR** 4. #2580 — Move the Step family (step3p5, shared by step3p7) The DBRX slice (#2579) was dropped: #2649 removed DBRX support, so there was nothing left to move. Likewise, the Qwen3-VL-MoE half of [3/4] was dropped: #2670 removed its pre-5.12 wrapper. ### This PR [3/4] Moves `_QuantGptOssExperts` and its registration into `gpt_oss/modeling_ptq.py`, adds `gpt_oss` to the import loop, and extends the registration test to it. This slice originally also moved the pre-5.12 Qwen3-VL-MoE wrapper; #2670 removed that wrapper, so that half was dropped. ### Usage No API change. To add PTQ support for a new model, put its wrapper in `modelopt/torch/models/<model_type>/modeling_ptq.py` and add `<model_type>` to the import loop at the end of `quantization/plugins/huggingface.py`. ### Testing Each branch of the stack checked out and tested on `main` at 90ba9fb, CPU, torch 2.11. With transformers 5.18: `pytest tests/unit/torch/models tests/unit/torch/quantization tests/unit/torch/export tests/unit/recipe`, 2371 passed, 1 skipped. With transformers 5.5 (the new floor): `tests/unit/torch/models` plus `test_moe_linear.py` and `test_huggingface.py`, 206 passed, 1 skipped. GPU tests are left to CI. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ Only private (`_`-prefixed) names change module; public API and quantization behavior are unchanged. - 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?: N/A (verbatim move; existing tests re-pointed) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (internal refactor) - Did you get Claude approval on this PR?: ❌ ### Additional Information Part of a 5-PR stack; see the series list above. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added post-training quantization support for GPT-OSS models. * Added support for quantizing legacy Qwen3-VL-MoE layouts, while retaining support for newer fused layouts. * Quantization support now accommodates both legacy and newer Qwen3-VL-MoE model layouts. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…dels [4/4] (#2580) ### What does this PR do? Type of change: Refactor ### Summary of the series Model-specific PTQ modeling moves out of `modelopt/torch/quantization/plugins/huggingface.py` into the per-model packages under `modelopt/torch/models/<model_type>/modeling_ptq.py`, next to each model's `specs.py`. By the end of the series `huggingface.py` shrinks from ~1880 to ~1500 lines and keeps only the code shared across models: sequential/fused MoE auto-detection, attention, `FP8Linear`, `CompressedLinear`, parallel linears, and the transposed-quantization helpers that gpt_oss and llama4 share. Model types that gain PTQ modeling but had no `ModelSpec` get one. - Class and function bodies move verbatim. The only rewrites: on-the-fly callbacks are added to `CUSTOM_MODEL_PLUGINS` by their own module, and `is_homogeneous_hf_model` imports `is_nemotron_h_model` lazily. - `huggingface.py` imports every `modeling_ptq` from an explicit list. The list sits after the generic wrappers these modules build on and before the homogeneous decoder discoverer is registered, so Nemotron-H's more specific discoverer still matches first. - Each package's `__init__.py` keeps importing only `specs`, so `import modelopt.torch.models` does not pull in quantization or transformers. - No public API or quantization behavior changes; only private (`_`-prefixed) names change module. Merge in order. #2576, #2577 and #2578 have merged; this last slice, #2580, now targets `main`. 1. #2576 — Add model specs for llama4, qwen3_vl_moe, falcon, step3p5 and step3p7 (merged) 2. #2577 — Add the `modeling_ptq.py` convention; move Nemotron-H, Falcon and Llama4 (merged) 3. #2578 — Move GPT-OSS (merged) 4. #2580 — Move the Step family (step3p5, shared by step3p7) ← **this PR** The DBRX slice (#2579) was dropped: #2649 removed DBRX support, so there was nothing left to move. Likewise, the Qwen3-VL-MoE half of [3/4] was dropped: #2670 removed its pre-5.12 wrapper. ### This PR [4/4] Moves `_QuantMoELinear`, `_is_expert_indexed_moe_linear`, `_is_step_family_model`, `register_moe_linear_on_the_fly` and `_reconstruct_fused_moe_linear` into `step3p5/modeling_ptq.py`, which also serves `step3p7`. This carries over #2569's switch of `_is_step_family_model` to `hf_model_type`; `huggingface.py` drops its now-unused `hf_model_type` and `re` imports. Export (`hf_export_handlers.py`, `layerwise_export.py`, `unified_export_hf.py`, `unified_export_hf_streaming.py`) imports these names from the new module. The imports stay deferred, because `modelopt.torch.quantization` imports `modelopt.torch.export` and importing at module scope would risk an import cycle; the comments now give that reason. The PTQ skill references (`unsupported-models.md`, `checkpoint-validation.md`) now point agents at `modelopt/torch/models/<model_type>/modeling_ptq.py` for model-specific patches. Rebasing onto #2649 also replaces the README's DBRX example of a model-specific wrapper with Llama4's fused BMM experts. `step3p7` also gets its own `modeling_ptq.py`, which imports the Step-3.5 module: Step-3.7 remote-code checkpoints use the same `MoELinear`, and today they are covered only because the HF plugin imports every `modeling_ptq`. With its own module, a future loader that imports PTQ modeling by model type finds it too. ### Usage No API change. To add PTQ support for a new model, put its wrapper in `modelopt/torch/models/<model_type>/modeling_ptq.py` and add `<model_type>` to the import loop at the end of `quantization/plugins/huggingface.py`. ### Testing Each branch of the stack checked out and tested on `main` at 90ba9fb, CPU, torch 2.11. With transformers 5.18: `pytest tests/unit/torch/models tests/unit/torch/quantization tests/unit/torch/export tests/unit/recipe`, 2372 passed, 1 skipped. With transformers 5.5 (the new floor): `tests/unit/torch/models` plus `test_moe_linear.py` and `test_huggingface.py`, 207 passed, 1 skipped. GPU tests are left to CI. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ Only private (`_`-prefixed) names change module; public API and quantization behavior are unchanged. - 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?: N/A (verbatim move; existing tests re-pointed) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (internal refactor) - Did you get Claude approval on this PR?: ❌ ### Additional Information Part of a 5-PR stack; see the series list above. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added post-training quantization support for Step-family models with expert-indexed MoE layers. * Checkpoint exports preserve the original layout of quantized expert weights and scales. * **Documentation** * Updated MoE architecture examples and guidance for configuring model-specific quantization support, including Step model revisions. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
What does this PR do?
Type of change: Refactor
Summary of the series
Model-specific PTQ modeling moves out of
modelopt/torch/quantization/plugins/huggingface.pyinto the per-model packages undermodelopt/torch/models/<model_type>/modeling_ptq.py, next to each model'sspecs.py. By the end of the serieshuggingface.pyshrinks from ~2025 to ~1525 lines and keeps only the code shared across models: sequential/fused MoE auto-detection, attention,FP8Linear,CompressedLinear, parallel linears, and the transposed-quantization helpers that gpt_oss and llama4 share. Model types that gain PTQ modeling but had noModelSpecget one.CUSTOM_MODEL_PLUGINSby their own module, andis_homogeneous_hf_modelimportsis_nemotron_h_modellazily.huggingface.pyimports everymodeling_ptqfrom an explicit list. The list sits after the generic wrappers these modules build on and before the homogeneous decoder discoverer is registered, so Nemotron-H's more specific discoverer still matches first.__init__.pykeeps importing onlyspecs, soimport modelopt.torch.modelsdoes not pull in quantization or transformers._-prefixed) names change module.Merge in order. #2576 targets
main; each later PR targets the one before it.modeling_ptq.pyconvention; move Nemotron-H, Falcon and Llama4The DBRX slice (#2579) was dropped: #2649 removed DBRX support, so there was nothing left to move.
This PR [1/4]
Registers a
ModelSpecfor each model type that gainsmodeling_ptq.pylater in the series but has no spec yet. Each spec records only facts checked against transformers 4.57 and 5.14:llama4andqwen3_vl_moe: fused MoE layout (gate_up_proj/down_proj), no gate/up pair, grouped export off. This is data only: both blocks were already detected as MoE (structurally and by name), and their expert containers are exported without name lookups.falcon: dense, no sections.step3p5/step3p7:modeling_source="remote_code"and intentionally noMoESpec. Step's expert projections sit directly on the MoE MLP with noexpertscontainer. Declaring the block would makeis_moeclaim it and send AWQ export intoget_experts_list, which does not support that layout.step3p7gets its own package, following thegemma4_text/gemma4precedent.The exhaustive spec tables in
tests/unit/torch/models/test_model_specs.py(EXPECTED_MOE_LAYOUTS,root_class_names) gain the two MoE rows.Usage
No API change. The new specs are read through the existing registry, e.g.
get_spec("llama4").moe_spec.Testing
Each branch of the stack checked out and tested with
pytest tests/unit/torch/models tests/unit/torch/quantization tests/unit/torch/export tests/unit/recipeon CPU (torch 2.11, transformers 5.14), onmainat 9f902ae (#2649): 2248 passed, 7 skipped. The new block names resolve intest_specs_vs_transformers.py, and the step3p5/step3p7 remote-code absence checks pass. Class names were also checked against the transformers 4.57 wheel. GPU tests and the transformers 4.57 matrix are left to CI.Before your PR is "Ready for review"
_-prefixed) names change module; public API and quantization behavior are unchanged.CONTRIBUTING.md: N/AEXPECTED_MOE_LAYOUTS,root_class_names).Additional Information
Part of a 5-PR stack; see the series list above.
🤖 Generated with Claude Code
Summary by CodeRabbit