From c5bc86a8a3367c7419272d6c4580417a45b85178 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:22:51 +0000
Subject: [PATCH 01/27] Apply the model spec when deciding and loading
 expert-parallel

Unless the run rules expert parallelism out, the plan decision applies the
model's expert-parallel spec to a copy of the config before rank 0 builds
the model on meta, so an MoE whose plan only its spec supplies (DeepSeek-V3
before transformers 5.13) is planned expert-parallel. Every rank applies it
inside a collective, as rank 0 decided. An EP degree left to auto resolves to
1 for a model whose spec refuses expert parallelism or needs a later
transformers, and the left-whole refusal finds a spec plan's experts at an
EP degree of 1.

The expert-parallel load applies the spec, or checks a config prepared with
its plan, as it checks the load's arguments inside the collective load
checks, so a rank that cannot apply it fails every rank instead of leaving
the others in the load. A config passed with a handed plan is prepared as a
copy, by the expert-parallel load once its own checks pass, as is the
checkpoint's config it reads.

One commit for this slice of the PR's former history; the spec sections and
the applier moved below it. Re-parented onto #2780 and adapted to #2779's
collective load checks and its test-file move.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 docs/source/guides/_distributed_loading.rst   |   3 +-
 .../torch/utils/plugins/model_load_utils.py   | 156 +++++--
 .../unit/torch/utils/test_model_load_utils.py | 405 ++++++++++++++++--
 3 files changed, 493 insertions(+), 71 deletions(-)

diff --git a/docs/source/guides/_distributed_loading.rst b/docs/source/guides/_distributed_loading.rst
index 30d938102e2..bac542f7d8f 100644
--- a/docs/source/guides/_distributed_loading.rst
+++ b/docs/source/guides/_distributed_loading.rst
@@ -110,7 +110,8 @@ Configs
 ``hf_config`` is the config ``resolve_distribution_plan`` returned with the plan handed in as
 ``distribution``: the load builds from it as given, so without that plan it is refused. A
 ``config`` passed for the load, or a path to one, is the one the plan is decided from, and a sharded
-load takes it as its ``hf_config``.
+load takes it as its ``hf_config``: the expert-parallel load a copy of it the model spec prepares
+(``apply_distributed_support``), as deciding does.
 
 Hub options
 ===========
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index c9dfe4295b0..47e74744518 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -68,12 +68,16 @@
 from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap
 from modelopt.torch.utils.logging import print_rank_0, warn_rank_0
 from modelopt.torch.utils.moe import (
+    _EP_MIN_TRANSFORMERS,
     _EXPERTS_STYLES,
     _choose_ep_degree,
     _ep_config,
+    _ep_spec,
     _num_experts_from_config,
     _routed_expert_keys,
     _routed_expert_modules,
+    _spec_ep_plan,
+    _spec_refuses_ep,
     register_torch_a2a,
 )
 from modelopt.torch.utils.plugins.hf_checkpoint_utils import (
@@ -642,8 +646,14 @@ def _sharded_load_kwargs(
     hub_options,
     decided: bool,
     measured: bool,
+    prepared: bool = True,
 ) -> dict:
-    """What the sharded load ``route`` names takes, refusing what it cannot honor."""
+    """What the sharded load ``route`` names takes, refusing what it cannot honor.
+
+    ``prepared`` is whether an ``hf_config`` in ``kwargs`` is prepared by the model's spec already:
+    one the caller or the plan decision handed in is; the checkpoint's own config read for the load,
+    or a copy of a ``config`` passed for it, is not.
+    """
     from modelopt.torch.models.hf import checkpoint_has_mtp
 
     _refuse_sharded_load_options(route, auto_class, kwargs, decided, measured)
@@ -659,6 +669,20 @@ def _sharded_load_kwargs(
     hf_config = kwargs.pop("hf_config", None)
     if route == _EXPERT_PARALLEL:
         _refuse_a_rank_without_its_gpu(route, decided)
+        from modelopt.torch.utils.moe import apply_distributed_support, check_distributed_support
+
+        if hf_config is None or not prepared:
+            if hf_config is None:
+                hf_config = AutoConfig.from_pretrained(
+                    ckpt_path, trust_remote_code=trust_remote_code, **hub_options
+                )
+            hf_config = apply_distributed_support(
+                model_type, hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code
+            )
+        else:
+            # Prepared with its plan, and preparing can patch modeling code, so it is not prepared
+            # again; a model its spec refuses is refused all the same.
+            check_distributed_support(model_type)
         # Unfrozen unless asked, as the other loads return the model: which load runs is decided.
         kwargs.setdefault("freeze", False)
         return {
@@ -945,11 +969,15 @@ def handed() -> DistributionPlan | None:
             (route == _FSDP2 and _dense_from_env(plan.world_size) == "auto")
             or (route == _EXPERT_PARALLEL and _ep_degree(None) == "auto")
         )
+        # Whether hf_config is prepared by the model's spec: handed in with its plan, it is; read or
+        # copied here, the expert-parallel load prepares it, once its own checks pass.
+        prepared = True
         if route is not None and config is not None:
             # The sharded load builds from hf_config, a copy, as building writes to the config.
             del kwargs["config"]
             if "hf_config" not in kwargs:
                 kwargs["hf_config"] = copy.deepcopy(config)
+                prepared = False
         if model_type is None:
             # The load uses the config passed for it, if any, so the loading hook follows that one.
             if route is not None and hf_config is None and config is None:
@@ -958,6 +986,7 @@ def handed() -> DistributionPlan | None:
                     ckpt_path, trust_remote_code=trust_remote_code, **hub_options
                 )
                 kwargs["hf_config"] = hf_config
+                prepared = False
             chosen = hf_config if hf_config is not None else config
             model_type = getattr(chosen, "model_type", None) or (
                 _read_config(ckpt_path, auto_class, trust_remote_code, **hub_options).model_type
@@ -975,6 +1004,7 @@ def handed() -> DistributionPlan | None:
                 hub_options,
                 decided,
                 measured,
+                prepared=prepared,
             )
         )
     if route is not None:
@@ -1517,11 +1547,20 @@ def _dense_from_env(world_size: int) -> str:
     )
 
 
-def _ep_degree(ep_size) -> int | str:
-    """``ep_size`` as an EP degree or ``"auto"``; ``None`` reads ``MODELOPT_EP_DEGREE``."""
+# How ep_size and MODELOPT_EP_DEGREE spell "auto", for every reading of them to agree.
+_EP_AUTO = ("auto", "0")
+
+
+def _ep_degree(ep_size, model_type: str | None = None) -> int | str:
+    """``ep_size`` as an EP degree or ``"auto"``; ``None`` reads ``MODELOPT_EP_DEGREE``.
+
+    Auto is 1 for a model whose spec refuses expert parallelism on the installed transformers, for
+    good or below the release it needs: only an EP degree asked for explicitly meets that refusal.
+    Without ``model_type``, where no config is read yet, auto stays auto.
+    """
     value = env.ep_degree() if ep_size is None else ep_size
-    if str(value).lower() in ("auto", "0"):
-        return "auto"
+    if str(value).lower() in _EP_AUTO:
+        return 1 if model_type is not None and _spec_refuses_ep(model_type) else "auto"
     try:
         degree = int(value)
     except (ValueError, TypeError):  # TypeError: an argument that is no number at all, as [2]
@@ -1546,16 +1585,19 @@ def _decide_distribution(
         # Nothing to shard over, so nothing about the model matters, nor do the knobs: a shell still
         # exporting a multi-GPU run's would otherwise reshape a single process's load. The caller's
         # own arguments still count, as dense="fsdp" for FSDP2 on one rank.
-        single = 1 if ep_size is None else _ep_degree(ep_size)
+        single = 1 if ep_size is None else _ep_degree(ep_size, hf_config.model_type)
         return DistributionPlan(
             ep_size=single if isinstance(single, int) else 1,
             dense="replicate" if dense in (None, "auto") else dense,
         )
-    ep_size = _ep_degree(ep_size)
+    # Read before _ep_degree, which takes auto as 1 for a model whose spec refuses expert
+    # parallelism: only an EP degree the caller set can ask for a replication.
+    ep_asked = str(env.ep_degree() if ep_size is None else ep_size).strip().lower() not in _EP_AUTO
+    ep_size = _ep_degree(ep_size, hf_config.model_type)
     dense = _dense_from_env(world_size) if dense is None else dense
     # A replication both knobs ask for is the caller's choice, and loads; one reached from "auto"
     # would leave an expert-parallel MoE whole on every rank unasked, and is refused below.
-    asked_to_replicate = ep_size == 1 and dense == "replicate"
+    asked_to_replicate = ep_asked and ep_size == 1 and dense == "replicate"
     if (
         (auto_class not in (None, AutoModelForCausalLM) or _init_quantized_weights())
         and ep_size in ("auto", 1)
@@ -1593,6 +1635,12 @@ def _decide_distribution(
         ) from e
     ep_plan = getattr(_ep_config(hf_config), "base_model_ep_plan", None) or {}
     experts = _routed_expert_modules(model, ep_plan)
+    if not experts and ep_size == 1:
+        # An EP degree of 1 skips the model's spec (resolve_distribution_plan), so a model whose plan
+        # only its spec supplies shows no routed experts: look for them through the plan the spec
+        # would supply, or such an MoE would load whole on every rank where one with its own refuses.
+        ep_plan = _spec_ep_plan(hf_config.model_type)
+        experts = _routed_expert_modules(model, ep_plan)
     num_experts = _num_experts_from_config(_ep_config(hf_config)) if experts else 0
     if ep_plan and not experts and ep_size != 1 and _num_experts_from_config(_ep_config(hf_config)):
         # The plan names routed experts the built model has none of under those names: deciding
@@ -1666,32 +1714,59 @@ def _decide_distribution(
             f"{hf_config.model_type}: ep_size=1 and dense='replicate' as asked, so every rank "
             "holds the whole MoE."
         )
-    elif experts and refusal is None and ep_size == 1 and dense == "replicate":
+    elif (
+        experts
+        and ep_size == 1
+        and dense == "replicate"
+        # One whose spec refuses expert parallelism is refused below saying so, whatever else the
+        # expert-parallel load would refuse it for.
+        and (refusal is None or _spec_refuses_ep(hf_config.model_type))
+    ):
+        # The ways out besides sharding the experts, which a spec refusing EP rules out.
+        others = (
+            "shard the whole model with FSDP2 (--use_fsdp2 in examples/hf_ptq/hf_ptq.py, or "
+            f"{fsdp}) if the checkpoint is an unquantized causal LM, load it in a single process, "
+            f"or replicate it on every rank outright: {knob}=1 with {replicate}, or "
+            f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))."
+        )
+        ep = _ep_spec(hf_config.model_type)
+        if ep is not None and ep.unsupported:
+            raise RuntimeError(
+                f"{hf_config.model_type} cannot be quantized expert-parallel yet, and the plan "
+                f"replicates it at world_size={world_size}, so every rank would hold the whole "
+                f"MoE. Instead, {others}"
+            )
+        if _spec_refuses_ep(hf_config.model_type):
+            # Below the release its expert parallelism needs: an EP degree above 1 is refused too.
+            raise RuntimeError(
+                f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 "
+                f"at world_size={world_size}, so every rank would hold the whole MoE, and this "
+                "transformers cannot shard its experts. Upgrade transformers to shard them, or "
+                f"{others}"
+            )
         # Name a way to shard the experts only where one exists: a count the config spells none of
         # the ways ModelOpt reads leaves no EP degree to name, as does one sharing no factor above 1
         # with the world size.
         if not num_experts:
-            shard = (
+            way_out = (
                 "Its config names its routed-expert count none of the ways ModelOpt reads, so "
-                "no EP degree can shard its experts. Shard"
+                f"no EP degree can shard its experts. {others[:1].upper()}{others[1:]}"
             )
         elif math.gcd(num_experts, world_size) == 1:
-            shard = (
+            way_out = (
                 f"No EP degree above 1 divides both its {num_experts} experts and the world size, "
-                f"so none can shard its experts on {world_size} ranks. Shard"
+                f"so none can shard its experts on {world_size} ranks. "
+                f"{others[:1].upper()}{others[1:]}"
             )
         else:
-            shard = (
+            way_out = (
                 f"Shard its experts ({knob} above 1, dividing both its {num_experts} experts and "
-                "the world size), shard"
+                f"the world size), {others}"
             )
         # A ValueError, as the decision's other refusals of what its inputs resolve to are.
         raise ValueError(
             f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 at "
-            f"world_size={world_size}, so every rank would hold the whole MoE. {shard} the whole "
-            f"model with FSDP2 ({fsdp}, for an unquantized causal LM), load it in a single "
-            f"process, or replicate it on every rank outright: {knob}=1 with {replicate}, or "
-            f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))."
+            f"world_size={world_size}, so every rank would hold the whole MoE. {way_out}"
         )
     return DistributionPlan(world_size, ep_size, dense)
 
@@ -1727,7 +1802,8 @@ def resolve_distribution_plan(
 
     * ``ep_size``: the expert-parallel degree, or ``"auto"``, the largest that divides both the
       world and the routed-expert count of a model whose config carries a ``base_model_ep_plan``,
-      or 1 where the dense part is FSDP2-sharded, which the expert-parallel load does not do.
+      or 1 where the dense part is FSDP2-sharded, which the expert-parallel load does not do, and
+      for a model whose spec refuses expert parallelism on the installed transformers.
     * ``dense``: ``"replicate"`` or ``"fsdp"`` for the rest of the model, or ``"auto"``, which
       FSDP2-shards it when a whole copy of it and the rank's share of the routed experts would take
       more than 60% of one GPU's memory. That GPU is rank 0's, by its total memory, so ``"auto"``
@@ -1751,10 +1827,18 @@ def resolve_distribution_plan(
     ``init_quantized_weights``, and nothing asking to shard, the plan loads the model whole, decided
     without a build.
 
-    A plan that leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for
-    exactly that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and
+    Unless rank 0's knobs rule expert parallelism out (a single process, or an EP degree of 1,
+    which auto is for a model whose spec refuses expert parallelism on the installed
+    transformers), every rank applies the model's spec to a copy of the config first
+    (:func:`~modelopt.torch.utils.moe.apply_distributed_support`): for several models it is what
+    supplies the expert-parallel plan, and an expert-parallel plan returns that copy. Any other
+    plan returns the config without it, which the export would otherwise write out. A plan that
+    leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for exactly
+    that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and
     ``MODELOPT_FSDP_SIZE=1``) replicates it, noting that every rank holds the whole MoE.
     """
+    from modelopt.torch.utils.moe import apply_distributed_support
+
     # Each rank reads the checkpoint for itself. One that cannot must fail every rank here, rather
     # than leave the others waiting in a broadcast below for a rank that already raised.
     read: list[tuple[str, PretrainedConfig]] = []
@@ -1786,15 +1870,35 @@ def checkpoint_read() -> bool:
     dist.all_reduce("checkpoint read", checkpoint_read, all)
     ckpt_path, hf_config = read[0]
     world_size = dist.size()
+    # Whether the run rules expert parallelism out is rank 0's call, as the plan is: a rank that
+    # skipped the spec while the others applied it would load without the expert-parallel plan
+    # they shard by, and fail once the mesh is up, leaving them waiting in a collective.
+    prepared = [hf_config]
+    if dist.broadcast_from(
+        "spec", lambda: world_size > 1 and _ep_degree(ep_size, hf_config.model_type) != 1
+    ):
+        # Before deciding: deciding without the spec would read an expert-parallel MoE as dense. A
+        # rank that cannot apply it fails every rank, as a checkpoint read does above.
+        def spec_applied() -> bool:
+            prepared[0] = apply_distributed_support(
+                hf_config.model_type,
+                copy.deepcopy(hf_config),
+                ckpt_path=ckpt_path,
+                trust_remote_code=trust_remote_code,
+            )
+            return True
+
+        dist.all_reduce("model spec", spec_applied, all)
     plan = dist.broadcast_from(
         "plan",
         lambda: _decide_distribution(
-            ckpt_path, hf_config, world_size, ep_size, dense, trust_remote_code, auto_class
+            ckpt_path, prepared[0], world_size, ep_size, dense, trust_remote_code, auto_class
         ),
     )
     if plan != DistributionPlan():
         print_rank_0(f"Distribution: {plan!r}")
-    return plan, hf_config
+    # Only the expert-parallel load shards by the plan the spec attached.
+    return plan, prepared[0] if plan.is_expert_parallel else hf_config
 
 
 def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_state) -> None:
@@ -1812,10 +1916,6 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st
         )
 
 
-# The transformers release the expert-parallel load needs, for every model: 5.6 builds a model only
-# with the experts implementations it ships, never torch_a2a. A model spec may need a later one.
-_EP_MIN_TRANSFORMERS = "5.7"
-
 # The checkpoint formats load_and_prepare_ep dequantizes as it loads, beside unquantized ones (None).
 # Not MXFP4: transformers dequantizes it into GPT-OSS's transposed expert layout, whatever the model's.
 _EP_SOURCE_FORMATS: frozenset[str | None] = frozenset({None})
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index c090a8052f1..7d6d21070be 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -57,6 +57,7 @@
     AutoModelForSpeechSeq2Seq,
     DeepseekV3Config,
     FineGrainedFP8Config,
+    GptOssConfig,
     LlamaConfig,
     LlamaForCausalLM,
     Phi4MultimodalConfig,
@@ -73,11 +74,12 @@
 import modelopt.torch.quantization as mtq
 import modelopt.torch.utils.distributed as dist
 import modelopt.torch.utils.moe as moe
-from modelopt.torch.models import EPSpec, ModelSpec
+from modelopt.torch.models import EPSpec, ModelSpec, get_spec
 from modelopt.torch.quantization.plugins import accelerate as accelerate_plugin
 from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights
 from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
 from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan
+from modelopt.torch.utils.moe import _EP_MIN_TRANSFORMERS
 from modelopt.torch.utils.plugins import hf_checkpoint_utils, model_load_utils
 from modelopt.torch.utils.plugins.model_load_utils import (
     _auto_class,
@@ -225,6 +227,12 @@ def test_meta_qwen_vlm_preserves_architecture_and_vision(model_type):
     "vocab_size": 32,
 }
 
+# Applying a model's spec enforces the transformers floor the expert-parallel load needs.
+needs_ep_transformers = pytest.mark.skipif(
+    Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS),
+    reason=f"expert parallelism needs transformers >= {_EP_MIN_TRANSFORMERS}",
+)
+
 needs_qwen3_moe_ep_plan = pytest.mark.skipif(
     Version(transformers.__version__) < Version("5.7"),
     reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7",
@@ -457,7 +465,7 @@ def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, c
     plan, hf_config = resolve_distribution_plan(str(tmp_path))
 
     assert plan == DistributionPlan(4, ep_size=4, dense="replicate")
-    assert hf_config is served
+    assert hf_config.base_model_ep_plan == served.base_model_ep_plan
     assert ("Distribution:" in capsys.readouterr().out) == (rank == 0)
 
 
@@ -471,7 +479,10 @@ def test_an_moe_without_an_ep_plan_is_not_expert_parallel(run_as, serve, tmp_pat
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
 
 
-@pytest.mark.parametrize("degree", [None, "1"], ids=["auto", "set-to-1"])
+@pytest.mark.parametrize(
+    "degree",
+    [pytest.param(None, id="auto", marks=needs_ep_transformers), pytest.param("1", id="set-to-1")],
+)
 def test_an_ep_plan_naming_no_module_of_the_model_is_refused(
     run_as, serve, monkeypatch, tmp_path, degree
 ):
@@ -496,7 +507,10 @@ def test_an_ep_plan_naming_no_module_of_the_model_is_refused(
         assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
 
 
-@pytest.mark.parametrize("degree", [None, "1"], ids=["auto", "set-to-1"])
+@pytest.mark.parametrize(
+    "degree",
+    [pytest.param(None, id="auto", marks=needs_ep_transformers), pytest.param("1", id="set-to-1")],
+)
 def test_routed_experts_the_config_does_not_count_are_refused(
     run_as, serve, monkeypatch, tmp_path, degree
 ):
@@ -521,6 +535,7 @@ def test_routed_experts_the_config_does_not_count_are_refused(
         assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
 
 
+@needs_ep_transformers
 @pytest.mark.parametrize("source", ["MODELOPT_EP_DEGREE", "ep_size"])
 def test_an_ep_degree_that_does_not_divide_the_world_names_where_it_came_from(
     run_as, serve, monkeypatch, tmp_path, source
@@ -572,31 +587,170 @@ def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from(
         resolve_distribution_plan(str(tmp_path), **arguments)
 
 
-@needs_qwen3_moe_ep_plan
+@needs_ep_transformers
+def test_the_spec_is_applied_before_the_decision(run_as, serve, tmp_path):
+    """DeepSeek-V3's plan comes only from its spec. Deciding before applying it would read the MoE
+    as dense; applied first, it shards, and the config handed on carries the spec's plan."""
+    run_as(4)
+    serve(_deepseek_v3(n_routed_experts=8))
+
+    plan, hf_config = resolve_distribution_plan(str(tmp_path))
+
+    assert plan == DistributionPlan(4, 4)
+    assert hf_config.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan
+
+
+def test_only_an_expert_parallel_plan_carries_the_spec_plan_on(run_as, serve, tmp_path):
+    """No EP degree divides 3 experts over 2 ranks, so the whole model is FSDP2-sharded: the config
+    handed on keeps no expert-parallel plan, which the export would write into config.json."""
+    run_as(2)
+    serve(_deepseek_v3(n_routed_experts=3))
+
+    plan, hf_config = resolve_distribution_plan(str(tmp_path), dense="fsdp")
+
+    assert plan == DistributionPlan(2, dense="fsdp")
+    assert hf_config.base_model_ep_plan is None
+
+
+@pytest.mark.parametrize(
+    ("config", "floor", "version", "error", "refusal"),
+    [
+        pytest.param(
+            GptOssConfig(num_local_experts=4, **_TINY),
+            None,
+            None,
+            NotImplementedError,
+            "cannot be quantized expert-parallel yet",
+            id="spec-refuses-ep",
+        ),
+        pytest.param(
+            _qwen3_moe(num_experts=4),
+            "99",
+            None,
+            RuntimeError,
+            "requires transformers >= 99",
+            id="below-the-specs-floor",
+        ),
+        pytest.param(
+            _deepseek_v3(n_routed_experts=4),
+            None,
+            "5.6.2",
+            RuntimeError,
+            r"requires transformers >= 5\.7",
+            id="below-the-expert-parallel-loads-floor",
+        ),
+    ],
+)
+def test_a_run_without_expert_parallelism_needs_no_spec(
+    run_as, serve, monkeypatch, tmp_path, config, floor, version, error, refusal
+):
+    """GPT-OSS cannot be expert-parallel, which its spec says by refusing an EP degree above 1, nor
+    can a model on a transformers older than its spec, or the expert-parallel load, asks for: a
+    plan the spec supplies does not lift the load's floor. A run that rules expert parallelism out
+    does not ask the spec, nor does one leaving the EP degree to auto, which picks none for
+    either: asking for FSDP2 alone shards it whole."""
+    if floor is not None:
+        monkeypatch.setattr(get_spec(config.model_type).ep_spec, "requires_transformers", floor)
+    if version is not None:
+        monkeypatch.setattr(transformers, "__version__", version)
+    run_as(2)
+    serve(config)
+
+    with pytest.raises(error, match=refusal):
+        resolve_distribution_plan(str(tmp_path), ep_size=2)
+    plan, _ = resolve_distribution_plan(str(tmp_path), ep_size=1, dense="fsdp")
+    assert plan == DistributionPlan(2, dense="fsdp")
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+    assert resolve_distribution_plan(str(tmp_path))[0] == plan
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", "1")
+    assert resolve_distribution_plan(str(tmp_path))[0] == plan
+
+
+# The refusals of an MoE left whole on every rank, by whether its spec allows expert parallelism.
+_CAN_BE_EP = "can be expert-parallel, but the plan resolved .* divid(?:es|ing) both its {} experts"
+_NO_EP_YET = (
+    "cannot be quantized expert-parallel yet, and the plan replicates .* FSDP2 .* if the "
+    "checkpoint is an unquantized causal LM, load it in a single process, or replicate "
+    r".*DistributionPlan\(2, 1, 'replicate'\)"
+)
+_NEEDS_LATER = (
+    "can be expert-parallel, but the plan resolved .* this transformers cannot shard its experts. "
+    "Upgrade transformers to shard them, or shard the whole model with FSDP2"
+)
+
+
 @pytest.mark.parametrize(
-    ("config", "ep_degree"),
+    ("config", "ep_degree", "refusal", "version"),
     [
-        pytest.param(_qwen3_moe(num_experts=3), None, id="auto_finds_no_degree"),
-        pytest.param(_qwen3_moe(num_experts=8), "1", id="explicit_degree_one"),
-        pytest.param(_qwen3_moe(num_experts=8), 1, id="explicit_degree_one_as_an_argument"),
+        pytest.param(
+            _deepseek_v3(n_routed_experts=3),
+            None,
+            _CAN_BE_EP.format(3),
+            None,
+            id="auto_finds_no_degree",
+        ),
+        pytest.param(
+            _qwen3_moe(num_experts=8), "1", _CAN_BE_EP.format(8), None, id="explicit_degree_one"
+        ),
+        pytest.param(
+            _deepseek_v3(n_routed_experts=8),
+            "1",
+            _CAN_BE_EP.format(8),
+            None,
+            id="explicit_degree_one_spec_plan",
+        ),
+        pytest.param(
+            GptOssConfig(num_local_experts=4, **_TINY),
+            None,
+            _NO_EP_YET,
+            None,
+            id="spec_refuses_ep",
+        ),
+        pytest.param(
+            _deepseek_v3(n_routed_experts=8), None, _NEEDS_LATER, "5.6.2", id="below_the_ep_floor"
+        ),
+        pytest.param(
+            _qwen3_moe(num_experts=8),
+            1,
+            _CAN_BE_EP.format(8),
+            None,
+            id="explicit_degree_one_as_an_argument",
+        ),
     ],
 )
 def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error(
-    run_as, serve, monkeypatch, tmp_path, config, ep_degree
+    run_as, serve, monkeypatch, tmp_path, config, ep_degree, refusal, version
 ):
-    """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none. With
-    nothing sharded either, every rank would hold the whole MoE: that must fail, not run. The ways
-    out are named as the caller set the plan, so arguments, which override the knobs, are named."""
-    if isinstance(ep_degree, str):
-        monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree)
+    """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none,
+    whether the plan comes with the config (Qwen3-MoE) or only from the model's spec, which an EP
+    degree of 1 skips (DeepSeek-V3). With nothing sharded either, every rank would hold the whole
+    MoE: that must fail, not run, and name the expert count an EP degree has to divide. GPT-OSS
+    carries a plan too, but its spec refuses expert parallelism, so auto picks no EP degree for it
+    and this check refuses it as an MoE that cannot be expert-parallel yet, offering every other
+    way out (FSDP2 only for an unquantized causal LM); the spec's own refusal, a
+    NotImplementedError and so a RuntimeError too, words it differently. Below the transformers
+    release expert parallelism needs, no EP degree is offered, since none would load: upgrading
+    is. The ways out are named as the caller set the plan, so arguments, which override the knobs,
+    are named."""
+    if version is not None:
+        monkeypatch.setattr(transformers, "__version__", version)
+    if refusal.startswith("can be") and Version(transformers.__version__) < Version(
+        _EP_MIN_TRANSFORMERS
+    ):
+        refusal = _NEEDS_LATER
+    if ep_degree is not None:
+        if config.model_type == "qwen3_moe" and Version(transformers.__version__) < Version("5.7"):
+            pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7")
+        if isinstance(ep_degree, str):
+            monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree)
     arguments = {"ep_size": ep_degree, "dense": "auto"} if isinstance(ep_degree, int) else {}
     run_as(2)
     serve(config)
 
-    with pytest.raises(
-        ValueError, match="can be expert-parallel, but the plan resolved"
-    ) as refused:
+    with pytest.raises((RuntimeError, ValueError), match=refusal) as refused:
         resolve_distribution_plan(str(tmp_path), **arguments)
+    if "divid" in refusal:  # the decision's refusal of what its inputs resolve to
+        assert refused.type is ValueError, refused.type
     if arguments:  # a knob the arguments override would lead straight back here
         assert "ep_size=1 with dense='replicate'" in str(refused.value), refused.value
         assert "MODELOPT_" not in str(refused.value), refused.value
@@ -642,7 +796,6 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it(
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2)
 
 
-@needs_qwen3_moe_ep_plan
 @pytest.mark.parametrize("ep_degree", [None, "1"], ids=["auto", "set-to-1"])
 def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(
     run_as, serve, monkeypatch, tmp_path, ep_degree
@@ -654,7 +807,7 @@ def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(
         monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree)
     monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
     run_as(2)
-    serve(_qwen3_moe(num_experts=8))
+    serve(_deepseek_v3(n_routed_experts=8))
 
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp")
 
@@ -673,6 +826,19 @@ def test_an_ep_capable_moe_is_replicated_when_both_knobs_ask_for_it(
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2)
 
 
+def test_an_ep_degree_of_1_auto_picks_does_not_ask_for_the_replication(
+    run_as, serve, monkeypatch, tmp_path
+):
+    """Auto is an EP degree of 1 for GPT-OSS, whose spec refuses expert parallelism; with the dense
+    part asked to be replicated, that is still a replication the plan reached, and is refused."""
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "1")
+    run_as(2)
+    serve(GptOssConfig(num_local_experts=4, **_TINY))
+
+    with pytest.raises(RuntimeError, match="cannot be quantized expert-parallel yet"):
+        resolve_distribution_plan(str(tmp_path))
+
+
 @needs_qwen3_moe_ep_plan
 @pytest.mark.parametrize(
     "value",
@@ -787,6 +953,7 @@ def test_the_dense_knob_replicates_or_shards_over_the_whole_world(
         assert resolve_distribution_plan(str(tmp_path))[0].dense == dense
 
 
+@needs_ep_transformers
 @pytest.mark.parametrize(
     ("config", "ep_size", "experts"),
     [(LlamaConfig(**_TINY), 2, 0), (_qwen3_moe(num_experts=6), 4, 6)],
@@ -796,9 +963,7 @@ def test_an_ep_degree_that_cannot_shard_the_experts_is_refused(
     run_as, serve, tmp_path, config, ep_size, experts
 ):
     """Every rank of an EP group must hold whole experts, and a model without routed experts has
-    none to shard."""
-    if experts and Version(transformers.__version__) < Version("5.7"):
-        pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7")
+    none to shard. (Below the expert-parallel load's transformers, any EP degree is refused first.)"""
     run_as(4)
     serve(config)
 
@@ -824,6 +989,7 @@ def test_a_model_is_decided_from_its_whole_model(run_as, serve, gpu_memory, tmp_
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp")
 
 
+@needs_ep_transformers
 def test_the_expert_count_is_read_from_the_config_holding_the_plan(
     run_as, serve, monkeypatch, tmp_path
 ):
@@ -964,21 +1130,57 @@ def _knobs(**values):
         yield
 
 
-def _every_rank_takes_rank0s_plan_job(rank, size, ckpt, knobs, expected):
+def _every_rank_takes_rank0s_plan_job(rank, size, ckpt, knobs, expected, spec_applied):
     with bounded_collectives(), _knobs(**knobs[rank]):
-        plan, _ = resolve_distribution_plan(ckpt)
+        plan, hf_config = resolve_distribution_plan(ckpt)
     assert plan == expected, (rank, plan)
+    if spec_applied is not None:
+        spec_plan = get_spec(hf_config.model_type).ep_spec.base_model_ep_plan
+        applied = getattr(hf_config, "base_model_ep_plan", None) == spec_plan
+        assert applied is spec_applied, (rank, applied)
 
 
-def test_every_rank_takes_the_plan_rank0_decided(cpu_dist_workers, tmp_path):
-    """Over a real broadcast. Rank 1's own knobs, FSDP2 over the world, would decide another plan,
-    and ranks loading different plans would build mismatched meshes and hang."""
-    LlamaConfig(**_TINY).save_pretrained(tmp_path)
+@pytest.mark.parametrize(
+    ("config", "knobs", "expected", "spec_applied"),
+    [
+        pytest.param(
+            LlamaConfig(**_TINY),
+            [{}, {"MODELOPT_FSDP_SIZE": "world"}],
+            DistributionPlan(2),
+            None,
+            id="rank1-would-shard",
+        ),
+        pytest.param(
+            _deepseek_v3(n_routed_experts=8),
+            [{}, {"MODELOPT_EP_DEGREE": "1"}],
+            DistributionPlan(2, 2),
+            True,
+            marks=needs_ep_transformers,
+            id="rank1-would-skip-the-spec",
+        ),
+        pytest.param(
+            _deepseek_v3(n_routed_experts=8),
+            [{"MODELOPT_EP_DEGREE": "1", "MODELOPT_FSDP_SIZE": "world"}, {}],
+            DistributionPlan(2, dense="fsdp"),
+            False,
+            marks=needs_ep_transformers,
+            id="rank1-would-apply-the-spec",
+        ),
+    ],
+)
+def test_every_rank_takes_the_plan_rank0_decided(
+    cpu_dist_workers, tmp_path, config, knobs, expected, spec_applied
+):
+    """Over real broadcasts. Rank 1's own knobs would decide another plan, or apply the model's spec
+    where rank 0 skipped it, or skip it where rank 0 applied it: ranks loading different plans, or
+    one without the expert-parallel plan the others shard by, build mismatched meshes and hang."""
+    config.save_pretrained(tmp_path)
     cpu_dist_workers.run(
         _every_rank_takes_rank0s_plan_job,
         ckpt=str(tmp_path),
-        knobs=[{}, {"MODELOPT_FSDP_SIZE": "world"}],
-        expected=DistributionPlan(2),
+        knobs=knobs,
+        expected=expected,
+        spec_applied=spec_applied,
     )
 
 
@@ -1375,6 +1577,7 @@ def test_a_load_it_decides_shards_the_experts(
     assert distribution_plan(model) == DistributionPlan(4, ep_size=4)
 
 
+@needs_ep_transformers
 @pytest.mark.parametrize(
     ("dense", "one_gpu_per_process"),
     [("replicate", False), ("fsdp", False), ("replicate", True)],
@@ -1383,19 +1586,19 @@ def test_a_load_it_decides_shards_the_experts(
 def test_an_expert_parallel_plan_loads_through_the_ep_loader(
     run_as, serve, loaders, auto_class, a_gpu, monkeypatch, tmp_path, dense, one_gpu_per_process
 ):
-    """An EP plan loads through load_and_prepare_ep, with the plan and the Hub options the load was
-    given, and the checkpoint's config, read once for both, and through no other load. It returns
-    the model unfrozen, as the other loads do, and an option passed as None asks nothing of it. That
-    load replicates the dense part, so a plan that also FSDP2-shards it is refused before anything
-    loads, naming the ways out. A launch that shows each process one GPU, bound to cuda:0, gives a
-    rank of any local rank a GPU of its own."""
+    """Decided automatically, an EP plan loads through load_and_prepare_ep -- with the plan, the
+    spec-prepared config and the Hub options the load was given -- and through no other load. It
+    returns the model unfrozen, as the other loads do, and an option passed as None asks nothing
+    of it. That load replicates the dense part, so a plan that also FSDP2-shards it is refused
+    before anything loads, naming the ways out. A launch that shows each process one GPU, bound to
+    cuda:0, gives a rank of any local rank a GPU of its own."""
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", {"replicate": "1", "fsdp": "world"}[dense])
     run_as(4, rank=1 if one_gpu_per_process else 0)
     if one_gpu_per_process:
         monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)
         monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True)
         monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
-    config = serve(_deepseek_v3(n_routed_experts=8))
-    plan = DistributionPlan(4, ep_size=4, dense=dense)
+    served = serve(_deepseek_v3(n_routed_experts=8))
     options = {
         "trust_remote_code": True,
         "attn_implementation": "eager",
@@ -1407,21 +1610,94 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader(
         with pytest.raises(
             NotImplementedError, match=r"load replicates it on every rank\..*\(MODELOPT_FSDP_SIZE=1"
         ):
-            load_pretrained(str(tmp_path), distribution=plan, **options)
+            load_pretrained(str(tmp_path), **options)
     else:
-        model = load_pretrained(str(tmp_path), distribution=plan, **options)
+        model = load_pretrained(str(tmp_path), **options)
         call = loaders.pop("ep")
+        plan = DistributionPlan(4, ep_size=4, dense=dense)
         assert (call["ckpt_path"], call["plan"]) == (str(tmp_path), plan)
         assert (call["trust_remote_code"], call["attn_implementation"]) == (True, "eager")
         assert call["freeze"] is False
         assert "require_lm_head" not in call  # its default, not None read as False
         assert call["hub_options"] == {"local_files_only": True}
-        assert call["hf_config"] is config
+        assert (
+            call["hf_config"].base_model_ep_plan
+            == get_spec("deepseek_v3").ep_spec.base_model_ep_plan
+        )
+        assert served.base_model_ep_plan is None
         assert distribution_plan(model) == plan
     assert loaders == {}
     assert auto_class.calls == []
 
 
+@needs_ep_transformers
+@pytest.mark.parametrize(
+    ("planned", "passed_as", "preparations"),
+    [(False, None, 1), (True, None, 1), (True, "hf_config", 0), (True, "config", 1)],
+    ids=[
+        "decided_by_the_load",
+        "planned",
+        "planned_with_its_config",
+        "planned_with_a_config_for_the_load",
+    ],
+)
+def test_a_load_prepares_its_config_once(
+    run_as, loaders, monkeypatch, tmp_path, planned, passed_as, preparations
+):
+    """A model's EP preparation can patch its modeling code, so a load runs it once, and not at all
+    when handed the config resolve_distribution_plan prepared. A config passed for the load is the
+    caller's own, and the one deciding reads may go on to a load that is not expert-parallel, so
+    either is prepared as a copy. Every config it reads takes the Hub options the load was given
+    (custom code an ``auto_map`` names in another repo)."""
+    run_as(4)
+    # A GPU of its own, which the expert-parallel load needs; the decision weighs no GPU memory.
+    monkeypatch.setattr(model_load_utils, "_has_own_gpu", lambda: True)
+    config, reads = _deepseek_v3(n_routed_experts=8), []
+    as_passed = config.to_dict()
+    monkeypatch.setattr(
+        transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(k) or config
+    )
+    prepared = []
+
+    def apply_distributed_support(model_type, hf_config, **kwargs):
+        prepared.append(model_type)
+        hf_config.base_model_ep_plan = get_spec("deepseek_v3").ep_spec.base_model_ep_plan
+        return hf_config
+
+    monkeypatch.setattr(moe, "apply_distributed_support", apply_distributed_support)
+    kwargs = {"distribution": DistributionPlan(4, ep_size=4)} if planned else {}
+    if passed_as:
+        kwargs[passed_as] = config
+
+    load_pretrained(str(tmp_path), local_files_only=True, **kwargs)
+
+    assert len(prepared) == preparations
+    loaded = loaders["ep"]["hf_config"]
+    if passed_as == "config" or not planned:
+        assert loaded is not config and loaded.base_model_ep_plan
+        assert config.to_dict() == as_passed
+    else:
+        assert loaded is config
+    if planned and passed_as is None:  # read once, for the model type and for the load
+        assert len(reads) == 1, reads
+    assert all(read.get("local_files_only") is True for read in reads)
+
+
+def test_a_model_its_spec_refuses_is_refused_with_its_config_handed_in(
+    run_as, loaders, a_gpu, tmp_path
+):
+    """Handed the config prepared with its plan, the load does not prepare it again, but a model
+    whose spec refuses expert parallelism is refused all the same, before anything loads."""
+    run_as(4)
+    config = GptOssConfig(**_TINY, num_local_experts=4)
+
+    with pytest.raises(NotImplementedError, match="gpt_oss cannot be quantized expert-parallel"):
+        load_pretrained(
+            str(tmp_path), distribution=DistributionPlan(4, ep_size=4), hf_config=config
+        )
+    assert loaders == {}
+
+
 def _every_rank_loads_rank0s_plan_job(rank, size, ckpt, handed, placed):
     model, refused = None, None
     # Ranks without GPUs, whatever the host has: the pooled workers keep no CUDA context.
@@ -1643,6 +1919,51 @@ def test_a_placement_warning_raised_on_one_rank_fails_every_rank(cpu_dist_worker
     cpu_dist_workers.run(_one_rank_warns_as_an_error_job, ckpt=str(tmp_path))
 
 
+def _one_rank_cannot_apply_the_spec_job(rank, size, ckpt, step):
+    apply = moe.apply_distributed_support
+
+    def apply_on_rank0_only(*args, **kwargs):
+        if rank == 1:
+            raise OSError("this rank cannot import the checkpoint's code")
+        return apply(*args, **kwargs)
+
+    with ExitStack() as stack:
+        stack.enter_context(bounded_collectives())
+        stack.enter_context(
+            mock.patch.object(moe, "apply_distributed_support", apply_on_rank0_only)
+        )
+        if step == "load":
+            # Every rank has a GPU of its own, without which the load refuses it first.
+            stack.enter_context(mock.patch.object(torch.cuda, "is_available", return_value=True))
+            stack.enter_context(mock.patch.object(torch.cuda, "device_count", return_value=size))
+            stack.enter_context(
+                mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step)
+            )
+            stack.enter_context(
+                mock.patch.object(model_load_utils, "load_and_prepare_ep", _loaded_in_step)
+            )
+        try:
+            if step == "decision":
+                resolve_distribution_plan(ckpt)
+            else:
+                load_pretrained(ckpt, distribution=DistributionPlan(2, ep_size=2))
+        except (OSError, RuntimeError) as error:
+            raised = str(error)
+        else:
+            raised = "no error"
+    assert "cannot import the checkpoint's code" in raised, raised
+
+
+@needs_ep_transformers
+@pytest.mark.parametrize("step", ["decision", "load"])
+def test_a_rank_that_cannot_apply_the_spec_fails_every_rank(cpu_dist_workers, tmp_path, step):
+    """A spec's preparation reads the checkpoint and imports its code on each rank, which can fail
+    on one alone; the others must not wait for it, in the plan's broadcast when deciding, nor in the
+    load a handed expert-parallel plan starts."""
+    _deepseek_v3(n_routed_experts=8).save_pretrained(tmp_path)
+    cpu_dist_workers.run(_one_rank_cannot_apply_the_spec_job, ckpt=str(tmp_path), step=step)
+
+
 @pytest.mark.parametrize("sharding", ["expert-parallel", "fsdp2"])
 @pytest.mark.parametrize(
     ("change", "error", "match"),

From 1522f1731f8de0cc086cb12c29d631683c7d052a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 02:01:52 +0000
Subject: [PATCH 02/27] Import the spec applier at the top of model_load_utils

The decision and the load imported apply_distributed_support and
check_distributed_support inside the functions that call them, with no
circular import or optional dependency to justify it: model_load_utils
already imports from moe at module scope. The tests that stand in for
the applier now patch it where the load looks it up.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 6 ++----
 tests/unit/torch/utils/test_model_load_utils.py  | 6 +++---
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 47e74744518..0e408e0971d 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -78,6 +78,8 @@
     _routed_expert_modules,
     _spec_ep_plan,
     _spec_refuses_ep,
+    apply_distributed_support,
+    check_distributed_support,
     register_torch_a2a,
 )
 from modelopt.torch.utils.plugins.hf_checkpoint_utils import (
@@ -669,8 +671,6 @@ def _sharded_load_kwargs(
     hf_config = kwargs.pop("hf_config", None)
     if route == _EXPERT_PARALLEL:
         _refuse_a_rank_without_its_gpu(route, decided)
-        from modelopt.torch.utils.moe import apply_distributed_support, check_distributed_support
-
         if hf_config is None or not prepared:
             if hf_config is None:
                 hf_config = AutoConfig.from_pretrained(
@@ -1837,8 +1837,6 @@ def resolve_distribution_plan(
     that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and
     ``MODELOPT_FSDP_SIZE=1``) replicates it, noting that every rank holds the whole MoE.
     """
-    from modelopt.torch.utils.moe import apply_distributed_support
-
     # Each rank reads the checkpoint for itself. One that cannot must fail every rank here, rather
     # than leave the others waiting in a broadcast below for a rank that already raised.
     read: list[tuple[str, PretrainedConfig]] = []
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 7d6d21070be..058b09fb4e7 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -1664,7 +1664,7 @@ def apply_distributed_support(model_type, hf_config, **kwargs):
         hf_config.base_model_ep_plan = get_spec("deepseek_v3").ep_spec.base_model_ep_plan
         return hf_config
 
-    monkeypatch.setattr(moe, "apply_distributed_support", apply_distributed_support)
+    monkeypatch.setattr(model_load_utils, "apply_distributed_support", apply_distributed_support)
     kwargs = {"distribution": DistributionPlan(4, ep_size=4)} if planned else {}
     if passed_as:
         kwargs[passed_as] = config
@@ -1920,7 +1920,7 @@ def test_a_placement_warning_raised_on_one_rank_fails_every_rank(cpu_dist_worker
 
 
 def _one_rank_cannot_apply_the_spec_job(rank, size, ckpt, step):
-    apply = moe.apply_distributed_support
+    apply = model_load_utils.apply_distributed_support
 
     def apply_on_rank0_only(*args, **kwargs):
         if rank == 1:
@@ -1930,7 +1930,7 @@ def apply_on_rank0_only(*args, **kwargs):
     with ExitStack() as stack:
         stack.enter_context(bounded_collectives())
         stack.enter_context(
-            mock.patch.object(moe, "apply_distributed_support", apply_on_rank0_only)
+            mock.patch.object(model_load_utils, "apply_distributed_support", apply_on_rank0_only)
         )
         if step == "load":
             # Every rank has a GPU of its own, without which the load refuses it first.

From 5bfc000491115501852e77f79d87661b201f0e20 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 02:02:05 +0000
Subject: [PATCH 03/27] Say an MTP head the FSDP2 load cannot build rules FSDP2
 out too

The ways out of an MoE left whole on every rank offered FSDP2 for any
unquantized causal LM, but both FSDP2 loads refuse a Nemotron-H
checkpoint carrying MTP weights. Name that as well.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 5 +++--
 tests/unit/torch/utils/test_model_load_utils.py  | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 0e408e0971d..28a0197e8d9 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1725,8 +1725,9 @@ def _decide_distribution(
         # The ways out besides sharding the experts, which a spec refusing EP rules out.
         others = (
             "shard the whole model with FSDP2 (--use_fsdp2 in examples/hf_ptq/hf_ptq.py, or "
-            f"{fsdp}) if the checkpoint is an unquantized causal LM, load it in a single process, "
-            f"or replicate it on every rank outright: {knob}=1 with {replicate}, or "
+            f"{fsdp}) if the FSDP2 load supports the checkpoint (an unquantized causal LM with no "
+            "MTP head it cannot build, such as Nemotron-H's), load it in a single process, or "
+            f"replicate it on every rank outright: {knob}=1 with {replicate}, or "
             f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))."
         )
         ep = _ep_spec(hf_config.model_type)
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 058b09fb4e7..c584000b715 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -670,7 +670,7 @@ def test_a_run_without_expert_parallelism_needs_no_spec(
 _CAN_BE_EP = "can be expert-parallel, but the plan resolved .* divid(?:es|ing) both its {} experts"
 _NO_EP_YET = (
     "cannot be quantized expert-parallel yet, and the plan replicates .* FSDP2 .* if the "
-    "checkpoint is an unquantized causal LM, load it in a single process, or replicate "
+    r"FSDP2 load supports the checkpoint \(.*\), load it in a single process, or replicate "
     r".*DistributionPlan\(2, 1, 'replicate'\)"
 )
 _NEEDS_LATER = (

From 12201d19af142c52f43b1ed9e56d63b75e9d1bb5 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 07:05:24 +0000
Subject: [PATCH 04/27] Load a MoE only its spec makes expert-parallel on GPUs,
 and run it

The GPU test of an expert-parallel load_pretrained covered Qwen3-MoE,
whose config ships its own plan. It now also loads DeepSeek-V3 saved
with no plan of its own, so the plan comes from its model spec, as on a
transformers before 5.13. For both, each rank's batch runs through the
sharded model and must give the whole model's logits, so the test checks
the all-to-all forward as well as each rank's expert slices.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../gpu/torch/utils/test_model_load_utils.py  | 59 ++++++++++++++++---
 1 file changed, 51 insertions(+), 8 deletions(-)

diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py
index eb2685c68b2..e5be242cdd2 100644
--- a/tests/gpu/torch/utils/test_model_load_utils.py
+++ b/tests/gpu/torch/utils/test_model_load_utils.py
@@ -27,7 +27,7 @@
 from packaging.version import Version
 from safetensors.torch import load_file
 from torch.distributed.tensor import DTensor
-from transformers import AutoModelForCausalLM, LlamaConfig, Qwen3MoeConfig
+from transformers import AutoModelForCausalLM, DeepseekV3Config, LlamaConfig, Qwen3MoeConfig
 
 from modelopt.torch.export.unified_export_hf import export_hf_checkpoint
 from modelopt.torch.utils.distributed import (
@@ -36,6 +36,7 @@
     distribution_plan,
     set_distribution_plan,
 )
+from modelopt.torch.utils.moe import _EP_MIN_TRANSFORMERS
 from modelopt.torch.utils.plugins.model_load_utils import (
     load_pretrained,
     parallel_load_and_prepare_fsdp2,
@@ -216,7 +217,8 @@ def _load_and_check(rank, size, ckpt_dir, knobs, expected):
     Left to decide, ``load_pretrained`` shards an expert-parallel MoE's experts, FSDP2-shards a dense
     model when asked to, and otherwise gives every rank the whole model on its own GPU. Whichever
     load runs, every rank ends up with the checkpoint's weights and the load's record of what it
-    could not place.
+    could not place. An expert-parallel model also computes what the whole model does: each rank's
+    tokens reach the ranks holding their experts and come back.
     """
     device = torch.device(f"cuda:{rank}")
     os.environ.update(knobs)
@@ -256,6 +258,15 @@ def _load_and_check(rank, size, ckpt_dir, knobs, expected):
                 "expert-parallel": DistributionPlan(size, ep_size=size),
             }[expected]
         )
+        if expected == "expert-parallel":
+            # Every rank sends its own batch through the experts, as calibration does.
+            whole.to(device)
+            generator = torch.Generator().manual_seed(rank)
+            input_ids = torch.randint(0, _TINY["vocab_size"], (2, 16), generator=generator)
+            with torch.no_grad():
+                sharded = model(input_ids.to(device)).logits
+                reference = whole(input_ids.to(device)).logits
+            torch.testing.assert_close(sharded.float(), reference.float(), atol=2e-2, rtol=2e-2)
     finally:
         for knob in knobs:
             del os.environ[knob]
@@ -271,14 +282,46 @@ def test_a_dense_model_asked_to_shard_loads_fsdp2_sharded(need_2_gpus, dist_work
     dist_workers.run(_load_and_check, str(tmp_path), {"MODELOPT_FSDP_SIZE": "world"}, "fsdp2")
 
 
+def _qwen3_moe():
+    """Qwen3-MoE, whose config ships its own expert-parallel plan from transformers 5.7."""
+    config = Qwen3MoeConfig(moe_intermediate_size=32, num_experts=NUM_EXPERTS, **_TINY)
+    config.num_experts_per_tok = 2
+    return config
+
+
+def _deepseek_v3():
+    """DeepSeek-V3 with no plan of its own, as before transformers 5.13, so only its model spec
+    makes it expert-parallel. Every layer is an MoE layer, its attention shrunk to the tiny width."""
+    config = DeepseekV3Config(
+        n_routed_experts=NUM_EXPERTS,
+        num_experts_per_tok=2,
+        moe_intermediate_size=32,
+        first_k_dense_replace=0,
+        n_group=1,
+        topk_group=1,
+        q_lora_rank=32,
+        kv_lora_rank=16,
+        qk_rope_head_dim=8,
+        qk_nope_head_dim=8,
+        v_head_dim=16,
+        **_TINY,
+    )
+    config.base_model_ep_plan = None  # saved in config.json: no plan of its own on any transformers
+    return config
+
+
 @pytest.mark.skipif(
-    Version(transformers.__version__) < Version("5.7"),
-    reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7",
+    Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS),
+    reason=f"expert parallelism needs transformers >= {_EP_MIN_TRANSFORMERS}",
+)
+@pytest.mark.parametrize(
+    "config", [_qwen3_moe, _deepseek_v3], ids=["its-own-plan", "its-spec-plan"]
 )
-def test_an_expert_parallel_moe_loads_with_its_experts_sharded(need_2_gpus, dist_workers, tmp_path):
+def test_an_expert_parallel_moe_loads_with_its_experts_sharded(
+    need_2_gpus, dist_workers, tmp_path, config
+):
+    """Whether the plan comes with the model's config or only from its model spec."""
     if NUM_EXPERTS % dist_workers.world_size:
         pytest.skip(f"{NUM_EXPERTS} experts do not shard over {dist_workers.world_size} GPUs")
-    config = Qwen3MoeConfig(moe_intermediate_size=32, num_experts=NUM_EXPERTS, **_TINY)
-    config.num_experts_per_tok = 2
-    _save(config, tmp_path)
+    _save(config(), tmp_path)
     dist_workers.run(_load_and_check, str(tmp_path), {}, "expert-parallel")

From 46d4024725e0a9ba4e5ded8f7bc4bd37d6726620 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 07:27:22 +0000
Subject: [PATCH 05/27] Keep a spec-attached EP plan out of the loaded model's
 config

load_and_prepare_ep put the config the model keeps back to the plan the
checkpoint declares, which it inferred from the plan being an attribute
of the config it was handed. Since the decision and the load apply the
model spec, a plan the spec attached is such an attribute too, so the
kept config held ModelOpt's plan and an export would write it into
config.json: for Nemotron-H a style only ModelOpt registers, and so for
DeepSeek-V3 below transformers 5.13 or with a plan pinned to null.

The load now reads the declared plan from the checkpoint's config.json,
where the expert-parallel config sits in it, and the kept config gets
exactly that, or no plan of its own.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   | 33 +++++++++++++++++--
 .../unit/torch/utils/test_model_load_utils.py | 30 ++++++++++++-----
 2 files changed, 52 insertions(+), 11 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 28a0197e8d9..d8672885e5c 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -18,6 +18,7 @@
 import contextlib
 import copy
 import dataclasses
+import json
 import logging
 import math
 import os
@@ -70,6 +71,7 @@
 from modelopt.torch.utils.moe import (
     _EP_MIN_TRANSFORMERS,
     _EXPERTS_STYLES,
+    _SUB_CONFIG_ATTRS,
     _choose_ep_degree,
     _ep_config,
     _ep_spec,
@@ -2013,6 +2015,31 @@ def _ckpt_quant_method(cfg):
     return None
 
 
+_UNDECLARED = object()
+
+
+def _declared_ep_plan(ckpt_dir: str, hf_config, ep_cfg):
+    """The plan the checkpoint's ``config.json`` declares for ``ep_cfg``, or ``_UNDECLARED``.
+
+    ``ep_cfg`` is ``hf_config`` or one of its nested configs, and the plan is read from the same
+    place in the file. Read from the file, not the config: a plan the config carries besides is
+    ModelOpt's -- the one a model spec attaches
+    (:func:`~modelopt.torch.utils.moe.apply_distributed_support`) -- or the config class's, and an
+    export must write out neither.
+    """
+    try:
+        with open(os.path.join(ckpt_dir, "config.json")) as f:
+            source = json.load(f)
+    except (OSError, ValueError):
+        return _UNDECLARED
+    if ep_cfg is not hf_config:
+        attr = next((a for a in _SUB_CONFIG_ATTRS if getattr(hf_config, a, None) is ep_cfg), None)
+        source = source.get(attr) if attr else None
+    if isinstance(source, dict) and "base_model_ep_plan" in source:
+        return source["base_model_ep_plan"]
+    return _UNDECLARED
+
+
 def _transformers_mesh_builder():
     """The module whose ``initialize_distributed_mesh`` ``from_pretrained`` calls, or ``None``.
 
@@ -2226,8 +2253,8 @@ def gpu() -> tuple[bool, bool, str | None]:
     # some plans style only their parameters, grouped_gemm.
     swapped_plan.update(dict.fromkeys(_routed_expert_keys(ep_plan), "torch_a2a_experts"))
     from_pretrained_kwargs["experts_implementation"] = "torch_a2a"
-    # The checkpoint's own plan, for the config the model keeps (below); None: a class default.
-    declared_plan = vars(ep_cfg).get("base_model_ep_plan")
+    # The checkpoint's own plan, for the config the model keeps (below).
+    declared_plan = _declared_ep_plan(ckpt_path, hf_config, ep_cfg)
     ep_cfg.base_model_ep_plan = swapped_plan
     from_pretrained_kwargs["config"] = hf_config
     print_rank_0(f"[ep] ep_plan={swapped_plan} experts_implementation=torch_a2a")
@@ -2245,7 +2272,7 @@ def gpu() -> tuple[bool, bool, str | None]:
         kept_config = getattr(model, "config", None)
         if kept_config is not None:
             kept_ep_config = _ep_config(kept_config)
-            if declared_plan is None:
+            if declared_plan is _UNDECLARED:
                 vars(kept_ep_config).pop("base_model_ep_plan", None)
             else:
                 kept_ep_config.base_model_ep_plan = declared_plan
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index c584000b715..edbce55ccb9 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -25,6 +25,7 @@
 import copy
 import dataclasses
 import inspect
+import json
 import math
 import os
 import warnings
@@ -2724,13 +2725,19 @@ def from_pretrained(path, **kwargs):
         load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config)
 
 
-@pytest.mark.parametrize("declared", [True, False], ids=["on-the-config", "a-class-default"])
+@pytest.mark.parametrize(
+    ("declared", "on_the_config"),
+    [(True, True), (False, False), (False, True)],
+    ids=["by-its-checkpoint", "a-class-default", "attached-by-its-spec"],
+)
 def test_the_model_keeps_the_plan_its_checkpoint_declares(
-    ep_world_of_two, gloo_world_of_one, monkeypatch, declared
+    ep_world_of_two, gloo_world_of_one, monkeypatch, tmp_path, declared, on_the_config
 ):
     """The load shards by the torch_a2a styles, but the config the model keeps goes back to the
-    checkpoint's plan: an export writes it out, and nothing without ModelOpt knows those styles. A
-    plan the config class ships (transformers' own, e.g. Qwen3-MoE's) stays off the instance."""
+    plan its checkpoint's config.json declares: an export writes it out, and nothing without
+    ModelOpt knows those styles. A plan the config class ships (transformers' own, e.g.
+    Qwen3-MoE's) stays off the instance, and so does one the model's spec attached to the config,
+    which names styles only ModelOpt registers (Nemotron-H's nemotron_experts_ep)."""
 
     class _PastError(Exception):
         """Raised by the load's next check."""
@@ -2753,18 +2760,25 @@ def next_check(*args, **kwargs):
     monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None)
     monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check)
 
-    class _Config(SimpleNamespace):
-        model_type = _NO_SPEC
+    class _Config(transformers.PretrainedConfig):
+        model_type = "_probe"
         base_model_ep_plan = _EP_PLAN
 
+    (tmp_path / "config.json").write_text(
+        json.dumps({"base_model_ep_plan": _EP_PLAN} if declared else {})
+    )
     # 8 experts over 2 ranks: the stand-in model holds this rank's 4.
-    config = _Config(num_experts=8, **({"base_model_ep_plan": dict(_EP_PLAN)} if declared else {}))
+    plan = {"base_model_ep_plan": dict(_EP_PLAN)} if on_the_config else {}
+    config = _Config(num_experts=8, **plan)
 
     with pytest.raises(_PastError):
-        load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config)
+        load_and_prepare_ep(str(tmp_path), DistributionPlan(2, ep_size=2), hf_config=config)
 
     assert model.config.base_model_ep_plan == _EP_PLAN
     assert ("base_model_ep_plan" in vars(model.config)) is declared
+    model.config.save_pretrained(tmp_path / "saved")
+    written = json.loads((tmp_path / "saved" / "config.json").read_text())
+    assert ("base_model_ep_plan" in written) is declared
 
 
 def test_a_loaded_model_keeps_its_sharding_styles_out_of_its_saved_config(

From 4b9f10dfcf791ec7427aac6695a85dc5023bb082 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 08:05:35 +0000
Subject: [PATCH 06/27] Read the plan to keep inside the load's collective
 refusal

The expert-parallel load reads the plan the checkpoint declares from its
config.json. For a handed config whose plan sits on a nested config, a
config.json that is no JSON object made the nested lookup raise on that
rank alone, outside the load's collective refusal, so the other ranks
went on into the load and waited for it.

A file a rank cannot read as a JSON object now declares no plan, and the
read runs inside the expert-parallel load's refusal block, so anything
else it raises reaches every rank.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   |  9 +++--
 .../unit/torch/utils/test_model_load_utils.py | 34 ++++++++++++++++---
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index d8672885e5c..856c5a65661 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -2025,13 +2025,15 @@ def _declared_ep_plan(ckpt_dir: str, hf_config, ep_cfg):
     place in the file. Read from the file, not the config: a plan the config carries besides is
     ModelOpt's -- the one a model spec attaches
     (:func:`~modelopt.torch.utils.moe.apply_distributed_support`) -- or the config class's, and an
-    export must write out neither.
+    export must write out neither. A file this rank cannot read as a JSON object declares none.
     """
     try:
         with open(os.path.join(ckpt_dir, "config.json")) as f:
             source = json.load(f)
     except (OSError, ValueError):
         return _UNDECLARED
+    if not isinstance(source, dict):
+        return _UNDECLARED
     if ep_cfg is not hf_config:
         attr = next((a for a in _SUB_CONFIG_ATTRS if getattr(hf_config, a, None) is ep_cfg), None)
         source = source.get(attr) if attr else None
@@ -2253,14 +2255,15 @@ def gpu() -> tuple[bool, bool, str | None]:
     # some plans style only their parameters, grouped_gemm.
     swapped_plan.update(dict.fromkeys(_routed_expert_keys(ep_plan), "torch_a2a_experts"))
     from_pretrained_kwargs["experts_implementation"] = "torch_a2a"
-    # The checkpoint's own plan, for the config the model keeps (below).
-    declared_plan = _declared_ep_plan(ckpt_path, hf_config, ep_cfg)
     ep_cfg.base_model_ep_plan = swapped_plan
     from_pretrained_kwargs["config"] = hf_config
     print_rank_0(f"[ep] ep_plan={swapped_plan} experts_implementation=torch_a2a")
     # A rank failing in the load fails every rank here, before the warmup's collectives:
     # from_pretrained runs no collective, so the others finish it and meet the failure in the gather.
     with _refused_together("the expert-parallel load", alone=False):
+        # The checkpoint's own plan, for the config the model keeps (below): read from disk, so
+        # inside this block, where a rank failing to read it cannot leave the others in the load.
+        declared_plan = _declared_ep_plan(ckpt_path, hf_config, ep_cfg)
         # The class the plan's decision chose, with its LM head where an auto class builds one.
         loader = _auto_class(hf_config, require_lm_head, trust_remote_code=trust_remote_code)
         print_rank_0(f"[ep] loader={loader.__name__} (require_lm_head={require_lm_head})")
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index edbce55ccb9..620f7436ed3 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -3076,9 +3076,23 @@ def _meshes_in_step(*args, **kwargs):
     raise _PastTheMeshesError
 
 
+class _InTheLoadError(Exception):
+    """Raised by the weights' load stand-in: this rank got into the load."""
+
+
+def _load_in_step(*args, **kwargs):
+    """The weights' load, as the collective it is: a rank in it waits for every peer."""
+    torch.distributed.barrier()
+    raise _InTheLoadError
+
+
 def _one_rank_fails_job(rank, size, ckpt, failure, failing_rank, match):
     path, hf_config, patches = ckpt, None, []
-    if failure != "unreadable-config":
+    if failure == "unreadable-plan-provenance":
+        # A VLM's config, handed in: the plan to keep is read from where it sits in config.json.
+        plan = {"base_model_ep_plan": dict(_EP_PLAN), "num_experts": 4}
+        hf_config = SimpleNamespace(model_type=_NO_SPEC, text_config=SimpleNamespace(**plan))
+    elif failure != "unreadable-config":
         hf_config = SimpleNamespace(
             model_type=_NO_SPEC, base_model_ep_plan=dict(_EP_PLAN), num_experts=4
         )
@@ -3103,9 +3117,15 @@ def fail(*args, **kwargs):
     elif failure == "unseen-checkpoint" and fails_here:
         # Rank 0 makes the checkpoint local; a rank on a node without that storage cannot see it.
         patches.append(mock.patch.object(hf_checkpoint_utils.os.path, "isdir", return_value=False))
-    if failure in ("load", "record"):
+    elif failure == "unreadable-plan-provenance" and fails_here:
+        # This rank reads a config.json that is no JSON object: it declares no plan, not an error.
+        reads = SimpleNamespace(load=lambda f: [])
+        patches.append(mock.patch.object(model_load_utils, "json", reads))
+    if failure in ("load", "record", "unreadable-plan-provenance"):
         loaded = model_with_decoder_at("model.")
         load = fail if failure == "load" and fails_here else (lambda path, **kwargs: loaded)
+        if failure == "unreadable-plan-provenance":
+            load = _load_in_step
         loader = SimpleNamespace(__name__="Loader", from_pretrained=load)
         record = fail if failure == "record" and fails_here else None
         patches += [
@@ -3128,7 +3148,10 @@ def fail(*args, **kwargs):
     # The one rank's failure on every rank, before any rank went on alone into a collective. The GPU
     # check decides on every rank at once, from one gather, so no rank relays another's failure.
     assert refused is not None and match in refused, (rank, refused)
-    assert fails_here or failure == "no-gpu" or f"rank {failing_rank} failed" in refused, refused
+    if match != _InTheLoadError.__name__:  # else no rank failed, and every rank got into the load
+        assert fails_here or failure == "no-gpu" or f"rank {failing_rank} failed" in refused, (
+            refused
+        )
 
 
 @pytest.mark.skipif(
@@ -3144,6 +3167,7 @@ def fail(*args, **kwargs):
         ("no-gpu", 1, "rank(s) [1] have none while the others have theirs"),
         ("load", 1, "the GPU ran out of memory"),
         ("record", 1, "a weight file this rank cannot read"),
+        ("unreadable-plan-provenance", 1, _InTheLoadError.__name__),
     ],
 )
 def test_a_load_failing_on_one_rank_fails_every_rank(
@@ -3152,7 +3176,9 @@ def test_a_load_failing_on_one_rank_fails_every_rank(
     """One rank fails alone: reading the config, on its transformers, on its arguments, seeing the
     checkpoint rank 0 made local, without a GPU of its own where the others have theirs, or in the
     load. Every rank raises that failure, rather than leave the others waiting for it in the meshes
-    (stood in by a bounded barrier) or the warmup's collectives."""
+    (stood in by a bounded barrier) or the warmup's collectives. A config.json one rank cannot read
+    as a JSON object declares no plan to keep there, so every rank goes on into the load (stood in
+    by a bounded barrier too)."""
     Qwen3MoeConfig(num_experts=4, num_hidden_layers=1, hidden_size=8).save_pretrained(tmp_path)
     cpu_dist_workers.run(
         _one_rank_fails_job,

From b955be7c4b991dd333c40e9913203b7071fa9857 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 08:21:01 +0000
Subject: [PATCH 07/27] Warn when the load cannot read the plan its checkpoint
 declares

A config.json the expert-parallel load cannot read, or that is no JSON
object, declares no plan to keep, so the loaded model keeps none and an
export writes none, whatever the checkpoint declared. That loss is now
a warning instead of silence. The call-site comment says the read no
longer raises on the file, and why it stays in the collective block.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   | 18 +++++++++----
 .../unit/torch/utils/test_model_load_utils.py | 27 ++++++++++++-------
 2 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 856c5a65661..f6ee2353c46 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -2025,14 +2025,21 @@ def _declared_ep_plan(ckpt_dir: str, hf_config, ep_cfg):
     place in the file. Read from the file, not the config: a plan the config carries besides is
     ModelOpt's -- the one a model spec attaches
     (:func:`~modelopt.torch.utils.moe.apply_distributed_support`) -- or the config class's, and an
-    export must write out neither. A file this rank cannot read as a JSON object declares none.
+    export must write out neither. A file this rank cannot read as a JSON object declares none, with
+    a warning: what the checkpoint declared is lost to the export then.
     """
+    path = os.path.join(ckpt_dir, "config.json")
+    keeps_none = (
+        "so the loaded model keeps no expert-parallel plan of its own, and an export writes none"
+    )
     try:
-        with open(os.path.join(ckpt_dir, "config.json")) as f:
+        with open(path) as f:
             source = json.load(f)
-    except (OSError, ValueError):
+    except (OSError, ValueError) as e:
+        warn_rank_0(f"Cannot read {path} ({e}), {keeps_none}.")
         return _UNDECLARED
     if not isinstance(source, dict):
+        warn_rank_0(f"{path} is not a JSON object, {keeps_none}.")
         return _UNDECLARED
     if ep_cfg is not hf_config:
         attr = next((a for a in _SUB_CONFIG_ATTRS if getattr(hf_config, a, None) is ep_cfg), None)
@@ -2261,8 +2268,9 @@ def gpu() -> tuple[bool, bool, str | None]:
     # A rank failing in the load fails every rank here, before the warmup's collectives:
     # from_pretrained runs no collective, so the others finish it and meet the failure in the gather.
     with _refused_together("the expert-parallel load", alone=False):
-        # The checkpoint's own plan, for the config the model keeps (below): read from disk, so
-        # inside this block, where a rank failing to read it cannot leave the others in the load.
+        # The checkpoint's own plan, for the config the model keeps (below). The read raises on no
+        # failure of the file's, but sits in this block regardless: a rank failing there alone
+        # would leave the others in the load.
         declared_plan = _declared_ep_plan(ckpt_path, hf_config, ep_cfg)
         # The class the plan's decision chose, with its LM head where an auto class builds one.
         loader = _auto_class(hf_config, require_lm_head, trust_remote_code=trust_remote_code)
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 620f7436ed3..d073af6e1e0 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -29,7 +29,7 @@
 import math
 import os
 import warnings
-from contextlib import ExitStack, contextmanager
+from contextlib import ExitStack, contextmanager, nullcontext
 from types import SimpleNamespace
 from unittest import mock
 
@@ -2726,18 +2726,25 @@ def from_pretrained(path, **kwargs):
 
 
 @pytest.mark.parametrize(
-    ("declared", "on_the_config"),
-    [(True, True), (False, False), (False, True)],
-    ids=["by-its-checkpoint", "a-class-default", "attached-by-its-spec"],
+    ("source", "on_the_config"),
+    [
+        ({"base_model_ep_plan": _EP_PLAN}, True),
+        ({}, False),
+        ({}, True),
+        ('{"base_model_ep_plan": ', True),
+    ],
+    ids=["by-its-checkpoint", "a-class-default", "attached-by-its-spec", "in-an-unreadable-file"],
 )
 def test_the_model_keeps_the_plan_its_checkpoint_declares(
-    ep_world_of_two, gloo_world_of_one, monkeypatch, tmp_path, declared, on_the_config
+    ep_world_of_two, gloo_world_of_one, monkeypatch, tmp_path, source, on_the_config
 ):
     """The load shards by the torch_a2a styles, but the config the model keeps goes back to the
     plan its checkpoint's config.json declares: an export writes it out, and nothing without
     ModelOpt knows those styles. A plan the config class ships (transformers' own, e.g.
     Qwen3-MoE's) stays off the instance, and so does one the model's spec attached to the config,
-    which names styles only ModelOpt registers (Nemotron-H's nemotron_experts_ep)."""
+    which names styles only ModelOpt registers (Nemotron-H's nemotron_experts_ep). A config.json
+    the load cannot read declares none, with a warning, since an export then loses the plan."""
+    declared = isinstance(source, dict) and "base_model_ep_plan" in source
 
     class _PastError(Exception):
         """Raised by the load's next check."""
@@ -2764,14 +2771,14 @@ class _Config(transformers.PretrainedConfig):
         model_type = "_probe"
         base_model_ep_plan = _EP_PLAN
 
-    (tmp_path / "config.json").write_text(
-        json.dumps({"base_model_ep_plan": _EP_PLAN} if declared else {})
-    )
+    unreadable = isinstance(source, str)
+    (tmp_path / "config.json").write_text(source if unreadable else json.dumps(source))
     # 8 experts over 2 ranks: the stand-in model holds this rank's 4.
     plan = {"base_model_ep_plan": dict(_EP_PLAN)} if on_the_config else {}
     config = _Config(num_experts=8, **plan)
 
-    with pytest.raises(_PastError):
+    warns = pytest.warns(UserWarning, match="config.json") if unreadable else nullcontext()
+    with warns, pytest.raises(_PastError):
         load_and_prepare_ep(str(tmp_path), DistributionPlan(2, ep_size=2), hf_config=config)
 
     assert model.config.base_model_ep_plan == _EP_PLAN

From d6036640261674db47e49e0dd1493e598481f45b Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 08:21:10 +0000
Subject: [PATCH 08/27] Refuse an MoE whose spec refuses EP with
 NotImplementedError

Left whole on every rank, a model whose spec marks expert parallelism
unsupported was refused with a RuntimeError, while the load refuses the
same spec with check_distributed_support's NotImplementedError. Raise
the same type on both, which is still a RuntimeError.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 3 ++-
 tests/unit/torch/utils/test_model_load_utils.py  | 2 ++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index f6ee2353c46..cd915f793c3 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1734,7 +1734,8 @@ def _decide_distribution(
         )
         ep = _ep_spec(hf_config.model_type)
         if ep is not None and ep.unsupported:
-            raise RuntimeError(
+            # NotImplementedError, as check_distributed_support refuses the same spec.
+            raise NotImplementedError(
                 f"{hf_config.model_type} cannot be quantized expert-parallel yet, and the plan "
                 f"replicates it at world_size={world_size}, so every rank would hold the whole "
                 f"MoE. Instead, {others}"
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index d073af6e1e0..584f8a1f11a 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -750,6 +750,8 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error(
 
     with pytest.raises((RuntimeError, ValueError), match=refusal) as refused:
         resolve_distribution_plan(str(tmp_path), **arguments)
+    if refusal is _NO_EP_YET:  # as the load refuses the same spec
+        assert refused.type is NotImplementedError
     if "divid" in refusal:  # the decision's refusal of what its inputs resolve to
         assert refused.type is ValueError, refused.type
     if arguments:  # a knob the arguments override would lead straight back here

From 132152dd25656c1dbdcbd80cc69ed266c629b3a9 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 18:48:38 +0000
Subject: [PATCH 09/27] Warn on the rank that cannot read the plan its
 checkpoint declares

The warning for a config.json the load cannot read went through
warn_rank_0, so a rank other than 0 failing alone dropped the plan from
its model's config with no warning anywhere. Each rank now warns for
itself and names its rank.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 10 ++++++----
 tests/unit/torch/utils/test_model_load_utils.py  |  8 ++++++++
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index cd915f793c3..59e33824376 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -2027,20 +2027,22 @@ def _declared_ep_plan(ckpt_dir: str, hf_config, ep_cfg):
     ModelOpt's -- the one a model spec attaches
     (:func:`~modelopt.torch.utils.moe.apply_distributed_support`) -- or the config class's, and an
     export must write out neither. A file this rank cannot read as a JSON object declares none, with
-    a warning: what the checkpoint declared is lost to the export then.
+    a warning on this rank, which may be the only one that cannot: what the checkpoint declared is
+    lost to its model's config then.
     """
     path = os.path.join(ckpt_dir, "config.json")
     keeps_none = (
-        "so the loaded model keeps no expert-parallel plan of its own, and an export writes none"
+        f"so rank {dist.rank()}'s model keeps no expert-parallel plan of its own, and an export "
+        "from it writes none"
     )
     try:
         with open(path) as f:
             source = json.load(f)
     except (OSError, ValueError) as e:
-        warn_rank_0(f"Cannot read {path} ({e}), {keeps_none}.")
+        warnings.warn(f"Cannot read {path} ({e}), {keeps_none}.")
         return _UNDECLARED
     if not isinstance(source, dict):
-        warn_rank_0(f"{path} is not a JSON object, {keeps_none}.")
+        warnings.warn(f"{path} is not a JSON object, {keeps_none}.")
         return _UNDECLARED
     if ep_cfg is not hf_config:
         attr = next((a for a in _SUB_CONFIG_ATTRS if getattr(hf_config, a, None) is ep_cfg), None)
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 584f8a1f11a..80ea2bef784 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -3150,6 +3150,8 @@ def fail(*args, **kwargs):
     with bounded_collectives(), ExitStack() as stack:
         for patch in patches:
             stack.enter_context(patch)
+        caught = stack.enter_context(warnings.catch_warnings(record=True))
+        warnings.simplefilter("always")
         try:  # not pytest.raises: its failure would end the pooled worker instead of the job
             load_and_prepare_ep(path, DistributionPlan(2, ep_size=2), hf_config=hf_config)
         except Exception as error:
@@ -3161,6 +3163,12 @@ def fail(*args, **kwargs):
         assert fails_here or failure == "no-gpu" or f"rank {failing_rank} failed" in refused, (
             refused
         )
+    else:  # the plan the one rank could not read is lost to its model alone, which says so
+        warned = [
+            str(w.message) for w in caught if "keeps no expert-parallel plan" in str(w.message)
+        ]
+        assert bool(warned) == fails_here, (rank, warned)
+        assert not warned or f"rank {rank}'s model" in warned[0], warned
 
 
 @pytest.mark.skipif(

From 4bd04d0c65be865c93fdf18df828f9a3c5f5ed40 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 18:50:38 +0000
Subject: [PATCH 10/27] Name the transformers release the left-whole refusal
 asks for

Below the release a model's expert parallelism needs, an MoE left whole
on every rank was refused with "Upgrade transformers" and no release,
though the spec's own refusal names it. Both now read the release from
one place, moe._ep_floor: the expert-parallel load's own, or the spec's
when that is later.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/moe.py                    | 18 +++++++++++++-----
 .../torch/utils/plugins/model_load_utils.py    |  5 +++--
 .../unit/torch/utils/test_model_load_utils.py  |  2 +-
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index 5dd6ce764ea..a78dca19ccd 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -156,6 +156,18 @@ def _ep_spec(model_type: str):
     return getattr(get_spec(model_type), "ep_spec", None)
 
 
+def _ep_floor(model_type: str) -> str:
+    """The transformers release ``model_type``'s expert parallelism needs.
+
+    The expert-parallel load's own, or its spec's when that is later. A model whose spec states no
+    release, or that has no spec, still needs the load's: auto must pick no EP degree below it,
+    rather than reach the load and fail there.
+    """
+    ep = _ep_spec(model_type)
+    stated = ep.requires_transformers if ep is not None else None
+    return max(_EP_MIN_TRANSFORMERS, stated or _EP_MIN_TRANSFORMERS, key=parse)
+
+
 def check_distributed_support(model_type: str) -> None:
     """Refuse ``model_type``'s expert parallelism if its spec does, or below its transformers floor.
 
@@ -181,12 +193,8 @@ def check_distributed_support(model_type: str) -> None:
             f"{model_type} cannot be quantized expert-parallel yet: {ep.unsupported} Quantize it "
             f"{without_ep}."
         )
-    # A model whose spec states no release, or that has no spec, still needs the load's floor: auto
-    # must pick no EP degree below it, rather than reach the load and fail there.
-    stated = ep.requires_transformers if ep is not None else None
-    floor = max(_EP_MIN_TRANSFORMERS, stated or _EP_MIN_TRANSFORMERS, key=parse)
     try:
-        require_transformers(floor, f"{model_type} expert-parallel PTQ")
+        require_transformers(_ep_floor(model_type), f"{model_type} expert-parallel PTQ")
     except RuntimeError as e:
         raise RuntimeError(f"{e} Or quantize it {without_ep}.") from None
 
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 59e33824376..6a3a4afde17 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -74,6 +74,7 @@
     _SUB_CONFIG_ATTRS,
     _choose_ep_degree,
     _ep_config,
+    _ep_floor,
     _ep_spec,
     _num_experts_from_config,
     _routed_expert_keys,
@@ -1745,8 +1746,8 @@ def _decide_distribution(
             raise RuntimeError(
                 f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 "
                 f"at world_size={world_size}, so every rank would hold the whole MoE, and this "
-                "transformers cannot shard its experts. Upgrade transformers to shard them, or "
-                f"{others}"
+                "transformers cannot shard its experts. Upgrade transformers to "
+                f"{_ep_floor(hf_config.model_type)} or later to shard them, or {others}"
             )
         # Name a way to shard the experts only where one exists: a count the config spells none of
         # the ways ModelOpt reads leaves no EP degree to name, as does one sharing no factor above 1
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 80ea2bef784..6411e1111c7 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -676,7 +676,7 @@ def test_a_run_without_expert_parallelism_needs_no_spec(
 )
 _NEEDS_LATER = (
     "can be expert-parallel, but the plan resolved .* this transformers cannot shard its experts. "
-    "Upgrade transformers to shard them, or shard the whole model with FSDP2"
+    r"Upgrade transformers to 5\.7 or later to shard them, or shard the whole model with FSDP2"
 )
 
 

From 7812728d413d1bebefc150291cdfaf23d3ec8287 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 19:34:32 +0000
Subject: [PATCH 11/27] Offer the transformers upgrade only where an EP degree
 could shard

Below the release a model's expert parallelism needs, an MoE left whole
on every rank was told to upgrade transformers to shard its experts,
even when no EP degree divides both its experts and the world size, or
its config names no count: there an upgrade shards nothing, and the
next run meets the refusal that names the real obstacle. Offer the
upgrade only where a degree exists, and let the other cases name what
stands in the way.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py    |  8 +++++++-
 .../unit/torch/utils/test_model_load_utils.py  | 18 +++++++++++++-----
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 6a3a4afde17..d1ae0ccd230 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1741,8 +1741,14 @@ def _decide_distribution(
                 f"replicates it at world_size={world_size}, so every rank would hold the whole "
                 f"MoE. Instead, {others}"
             )
-        if _spec_refuses_ep(hf_config.model_type):
+        if (
+            _spec_refuses_ep(hf_config.model_type)
+            and num_experts
+            and math.gcd(num_experts, world_size) > 1
+        ):
             # Below the release its expert parallelism needs: an EP degree above 1 is refused too.
+            # Where no EP degree could shard the experts anyway, the arms below say so instead, as
+            # upgrading would not change it.
             raise RuntimeError(
                 f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 "
                 f"at world_size={world_size}, so every rank would hold the whole MoE, and this "
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 6411e1111c7..2c3d0296c3d 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -717,6 +717,13 @@ def test_a_run_without_expert_parallelism_needs_no_spec(
             None,
             id="explicit_degree_one_as_an_argument",
         ),
+        pytest.param(
+            _deepseek_v3(n_routed_experts=3),
+            None,
+            _CAN_BE_EP.format(3),
+            "5.6.2",
+            id="below_the_ep_floor_no_degree",
+        ),
     ],
 )
 def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error(
@@ -731,13 +738,14 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error(
     way out (FSDP2 only for an unquantized causal LM); the spec's own refusal, a
     NotImplementedError and so a RuntimeError too, words it differently. Below the transformers
     release expert parallelism needs, no EP degree is offered, since none would load: upgrading
-    is. The ways out are named as the caller set the plan, so arguments, which override the knobs,
-    are named."""
+    is, where an EP degree divides the experts; where none does, upgrading would not shard them.
+    The ways out are named as the caller set the plan, so arguments, which override the knobs, are
+    named."""
     if version is not None:
         monkeypatch.setattr(transformers, "__version__", version)
-    if refusal.startswith("can be") and Version(transformers.__version__) < Version(
-        _EP_MIN_TRANSFORMERS
-    ):
+    experts = getattr(config, "n_routed_experts", None) or getattr(config, "num_experts", 0)
+    below = Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS)
+    if refusal.startswith("can be") and below and experts % 2 == 0:  # an EP degree divides them
         refusal = _NEEDS_LATER
     if ep_degree is not None:
         if config.model_type == "qwen3_moe" and Version(transformers.__version__) < Version("5.7"):

From 0e5b9abea8d2aee76d8fd7987f3ae8932d96ab82 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 19:58:04 +0000
Subject: [PATCH 12/27] Name the EP degree to raise beside the upgrade when the
 caller set it

Below the release a model's expert parallelism needs, the left-whole
refusal tells the caller to upgrade transformers to shard the experts.
With the EP degree set to 1 by the caller, an upgrade alone leaves it at
1, and the next run meets the refusal again. Name raising the degree as
well there, with the knob the caller set it by.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   | 11 +++++++++--
 .../unit/torch/utils/test_model_load_utils.py | 19 +++++++++++++++++--
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index d1ae0ccd230..2803a820ed5 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1748,12 +1748,19 @@ def _decide_distribution(
         ):
             # Below the release its expert parallelism needs: an EP degree above 1 is refused too.
             # Where no EP degree could shard the experts anyway, the arms below say so instead, as
-            # upgrading would not change it.
+            # upgrading would not change it. An EP degree the caller set to 1 stays 1 after an
+            # upgrade, so it has to be raised as well.
+            degree = (
+                f" and set {knob} above 1, dividing both its {num_experts} experts and the world "
+                "size,"
+                if ep_asked
+                else ""
+            )
             raise RuntimeError(
                 f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 "
                 f"at world_size={world_size}, so every rank would hold the whole MoE, and this "
                 "transformers cannot shard its experts. Upgrade transformers to "
-                f"{_ep_floor(hf_config.model_type)} or later to shard them, or {others}"
+                f"{_ep_floor(hf_config.model_type)} or later{degree} to shard them, or {others}"
             )
         # Name a way to shard the experts only where one exists: a count the config spells none of
         # the ways ModelOpt reads leaves no EP degree to name, as does one sharing no factor above 1
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 2c3d0296c3d..c2463560e7d 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -678,6 +678,12 @@ def test_a_run_without_expert_parallelism_needs_no_spec(
     "can be expert-parallel, but the plan resolved .* this transformers cannot shard its experts. "
     r"Upgrade transformers to 5\.7 or later to shard them, or shard the whole model with FSDP2"
 )
+# The same, for an EP degree the caller set to 1, which an upgrade alone leaves at 1.
+_NEEDS_LATER_AND_A_DEGREE = (
+    "can be expert-parallel, but the plan resolved .* this transformers cannot shard its experts. "
+    r"Upgrade transformers to 5\.7 or later and set MODELOPT_EP_DEGREE above 1, dividing both its "
+    "8 experts and the world size, to shard them, or shard the whole model with FSDP2"
+)
 
 
 @pytest.mark.parametrize(
@@ -724,6 +730,13 @@ def test_a_run_without_expert_parallelism_needs_no_spec(
             "5.6.2",
             id="below_the_ep_floor_no_degree",
         ),
+        pytest.param(
+            _deepseek_v3(n_routed_experts=8),
+            "1",
+            _NEEDS_LATER_AND_A_DEGREE,
+            "5.6.2",
+            id="below_the_ep_floor_explicit_degree_one",
+        ),
     ],
 )
 def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error(
@@ -745,8 +758,10 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error(
         monkeypatch.setattr(transformers, "__version__", version)
     experts = getattr(config, "n_routed_experts", None) or getattr(config, "num_experts", 0)
     below = Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS)
-    if refusal.startswith("can be") and below and experts % 2 == 0:  # an EP degree divides them
-        refusal = _NEEDS_LATER
+    if version is None and refusal.startswith("can be") and below and experts % 2 == 0:
+        # On a transformers below the floor, upgrading is the way out where an EP degree divides
+        # the experts, with a degree set to 1 raised as well.
+        refusal = _NEEDS_LATER_AND_A_DEGREE if ep_degree is not None else _NEEDS_LATER
     if ep_degree is not None:
         if config.model_type == "qwen3_moe" and Version(transformers.__version__) < Version("5.7"):
             pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7")

From 97c23b3a9a53a8ea3d5cd04ffe468028d741dba7 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 20:27:13 +0000
Subject: [PATCH 13/27] Tell the decision's own refusal from the upgrade's by
 its pattern

The left-whole refusal test checked for a ValueError by the word
"divid" in the expected pattern, which the pattern for an EP degree to
raise beside an upgrade contains too, though that refusal stays a
RuntimeError. Check by the pattern of the decision's own refusal.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/utils/test_model_load_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index c2463560e7d..6314774209f 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -775,7 +775,7 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error(
         resolve_distribution_plan(str(tmp_path), **arguments)
     if refusal is _NO_EP_YET:  # as the load refuses the same spec
         assert refused.type is NotImplementedError
-    if "divid" in refusal:  # the decision's refusal of what its inputs resolve to
+    if refusal.startswith(_CAN_BE_EP[: _CAN_BE_EP.index("{")]):  # what its inputs resolve to
         assert refused.type is ValueError, refused.type
     if arguments:  # a knob the arguments override would lead straight back here
         assert "ep_size=1 with dense='replicate'" in str(refused.value), refused.value

From 6a819609fcee5a4239c1d905b249f6baf38efa21 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 20:48:51 +0000
Subject: [PATCH 14/27] Leave the model spec out of a single process's plan

With no ranks to shard over, the plan holds the whole model whatever
the model is, but the EP degree argument was still parsed with the model
type, which consults the model's spec and its transformers floor for a
result that cannot change. Parse it without.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py |  2 +-
 tests/unit/torch/utils/test_model_load_utils.py  | 10 +++++++---
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 2803a820ed5..4321c387121 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1588,7 +1588,7 @@ def _decide_distribution(
         # Nothing to shard over, so nothing about the model matters, nor do the knobs: a shell still
         # exporting a multi-GPU run's would otherwise reshape a single process's load. The caller's
         # own arguments still count, as dense="fsdp" for FSDP2 on one rank.
-        single = 1 if ep_size is None else _ep_degree(ep_size, hf_config.model_type)
+        single = 1 if ep_size is None else _ep_degree(ep_size)
         return DistributionPlan(
             ep_size=single if isinstance(single, int) else 1,
             dense="replicate" if dense in (None, "auto") else dense,
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 6314774209f..c09923018b5 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -437,9 +437,9 @@ class counts only when the checkpoint's code is trusted. An encoder-decoder mode
 def test_a_single_process_holds_the_whole_model(
     run_as, serve, monkeypatch, tmp_path, world_size, knob
 ):
-    """With no ranks to shard over, even an MoE loads whole, and no model is built to decide it. A
-    shell still exporting a multi-GPU run's knobs must not reshape that: only the caller's own
-    arguments ask for FSDP2 on one rank."""
+    """With no ranks to shard over, even an MoE loads whole, and no model is built, nor its spec
+    consulted, to decide it. A shell still exporting a multi-GPU run's knobs must not reshape that:
+    only the caller's own arguments ask for FSDP2 on one rank."""
     if world_size is not None:
         run_as(world_size)
     if knob is not None:
@@ -448,8 +448,12 @@ def test_a_single_process_holds_the_whole_model(
     monkeypatch.setattr(
         model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model")
     )
+    monkeypatch.setattr(
+        model_load_utils, "_spec_refuses_ep", lambda *a: pytest.fail("consulted the spec")
+    )
 
     assert resolve_distribution_plan(str(tmp_path)) == (DistributionPlan(), served)
+    assert resolve_distribution_plan(str(tmp_path), ep_size="auto")[0] == DistributionPlan()
     assert resolve_distribution_plan(str(tmp_path), dense="fsdp")[0] == DistributionPlan(
         dense="fsdp"
     )

From edba0818a78bf9e223df62f21e35e98dc8d7c6f0 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 20:59:45 +0000
Subject: [PATCH 15/27] Apply no model spec for a model class the sharded loads
 refuse

resolve_distribution_plan applied the model's spec whenever the world and
the EP degree left room for expert parallelism, whatever class the load
builds the model with. A class the sharded loads refuse is planned whole on
every rank, or refused, so it is never expert-parallel, yet every rank ran
the spec's prepare first: a prepare that raises failed a load that loaded
the whole model on every rank before, and one that patches the process did
so for a load that never uses it. Skip the spec for such a class, as the
decision does.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   | 19 ++++++++++++++-----
 .../unit/torch/utils/test_model_load_utils.py | 12 +++++++++---
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 4321c387121..d99d4825bc8 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1845,9 +1845,10 @@ def resolve_distribution_plan(
     ``init_quantized_weights``, and nothing asking to shard, the plan loads the model whole, decided
     without a build.
 
-    Unless rank 0's knobs rule expert parallelism out (a single process, or an EP degree of 1,
-    which auto is for a model whose spec refuses expert parallelism on the installed
-    transformers), every rank applies the model's spec to a copy of the config first
+    Unless the call or rank 0's knobs rule expert parallelism out (a single process, an
+    ``auto_class`` the sharded loads refuse, or an EP degree of 1, which auto is for a model whose
+    spec refuses expert parallelism on the installed transformers), every rank applies the model's
+    spec to a copy of the config first
     (:func:`~modelopt.torch.utils.moe.apply_distributed_support`): for several models it is what
     supplies the expert-parallel plan, and an expert-parallel plan returns that copy. Any other
     plan returns the config without it, which the export would otherwise write out. A plan that
@@ -1888,10 +1889,18 @@ def checkpoint_read() -> bool:
     world_size = dist.size()
     # Whether the run rules expert parallelism out is rank 0's call, as the plan is: a rank that
     # skipped the spec while the others applied it would load without the expert-parallel plan
-    # they shard by, and fail once the mesh is up, leaving them waiting in a collective.
+    # they shard by, and fail once the mesh is up, leaving them waiting in a collective. A model
+    # class the sharded loads refuse rules it out too: 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.
     prepared = [hf_config]
     if dist.broadcast_from(
-        "spec", lambda: world_size > 1 and _ep_degree(ep_size, hf_config.model_type) != 1
+        "spec",
+        lambda: (
+            world_size > 1
+            and auto_class in (None, AutoModelForCausalLM)
+            and _ep_degree(ep_size, hf_config.model_type) != 1
+        ),
     ):
         # Before deciding: deciding without the spec would read an expert-parallel MoE as dense. A
         # rank that cannot apply it fails every rank, as a checkpoint read does above.
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index c09923018b5..33cfd664101 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -1550,9 +1550,10 @@ def test_a_model_class_of_the_callers_own_loads_whole(
     """The sharded loads build their own model class and refuse another, so with nothing asking to
     shard, a class of the caller's own loads whole on every rank, and with a knob asking to shard
     it is refused. Neither builds the model on the meta device, which the decision's own pick of
-    class may not build (an encoder-decoder loaded through AutoModel). The decision reads the
-    checkpoint's config as the class does, through its own config class, which AutoConfig need not
-    know."""
+    class may not build (an encoder-decoder loaded through AutoModel), nor applies the model's spec,
+    whose prepare could fail every rank, or patch the process, for a load that is never
+    expert-parallel. The decision reads the checkpoint's config as the class does, through its own
+    config class, which AutoConfig need not know."""
     run_as(2)
     if own_config:
         _OwnConfig(**_TINY).save_pretrained(tmp_path)
@@ -1562,6 +1563,11 @@ class may not build (an encoder-decoder loaded through AutoModel). The decision
     monkeypatch.setattr(
         model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model")
     )
+    monkeypatch.setattr(
+        model_load_utils, "apply_distributed_support", lambda *a, **k: pytest.fail("applied a spec")
+    )
+    # As on a transformers new enough for expert parallelism, whose floor would rule it out here.
+    monkeypatch.setattr(model_load_utils, "_spec_refuses_ep", lambda model_type: False)
     if fsdp_size is None:
         model = load_pretrained(str(tmp_path), auto_class)
         assert distribution_plan(model) == DistributionPlan(2)

From a52d75add09e716ae75f4ecc3f0de50930f91873 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 21:32:18 +0000
Subject: [PATCH 16/27] Check that a rank without its GPU prepares no config
 passed in

With an expert-parallel plan handed in, the expert-parallel load prepares
its copy of a config= passed for the load once its own checks pass, the
rank's-own-GPU refusal among them, so a rank that cannot load runs no
spec prepare, which can import the checkpoint's code. A row of
test_a_sharded_load_refuses_what_it_cannot_honor holds it: its spec
applier fails if a rank without a GPU prepares anything.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../unit/torch/utils/test_model_load_utils.py | 28 +++++++++++++++----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 33cfd664101..990954aa3c6 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -2023,28 +2023,44 @@ def test_a_rank_that_cannot_apply_the_spec_fails_every_rank(cpu_dist_workers, tm
             "config and hf_config",
         ),
         ({"gpus": 0}, RuntimeError, "this rank has none"),
+        ({"gpus": 0, "config": True}, RuntimeError, "this rank has none"),
+    ],
+    ids=[
+        "placement",
+        "dtype",
+        "torch_dtype",
+        "model_class",
+        "mtp_head",
+        "two_configs",
+        "no_gpu",
+        "no_gpu_config_passed",
     ],
-    ids=["placement", "dtype", "torch_dtype", "model_class", "mtp_head", "two_configs", "no_gpu"],
 )
 def test_a_sharded_load_refuses_what_it_cannot_honor(
     run_as, serve, loaders, process_group, monkeypatch, tmp_path, sharding, change, error, match
 ):
     """Placement, a model class, a head only a whole model builds, two configs to load with, a rank
     without a GPU of its own: none reaches a sharded load, which fails before it starts rather than
-    load something other than what was asked."""
+    load something other than what was asked. A rank without a GPU is refused before the model's
+    spec prepares its config, which can import the checkpoint's code, whichever config it loads
+    with."""
     run_as(2)
     expert_parallel = sharding == "expert-parallel"
-    serve(_deepseek_v3() if expert_parallel else LlamaConfig(**_TINY))
+    config = serve(_deepseek_v3() if expert_parallel else LlamaConfig(**_TINY))
     if change.get("mtp"):
         monkeypatch.setattr(models_hf, "checkpoint_has_mtp", lambda *a: True)
     if "gpus" in change:  # a host that shows this rank no GPU
         monkeypatch.setattr(torch.cuda, "device_count", lambda: change["gpus"])
+        monkeypatch.setattr(
+            model_load_utils,
+            "apply_distributed_support",
+            lambda *a, **k: pytest.fail("prepared a rank that cannot load"),
+        )
     plan = DistributionPlan(2, ep_size=2) if expert_parallel else DistributionPlan(2, dense="fsdp")
+    kwargs = {**change.get("kwargs", {}), **({"config": config} if change.get("config") else {})}
 
     with pytest.raises(error, match=match) as refused:
-        load_pretrained(
-            str(tmp_path), change.get("auto_class"), distribution=plan, **change.get("kwargs", {})
-        )
+        load_pretrained(str(tmp_path), change.get("auto_class"), distribution=plan, **kwargs)
     if "device_map" in change.get("kwargs", {}) or "gpus" in change:
         # The way to a whole model it names, for the plan handed in: an expert-parallel MoE's is a
         # plan with an EP degree of 1.

From f11eaa7ee2d7651a2ddb54835c76b91838d66aa3 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 21:32:41 +0000
Subject: [PATCH 17/27] Say a plan set on a config passed to the load is not
 kept either

The expert-parallel load keeps on the loaded model's config the plan the
checkpoint's config.json declares, read from the file. A plan set on a
config passed for the load, and not declared there, is not kept, so an
export writes out what the checkpoint declares; the docstring said only
the spec's and the config class's plans were left out.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index d99d4825bc8..50a95a9520b 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -2048,10 +2048,11 @@ def _declared_ep_plan(ckpt_dir: str, hf_config, ep_cfg):
     ``ep_cfg`` is ``hf_config`` or one of its nested configs, and the plan is read from the same
     place in the file. Read from the file, not the config: a plan the config carries besides is
     ModelOpt's -- the one a model spec attaches
-    (:func:`~modelopt.torch.utils.moe.apply_distributed_support`) -- or the config class's, and an
-    export must write out neither. A file this rank cannot read as a JSON object declares none, with
-    a warning on this rank, which may be the only one that cannot: what the checkpoint declared is
-    lost to its model's config then.
+    (:func:`~modelopt.torch.utils.moe.apply_distributed_support`) -- the config class's, or one set
+    on a config passed to the load, and the loaded model keeps none of them: an export writes out
+    what the checkpoint declares, if anything. A file this rank cannot read as a JSON object
+    declares none, with a warning on this rank, which may be the only one that cannot: what the
+    checkpoint declared is lost to its model's config then.
     """
     path = os.path.join(ckpt_dir, "config.json")
     keeps_none = (

From 0090fdfb08f0702ec77946817c123c232aacc66a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 21:35:35 +0000
Subject: [PATCH 18/27] Word the FSDP2 way around expert parallelism once

check_distributed_support and the refusal of an MoE left whole on every
rank both offered FSDP2 with the same caveat, what the FSDP2 load takes,
in two wordings that each round had to keep in step. Both now take it
from one helper in moe, with the knob spelled as the caller set it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/moe.py                   | 20 ++++++++++++-------
 .../torch/utils/plugins/model_load_utils.py   |  7 +++----
 .../unit/torch/utils/test_model_load_utils.py |  4 ++--
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index a78dca19ccd..20d5ef7be04 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -168,6 +168,18 @@ def _ep_floor(model_type: str) -> str:
     return max(_EP_MIN_TRANSFORMERS, stated or _EP_MIN_TRANSFORMERS, key=parse)
 
 
+def _fsdp2_way_out(fsdp: str = "MODELOPT_FSDP_SIZE=world") -> str:
+    """FSDP2 as a refusal names it, the way around expert parallelism ``fsdp`` asks for."""
+    # The FSDP2 load builds a causal LM and reads its tensors as stored, so a quantized or
+    # multimodal checkpoint (GPT-OSS's MXFP4, Llama-4's) needs a single process, as does one with
+    # an MTP head the load cannot build (--use_fsdp2 refuses Nemotron-H's).
+    return (
+        f"(--use_fsdp2 in examples/hf_ptq/hf_ptq.py, or {fsdp}) if the FSDP2 load supports the "
+        "checkpoint, which takes an unquantized causal LM with no MTP head it cannot build (such "
+        "as Nemotron-H's)"
+    )
+
+
 def check_distributed_support(model_type: str) -> None:
     """Refuse ``model_type``'s expert parallelism if its spec does, or below its transformers floor.
 
@@ -179,14 +191,8 @@ def check_distributed_support(model_type: str) -> None:
     expert-parallel plan handed in, meets these errors; both name the way around expert parallelism.
     """
     ep = _ep_spec(model_type)
-    # The FSDP2 load builds a causal LM and reads its tensors as stored, so a quantized or
-    # multimodal checkpoint (GPT-OSS's MXFP4, Llama-4's) needs a single process, as does one with
-    # an MTP head the load cannot build (--use_fsdp2 refuses Nemotron-H's).
     without_ep = (
-        "without expert parallelism: in a single process, or FSDP2-sharded (--use_fsdp2 in "
-        "examples/hf_ptq/hf_ptq.py, or MODELOPT_FSDP_SIZE=world) if the FSDP2 load supports the "
-        "checkpoint, which takes an unquantized causal LM with no MTP head it cannot build (such "
-        "as Nemotron-H's)"
+        f"without expert parallelism: in a single process, or FSDP2-sharded {_fsdp2_way_out()}"
     )
     if ep is not None and ep.unsupported:
         raise NotImplementedError(
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 50a95a9520b..29111fd242a 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -76,6 +76,7 @@
     _ep_config,
     _ep_floor,
     _ep_spec,
+    _fsdp2_way_out,
     _num_experts_from_config,
     _routed_expert_keys,
     _routed_expert_modules,
@@ -1727,10 +1728,8 @@ def _decide_distribution(
     ):
         # The ways out besides sharding the experts, which a spec refusing EP rules out.
         others = (
-            "shard the whole model with FSDP2 (--use_fsdp2 in examples/hf_ptq/hf_ptq.py, or "
-            f"{fsdp}) if the FSDP2 load supports the checkpoint (an unquantized causal LM with no "
-            "MTP head it cannot build, such as Nemotron-H's), load it in a single process, or "
-            f"replicate it on every rank outright: {knob}=1 with {replicate}, or "
+            f"shard the whole model with FSDP2 {_fsdp2_way_out(fsdp)}, load it in a single "
+            f"process, or replicate it on every rank outright: {knob}=1 with {replicate}, or "
             f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))."
         )
         ep = _ep_spec(hf_config.model_type)
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 990954aa3c6..7350cfc8ce9 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -675,8 +675,8 @@ def test_a_run_without_expert_parallelism_needs_no_spec(
 _CAN_BE_EP = "can be expert-parallel, but the plan resolved .* divid(?:es|ing) both its {} experts"
 _NO_EP_YET = (
     "cannot be quantized expert-parallel yet, and the plan replicates .* FSDP2 .* if the "
-    r"FSDP2 load supports the checkpoint \(.*\), load it in a single process, or replicate "
-    r".*DistributionPlan\(2, 1, 'replicate'\)"
+    "FSDP2 load supports the checkpoint, which takes .*, load it in a single process, or "
+    r"replicate .*DistributionPlan\(2, 1, 'replicate'\)"
 )
 _NEEDS_LATER = (
     "can be expert-parallel, but the plan resolved .* this transformers cannot shard its experts. "

From e70d38d72122791b1d1ebfa9571a6983b195865e Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 21:55:00 +0000
Subject: [PATCH 19/27] State that a spec's prepare may run for a plan that is
 not expert-parallel

The plan decision applies the model's spec, prepare included, before it
knows the plan, wherever the plan could still be expert-parallel, so a
run whose auto EP degree comes out at 1 has run prepare too. That is now
the documented contract of EPSpec.prepare: it must leave a whole-model or
FSDP2 load as it would be without it, and may run more than once.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/models/specs.py                   | 6 +++++-
 modelopt/torch/utils/plugins/model_load_utils.py | 4 +++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py
index 4599813e12a..fba161de33d 100644
--- a/modelopt/torch/models/specs.py
+++ b/modelopt/torch/models/specs.py
@@ -211,7 +211,11 @@ class EPSpec(SpecSection):
     config to load. Reserved for checkpoints carrying their own modeling code, where the work is
     patching classes or installing on-the-fly weight decompression. Runs after the declarative
     fields above have been applied, so a ``prepare`` that returns another config than the one it
-    was handed must carry over the ``base_model_ep_plan`` they attached."""
+    was handed must carry over the ``base_model_ep_plan`` they attached.
+
+    It runs before the run's plan is decided, wherever the plan could still be expert-parallel, so
+    a run whose EP degree comes out at 1 has run it too: it must leave a load of the whole model,
+    on every rank or FSDP2-sharded, as it would be without it, and may run more than once."""
 
     prepared_ep_plan: dict[str, str] | None = None
     """The ``base_model_ep_plan`` ``prepare`` attaches, for a model whose plan only ``prepare``
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 29111fd242a..c162c1f86b1 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1850,7 +1850,9 @@ def resolve_distribution_plan(
     spec to a copy of the config first
     (:func:`~modelopt.torch.utils.moe.apply_distributed_support`): for several models it is what
     supplies the expert-parallel plan, and an expert-parallel plan returns that copy. Any other
-    plan returns the config without it, which the export would otherwise write out. A plan that
+    plan returns the config without it, which the export would otherwise write out; the spec's
+    ``prepare`` has run for it all the same, which its contract allows
+    (:attr:`~modelopt.torch.models.EPSpec.prepare`). A plan that
     leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for exactly
     that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and
     ``MODELOPT_FSDP_SIZE=1``) replicates it, noting that every rank holds the whole MoE.

From c5db496c553962d75e3c4e16c387c57372d5e00c Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 23:08:21 +0000
Subject: [PATCH 20/27] Take a config left unmarked as not yet prepared by the
 model spec

_sharded_load_kwargs took an hf_config as prepared by the model's spec
unless told otherwise, so a caller that left prepared out would skip the
spec for a config that needs it: DeepSeek-V3 below transformers 5.13
would load sharding nothing. Left out, it is now taken as not prepared;
preparing again is harmless by EPSpec.prepare's contract, skipping it is
not.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index c162c1f86b1..4096ae6e4e8 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -652,13 +652,14 @@ def _sharded_load_kwargs(
     hub_options,
     decided: bool,
     measured: bool,
-    prepared: bool = True,
+    prepared: bool = False,
 ) -> dict:
     """What the sharded load ``route`` names takes, refusing what it cannot honor.
 
     ``prepared`` is whether an ``hf_config`` in ``kwargs`` is prepared by the model's spec already:
     one the caller or the plan decision handed in is; the checkpoint's own config read for the load,
-    or a copy of a ``config`` passed for it, is not.
+    or a copy of a ``config`` passed for it, is not. Left out, it is taken as not, as preparing
+    again is harmless (:attr:`~modelopt.torch.models.EPSpec.prepare`) and skipping it is not.
     """
     from modelopt.torch.models.hf import checkpoint_has_mtp
 

From b8a79e972de426976f3fca58a37526f889f35d37 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 23:12:06 +0000
Subject: [PATCH 21/27] Ask for the EP degree where the routing test expects
 the FSDP2 refusal

The plan decision now gives an EP degree it chose way to FSDP2 for the
whole model when the dense part is FSDP2-sharded, so the routing test's
case for the refusal of an expert-parallel plan that FSDP2-shards the
dense part asks for the EP degree, as a caller meeting that refusal has.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/utils/test_model_load_utils.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 7350cfc8ce9..a40367ba216 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -1625,10 +1625,13 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader(
     """Decided automatically, an EP plan loads through load_and_prepare_ep -- with the plan, the
     spec-prepared config and the Hub options the load was given -- and through no other load. It
     returns the model unfrozen, as the other loads do, and an option passed as None asks nothing
-    of it. That load replicates the dense part, so a plan that also FSDP2-shards it is refused
-    before anything loads, naming the ways out. A launch that shows each process one GPU, bound to
-    cuda:0, gives a rank of any local rank a GPU of its own."""
+    of it. That load replicates the dense part, so a plan that asks for expert parallelism and
+    FSDP2-shards the dense part too is refused before anything loads, naming the ways out. A launch
+    that shows each process one GPU, bound to cuda:0, gives a rank of any local rank a GPU of its
+    own."""
     monkeypatch.setenv("MODELOPT_FSDP_SIZE", {"replicate": "1", "fsdp": "world"}[dense])
+    if dense == "fsdp":  # left to the decision, the EP degree would give way to FSDP2
+        monkeypatch.setenv("MODELOPT_EP_DEGREE", "4")
     run_as(4, rank=1 if one_gpu_per_process else 0)
     if one_gpu_per_process:
         monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)

From 5c3258c3cd96f38440e494de9a29380ebc60ed61 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 31 Oct 2026 00:13:24 +0000
Subject: [PATCH 22/27] Say why the spec rules expert parallelism out in the
 left-whole refusal

The plan decision's refusal of an MoE its spec rules out read the spec's
unsupported text to decide, then left it out, so a run that set nothing,
the one most users reach as auto resolves to 1 for these models, learned
only that the model cannot be expert-parallel. It now says why, as the
load's own refusal of the same spec does.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 6 +++---
 tests/unit/torch/utils/test_model_load_utils.py  | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 4096ae6e4e8..eb1a0abec35 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1737,9 +1737,9 @@ def _decide_distribution(
         if ep is not None and ep.unsupported:
             # NotImplementedError, as check_distributed_support refuses the same spec.
             raise NotImplementedError(
-                f"{hf_config.model_type} cannot be quantized expert-parallel yet, and the plan "
-                f"replicates it at world_size={world_size}, so every rank would hold the whole "
-                f"MoE. Instead, {others}"
+                f"{hf_config.model_type} cannot be quantized expert-parallel yet: {ep.unsupported} "
+                f"The plan replicates it at world_size={world_size}, so every rank would hold the "
+                f"whole MoE. Instead, {others}"
             )
         if (
             _spec_refuses_ep(hf_config.model_type)
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index a40367ba216..d147f0cafa8 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -674,7 +674,7 @@ def test_a_run_without_expert_parallelism_needs_no_spec(
 # The refusals of an MoE left whole on every rank, by whether its spec allows expert parallelism.
 _CAN_BE_EP = "can be expert-parallel, but the plan resolved .* divid(?:es|ing) both its {} experts"
 _NO_EP_YET = (
-    "cannot be quantized expert-parallel yet, and the plan replicates .* FSDP2 .* if the "
+    "cannot be quantized expert-parallel yet: .* The plan replicates .* FSDP2 .* if the "
     "FSDP2 load supports the checkpoint, which takes .*, load it in a single process, or "
     r"replicate .*DistributionPlan\(2, 1, 'replicate'\)"
 )

From a8afcfc92fdfe9bb089d6aeecbbd6010f3e16646 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 31 Oct 2026 00:16:58 +0000
Subject: [PATCH 23/27] Expect the spec's floor refusal for an MoE below it,
 left to decide

With the model spec applied, an EP degree left to auto resolves to 1 for
an MoE below the transformers floor its spec states, and the plan
decision refuses to leave it whole, naming the upgrade that would shard
it, before the load's give-way could plan it whole. The give-way test's
below-the-floor case expects that refusal here.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/utils/test_model_load_utils.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index d147f0cafa8..deaf32e448d 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -801,7 +801,9 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it(
     later one the spec states (MiniMax-M3-VL's) -- is planned without expert parallelism rather
     than committed to a load that refuses it, for some reasons only after the download. Nor is it
     refused for being left whole, whether the decision or the caller made its EP degree 1: it
-    cannot be expert-parallel here."""
+    cannot be expert-parallel here. Below the floor its spec states, though, the decision's own
+    refusal names the upgrade, as one would shard it, and an EP degree the caller set to 1 to
+    raise."""
     if ep_degree is not None:
         monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree)
     run_as(2)
@@ -823,6 +825,11 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it(
     serve(config)
     monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None)
 
+    if refused_by == "its-transformers":
+        refusal = _NEEDS_LATER if ep_degree is None else _NEEDS_LATER_AND_A_DEGREE
+        with pytest.raises(RuntimeError, match=refusal):
+            resolve_distribution_plan(str(tmp_path))
+        return
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2)
 
 

From 36760db9a6201591abf1aef099fa1596cd404d15 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 31 Oct 2026 01:29:16 +0000
Subject: [PATCH 24/27] Leave the spec's expert-parallel refusals to the
 spec-aware decision

The plan decision asks the model spec (_spec_refuses_ep) before the
expert-parallel load's refusal: auto resolves to an EP degree of 1 for a
model whose spec refuses expert parallelism or needs a later transformers,
and the left-whole refusal says why. The load refusal's own spec and floor
checks could no longer change a plan, so they are gone, with the floor they
spelled out beside _ep_floor's.

The give-way test's model-spec case patched a spec lookup the decision no
longer reads; GPT-OSS's left-whole case covers that refusal, now with the
MXFP4 checkpoint format its releases have, which the load refuses as well.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   | 20 ++-------
 .../unit/torch/utils/test_model_load_utils.py | 42 ++++++++-----------
 2 files changed, 21 insertions(+), 41 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index eb1a0abec35..70153929051 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -34,7 +34,6 @@
 import torch.nn as nn
 import transformers.distributed as hf_distributed
 from huggingface_hub import snapshot_download
-from packaging.version import parse
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -55,7 +54,6 @@
 from transformers.conversion_mapping import get_model_conversion_mapping
 from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key
 
-from modelopt.torch.models.specs import get_spec
 from modelopt.torch.utils import distributed as dist
 from modelopt.torch.utils import env, moe
 from modelopt.torch.utils.distributed import (
@@ -2002,21 +2000,11 @@ def _ep_plan_refusal(ep_plan: dict[str, str]) -> str | None:
 def _ep_load_refusal(hf_config, ep_plan: dict[str, str]) -> Exception | None:
     """The error the expert-parallel load refuses this model with, by its config, or ``None``.
 
-    A model spec that rules it out (``NotImplementedError``), the load's transformers floor or the
-    spec's later one (``RuntimeError``), a checkpoint format the load does not take
-    (``NotImplementedError``), and the plan's styles (``ValueError``); the rest shows only as the
-    experts load.
+    A checkpoint format the load does not take (``NotImplementedError``) and the plan's styles
+    (``ValueError``); the rest shows only as the experts load. A model spec ruling expert
+    parallelism out, or the transformers release it needs, is not asked here: the plan decision asks
+    the spec first (``_spec_refuses_ep``).
     """
-    spec = get_spec(hf_config.model_type)
-    ep = spec.ep_spec if spec is not None else None
-    if ep is not None and ep.unsupported:
-        return NotImplementedError(ep.unsupported)
-    stated = ep.requires_transformers if ep is not None else None
-    floor = max(_EP_MIN_TRANSFORMERS, stated or _EP_MIN_TRANSFORMERS, key=parse)
-    try:
-        moe.require_transformers(floor, f"{hf_config.model_type}'s expert-parallel load")
-    except RuntimeError as e:
-        return e
     if (method := _ckpt_quant_method(hf_config)) not in _EP_SOURCE_FORMATS:
         return NotImplementedError(
             f"The expert-parallel load cannot dequantize a {method!r} checkpoint."
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index deaf32e448d..213dbe9d453 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -75,7 +75,7 @@
 import modelopt.torch.quantization as mtq
 import modelopt.torch.utils.distributed as dist
 import modelopt.torch.utils.moe as moe
-from modelopt.torch.models import EPSpec, ModelSpec, get_spec
+from modelopt.torch.models import get_spec
 from modelopt.torch.quantization.plugins import accelerate as accelerate_plugin
 from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights
 from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
@@ -711,7 +711,10 @@ def test_a_run_without_expert_parallelism_needs_no_spec(
             id="explicit_degree_one_spec_plan",
         ),
         pytest.param(
-            GptOssConfig(num_local_experts=4, **_TINY),
+            # MXFP4, as GPT-OSS's checkpoints are, which the expert-parallel load refuses too.
+            GptOssConfig(
+                num_local_experts=4, quantization_config={"quant_method": "mxfp4"}, **_TINY
+            ),
             None,
             _NO_EP_YET,
             None,
@@ -751,8 +754,9 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error(
     degree of 1 skips (DeepSeek-V3). With nothing sharded either, every rank would hold the whole
     MoE: that must fail, not run, and name the expert count an EP degree has to divide. GPT-OSS
     carries a plan too, but its spec refuses expert parallelism, so auto picks no EP degree for it
-    and this check refuses it as an MoE that cannot be expert-parallel yet, offering every other
-    way out (FSDP2 only for an unquantized causal LM); the spec's own refusal, a
+    and this check refuses it as an MoE that cannot be expert-parallel yet, whatever else the
+    expert-parallel load refuses it for, offering every other way out (FSDP2 only for an
+    unquantized causal LM); the spec's own refusal, a
     NotImplementedError and so a RuntimeError too, words it differently. Below the transformers
     release expert parallelism needs, no EP degree is offered, since none would load: upgrading
     is, where an EP degree divides the experts; where none does, upgrading would not shard them.
@@ -788,22 +792,18 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error(
 
 @needs_qwen3_moe_ep_plan
 @pytest.mark.parametrize("ep_degree", [None, "1"], ids=["auto", "set-to-1"])
-@pytest.mark.parametrize(
-    "refused_by",
-    ["its-plan", "its-format", "its-spec", "its-transformers", "its-specs-transformers"],
-)
+@pytest.mark.parametrize("refused_by", ["its-plan", "its-format", "its-transformers"])
 def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it(
     run_as, serve, monkeypatch, tmp_path, refused_by, ep_degree
 ):
-    """Left to the decision, an MoE the expert-parallel load would refuse -- by a plan that
-    shards no expert weights through grouped_gemm, a checkpoint format it does not load, a model
-    spec ruling it out (Llama-4's, GPT-OSS's), or a transformers below its floor or below the
-    later one the spec states (MiniMax-M3-VL's) -- is planned without expert parallelism rather
-    than committed to a load that refuses it, for some reasons only after the download. Nor is it
-    refused for being left whole, whether the decision or the caller made its EP degree 1: it
-    cannot be expert-parallel here. Below the floor its spec states, though, the decision's own
-    refusal names the upgrade, as one would shard it, and an EP degree the caller set to 1 to
-    raise."""
+    """Left to the decision, an MoE the expert-parallel load would refuse -- by a plan that shards no
+    expert weights through grouped_gemm or a checkpoint format it does not load -- is planned
+    without expert parallelism rather than committed to a load that refuses it, for some reasons
+    only after the download. Nor is it refused for being left whole, whether the decision or the
+    caller made its EP degree 1: it cannot be expert-parallel here. Below the floor its spec
+    states, though, the decision's own refusal names the upgrade, as one would shard it, and an EP
+    degree the caller set to 1 to raise; a spec ruling expert parallelism out has a refusal of its
+    own (GPT-OSS's, in test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error)."""
     if ep_degree is not None:
         monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree)
     run_as(2)
@@ -812,14 +812,6 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it(
         config.base_model_ep_plan = {"layers.*.mlp.experts": "moe_tp_experts"}
     elif refused_by == "its-format":
         config.quantization_config = {"quant_method": "fp8"}
-    elif refused_by == "its-spec":
-        spec = ModelSpec(model_type="qwen3_moe", ep_spec=EPSpec(unsupported="Not yet."))
-        monkeypatch.setattr(model_load_utils, "get_spec", lambda model_type: spec)
-    elif refused_by == "its-specs-transformers":
-        # Above the load's own 5.7 by version, as 5.10 is not, though not as strings compare.
-        spec = ModelSpec(model_type="qwen3_moe", ep_spec=EPSpec(requires_transformers="5.13"))
-        monkeypatch.setattr(model_load_utils, "get_spec", lambda model_type: spec)
-        monkeypatch.setattr(transformers, "__version__", "5.10.0")
     else:
         monkeypatch.setattr(transformers, "__version__", "5.6.2")
     serve(config)

From 6c21c77f2f072766489e5a2664917aecb1349ee8 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 31 Oct 2026 03:57:52 +0000
Subject: [PATCH 25/27] Refuse a model its spec refuses in the expert-parallel
 load itself

and floor checks this PR leaves to the spec-aware decision. A plan handed
to the load reaches it without that decision, so the load asks the spec
itself, as load_pretrained's expert-parallel route does, and refuses with
the same NotImplementedError or RuntimeError.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 5 +++--
 tests/unit/torch/utils/test_model_load_utils.py  | 7 +++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 70153929051..835b1180253 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -2180,9 +2180,10 @@ def load_and_prepare_ep(
         ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {})
         if not ep_plan:
             raise ValueError("expert-parallel PTQ requires the model to declare base_model_ep_plan")
-        # The refusals the plan decision gives way for, which an EP degree asked for, or a plan
+        # What the plan decision refuses or gives way for, which an EP degree asked for, or a plan
         # handed in, meets here: a model spec ruling expert parallelism out or stating a later
-        # transformers, and the plan's styles.
+        # transformers (refused as load_pretrained refuses it), and the plan's styles.
+        check_distributed_support(hf_config.model_type)
         if (refusal := _ep_load_refusal(hf_config, ep_plan)) is not None:
             raise refusal
         # The all-to-all transport needs every rank to hold as many experts, counted as the plan's
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 213dbe9d453..9926d736396 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -2709,7 +2709,7 @@ def resolve(path, **hub_options):
         (
             SimpleNamespace(model_type="gpt_oss", base_model_ep_plan=_EP_PLAN),
             NotImplementedError,
-            "the expert-parallel export cannot write that layout",
+            "gpt_oss cannot be quantized expert-parallel yet: .* cannot write that layout",
         ),
     ],
     ids=[
@@ -2733,9 +2733,8 @@ def test_a_config_the_load_cannot_take_is_refused_before_anything_loads(
     that leaves the experts whole on every rank: one naming no routed experts, or none of
     whose entries shards their weights by grouped_gemm, the only expert-weight style the load
     shards by, or one that styles some of their weights otherwise. So is a model whose spec rules
-    expert parallelism out (GPT-OSS's), which the plan decision plans without it: an EP degree asked
-    for, or a plan handed in, meets that refusal here. The load refuses each before it reads
-    anything."""
+    expert parallelism out (GPT-OSS's), as load_pretrained refuses it: a plan handed to this load
+    meets that refusal here. The load refuses each before it reads anything."""
     with pytest.raises(error, match=match):
         load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config)
     assert ep_world_of_two == []

From 4258beae36970c7cb1e8e7e4dafb1c7c51356951 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 31 Oct 2026 03:58:09 +0000
Subject: [PATCH 26/27] Read the checkpoint's config.json as UTF-8

transformers writes config.json as UTF-8. Read in the locale's encoding,
a non-ASCII byte under a non-UTF-8 locale raised UnicodeDecodeError, which
the plan-provenance read took for an unreadable file and dropped the plan
the checkpoint declares.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 835b1180253..a0176728240 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -2050,7 +2050,7 @@ def _declared_ep_plan(ckpt_dir: str, hf_config, ep_cfg):
         "from it writes none"
     )
     try:
-        with open(path) as f:
+        with open(path, encoding="utf-8") as f:
             source = json.load(f)
     except (OSError, ValueError) as e:
         warnings.warn(f"Cannot read {path} ({e}), {keeps_none}.")

From 997f9a6305415883b27c81ed15c310bd2d3e9c60 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 31 Oct 2026 04:29:46 +0000
Subject: [PATCH 27/27] Ask the model spec first in the expert-parallel load

The load asked the spec only past its format check and its no-plan guard,
so a direct call for a model whose plan only its spec supplies, or whose
MXFP4 checkpoint the load refuses as well (GPT-OSS's), reported those
instead of the spec's reason, whose remedy is the one that loads. It now
asks the spec as soon as it has the config, as load_pretrained does.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 13 +++++++------
 tests/unit/torch/utils/test_model_load_utils.py  | 12 +++++++-----
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index a0176728240..122455daaa0 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -2002,8 +2002,8 @@ def _ep_load_refusal(hf_config, ep_plan: dict[str, str]) -> Exception | None:
 
     A checkpoint format the load does not take (``NotImplementedError``) and the plan's styles
     (``ValueError``); the rest shows only as the experts load. A model spec ruling expert
-    parallelism out, or the transformers release it needs, is not asked here: the plan decision asks
-    the spec first (``_spec_refuses_ep``).
+    parallelism out, or the transformers release it needs, is not asked here: the plan decision and
+    the load ask the spec first (:func:`~modelopt.torch.utils.moe.check_distributed_support`).
     """
     if (method := _ckpt_quant_method(hf_config)) not in _EP_SOURCE_FORMATS:
         return NotImplementedError(
@@ -2163,6 +2163,10 @@ def load_and_prepare_ep(
             if attn_implementation is not None:
                 config_kwargs["attn_implementation"] = attn_implementation
             hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs)
+        # A model spec ruling expert parallelism out, or stating a later transformers, refuses it
+        # first, as load_pretrained does, for an EP degree asked for or a plan handed in: no other
+        # checkpoint of the model, nor a plan of its own, would load.
+        check_distributed_support(hf_config.model_type)
         quant_method = _ckpt_quant_method(hf_config)
         if quant_method not in _EP_SOURCE_FORMATS:
             # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights.
@@ -2180,10 +2184,7 @@ def load_and_prepare_ep(
         ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {})
         if not ep_plan:
             raise ValueError("expert-parallel PTQ requires the model to declare base_model_ep_plan")
-        # What the plan decision refuses or gives way for, which an EP degree asked for, or a plan
-        # handed in, meets here: a model spec ruling expert parallelism out or stating a later
-        # transformers (refused as load_pretrained refuses it), and the plan's styles.
-        check_distributed_support(hf_config.model_type)
+        # The plan's styles, which the plan decision gives way for too.
         if (refusal := _ep_load_refusal(hf_config, ep_plan)) is not None:
             raise refusal
         # The all-to-all transport needs every rank to hold as many experts, counted as the plan's
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 9926d736396..63894086116 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -2673,13 +2673,13 @@ def resolve(path, **hub_options):
     [
         *(
             (
-                SimpleNamespace(quantization_config={"quant_method": method}),
+                SimpleNamespace(model_type=_NO_SPEC, quantization_config={"quant_method": method}),
                 NotImplementedError,
                 f"cannot dequantize a '{method}'",
             )
             for method in ("gptq", "modelopt", "compressed-tensors", "mxfp4")
         ),
-        (SimpleNamespace(), ValueError, "declare base_model_ep_plan"),
+        (SimpleNamespace(model_type=_NO_SPEC), ValueError, "declare base_model_ep_plan"),
         (
             SimpleNamespace(
                 model_type=_NO_SPEC, base_model_ep_plan={"layers.*.self_attn.q_proj": "colwise"}
@@ -2707,7 +2707,8 @@ def resolve(path, **hub_options):
             r"styles layers\.\*\.mlp\.experts\.down_proj by a style this load drops",
         ),
         (
-            SimpleNamespace(model_type="gpt_oss", base_model_ep_plan=_EP_PLAN),
+            # An MXFP4 checkpoint with no plan, refused by its spec before either.
+            SimpleNamespace(model_type="gpt_oss", quantization_config={"quant_method": "mxfp4"}),
             NotImplementedError,
             "gpt_oss cannot be quantized expert-parallel yet: .* cannot write that layout",
         ),
@@ -2733,8 +2734,9 @@ def test_a_config_the_load_cannot_take_is_refused_before_anything_loads(
     that leaves the experts whole on every rank: one naming no routed experts, or none of
     whose entries shards their weights by grouped_gemm, the only expert-weight style the load
     shards by, or one that styles some of their weights otherwise. So is a model whose spec rules
-    expert parallelism out (GPT-OSS's), as load_pretrained refuses it: a plan handed to this load
-    meets that refusal here. The load refuses each before it reads anything."""
+    expert parallelism out (GPT-OSS's), as load_pretrained refuses it, before its checkpoint's
+    format or plan is looked at: a plan handed to this load meets that refusal here. The load
+    refuses each before it reads anything."""
     with pytest.raises(error, match=match):
         load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config)
     assert ep_world_of_two == []