[shengliangxu]

What does this PR do?

Type of change: New feature

Before it can shard a model's experts, the expert-parallel path needs three facts about the model:

• which modules are the routed experts, and how they shard;
• which transformers release first ships that information;
• for a VLM, which nested config the plan must be attached to.

All three are facts about the architecture, so they belong with the model's other facts in modelopt.torch.models.

• EPSpec, a new ModelSpec section, and apply_distributed_support, the one applier that reads it. Most models become data in the spec package they already had:
  • deepseek_v3 carries the plan stock transformers ships only from 5.13;
  • qwen3_moe and minimax_m3_vl carry only the transformers release their own plan needs.
  • An existing plan on the config always wins, since overwriting it would substitute our idea of the layout for the model's.
• One thing stays code: EPSpec.prepare, the escape hatch for checkpoints that ship their own modeling code ([18/21] to [20/21]). A spec whose plan only prepare attaches states that plan too (EPSpec.prepared_ep_plan), so it can be read without running prepare, which patches the model's code. Nemotron-H needs none. Its gate emits logits and the block runs its own top-k, so the experts receive global expert ids, which the all-to-all experts forward ([1/21]) routes on directly.
• Nemotron-H-Omni is registered separately. Its experts live in the language sub-model, and transformers aggregates a sub-model's plan up to the root with its module prefix, so a plan on the VLM root would match nothing, silently. A test pins the placement and that the two spellings share one plan. transformers ships the Omni model itself from 5.18, which its spec records. The plan goes on the language-model config, which the checkpoint's own code names llm_config and transformers' 5.18 port text_config; the spec lists both, text_config first as the load reads them, and the first one a config has takes the plan. On a tiny port built on CPU, the plan on text_config names its NemotronHExperts parameters; the port has not been run expert-parallel on GPUs. Its root type registers only the expert-parallel plan, so is_moe by the root type is False for its NemotronHMoE blocks. No export depends on the missing MoE and export sections: on transformers 5 its experts are one fused container, exported by its own handler (checked locally, see Testing).
• A plan goes where the load reads it. apply_distributed_support refuses a spec whose plan_config_attr is neither text_config nor llm_config, the two nested configs the expert-parallel load reads a plan and the expert count from (one module-level tuple), instead of attaching a plan that would be read as none and plan the MoE as dense. plan_config_attr may list several spellings, the first one the config has taking the plan; each must be one the load reads, and a config with none of them is refused, naming them all. The spellings are checked for every spec, whether the spec or its prepare supplies the plan, before anything is attached or prepared.
• EPSpec.unsupported refuses a model whose own plan would route it into the expert-parallel path before that path can handle it. The refusal comes before anything is loaded and names a single process first, and FSDP2 (--use_fsdp2 or MODELOPT_FSDP_SIZE=world) only for an unquantized causal LM: GPT-OSS ships MXFP4, which the FSDP2 load refuses, and Llama-4's checkpoint is multimodal. A spec's transformers-floor refusal says the same.
  • GPT-OSS and Llama-4 store their experts transposed, which the fused-experts export cannot write (GPT-OSS also interleaves gate and up).
  • Llama-4's MoE block also routes tokens itself rather than through the experts interface.
• The decision and the load apply the spec. [12/21] carried this before the series was reordered.
  • Unless the run rules expert parallelism out (a single process, or an EP degree of 1), resolve_distribution_plan applies the model's spec to a copy of the config before rank 0 builds the model on meta; without it, DeepSeek-V3 before 5.13 would show no experts to shard. Every rank applies it under a collective, so a rank that cannot fails every rank. An expert-parallel plan returns the prepared copy; any other plan returns the config as read, which the export would otherwise write the spec's plan into. Neither the spec's plan nor the load's torch_a2a swap of it reaches an exported config.json: this PR adds no export path, and the expert-parallel export ([17/21]) writes each config's base_model_ep_plan as the source config.json declares it (none for DeepSeek-V3 before 5.13, Nemotron-H or Nemotron-H-Omni).
  • A load handed that config does not prepare it again, since a model's preparation can patch its modeling code; a load that decides for itself prepares it once.
  • A handed config still meets its spec's refusal and transformers floor (check_distributed_support). A config= passed for the load with a plan handed in is the caller's own, so an expert-parallel load prepares a copy of it once, as deciding does.
  • An EP degree left to auto resolves to 1 for a model whose spec refuses expert parallelism or needs a later transformers than the one installed (one predicate, _spec_refuses_ep, asks check_distributed_support), so its spec is skipped and only an explicit EP degree above 1 (or a handed expert-parallel plan) meets the refusal: MODELOPT_FSDP_SIZE=world alone plans FSDP2 for GPT-OSS, and for Qwen3-MoE on transformers 5.5.
  • At an EP degree of 1, which skips the spec, the check that refuses an MoE held whole on every rank ([12/21]) still finds a spec-supplied plan's experts. It reads the plan the spec would supply, or the one it states prepare attaches, without applying anything. DeepSeek-V3 at MODELOPT_EP_DEGREE=1, replicated, is refused like a native-plan MoE instead of loading whole on every rank, and the refusal reports the expert count found through the spec's plan, which an EP degree has to divide.
  • For an MoE whose spec refuses expert parallelism (GPT-OSS, Llama-4), that refusal offers the ways out its sibling for an expert-parallel-capable MoE offers besides sharding the experts, from one shared text: FSDP2 (--use_fsdp2 in hf_ptq, or MODELOPT_FSDP_SIZE=world) if the checkpoint is an unquantized causal LM, a single process, or a replicated plan handed in (DistributionPlan(world_size, 1, 'replicate')).
  • Rank 0's knobs decide whether the spec applies, and rank 0 sends that decision as it does the plan: a rank whose own environment skipped the spec would load without the expert-parallel plan the others shard by, and fail once the mesh is up. The gate reads the EP knob through the decision's own _ep_degree, so a malformed knob fails every rank before any spec runs.

Usage

from modelopt.torch.utils.moe import apply_distributed_support

hf_config = apply_distributed_support(hf_config.model_type, hf_config, ckpt_path=path)
# raises NotImplementedError for gpt_oss / llama4 / llama4_text, before any weight loads

Testing

• tests/unit/torch/models/test_ep_specs.py:
  • every registered plan is well formed, and every EP section declares something;
  • the Omni placement and plan sharing, the plan landing under either name (llm_config or text_config), and on text_config for a config with both;
  • the applier's version floor, plan attachment, existing-plan precedence and prepare ordering;
  • a spec attaching its plan where the load reads none refused with the config untouched, whether the spec or its prepare supplies the plan;
  • the refusals: GPT-OSS, Llama-4 and Llama-4-text are refused with their configs untouched, naming a single process and FSDP2 only for an unquantized causal LM, and a refusal comes before the rest of the section.
• tests/unit/torch/utils/test_resolve_distribution_plan.py and test_load_pretrained_distribution.py gain the spec's cases, which [12/21] carried before the reorder:
  • the spec applied before the decision, DeepSeek-V3's plan coming only from its spec, and the returned config carrying it for an expert-parallel plan only;
  • on two real gloo ranks, a rank that cannot apply the spec failing every rank;
  • no spec for a run that rules expert parallelism out, by argument or by MODELOPT_EP_DEGREE=1, so GPT-OSS's spec does not refuse an FSDP2 run;
  • auto picks no EP degree for GPT-OSS: MODELOPT_FSDP_SIZE=world alone plans FSDP2, and an explicit ep_size=2 meets the spec's refusal;
  • a Qwen3-MoE whose spec floor is unmet plans FSDP2 under auto and is refused at an explicit ep_size=2; the decision test of an MoE without a plan runs on transformers 5.5 again;
  • on two real gloo ranks (cpu_dist_workers), every rank applying the spec, or not, as rank 0 decided, whatever its own knob says;
  • DeepSeek-V3 refused when left whole on every rank, at an explicit EP degree of 1 too, as a native-plan MoE is, and allowed FSDP2 instead; GPT-OSS left whole refused as an MoE that cannot be expert-parallel yet, by the decision rather than by its spec, naming a single process and FSDP2 only for an unquantized causal LM; the left-whole refusal names the expert count, DeepSeek-V3's at an EP degree of 1 included;
  • the plan a spec would supply, read without running prepare: the spec's own plan, or the one it states prepare attaches, and nothing for a model with no spec or EP section, or whose spec refuses expert parallelism, and a spec below its transformers floor keeping its plan;
  • a decided expert-parallel load handed the spec-prepared config, and the config prepared once, not at all when handed, and on a copy when passed as config= with a handed plan;
  • GPT-OSS refused with its config handed in, without preparing it.
• Mutation-checked: skipping the spec or applying it without expert parallelism, ignoring the EP knob, letting each rank decide the spec for itself, preparing a handed config twice or not at all, ignoring the spec's plan at an EP degree of 1, ignoring the plan a spec states for prepare, letting auto pick an EP degree for a model whose spec refuses expert parallelism or below a spec's floor, leaving the spec plan's experts out of the refusal's count, attaching a plan where the load reads none, dropping the single process or FSDP2's unquantized-causal-LM condition from the left-whole and spec refusals, reading the Omni plan off llm_config alone, trying llm_config before text_config, checking a plan spelling only when the spec supplies the plan itself, hiding the plan of a spec below its transformers floor, and dropping the EP-unsupported branch each fail a test.
• Tests the applier and decision tests already cover are dropped, and the refusal test names its three models rather than looping over whichever specs refuse. Each drop is checked by mutation.
• Checked locally, not as a unit test: an FP8, INT4-AWQ or W4A8-AWQ export of a tiny NemotronH writes byte-identical tensors with nemotron_h's MoE and export sections and without them (transformers 5.18 and 5.5.4 on torch 2.14, 5.18 on torch 2.8).
• Checked locally, at the stack's tip: DeepSeek-V3, Nemotron-H and Nemotron-H-Omni configs prepared by their specs and swapped to the torch_a2a styles as the expert-parallel load does are written by the expert-parallel export with no base_model_ep_plan, as their sources have none (transformers 5.18 on torch 2.14 and 2.8, and 5.5.4); a bare save_pretrained of the same configs writes the swapped plan.
• Locally on this revision, the stack rebased on main 0638a27: pre-commit over the whole PR against main, at every PR of the stack; each fix's own tests at its PR; and at the stack's tip the unit tests of the changed areas (tests/unit/torch/utils, export, models, the fused-experts quantization tests and the hf_ptq example tests) on transformers 5.18 and 5.5 and on torch 2.8. They pass apart from the diffusers-export and GGML tests, which fail the same way on main in that environment. CI runs the full unit suite.
• The GPU results here were measured on earlier revisions of this branch; this revision's GPU tests run in CI only.
• On GPUs (1 node, EP4): gpt-oss-20b under torchrun stopped at the plan decision with the spec's refusal message, before any weights load. That run predates auto picking no EP degree for GPT-OSS: left to auto, a replicated run now stops at the decision's refusal ("cannot be quantized expert-parallel yet, and the plan replicates it ..."), still before any weights load; this was not re-run on GPU. MiniMax-M3 runs on stock transformers 5.13 code (see [21/21]).

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?: N/A. The feature's entry lands with [21/21].
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 13 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [12/21]. The branch's history adds a Llama-4 and a MiniMax-M3 adapter and then removes them. The Llama-4 adapter's routing was numerically wrong (it scaled expert outputs where Llama-4 scales inputs), and the M3 one only repeated what transformers 5.13 ships. Review the net diff.

Summary by CodeRabbit

• New Features
  • Added expert-parallel configuration support for DeepSeek-V3, Nemotron-H, Nemotron-H-Omni, and Qwen3-MoE.
  • Added model registrations for MiniMax-M3-VL and Llama-4 Text.
  • Model loading now applies model-specific expert-parallel settings, including nested configurations, when selecting distribution plans.
• Compatibility
  • GPT-OSS and Llama-4 variants are identified as unsupported for expert-parallel loading.
  • Expert-parallel support requires minimum Transformers versions for applicable models.

[copy-pr-bot]

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.

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 3aaff7f
• Latest commit, not reviewed: e2db64b

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

The change adds expert-parallel specifications for model plans, version requirements, and unsupported cases. It applies these specifications to Hugging Face configs during distributed plan resolution and expert-parallel loading.

Changes: Expert-parallel support

Layer / File(s)	Summary
Expert-parallel contracts and model registrations modelopt/torch/models/specs.py, modelopt/torch/models/__init__.py, modelopt/torch/models/*/specs.py, modelopt/torch/models/llama4_text/*, modelopt/torch/models/minimax_m3_vl/*, modelopt/torch/models/nemotron_h_omni/*	Adds EPSpec to ModelSpec and registers expert-parallel plans, Transformers requirements, or unsupported declarations for DeepSeek-V3, GPT-OSS, Llama-4, Llama-4 text, MiniMax-M3-VL, Nemotron-H, Nemotron-H-Omni, and Qwen3-MoE.
Apply expert-parallel specifications modelopt/torch/utils/moe.py, tests/unit/torch/models/test_ep_specs.py	Adds helpers to retrieve and apply model specifications. The application checks model support and version requirements, attaches plans when absent, and runs optional preparation callbacks. Tests cover these behaviors and error cases.
Use specifications in distributed loading modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_load_pretrained_distribution.py, tests/unit/torch/utils/test_resolve_distribution_plan.py	Applies model specifications during plan resolution and prepares configs for expert-parallel loading when needed. Tests cover plan selection, config forwarding, and rank consistency.

Priority: Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Resolver as resolve_distribution_plan
  participant Ranks as Distributed ranks
  participant Support as apply_distributed_support
  participant Registry as Model spec registry
  Resolver->>Ranks: Broadcast whether to apply distributed support
  Resolver->>Support: Apply model spec to config
  Support->>Registry: Retrieve model expert-parallel spec
  Support-->>Resolver: Return prepared config
  Resolver->>Ranks: Resolve distribution plan

Loading

Merge Risk: Low · up to 3aaff

Default multi-GPU Qwen3-MoE loads can fail on Transformers 5.5 or 5.6. Users can set the EP degree to one, but the automatic planning regression should be fixed before merging.

Pre-merge checks: 6

Check name	Status	Explanation
Title check		The title names the main change: adding expert-parallel (EP) descriptions to the model specs. The raw summary confirms this. It adds EPSpec, per-model EP plans, Transformers version floors, and unsu…
Docstring Coverage		Docstring coverage is 82.93% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 41 functions across 18 files.
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		PASS. The PR adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), eval, exec, or # nosec usage. New Transformers loading paths propagate trust_remote_code from …
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Title check

Explanation

The title names the main change: adding expert-parallel (EP) descriptions to the model specs. The raw summary confirms this. It adds EPSpec, per-model EP plans, Transformers version floors, and unsupported-model refusals to the specs, and the PR objectives match this scope. The bracketed ticket and stack-position prefix are noise but do not make the title misleading. The title is specific and accurate, so it passes.

Full details: Security Anti-Patterns

Explanation

PASS. The PR adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), eval, exec, or # nosec usage. New Transformers loading paths propagate trust_remote_code from the caller; the new API defaults it to False. The only added trust_remote_code=True instances are test arguments or documentation text, not hardcoded model/tokenizer loading. No dependency files changed.

✨ Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

📝 Generate docstrings

• 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.}

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2746/
Built to branch gh-pages at 2026-10-25 01:02 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[shengliangxu]

/ok to test 2e134ec

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 61.75%. Comparing base (e2a88f1) to head (e2db64b).

Additional details and impacted files

@@                           Coverage Diff                           @@
##           shengliangx/hfep-9b-ep-load-routing    #2746      +/-   ##
=======================================================================
+ Coverage                                61.70%   61.75%   +0.04%     
=======================================================================
  Files                                      649      655       +6     
  Lines                                    72892    72990      +98     
=======================================================================
+ Hits                                     44976    45072      +96     
- Misses                                   27916    27918       +2     

Flag	Coverage Δ
unit	61.75% <100.00%> (+0.04%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shengliangxu]

/ok to test fdd217a

[shengliangxu]

/ok to test 67f1c8f

[shengliangxu]

/ok to test 2f18815

[shengliangxu]

/ok to test e2db64b