[shengliangxu]

Note

Slice [3/3] of this PR's split (2026-10-30), so the model-spec work is reviewed and merged in pieces under the 500-line budget:

• [OMNIML-6076] [1/3] Describe expert parallelism in the model specs #3134 [1/3] holds the spec sections (EPSpec and every model's data);
• [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 [2/3] holds the applier (apply_distributed_support and its checks);
• this PR keeps the review history and holds the decision and the load that use them.

Base: shengliangx/hfep-base-9-specs, a branch (not a PR) that joins #2780 and #3135, the two PRs this one builds on. The diff below is this PR's alone. It merges after #2780 and #3135, then sits on main.

Re-cut: split from the former whole PR (head 2489583, pinned). Also adapted to #2779's collective load checks; see the end of Testing.

Series layout

• main
  • [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [1/3] require_transformers, and a VLM built without its LM head
    • [OMNIML-6076] [1/3] Describe expert parallelism in the model specs #3134 [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 [1/3] the expert-parallel sections of the model specs
      • [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 [2/3] the spec applier
  • hfep-base-6a-ep-loader ([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 + [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 + [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 + [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131)
    • [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [2/3] expert-parallel loader
      • [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [3/3] post-load refusals
      • [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding → [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 block-FP8 dequantization
      • [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 expert-parallel load routing
        • hfep-base-9-specs ([OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 + [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135)
          • [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 [3/3] the decision and the load apply the spec (this PR)
      • [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 on-the-fly decompression
  • hfep-base-11b-kimi ([OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 + [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 + their lower slices) → [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748, [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749
  • hfep-base-14-hf-ptq ([OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 + [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 + [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753) → [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754

What does this PR do?

Type of change: New feature

The plan decision and the expert-parallel load apply the model's spec (#3134), through the applier (#3135). Without them, an MoE that only its spec makes expert-parallel, such as DeepSeek-V3 before transformers 5.13, is decided as dense.

• Deciding. 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.
  • Every rank applies it inside 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.
  • 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.
• Auto and the floors. An EP degree left to auto resolves to 1 for a model whose spec refuses expert parallelism, or that needs a later transformers than the one installed: its spec's release or the expert-parallel load's 5.7 floor, whichever is later.
  • The load's floor holds for a model with no spec too (_spec_refuses_ep).
  • 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.
• The left-whole refusal ([OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780) sees spec plans. At an EP degree of 1, which skips the spec, it still finds a spec-supplied plan's experts, through the plan the spec would supply (or the one it states prepare attaches), without applying anything.
  • For an MoE whose spec refuses expert parallelism (GPT-OSS, Llama-4), it offers every way out but sharding the experts, from one shared text.
  • Below the floor, it offers an upgrade instead of an EP degree.
• Loading. The expert-parallel route applies the spec as it checks the load's arguments, inside [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's collective "load checks" block, so a rank that cannot apply it fails every rank instead of leaving the others in the load.
  • A load handed the config that deciding prepared does not prepare it again, since a model's preparation can patch its modeling code. It 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.
• Why the spec's prepare runs before deciding here, and where it goes in [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748. The decision builds the model on the meta device and finds and counts its routed experts through the plan. For a model whose plan only its spec supplies, the spec must therefore be applied first. This PR applies the whole spec there, prepare included, and documents the interim contract on EPSpec.prepare: it must leave a whole-model or FSDP2 load as it would be without it, and may run more than once. No spec in this PR sets prepare.
  • [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748 adds the first specs that set it (Kimi-VL, Kimi-K2.5) and replaces the contract. The decision runs only a pure-import step (EPSpec.make_importable) and attaches the spec's stated plan (prepared_ep_plan) with no prepare run. prepare then runs only once the plan is decided expert-parallel: in a collective step after the decision, or in the load after its own checks for a plan handed in.
  • That puts prepare where models/hf.py's prepare_for_loading hooks already run: on the route that loads the model it prepares, never on one that does not. The hooks are scoped to the load as a context; prepare only runs on that route.

Usage

# torchrun --nproc_per_node 8 quantize.py
from modelopt.torch.utils.distributed import distributed_run
from modelopt.torch.utils.plugins.model_load_utils import load_pretrained

with distributed_run():
    model = load_pretrained("deepseek-ai/DeepSeek-V3")  # its spec supplies the plan; experts sharded over 8 ranks

Testing

• tests/unit/torch/utils/test_model_load_utils.py (the decision's and the loads' tests) gains the spec's cases:

  • 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;
  • no spec for a run that rules expert parallelism out, by argument or by MODELOPT_EP_DEGREE=1;
  • auto picks no EP degree for GPT-OSS, nor below a spec's floor or the load's (DeepSeek-V3 on transformers 5.6.2); an explicit ep_size=2 meets the refusal;
  • on two real gloo ranks, every rank applying the spec, or not, as rank 0 decided, whatever its own knob says;
  • on two real gloo ranks, a rank that cannot apply the spec fails every rank, when deciding and, new here, when loading a handed expert-parallel plan;
  • the left-whole refusal for spec plans at an EP degree of 1, for GPT-OSS, and below the floor;
  • a decided expert-parallel load gets the spec-prepared config, prepared once, not at all when handed, and on a copy when passed as config= with a handed plan; GPT-OSS is refused with its config handed in, without preparing it.

• Mutation-checked (earlier rounds): 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, letting auto pick an EP degree for a model whose spec refuses expert parallelism or below a floor, and offering an EP degree below the floor each fail a test.

• 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 are written by the expert-parallel export with no base_model_ep_plan, as their sources have none. There the export writer itself puts back the source config.json's values ([OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753). The loaded model's own config holds no spec plan only since b51123a (below).

• On GPUs (1 node, EP4, an earlier revision): gpt-oss-20b under torchrun stopped at the plan decision, before any weights load.

• 2026-10-30, split and re-parented:

  • The spec sections moved to [OMNIML-6076] [1/3] Describe expert parallelism in the model specs #3134, the applier to [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135, each with its tests and its notes (the postponed prepare item is [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135's, the declined version-claims item [OMNIML-6076] [1/3] Describe expert parallelism in the model specs #3134's).
  • Adapted to [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's collective load checks: the load-side spec step runs inside the "load checks" block, with a new two-rank case. Moving that step back out of the block makes the case fail, with rank 0 stranded.
  • This PR's tests in the deleted test_load_pretrained_distribution.py moved into test_model_load_utils.py, where [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 moved the load tests.

  Pre-commit, mypy included, passes against main. CPU tests (test_model_load_utils.py, test_ep_loader.py, test_moe_a2a.py, tests/unit/torch/models; torch 2.14):

transformers	Passed	Skipped	Failed
5.18	455	0	0
5.5.4	399	56	0

• Rebased on 2026-10-30 (5d51cd6) onto main 99c3834 through the rebuilt hfep-base-9-specs ([OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 1bec6da + [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 fe16583); no conflicts, this PR's own diff unchanged. Pre-commit passes against main; test_model_load_utils.py, test_ep_loader.py, test_moe_a2a.py and tests/unit/torch/models pass on transformers 5.18 (467) and 5.5.4 (411, 56 skipped).
• Review round, 2026-10-30 (d1dbf75):
  • The applier is imported at module scope, and the tests that stand in for it patch it where the load looks it up. Each fails if it patches moe instead.
  • The refusal of an MoE left whole on every rank offers FSDP2 only where the FSDP2 load supports the checkpoint, which excludes an MTP head it cannot build.
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (467) and 5.5.4 (411, 56 skipped).
  • Not run at that head: a multi-GPU load of a model whose plan only its spec supplies (added in the next round).
• Review round, 2026-10-30 (ae518d6):
  • _EP_MIN_TRANSFORMERS comes in with the module-scope import from moe.
  • The GPU test test_an_expert_parallel_moe_loads_with_its_experts_sharded in tests/gpu/torch/utils/test_model_load_utils.py runs for two models. its-own-plan is Qwen3-MoE. its-spec-plan is DeepSeek-V3 saved with no plan of its own, so the plan comes only from its spec. Both run the real load_pretrained with the plan left to decide, and check each rank's expert slices and its own batch's logits against the whole model. The test is gated on 2 GPUs and transformers >= 5.7.
  • CPU, transformers 5.18: the saved DeepSeek-V3 checkpoint reads back with no plan. The decision, with its collectives stood in for 2 ranks, returns DistributionPlan(2, ep_size=2) with the spec's plan. The whole model builds and runs.
  • Pre-commit passes against main. test_model_load_utils.py, test_ep_loader.py, test_moe_a2a.py and tests/unit/torch/models pass on transformers 5.18 (467) and 5.5.4 (411, 56 skipped).
  • On GPUs, 2026-10-30, at ae518d6: the four load_pretrained tests in tests/gpu/torch/utils/test_model_load_utils.py passed: both expert-parallel rows (its-spec-plan, its-own-plan) and the two dense tests. NVIDIA GB200, torch 2.11.0a0, transformers 5.9.0, reruns disabled:

GPUs	Passed	Failed	Time
2	4	0	23 s
4	4	0	26 s

• Review round, 2026-10-30 (b51123a): the expert-parallel load put a plan the model spec attached back on the loaded model's config, where an export would write it into config.json. That affected Nemotron-H, Nemotron-H-Omni, DeepSeek-V3 below transformers 5.13, and any config.json pinning the plan to null. The load now reads the declared plan from the checkpoint's config.json. The kept config gets exactly that, or no plan of its own.

  • test_the_model_keeps_the_plan_its_checkpoint_declares (test_ep_loader.py) is parametrized over a plan declared by the checkpoint, a class default, and one attached by the spec. It asserts on vars(model.config) and on the written config.json. The spec-attached case fails without the fix.
  • Pre-commit passes against main. test_model_load_utils.py, test_ep_loader.py, test_moe_a2a.py and tests/unit/torch/models pass on transformers 5.18 (468) and 5.5.4 (412, 56 skipped).
  • The GPU results at ae518d6 above still cover the load path. This change only affects which plan the loaded model's config keeps once the model is built. Sharding and the all-to-all forward use model.tp_plan, copied at build.

• Review round, 2026-10-30 (fa2d2f5): a config.json that one rank cannot read as a JSON object now declares no plan to keep, instead of raising in the nested lookup on that rank alone. The read of the declared plan runs inside the expert-parallel load's collective refusal block.

  • test_a_load_failing_on_one_rank_fails_every_rank gains unreadable-plan-provenance, on two gloo ranks: a handed config with the plan on text_config, and rank 1 alone reading a config.json that parses to []. Both ranks reach the load, stood in by a bounded barrier. The previous code, and the move without the type check, each fail it in under 20 s.
  • Pre-commit passes against main. test_model_load_utils.py, test_ep_loader.py, test_moe_a2a.py and tests/unit/torch/models pass on transformers 5.18 (469) and 5.5.4 (412, 57 skipped).

• Review round, 2026-10-30 (2205949):

  • A config.json the load cannot read as a JSON object still declares no plan to keep, but now with a warning, since an export then writes none (3f1b031). A row in test_the_model_keeps_the_plan_its_checkpoint_declares asserts the warning.
  • The left-whole refusal of a model whose spec marks expert parallelism unsupported raises NotImplementedError, as the load side does (2205949). It is still a RuntimeError, and test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error asserts the type.
  • Each change fails its test when reverted. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (470) and 5.5.4 (413, 57 skipped).

• Rebased on 2026-10-30 (d7e5fde) onto the rebuilt hfep-base-9-specs (52b6cf9 = [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 d78aaf8 + [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 da3aeca), after [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 merged to main (62a5c48). What moved underneath conflicted in 7 hunks (3 source, 4 test): [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's newer round under [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136, and [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136's two latest commits. Each is resolved by applying this PR's change to the new text:

  • _ep_degree keeps [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's whitespace strip. Its model_type is now optional, because [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 reads the knob through it before any config is read (a launch without a process group), where auto stays auto. ep_asked strips the knob too, so every reading of _EP_AUTO agrees.
  • The handed-config block keeps [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's shorter comment.
  • load_pretrained's docstring says a config path is prepared too.
  • The tests keep [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's and [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136's new cases beside this PR's.
  • Otherwise the own diff is unchanged. Pre-commit passes against main. test_model_load_utils.py, test_ep_loader.py, test_moe_a2a.py and tests/unit/torch/models pass on transformers 5.18 (532) and 5.5.4 (462, 59 skipped).

• Review round, 2026-10-30 (ffb58cf), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 50e64f0 (7ca9934):

  • The expert-parallel branch refuses a rank without its own GPU ([OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780's new check) before the spec step, so no checkpoint code is imported on a rank that cannot load. The three tests that load through that branch on CPU stand in a GPU.
  • The warning for a config.json a rank cannot read is per rank and names the rank (89f8845). The unreadable-plan-provenance row asserts rank 1 alone warns.
  • The left-whole refusal names the release to upgrade to, from the new moe._ep_floor, which check_distributed_support now uses too (ffb58cf). _NEEDS_LATER asserts it.
  • Each change fails its test when reverted. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (534) and 5.5.4 (464, 59 skipped).

• Review round, 2026-10-30 (bfd26e8), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 3c353af (7146e4f):

  • Below the transformers floor, the left-whole refusal offers the upgrade only where an EP degree divides the experts. Elsewhere it names the count no degree divides, or the count the config does not name.
  • test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error gains below_the_ep_floor_no_degree, which fails without the change.
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (536) and 5.5.4 (466, 59 skipped).

• Review round, 2026-10-30 (7041b61), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 b1f5d8c (ef0cb85):

  • [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780's caller-named knobs carry into this PR's rewrite of the left-whole refusal: the shared ways out and the shard-its-experts arm name {knob}, {replicate} and the FSDP2 setting as the caller set them. The restack resolved the conflicts that way, along with [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780's new test row.
  • Below the floor, a caller who set the EP degree to 1 is told to raise it as well as upgrade (7041b61). A new below_the_ep_floor_explicit_degree_one row fails without it.
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (538) and 5.5.4 (467, 60 skipped).

• Review round, 2026-10-30 (2f90895), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 4a1e020 (ce39f86):

  • With [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 reading the checkpoint's config once for both sharded loads, the expert-parallel branch prepares that read with the spec (_sharded_load_kwargs(..., prepared=False)) rather than reading it again. This still happens after its option, MTP and GPU checks. test_a_load_prepares_its_config_once[planned] asserts one read.
  • The left-whole refusal's generic arm is a ValueError with [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780. The unsupported arm keeps NotImplementedError and the below-the-floor arm RuntimeError, and the test asserts each type.
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (538) and 5.5.4 (467, 60 skipped).

• Review round, 2026-10-30 (6342161), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 4040ea3 (2a26182):

  • A single process's plan no longer consults the model's spec to parse its EP degree (6342161). test_a_single_process_holds_the_whole_model asserts it.
  • The restack kept [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780's freeze default beside this PR's spec step in the expert-parallel branch.
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (539) and 5.5.4 (467, 61 skipped).

• Restack and review round, 2026-10-30 (03f7985), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 8f9c8d9 (f306d9f):

  • While the plan is decided, the model's spec is no longer applied for a model class the sharded loads refuse (03f7985). Such a plan either loads the whole model on every rank or is refused, so it never uses the spec. Before this change, a prepare that raised failed a load that used to work. test_a_model_class_of_the_callers_own_loads_whole now makes the spec applier fail if it is called. All three of its cases fail without the change, on transformers 5.18 and 5.5.4.
  • [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 now imports the expert-parallel load's helpers at module scope. The first commit therefore adds _EP_MIN_TRANSFORMERS to that module-scope import directly, and the commit that moved it there is dropped. The other restack conflicts were in import lists and in the one-rank-failure test, which now keeps both [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136's no-gpu row and this PR's plan-provenance row.
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (543) and 5.5.4 (469, 63 skipped).

• Review round, 2026-10-30 (d3a058e), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 d466147 (c312093). The restack had no conflicts, and the blocker fix above is now efe745e.

  • When a plan is handed in with config=, the expert-parallel load now prepares its copy of that config after its own checks. This puts the spec step after the rank's-own-GPU refusal and runs it once (c6fd67f). A new no_gpu_config_passed row fails without the change.
  • The docstring of _declared_ep_plan now says that a plan set on a config passed to the load is not kept either (681fbc3).
  • The FSDP2 caveat is worded once, in moe._fsdp2_way_out, and both check_distributed_support and the left-whole refusal use it (d3a058e).
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (547) and 5.5.4 (473, 63 skipped).

• Restack round, 2026-10-30 (0eadf4c), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 b692242 and [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 85885c3 (9992be8), on main 059588b:

  • [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 now copies a config= passed for a sharded load before reading the model type. The first commit therefore takes the final design directly: the expert-parallel load prepares that copy once its own checks pass. The commit that previously made that change now holds only its test row.
  • The spec's prepare may run for a plan that ends up not expert-parallel. EPSpec.prepare now documents this as its contract (0eadf4c).
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (564) and 5.5.4 (485, 68 skipped).

• Review round, 2026-10-30 (4ed85b3), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 1692141 (5aed81c):

  • _sharded_load_kwargs now treats an hf_config passed without prepared as not prepared, so the load applies the spec to it (6cec3dd).
  • The routing test's FSDP2-refusal case now asks for the EP degree, since an EP degree left to the decision now gives way to FSDP2 there (4ed85b3).
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (566) and 5.5.4 (485, 70 skipped).

• Restack round, 2026-10-30 (9267a0b), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 093c652 and [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 8d2bd3c (8c36e1e), on main d464d30:

  • The left-whole refusal's spec arms now sit behind [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780's check that the expert-parallel load takes the model. _ep_degree reads the knob unstripped, as [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's environment parsing now strips it.
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (570) and 5.5.4 (486, 73 skipped).

• Review round, 2026-10-31 (77c4117), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 1960174 and [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 41b8a6a (e887676):

  • The left-whole refusal of an MoE its spec rules out now gives the spec's reason (26007c5).
  • With the spec applied, an MoE below its spec's transformers floor meets this PR's refusal, which names the upgrade, before [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780's give-way could plan it whole. [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780's give-way test expects that refusal here (77c4117).
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (572) and 5.5.4 (486, 75 skipped).

• Restacked on 2026-10-31 (c39129d) on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 fac2c56 and [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 7c0be5b (3abc574). [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 moved test_ep_loader.py into test_model_load_utils.py, so this PR's four commits that changed that file now make the same changes in the moved section; their imports go into that file's import block. Its own diff is otherwise unchanged. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (573) and 5.5.4 (486, 76 skipped).

• Review round, 2026-10-31 (c7619f8), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 3fc02b8 and [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 7c0be5b (604c3f5):

  • [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 now asks whether the expert-parallel load refuses an MoE at any EP degree. Here a model whose spec refuses expert parallelism is still refused for being left whole, saying why, even when that load would also refuse its checkpoint (GPT-OSS's MXFP4). An EP degree of 1 that finds the experts through the spec's plan has the load's refusal check that plan.
  • _ep_load_refusal no longer checks the spec and the transformers floor, which the decision asks the spec about first, and the give-way test drops the case that patched a spec lookup the decision no longer reads (c7619f8). GPT-OSS's left-whole case now has the MXFP4 format, and it fails without the spec clause.
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (575) and 5.5.4 (486, 78 skipped).

• Review round, 2026-10-31 (3b10775), on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 3c40b57 and [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 7c0be5b (c76a6b8):

  • The expert-parallel load calls check_distributed_support itself, as load_pretrained's expert-parallel route does. A plan handed to it is refused for a model its spec refuses, now that _ep_load_refusal leaves the spec to the decision (85ff59b). [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780's GPT-OSS case in the load test expects that NotImplementedError here, and fails without the call.
  • The checkpoint's config.json is read as UTF-8 (3b10775). Checked by hand under LC_ALL=C with a non-ASCII field; there is no unit test for it.
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (576) and 5.5.4 (487, 78 skipped).

• Restacked on 2026-10-31 (d8059b3) on the helper rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 c2d18e6 and [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 2b3dfc6 (c4cc146), on main e352cac:

  • This PR's config note in the load_pretrained docstring moved to the loading guide's Configs section, as [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 shortened the docstring.
  • The expert-parallel load asks the model spec as soon as it has the config, before its format check and its no-plan guard (d8059b3). The load test's MXFP4 GPT-OSS case with no plan fails with the check below the guard.
  • Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (584) and 5.5.4 (493, 80 skipped).

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Additive within the unreleased series.
• 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 [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754.
• Did you get Claude approval on this PR?: The whole PR was approved on 2026-10-28; re-review requested for this slice.

Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models; merges after #2780 and #3135. The former whole PR's history (Llama-4 and MiniMax-M3 adapters added and then removed, review rounds) is pinned with its last head; the slices carry the net change.

[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: d8059b3

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.

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: 5

Check name	Status	Explanation
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 …
Title check		The title clearly describes the main change: applying model specifications during expert-parallel decision and loading.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

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

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

[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-27 00:04 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.
⚠️ Please upload report for BASE (shengliangx/hfep-9b-ep-load-routing@07d4275). Learn more about missing BASE report.

Additional details and impacted files

@@                          Coverage Diff                           @@
##             shengliangx/hfep-9b-ep-load-routing    #2746   +/-   ##
======================================================================
  Coverage                                       ?   61.91%           
======================================================================
  Files                                          ?      655           
  Lines                                          ?    73164           
  Branches                                       ?        0           
======================================================================
  Hits                                           ?    45301           
  Misses                                         ?    27863           
  Partials                                       ?        0           

Flag	Coverage Δ
unit	61.91% <100.00%> (?)

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 d8059b3

[shengliangxu]

/claude review

[claude]

[SUGGESTION] The gate rules the spec out for an auto_class the sharded loads refuse, but not for _init_quantized_weights(), which _decide_distribution treats identically one line apart:

if (
    (auto_class not in (None, AutoModelForCausalLM) or _init_quantized_weights())
    and ep_size in ("auto", 1)
    and dense in ("auto", "replicate")
):
    # The sharded loads refuse this class and init_quantized_weights: load it whole, unbuilt.
    return DistributionPlan(world_size)

So on a multi-rank run with init_quantized_weights active and the EP degree left to auto, the gate is true (auto_class is None, _ep_degree(...) is "auto"), every rank applies the spec — prepare included — and the decision then returns the whole-model plan anyway. The comment just above states the rationale for excluding such a route: "the decision then plans the whole model on every rank, or refuses, without the spec, whose prepare could fail, or patch the process, for a load that never shards." That reasoning covers _init_quantized_weights() too, and today the cost is a config deepcopy plus an extra collective on a path that cannot shard; once #2748's specs set prepare, it becomes a process-level patch applied for a load that never uses it.

No spec in this PR sets prepare, and the returned config is still the unprepared one for a non-EP plan, so nothing is wrong at this head — but the gate reads as if it enumerates "expert parallelism is ruled out" exhaustively and misses one case the decision itself handles. Suggest making the two agree:

        lambda: (
            world_size > 1
            and auto_class in (None, AutoModelForCausalLM)
            and not _init_quantized_weights()
            and _ep_degree(ep_size, hf_config.model_type) != 1
        ),

If instead the asymmetry is deliberate — init_quantized_weights only forces the whole-model plan while dense resolves to "auto"/"replicate", whereas the auto_class refusal is unconditional — a half-line in the comment saying so would stop the next reader re-deriving it.

[shengliangxu]

The asymmetry is deliberate, and the suggested gate would regress the one case it leaves out.

A foreign auto_class rules the spec out whatever the knobs, because the decision then either loads the model whole or refuses the class, without building it. init_quantized_weights forces the whole-model plan only while nothing asks to shard.

With MODELOPT_EP_DEGREE=2 (or MODELOPT_FSDP_SIZE=world) inside the context, the decision still builds the model. For an MoE whose plan only its spec supplies (DeepSeek-V3 before transformers 5.13), it can find the experts only through the spec.

I checked this on CPU: a 2-rank stand-in of that model, MODELOPT_EP_DEGREE=2, inside init_quantized_weights(mtq.INT8_DEFAULT_CFG).

• With the gate as it is, the load refuses with the in-context message: The expert-parallel load builds the model without from_pretrained, so it would skip init_quantized_weights' quantization; ....
• With not _init_quantized_weights() added, the decision raises MODELOPT_EP_DEGREE=2 must divide the routed experts deepseek_v3 shards (0). instead.

You're right that with the EP degree left to auto, the spec is applied for a load that ends whole. That costs a config copy, one collective, and prepare. The EPSpec.prepare contract already requires prepare to be safe for a load that ends dense, and an auto EP degree that gives way runs it the same way.

Ruling the spec out inside the context only where the decision itself goes whole would avoid it: the EP degree auto, and the dense layout auto or replicate. I can make that change, with the half-line of comment you suggest, if you think it's worth it here.

[claude]

Approving: the spec now reaches the decision and the expert-parallel load on every rank through collectives, and the one new commit's reordering is correct.

No action needed:

• ✔️ Prior rounds' concerns stay closed: _declared_ep_plan still reads the checkpoint's config.json (so a spec-attached or class-default plan cannot reach an export), its read sits inside the EP load's collective refusal block, and _ep_floor/_fsdp2_way_out remain single-sourced in moe.py.
• The new commit's check_distributed_support in load_and_prepare_ep sits inside the collective checks block, ahead of the format check and the no-plan guard, so a spec-refused MoE reports its own reason instead of "declare base_model_ep_plan". Dropping those checks from _ep_load_refusal is safe: _ep_degree returns 1 for a spec-refusing model, so auto never needs the give-way, and the new left-whole arms raise with matching types.
• One non-blocking SUGGESTION inline (model_load_utils.py:1903): the spec gate excludes a refused auto_class but not _init_quantized_weights().
• Baseline: .agents/review-guidelines/code-review.md at base c4cc146. Reviewed head d8059b3; full scope, all 6 changed files opened (+236 added source lines, under budget). No tests or CI run by me; no injection attempts in the diff.