Repository navigation
[OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder - #3067
Open
shengliangxu wants to merge 13 commits into
Conversation
shengliangxu
requested review from
chadvoegele
and removed request for
a team
October 27, 2026 20:37
Contributor
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review: 3
Please check the settings in the CodeRabbit UI or the
 |
| Layer / File(s) | Summary |
|---|---|
Fused-expert wrapper dispatch modelopt/torch/quantization/plugins/huggingface.py, modelopt/torch/export/moe_utils.py, tests/unit/torch/quantization/plugins/test_fused_experts.py |
The generic wrapper no longer detects non-gated experts by matching an up_proj shape. Module detection now recognizes registered wrappers that subclass _QuantFusedExperts. Documentation describes non-gated subclasses. Tests remove the synthetic non-gated wrapper suite and fixtures, and run the eager native FP8 test for gated and non-gated configurations. |
Nemotron-H wrapper and validation modelopt/torch/models/nemotron_h/modeling_ptq.py, tests/unit/torch/models/test_nemotron_h_modeling_ptq.py, tests/gpu/torch/models/test_modeling_ptq_registration.py |
NemotronHExperts is registered with a wrapper that uses up_proj as its first projection and marks the layout non-gated. Tests check conversion, projection quantizers, eager and grouped_mm output equivalence, calibration statistics, GPTQ Hessians, export, implementation normalization, and conditional registry availability. |
Priority:
Estimated code review effort: 3 (Moderate) | ~20 minutes
Change: Feature
Sequence Diagram(s)
sequenceDiagram
participant NemotronHExperts
participant QuantModuleRegistry
participant QuantNemotronHExperts
participant Projections
NemotronHExperts->>QuantModuleRegistry: Resolve registered wrapper
QuantModuleRegistry->>QuantNemotronHExperts: Select Nemotron-H wrapper
QuantNemotronHExperts->>Projections: Configure up_proj and down_proj quantizers
Merge Risk: 
70994
-
No blocking issue was established for Nemotron-H expert quantization; the PR is mergeable after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: adding Nemotron-H expert quantization for grouped_mm with a model-specific wrapper. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns No listed security anti-pattern was introduced. The PR adds no torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/…Full details: Security Anti-Patterns
-
Explanation
No listed security anti-pattern was introduced. The PR adds no
torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# nosecline. The changed-file inventory includes nopyproject.tomlor requirements file, so no dependency was added. The existingtrust_remote_code=Truereference inhuggingface.pyis unchanged by this PR.
-
Finishing Touches: 💡 2
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
shengliangxu
requested review from
kevalmorabia97 and
meenchen
and removed request for
a team
October 27, 2026 20:37
This was referenced Oct 27, 2026
Merged
[OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank
#2751
Open
Collaborator
Author
|
/ok to test 23ab7f4 |
shengliangxu
force-pushed
the
shengliangx/hfep-8a-grouped-mm-calibration
branch
from
October 28, 2026 23:46
eac1a35 to
df2e3fe
Compare
Collaborator
Author
|
On CodeRabbit's nitpick in its review of
|
Collaborator
Author
|
On the fourth suggestion in the Claude review summary of |
cjluo-nv
reviewed
Oct 29, 2026
cjluo-nv
left a comment
cjluo-nv
left a comment
Collaborator
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.
LGTM: the earlier threads are resolved, and the newest change (dropping the _fused_experts_wrapper fallback in eager_only) gives the same result as before. A human still needs to sign off on the deleted tests.
Needs action:
- Confirm the test deletions as the human reviewer.
TestNonGatedFusedExperts, the "not registered yet" test and the non-gated parametrizations exercised the removed_QuantNonGatedFusedExpertsstand-in. Those cases now run on the realNemotronHExpertsintest_nemotron_h_modeling_ptq.py, and the FP8 force-eager test is parametrized overhas_gate.
No action needed:
- ✔️ Resolved since the last review: the two stale comments and the dead
or _fused_experts_wrapper(module)arm innormalize_experts_impl_on_the_fly. CodeRabbit's rename nit was declined because it would clash with the flattest_<model>_modeling_ptq.pynaming. - The author declined the warning for unregistered non-gated containers, the
_setupshape assert and theHAS_GROUPED_MMguards, giving reasons from the code. I agree with those reasons. - The new test file's header matches
LICENSE_HEADER.
shengliangxu
force-pushed
the
shengliangx/hfep-8a-grouped-mm-calibration
branch
from
October 29, 2026 07:09
df2e3fe to
d50b2c3
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-8b-nemotron-h-experts
branch
from
October 29, 2026 07:09
32869b0 to
15f3fa9
Compare
Collaborator
Author
|
/ok to test 15f3fa9 |
shengliangxu
force-pushed
the
shengliangx/hfep-8a-grouped-mm-calibration
branch
from
October 29, 2026 16:21
d50b2c3 to
ec9e09c
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-8b-nemotron-h-experts
branch
from
October 29, 2026 16:21
15f3fa9 to
03c4da1
Compare
Collaborator
Author
|
/ok to test 03c4da1 |
shengliangxu
force-pushed
the
shengliangx/hfep-8a-grouped-mm-calibration
branch
from
October 29, 2026 17:56
ec9e09c to
8eaf5a4
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-8b-nemotron-h-experts
branch
from
October 29, 2026 17:56
03c4da1 to
87b2e8b
Compare
Collaborator
Author
|
/ok to test 87b2e8b |
… its model folder Nemotron-H's non-gated experts get their own wrapper, _QuantNemotronHExperts, in modelopt/torch/models/nemotron_h/modeling_ptq.py, registered on transformers' NemotronHExperts (every supported release ships it) in place of the generic plugin's _QuantNonGatedFusedExperts. It is not eager-only: the grouped_mm interception reaches the non-gated layout through _first_proj_attr, so these experts calibrate on the backend transformers selects, as the gated ones do. The generic structural detection now claims only the gated layout. _is_fused_experts_module also counts a class registered with a _QuantFusedExperts subclass, so the backend normalization still finds Nemotron-H's experts, and transformers' non-gated FP8 experts, which stay on eager. The non-gated tests move to tests/unit/torch/models and run the real NemotronHExperts. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Nemotron-H's experts register at import here, so no class the on-the-fly registration claims is eager-only, and the stand-in the test needed is gone with the generic non-gated detection. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ed layout Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…near The grouped-vs-eager comparisons would pass just as well if the grouped forward fell back to per-expert F.linear calls, which the eager hook intercepts with the same values. The interception count tells the two hooks apart: one _grouped_linear call per projection, against one F.linear call per projection and expert. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The test this one replaced checked it: for the non-gated layout up_proj is the whole fused weight, so an export that kept it would hold every expert twice, the fused original beside the per-expert copies. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_is_fused_experts_module found Nemotron-H's experts through their class's registration, which the class quantize() converts them to does not carry, so a second quantize() over the converted model skipped the backend normalization for them. Gated experts are found by their shapes either way. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
up_proj's input is the hidden states, which the grouped forward hands _grouped_linear gathered by expert: the same rows eager feeds per expert, so their amax matches exactly. Only down_proj's input passes through a GEMM and keeps the tolerance. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Its one caller, normalize_experts_impl_on_the_fly, runs inside replace_quant_module, which refuses a model that is already quantized, so no converted experts module reaches it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…Experts The non-gated cases of these fused-experts tests went with the stand-in when the non-gated tests moved here. They run on the real class again: the eager forward recovers each up_proj slice's expert, and the quantizer-only iterator yields the weight iterator's quantizers in its order without indexing the fused weight; that test fails if the iterator names the gated projection. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
With the non-gated wrapper gone, register_fused_experts_on_the_fly registers only the generic wrapper, which is not eager-only, so looking it up for a class not registered yet cannot change which experts are forced onto eager: those wrappers (FP8's, GPT-OSS's, Nemotron-H's) register when their modules are imported. The transposed check's comment no longer speaks of non-gated layouts. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
force-pushed
the
shengliangx/hfep-8a-grouped-mm-calibration
branch
from
October 29, 2026 18:14
8eaf5a4 to
f688379
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-8b-nemotron-h-experts
branch
from
October 29, 2026 18:14
87b2e8b to
da59231
Compare
Collaborator
Author
|
/ok to test da59231 |
shengliangxu
added a commit
that referenced
this pull request
Oct 29, 2026
… on the meta device (#2737) > [!NOTE] > **The expert-parallel PTQ series is now a tree instead of one stack** (re-cut 2026-10-27, from main 1521053), so PRs that do not depend on each other can be reviewed and merged in parallel. This PR was [5/21]. Rebased onto main 7c3c868 later the same day, after #2739 merged, and onto main f03651e that night. > > **Base:** `main`, now that #2736 has merged (2026-10-28). The diff below is this PR's alone. > > **Re-cut:** this PR's diff is unchanged; only its base moved. Review fixes since are listed at the end of Testing. <details><summary>Series layout</summary> - `main` - #2745 fused experts under grouped_mm (merged) - #3074 calibration on grouped_mm - #3067 Nemotron-H's experts under grouped_mm - #2740 calibration on per-rank data - #2739 meta-DCP checkpoint load (merged) - #2736 distribution plan, run group, collectives (merged) - #2750 keep-fused fold - #2751 distributed-writer pieces - #2752 no-gather export (on #2751, #2750, #2779 and #2740) - #2753 expert-parallel export - **#2737 plan decision (this PR)** - #2779 load routing - #2741 expert-parallel loader (on #2779, #3067 and #2740) - #2742 dense sharding - #2743 block-FP8 dequantization - #2780 expert-parallel load routing - #2746 expert parallelism in the model specs - #2747 on-the-fly decompression - #2748 Kimi-K2.5 and Kimi-VL adapters - #2749 Kimi-K3 adapter - #2754 expert-parallel PTQ from hf_ptq.py (on #2746 and #2753) </details> ### What does this PR do? Type of change: New feature `resolve_distribution_plan(ckpt_path)` decides how a run spreads a model over its ranks. It returns a `DistributionPlan` (#2736) for every model, MoE or dense, and for a single process too. It is the one place a run decides this: #2779 loads the model by the plan and records the plan on it, and every later step reads it from there. - **From the model, not estimates.** Rank 0 builds the model on the meta device, using the class that loads it whole: a causal LM, or for a VLM the class that carries its LM head. A trusted remote VLM whose `auto_map` registers `AutoModelForImageTextToText` builds through that class, with its LM head. It decides from the model's parameters and then discards it, and leaves the auto class it built through as it was: building a remote-code model through an auto class registers its class there, as it is, and rank 0's later load through that class would then build it without the `GenerationMixin` every other rank's first load adds. Rank 0's Kimi-K3 lacked `generate` while the other ranks' had it, and an expert-parallel preview deadlocked. - **EP degree** (`ep_size`, or `MODELOPT_EP_DEGREE`). - `auto` takes the largest degree that divides both the world and the routed-expert count. - The routed experts are the ones the config's `base_model_ep_plan` names. They are found by the rule the expert-parallel loader (#2741) uses, so the decision and the load cannot disagree. - An explicit degree must divide the routed-expert count. - **Dense layout** (`dense`, or `MODELOPT_FSDP_SIZE`, which takes `auto`, `1` to replicate, or `world` or the world size to shard). - `auto` FSDP2-shards the dense part when a whole copy of it plus the rank's share of the routed experts would take more than 60% of one GPU's memory. - Sharding over part of the world is refused. - **Every rank reads the checkpoint first.** Each rank makes the checkpoint local and reads its config inside the failure-carrying `all_reduce` (#2736), before the plan's broadcast. A rank that cannot see the checkpoint fails every rank with its own error, instead of leaving rank 0 waiting in the broadcast. - **The checkpoint the load reads.** `resolve_distribution_plan` takes `revision`, `token`, `cache_dir`, `local_files_only` and `subfolder` as `load_pretrained` does, so the plan and its config are for the checkpoint the load reads. - **Rank 0 decides, and every rank gets its plan** (`broadcast_from`, #2736). The inputs are read locally: the knobs, this GPU's memory, the files on this node. Ranks deciding for themselves could disagree, and ranks loading different plans hang. If rank 0 fails, every rank raises. - **A single process** decides without building anything, and reads no knob: a shell still exporting a multi-GPU run's `MODELOPT_EP_DEGREE` or `MODELOPT_FSDP_SIZE` must not reshape it. It holds the whole model, unless the caller's own `dense="fsdp"` asks for FSDP2 on one rank. Measuring the meta model avoids the gaps of estimating from config fields, which this PR's earlier revision did: it ignored GQA and tied embeddings, assumed bf16, and treated a config missing the fields it read as one that fits a GPU. **Why not hf_ptq's `get_model`, or Accelerate's sizing.** The decision builds through the library's own meta build, `build_meta_causal_lm`, which main's FSDP2 load uses too. `get_model` lives in the hf_ptq example, which the library cannot import, and is a single-process `device_map` loader: it requires `architectures` and takes its class from the first entry, beside the example's own speculative, pack-quantized and MXFP4 branches. Accelerate's `infer_auto_device_map` and `compute_module_sizes` say which device holds each module of such a load. The plan needs one number instead: what one rank holds when expert parallelism splits the routed experts `1/ep_size` and keeps the rest whole, a single sum over the meta model's parameters (re-tied first, as `init_empty_weights` unties them). The class that builds a VLM with its head takes the checkpoint's remote class only under `trust_remote_code`, a `MODELOPT_EP_DEGREE` that is not a number is refused by name, and a single process parses its own `ep_size` as a distributed run does, refusing a bad one the same way. ### Usage ```python # torchrun --nproc_per_node 8 ... from modelopt.torch.utils.plugins.model_load_utils import resolve_distribution_plan plan, hf_config = resolve_distribution_plan("Qwen/Qwen3-30B-A3B") # DistributionPlan(world_size=8, ep_size=8, dense='replicate') ``` ### Testing - Tests in `tests/unit/torch/utils/test_model_load_utils.py` (the decision) and `tests/unit/torch/utils/test_moe_a2a.py` (the MoE helpers), the files of the modules they cover. They run on CPU, build tiny real configs on the meta device, and answer rank and size as torchrun would. They cover: - the knobs, read at call time; expert counting; EP-degree selection; - routed experts found under any model prefix, named by their style or as their grouped-GEMM parameters' parent; - the class that builds a whole model, a VLM's with its LM head; - the decision: - a single process, whatever knobs its shell exports; an EP-capable MoE; an MoE without an EP plan; - the dense fit from exact parameter bytes, counting what expert parallelism leaves whole and the rank's share of the routed experts, so an expert-dominated MoE shards its dense part; - knob and argument overrides; - refused partial sharding and indivisible EP degrees; - a VLM decided from its whole model; - a model that cannot be built on meta, named in the error; - the meta build leaving the auto class's registrations as it found them (fails without the fix); - on two real gloo ranks (`cpu_dist_workers`): both take rank 0's plan, though rank 1 exports a different dense layout, and rank 1 decides nothing; a rank that cannot see the checkpoint fails both ranks with its error, before the broadcast. Rank 0's own failures are `broadcast_from`'s, tested in #2736. - On 2 GPUs with transformers 5.18, a tiny Kimi-K3 at EP 2 with previews on: before this fix rank 0's model alone had no `generate`, so rank 0 moved on to calibration while rank 1 decoded in lockstep, and the run deadlocked (`TORCH_DISTRIBUTED_DEBUG=DETAIL`: an all-reduce on one rank against an all-to-all on the other); with it, both ranks preview, calibrate and export. - Mutation-checked: 30 mutants, each failing a test; and, since, a rank deciding for itself, the plan sent from the wrong rank, a single process reading the knobs, the checkpoint read outside the `all_reduce`, the experts' share left out of the dense fit, the Hub options not passed on, a single process's `ep_size` not parsed, and a trusted remote VLM's registered image-text class not taken. - Near-duplicate tests are merged into one test per behavior, and the env-knob, `ep * dp == world` and no-EP-plan tests, which the decision tests cover, are dropped. Each drop is checked by mutation. - The auto-class table has an untrusted remote VLM and a trusted remote image-text VLM, and a non-numeric `MODELOPT_EP_DEGREE` names the knob, as a bad `ep_size` in a single process names the argument. - The Hub options locate the checkpoint against an offline Hugging Face cache with two revisions and a subfolder. - Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full `tests/unit` on transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite. - Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite. - Review fixes, 2026-10-27: an encoder-decoder model builds through its seq2seq class; `MODELOPT_EP_DEGREE` reads auto in any case; a hybrid plan names its routed experts by both styles; the expert count is read from the config holding the plan; a docstring note on a rank 0 without a GPU. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head. - Rebased on 2026-10-27 onto main 7c3c868, after #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8). - Review suggestions, 2026-10-27 evening: the refusal of an EP degree that cannot shard the experts names the knob (`MODELOPT_EP_DEGREE`) or the argument the degree came from. - Review fixes, 2026-10-27 night: the meta build leaves `AutoModel`'s and `AutoModelForCausalLM`'s registrations as it found them; a routed-expert config whose expert-parallel plan names no module of the model is refused unless the degree is 1; an explicit degree that does not divide the world names the knob or argument it came from. Each new test fails without its change. - Review suggestions, 2026-10-27, last round: the routed-expert match computes each module name's wildcard form (every layer number replaced) once instead of once per key, with the same matches; `build_meta_causal_lm`'s docstring states the registrations its restore assumes, and `resolve_distribution_plan`'s that a caller passes the local checkpoint directory. - Review fix, 2026-10-27, last round: a VLM-shaped config that `AutoModelForCausalLM` recognizes (Phi-4-multimodal, the only such config in transformers 5.5.4 or 5.18) builds through it, head included, rather than through `AutoModel`. The new table row fails without the change. - Review suggestion, 2026-10-27, last round: `resolve_distribution_plan`'s docstring says each rank returns its own copy of the config, rank 0's deciding the plan. - Review suggestion, 2026-10-28: `_choose_ep_degree`'s docstring says why the automatic EP degree is the largest: it holds the fewest experts per rank, leaving the most memory for activations, and every degree above 1 runs the same two all-to-alls per MoE layer. - Declined review findings, 2026-10-27: - The plan decision builds the model on meta even when `ep_size=1` and `dense` are both given. Rank 0 takes 0.54-0.65 s and 8-208 MB of host memory on 232B-1T configs, the FSDP2 load builds the same meta model anyway, and a model that cannot be built on meta loads with `distribution=`, which skips the decision. - `_auto_class` falls back to `AutoModel` for a remote VLM whose `AutoModelForCausalLM` class is named `*ForCausalLM`. None of the remote checkpoints the series supports is one: Kimi-VL, Kimi-K2.5, Kimi-K3 and Step-3.7-Flash register `*ForConditionalGeneration`, and DeepSeek-VL2 registers no auto class. Qwen3-Omni, which `AutoModel` cannot build, is not supported on main. - Review suggestion, 2026-10-28, second round: the plan decision's MoE helpers and the three auto classes are imported at module top, nothing having deferred them. - Review suggestion, 2026-10-28, fifth round: `modelopt.torch.utils.distributed` is imported at module top as `dist`; #2779 drops its four function-local imports of it. - Rebased on 2026-10-28, late night, onto #2736 at `cf4597dcbeb`, on main `1ca8bd1fc4d` (after #3088 and #2745 merged). The one conflict was with main's #2676: `build_meta_causal_lm` keeps this PR's `auto_class` parameter, now defaulting to `None`, so a call without one still gets main's text-only-VLM factory. The plan decision passes `_auto_class(...)` as before, and the registry restore covers whichever class builds. This PR's tests and `test_model_load_utils.py` pass on transformers 5.18 (90) and 5.5.4, and pre-commit passes against main. - Rebased on 2026-10-28, late night, onto main `bdc45368d10` after #2736 merged; #2736's squash has the tree of its last head, so this PR's tree is unchanged. - Review fixes, 2026-10-29: the plan decision re-ties the meta model before sizing it, so a tied embedding and LM head count once (`init_empty_weights` unties them); the dense-fit test is sized from a normally built model, tied and untied, on either side of the 60% share. An encoder-decoder config that names no architecture (a bare `T5Config` or `WhisperConfig`) builds through its seq2seq class, checked through the public resolver too. Each fails its test without the change (mutation-checked). The routed-expert match's local transformers import names its reason (the optional `[hf]` extra). Pre-commit passes against main; this PR's tests, `test_model_load_utils.py` and `test_moe_a2a.py` pass on CPU with transformers 5.18 and 5.5.4. - Review fix, 2026-10-29: the tests moved into the files of the modules they cover (CONTRIBUTING.md): the plan decision's into `test_model_load_utils.py`, `moe.py`'s helpers' into `test_moe_a2a.py` (a `test_moe.py` in `tests/unit` already belongs to `modelopt/torch/models/moe.py`, and pytest refuses two test modules of one basename in a run), and the plain-module MoE they share into `tests/_test_utils/torch/moe_models.py`; `test_resolve_distribution_plan.py` is gone. Test bodies are unchanged apart from the helpers' names. Pre-commit passes against main; the moved tests pass on CPU with transformers 5.18 and 5.5.4. - Review suggestions, 2026-10-29, Claude's approving round: `build_meta_causal_lm`'s dead local import of `AutoModelForImageTextToText` is gone; routed experts the built model has but the config does not count (under `num_experts`, `num_local_experts`, `n_routed_experts` or `num_routed_experts`) are refused unless the degree is 1, as a plan naming no module of the model is, instead of being replicated on every rank silently. The new test fails without the refusal (mutation-checked). Pre-commit passes against main; this PR's tests pass on CPU with transformers 5.18 and 5.5.4. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ Additive. - 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A. The feature's entry lands with #2754. - Did you get Claude approval on this PR?: ✅ Approved 2026-10-29 at `1f9b0ee7e98`; its two suggestions are taken since. ### Additional Information Part 5 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2736. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added automatic model distribution planning, including expert- and data-parallel degree selection and dense-parameter sharding based on GPU memory. - Added support for environment-configured parallelism settings and model selection and sizing for vision-language models. - Added validation for model configurations and parallelism choices. - **Bug Fixes** - Distribution plans are now consistent across ranks, and planning failures are propagated to all ranks. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
This branch has not been deployed
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
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.
Note
Split out of #2745 on 2026-10-27, stacked on it, so that the Nemotron-H part is reviewed on its own. Later that night #2745 split again; this PR now stacks on #3074, its second half.
Base:
shengliangx/hfep-8a-grouped-mm-calibration(#3074), which sits on #3122 (its expert-parallel fix, split out on 2026-10-29), on main. It merges after #3074.In the tree: #2741's helper base merges this PR instead of #2745 and #3074. The expert-parallel chain quantizes Nemotron-H's experts on the
torch_a2aforward, which refuses an eager-only wrapper.Series layout
main, with [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 (FSDP2 weight-access window,_fsdp2_owned_params) and [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 (fused experts under grouped_mm), both merged 2026-10-28What does this PR do?
Type of change: New feature
#2745 quantizes fused experts under
grouped_mm, the experts backend transformers selects by default, in the gated layout (gate_up_proj). Nemotron-H's experts are non-gated,down_proj(act_fn(up_proj(x))), with a single 3-Dup_proj. On #2745 and #3074 they keep main's generic_QuantNonGatedFusedExpertsand calibrate on the eager forward, which main forces and #3074 keeps for them._QuantNemotronHExpertslives inmodelopt/torch/models/nemotron_h/modeling_ptq.py. It registers on transformers'NemotronHExperts, which every supported release (5.5 on) ships, asmodelopt/torch/models/README.mdasks of wrappers only one model needs. The generic plugin's_QuantNonGatedFusedExpertsand itsup_proj-only structural detection are removed: no other transformers experts class has that layout.grouped_mm. The wrapper is not eager-only: [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's grouped interception reaches the non-gated layout through_first_proj_attr, and their quantizer state matches the eager forward's. Under expert parallelism ([OMNIML-6060] Load an MoE expert-parallel: shard the experts at load time #2741 on), they run thetorch_a2aforward, which refuses an eager-only wrapper._is_fused_experts_modulealso counts a class registered with a_QuantFusedExpertssubclass. So a backend the quantizers cannot see (batched_mm) is still redirected for Nemotron-H. Transformers' non-gated FP8 experts (an FP8 Nemotron-H's) are still forced onto eager.Behavior change. Nemotron-H's MoE calibration runs on
grouped_mmwhere transformers selects it, so some of its activation scales can differ from main's, as #3074's changelog entry says for fused experts. Loading withexperts_implementation="eager"calibrates as before.Usage
No API change. To calibrate exactly as earlier releases did:
Testing
tests/unit/torch/models/test_nemotron_h_modeling_ptq.py(new) runs transformers' realNemotronHExperts. It covers:up_projquantizers;grouped_mm;grouped_mmcalibration and GPTQ Hessians equal to eager's, for every expert;up_projanddown_proj;main's tests of the non-gated layout ran a hand-written stand-in, which the registration does not reach, so they move here.
test_fused_experts.py's FP8 force-eager test also runs non-gated FP8 experts. It and the backend-redirect test fail with the registry lookup in_is_fused_experts_modulereverted.tests/gpu/torch/models/test_modeling_ptq_registration.pychecks thatNemotronHExpertsis registered whichever module is imported first.Locally at this head: pre-commit, mypy included;
tests/unit/torch/{models,export},test_fused_experts.pyandtest_huggingface.pyon transformers 5.18 and 5.5.4 and on torch 2.8, with the same failures as main in each environment (diffusers). The series' PRs merged together add no unit-test failure over main.GPU, NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 (a BF16 checkpoint, no
quantization_config), one GPU, transformers 5.18.0,general/ptq/nvfp4_default-kv_fp8, 64 calibration samples of 512 tokens, exported tensors compared one by one:grouped_mm, twice (run-to-run floor)grouped_mmvsexperts_implementation="eager"The weights and weight scales do not depend on the backend; activation scales move with it, in the MoE layers and in the layers after them, as [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's changelog entry says for fused experts.
tests/gpu/torch/models/test_modeling_ptq_registration.pypasses on 1 GPU.Re-cut onto [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's last review fixes (2026-10-27): the plugin-order test of the non-gated stand-in is dropped with the stand-in. Pre-commit and this PR's unit tests pass on transformers 5.18 and 5.5.4. A remote-code Nemotron-H keeps its experts as an
nn.ModuleListofnn.LinearMLPs, which the fused detection never claimed, and the remote-code VL model builds transformers'NemotronHForCausalLM, whose experts this PR registers.Review suggestion, 2026-10-27, last round: the wrapper's docstring says the explicit registration skips the transposed-layout check and relies on transformers'
(num_experts, out_dim, in_dim)layout, which every supported release uses.Rebased with [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 onto main f03651e (2026-10-27, night); this PR's own diff is unchanged.
Review suggestions, 2026-10-27, last round: the test pins that
grouped_mmwent through_grouped_linear, two calls rather than oneF.linearcall per expert and projection, and that the export releases the fusedup_projanddown_proj. Re-cut onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074, the second half of [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's split; this PR's own diff is otherwise unchanged.Review suggestions, 2026-10-28: an already-converted Nemotron-H experts module still counts as fused experts for a second
quantize(), pinned by the registration test, which fails without the change;up_proj's input amax is compared exactly, and onlydown_proj's keeps a tolerance.Review suggestion, 2026-10-28, second round: the converted-experts check reuses
_is_quant_fused_experts_module.Review suggestion, 2026-10-28, third round:
_is_fused_experts_moduleno longer counts already-converted experts, undoing the first round's change: its one caller runs insidereplace_quant_module, which refuses a quantized model, so none reaches it.Declined review suggestion, 2026-10-28, third round, raised before: a warning for unregistered non-gated fused experts. transformers'
NemotronHExperts, which is registered, is the only class with the layout, and Nemotron-H's remote code builds per-expert MLPs.Declined review suggestions, 2026-10-28, fourth round: skip guards on the grouped tests, since
NemotronHExpertsand_grouped_linearare in every release the pin allows; and a shape assert in_setup, which cannot tell a consistently transposed pair from the stored one, while the export test built on transformers' ownNemotronHExpertspins the per-expert shapes.Declined review suggestion, 2026-10-28, fifth round (CodeRabbit): resolving a converted experts module's original class in
_is_fused_experts_module; its one caller runs before any module is converted.Rebased on 2026-10-28, late night, onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 on main
1ca8bd1fc4d(after [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 and [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 merged), so the diff shows this PR's own slice only (+292/−362 in 6 files). Conflicts were with [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's later changes:_is_fused_experts_module's docstring keeps [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's GPT-OSS sentence, the FP8 force-eager test runs both [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's backend cases and this PR's gated and non-gated experts, and the FP8 wrapper's comment no longer compares it with the removed non-gated wrapper.Review suggestion (cjluo-nv's bot), same night: the non-gated cases that went with the stand-in run on the real
NemotronHExpertsagain: the eager forward recovers eachup_projslice's expert, and the quantizer-only iterator yields the weight iterator's quantizers in its order without indexing the fused weight (it fails if the iterator names the gated projection). pre-commit passes against main;test_fused_experts.pyand this PR's tests pass on transformers 5.18 and 5.5.4.Rebased on 2026-10-28, late night, onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 on main
bdc45368d10([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 merged), with its own diff unchanged. Tests and pre-commit pass as for [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074.Restacked on 2026-10-29 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 after its expert-parallel fix moved to [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 underneath it, with its own diff unchanged (one import-list conflict with [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's new
assert_fused_experts_interceptedimport). Its tests,test_nemotron_h_modeling_ptq.pyandtest_moe_a2a.pypass (185), and pre-commit passes against main.Before your PR is "Ready for review"
grouped_mm, so some activation scales differ from earlier releases;experts_implementation="eager"restores the earlier calibration.CONTRIBUTING.md: N/Agrouped_mm.Additional Information
Part of expert-parallel PTQ for Hugging Face MoE models; merges after #2745.
Summary by CodeRabbit