Repository navigation
[OMNIML-3817] Move model-specific PTQ modeling into modelopt/torch/models [4/4] - #2580
Merged
Merged
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. |
Contributor
|
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 (7)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review. 📝 WalkthroughWalkthroughStep-family expert-indexed MoE quantization now resides in ChangesStep MoE PTQ
Priority: ⬇️ Low Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Refactor Sequence Diagram(s)sequenceDiagram
participant StepModel
participant register_moe_linear_on_the_fly
participant QuantMoELinear
StepModel->>register_moe_linear_on_the_fly: inspect model and module types
register_moe_linear_on_the_fly->>StepModel: register matching types with QuantMoELinear
StepModel->>QuantMoELinear: call forward with x and expert_id
Suggested reviewers: Merge Risk: ⚪ Minimal · up to This refactor relocates Step-family MoE quantization code without changing behavior. Registration and export paths continue to work through the new module. It is ready to merge. 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings
🧪 Generate unit tests (beta)
🛠️ Fix failing CI checks 💡
Comment |
This was referenced Sep 28, 2026
shengliangxu
added this pull request to stack #2631
October 1, 2026 22:32
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2580 +/- ##
==========================================
+ Coverage 71.69% 78.16% +6.46%
==========================================
Files 642 644 +2
Lines 71285 71886 +601
==========================================
+ Hits 51109 56187 +5078
+ Misses 20176 15699 -4477
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:
|
shengliangxu
force-pushed
the
shengliangx/move-modeling_5
branch
from
October 2, 2026 01:12
de350ca to
824518b
Compare
Contributor
|
shengliangxu
force-pushed
the
shengliangx/move-modeling_5
branch
from
October 2, 2026 17:13
824518b to
b29aa6b
Compare
shengliangxu
force-pushed
the
shengliangx/move-modeling_5
branch
from
October 2, 2026 17:17
b29aa6b to
88c315f
Compare
shengliangxu
marked this pull request as ready for review
October 2, 2026 17:19
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
shengliangxu
force-pushed
the
shengliangx/move-modeling_5
branch
from
October 2, 2026 17:37
88c315f to
1702932
Compare
h-guo18
added a commit
that referenced
this pull request
Oct 3, 2026
…in fp32 Two unit tests failed in CI on #2279 (run 36999463172). - TestDFlash2Forward.test_overfits_a_single_batch hit the 60 s cap in the py3.11 and py3.14 torch 2.14 jobs. The tiny base is bf16, and the draft keeps Qwen3Config's default 22016-wide MLP and 32x128 attention. On a CPU without AVX-512, bf16 matmul has no fast kernel: on 4 Zen3 cores a [104, 22016] @ [22016, 32] mm takes 57 ms in bf16 and 0.93 ms in fp32. For this test the CI runners split into a fast group (2-9 s) and a slow one (51-58 s in other PRs' runs); the logs carry no CPU model, but the slow group behaves like such a CPU. The test had already timed out once on code that does not touch DFlash (#2580, unit-3.12 torch_212, 60.14 s). The pinned anchor count added one masked block here (13 instead of 12 at SEQ_LEN=16, block 4), 8% more rows in eager mode, and moved this PR's slow jobs to 58.5 s and past the cap. In fp32 the test runs in about 5 s on 4 cores, as the LiLiCorr overfit test already does. The test also asserts selector coverage now, which is what guards the target alignment in fp32. With the selector's gold token one position early, bf16 training ends at train_acc 0.847, but fp32 ends at 0.903 and slips past 0.9. Coverage separates the two cleanly: over 71 correct trajectories (batch content, draft init or base init reseeded) it ends at 0.986 or above, and over 112 trajectories with the gold token one position early or late at 0.764 or below. Of six selector misalignments, the fp32 test catches the same two the bf16 test did; shifting the predecessor or the hidden state goes unnoticed in both, as before. - TestDraftStackCompile.test_compiled_matches_eager failed on torch 2.10 with 2/256 elements over atol=2e-2. In bf16, Inductor and eager round at different points: the worst element already used 82% of the tolerance locally, and 3 of 61 random inits and inputs went over. In fp32 the difference stays below 2e-6 across torch 2.8-2.14, so the check runs in fp32 at 1e-4. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
shengliangxu
force-pushed
the
shengliangx/move-modeling_5
branch
2 times, most recently
from
October 6, 2026 06:17
ee38d2e to
d57e5b6
Compare
shengliangxu
removed this pull request from stack #2631
October 6, 2026 09:15
shengliangxu
force-pushed
the
shengliangx/move-modeling_5
branch
from
October 8, 2026 01:03
7a909ac to
71dbbca
Compare
shengliangxu
force-pushed
the
shengliangx/move-modeling_5
branch
2 times, most recently
from
October 8, 2026 16:47
3e91098 to
702f928
Compare
shengliangxu
added a commit
that referenced
this pull request
Oct 8, 2026
) ### 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>
shengliangxu
force-pushed
the
shengliangx/move-modeling_5
branch
from
October 8, 2026 20:10
702f928 to
dce1662
Compare
shengliangxu
changed the title
shengliangxu
force-pushed
the
shengliangx/move-modeling_5
branch
from
October 9, 2026 00:10
dce1662 to
dd01678
Compare
shengliangxu
added a commit
that referenced
this pull request
Oct 9, 2026
…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>
Move _QuantMoELinear, its detection and registration, and _reconstruct_fused_moe_linear verbatim into step3p5/modeling_ptq.py, shared by step3p7. Export keeps importing _reconstruct_fused_moe_linear and _QuantMoELinear lazily, now from the new module, since modelopt.torch.quantization imports modelopt.torch.export. The PTQ skill references now point agents at modelopt/torch/models for model-specific patches. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
DBRX support was removed in #2649; use Llama4's fused BMM experts as the example instead. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Import step3p5 modeling_ptq first too, and check the Step on-the-fly callback is installed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
force-pushed
the
shengliangx/move-modeling_5
branch
from
October 9, 2026 06:24
dd01678 to
c23979a
Compare
cjluo-nv
reviewed
Oct 9, 2026
cjluo-nv
left a comment
cjluo-nv
left a comment
Collaborator
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Nudge: the refactor appears behavior-preserving, but the new file's copyright year differs from LICENSE_HEADER and needs confirmation.
Needs action:
- Confirm the 2024 copyright in
modelopt/torch/models/step3p5/modeling_ptq.pyis intentional, or update it to the canonical 2026 header; obtain human licensing sign-off if retaining the difference. - Move
init_empty_weightsto module-level imports instep3p5/modeling_ptq.py, or document why Accelerate must remain an optional, deferred dependency at_setup. - Relocate
test_moe_linear.pyto the Step model tests and name ittest_modeling_ptq.py, followingCONTRIBUTING.md's requirement that tests track renamed or moved production modules.
No action needed:
- Design rationale is sufficient: this extends the existing per-model convention and callback registry using stdlib
importlib, rather than introducing another system or forcing Step's layout into generic MoE wrappers. - Test edits preserve existing assertions and add Step registration/import-order coverage; their justification is visible in the move.
test_moe_linear.py covers _QuantMoELinear and its detection, registration and export, which now live in modelopt/torch/models/step3p5/modeling_ptq.py. Rename it to tests/unit/torch/models/test_step3p5_modeling_ptq.py, following CONTRIBUTING.md's rule that test filenames track the module they cover and the test_nemotron_h_mtp.py precedent. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
accelerate is a separate optional dependency, loaded by its own ModelOpt plugin; the HF plugin, which imports every modeling_ptq, does not require it, so the import stays inside _QuantMoELinear._setup as it was in quantization/plugins/huggingface.py. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The module is new in this PR; give it the year in LICENSE_HEADER (2026). Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Collaborator
Author
|
@cjluo-nv re: the three items in your review — all addressed:
|
cjluo-nv
reviewed
Oct 9, 2026
cjluo-nv
left a comment
cjluo-nv
left a comment
Collaborator
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: all prior concerns are resolved; the behavior-preserving refactor is ready for human sign-off on the justified test relocation and registration-test extension.
Needs action:
- Confirm human sign-off on the Step test relocation and added import-order/callback coverage; existing assertions remain intact and track the moved production code.
No action needed:
- ✔️ Resolved since the last review: the header matches
LICENSE_HEADER, the deferred Accelerate import has an optional-dependency justification, and the Step tests moved to the model-test directory. - Design remains settled: moving Step-specific code reduces the shared HF plugin’s model-specific burden. The PR reuses the existing per-model convention, callback registry, and stdlib
importlib, rather than retaining centralized wrappers or extending generic MoE detection for Step’s distinct layout.
cjluo-nv
reviewed
Oct 9, 2026
Step-3.7 remote-code checkpoints use the same expert-indexed MoELinear as Step-3.5, handled by step3p5/modeling_ptq.py. It was covered only because the HF plugin imports every modeling_ptq; a loader that imports PTQ modeling by model type would have found nothing for step3p7. Add step3p7/modeling_ptq.py importing the Step-3.5 module, list step3p7 in the plugin's import loop and the registration test, and say so in the models README (also correcting its stale try/except wording from before #2670). Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
h-guo18
added a commit
that referenced
this pull request
Oct 10, 2026
…in fp32 Two unit tests failed in CI on #2279 (run 36999463172). - TestDFlash2Forward.test_overfits_a_single_batch hit the 60 s cap in the py3.11 and py3.14 torch 2.14 jobs. The tiny base is bf16, and the draft keeps Qwen3Config's default 22016-wide MLP and 32x128 attention. On a CPU without AVX-512, bf16 matmul has no fast kernel: on 4 Zen3 cores a [104, 22016] @ [22016, 32] mm takes 57 ms in bf16 and 0.93 ms in fp32. For this test the CI runners split into a fast group (2-9 s) and a slow one (51-58 s in other PRs' runs); the logs carry no CPU model, but the slow group behaves like such a CPU. The test had already timed out once on code that does not touch DFlash (#2580, unit-3.12 torch_212, 60.14 s). The pinned anchor count added one masked block here (13 instead of 12 at SEQ_LEN=16, block 4), 8% more rows in eager mode, and moved this PR's slow jobs to 58.5 s and past the cap. In fp32 the test runs in about 5 s on 4 cores, as the LiLiCorr overfit test already does. The test also asserts selector coverage now, which is what guards the target alignment in fp32. With the selector's gold token one position early, bf16 training ends at train_acc 0.847, but fp32 ends at 0.903 and slips past 0.9. Coverage separates the two cleanly: over 71 correct trajectories (batch content, draft init or base init reseeded) it ends at 0.986 or above, and over 112 trajectories with the gold token one position early or late at 0.764 or below. Of six selector misalignments, the fp32 test catches the same two the bf16 test did; shifting the predecessor or the hidden state goes unnoticed in both, as before. - TestDraftStackCompile.test_compiled_matches_eager failed on torch 2.10 with 2/256 elements over atol=2e-2. In bf16, Inductor and eager round at different points: the worst element already used 82% of the tolerance locally, and 3 of 61 random inits and inputs went over. In fp32 the difference stays below 2e-6 across torch 2.8-2.14, so the check runs in fp32 at 1e-4. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 ~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 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, #2577 and #2578 have merged; this last slice, #2580, now targets
main.modeling_ptq.pyconvention; move Nemotron-H, Falcon and Llama4 (merged)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_flyand_reconstruct_fused_moe_linearintostep3p5/modeling_ptq.py, which also servesstep3p7. This carries over #2569's switch of_is_step_family_modeltohf_model_type;huggingface.pydrops its now-unusedhf_model_typeandreimports.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, becausemodelopt.torch.quantizationimportsmodelopt.torch.exportand 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 atmodelopt/torch/models/<model_type>/modeling_ptq.pyfor model-specific patches.Rebasing onto #2649 also replaces the README's DBRX example of a model-specific wrapper with Llama4's fused BMM experts.
step3p7also gets its ownmodeling_ptq.py, which imports the Step-3.5 module: Step-3.7 remote-code checkpoints use the sameMoELinear, and today they are covered only because the HF plugin imports everymodeling_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.pyand add<model_type>to the import loop at the end ofquantization/plugins/huggingface.py.Testing
Each branch of the stack checked out and tested on
mainat 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/modelsplustest_moe_linear.pyandtest_huggingface.py, 207 passed, 1 skipped. GPU tests 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/AAdditional Information
Part of a 5-PR stack; see the series list above.
🤖 Generated with Claude Code
Summary by CodeRabbit