From c22907ba5cb708e8d7d8cf4610d9c6e0b3867513 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 21:22:14 +0000
Subject: [PATCH 01/44] Decide the expert-parallel plan from the model and the
 knobs

decide_distribution_plan turns a model's config and the world size into a
DistributionPlan. Each degree is an explicit value or "auto", and explicit
always wins.

- Capability: a model is expert-parallel capable when transformers ships a
  base_model_ep_plan for it, found on the root config or a VLM's nested text
  config, and it has experts under any of the spellings configs use.
- EP degree: it must divide both the world and the expert count, so every
  rank holds whole experts; auto takes the largest such degree.
- Dense: replicated when it fits one GPU with headroom, else FSDP2-sharded
  across the world.

The decision runs on rank 0 alone and is broadcast; the other ranks adopt
it. Its inputs are read locally -- the knobs, this GPU's memory, the config
on this rank's filesystem -- so ranks deciding independently could disagree,
and a disagreement does not raise: they build meshes of different shapes and
hang at mesh creation or the first collective. If rank 0 fails to decide, it
still reaches the broadcast and sends the failure, so every rank raises
instead of waiting.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/moe.py                   | 222 +++++++++++++++++-
 .../unit/torch/utils/test_ep_distribution.py  | 196 ++++++++++++++++
 2 files changed, 417 insertions(+), 1 deletion(-)
 create mode 100644 tests/unit/torch/utils/test_ep_distribution.py

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index fc62fe9f71c..d922ee54603 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -15,7 +15,8 @@
 
 """HuggingFace-coupled MoE expert-parallel (EP) helpers.
 
-The DP=EP all-to-all expert transport (``torch_a2a``).
+Capability detection and EP degree selection: deciding how a model distributes, from the model,
+this GPU's memory and the knobs.
 
 The all-to-all transport is plain ``torch.distributed.all_to_all_single`` (NCCL) wrapped around
 transformers' ``_grouped_linear``. ``register_torch_a2a()`` registers it into
@@ -29,6 +30,225 @@
 import torch.distributed as dist
 from torch.distributed.nn.functional import all_to_all_single as autograd_all_to_all_single
 
+
+def require_transformers(min_version: str, feature: str) -> None:
+    """Fail loudly if the installed transformers is older than ``min_version`` (the floor for ``feature``).
+
+    Callers enforce the version where the needed APIs first exist, so a too-old transformers errors
+    clearly instead of silently mis-behaving (e.g. transformers 5.5 silently shards nothing for EP).
+    See each call site for the empirically-determined floor.
+    """
+    import transformers
+    from packaging.version import parse
+
+    have = transformers.__version__
+    if parse(have) < parse(min_version):
+        raise RuntimeError(
+            f"{feature} requires transformers >= {min_version}, but {have} is installed. "
+            f"Upgrade with: pip install -U 'transformers>={min_version}'."
+        )
+
+
+def _ep_config(hf_config):
+    """Return the config that carries the EP plan / expert count.
+
+    For a plain MoE causal LM this is the root config. For a VLM whose nested text submodel
+    is the MoE (e.g. ``KimiK25.text_config`` -> deepseek_v3, or ``NemotronH_Omni.llm_config``
+    -> nemotron_h), this is that nested config.
+    """
+    if getattr(hf_config, "base_model_ep_plan", None):
+        return hf_config
+    for _attr in ("text_config", "llm_config"):
+        sub = getattr(hf_config, _attr, None)
+        if sub is not None and getattr(sub, "base_model_ep_plan", None):
+            return sub
+    # Fall through: no plan anywhere; let callers decide.
+    return hf_config
+
+
+def num_experts_from_config(hf_config) -> int:
+    """Return the MoE expert count from an HF config (0 if the model is not MoE).
+
+    Recurses into ``.text_config`` / ``.llm_config`` so VLMs whose nested text submodel is the MoE
+    (Kimi-K2.7, Llama-4 multimodal, NemotronH-Omni, ...) report the LM's expert count.
+    """
+    for cfg in (
+        hf_config,
+        getattr(hf_config, "text_config", None),
+        getattr(hf_config, "llm_config", None),
+    ):
+        if cfg is None:
+            continue
+        for attr in ("num_experts", "num_local_experts", "n_routed_experts", "num_routed_experts"):
+            value = getattr(cfg, attr, None)
+            if value:
+                return int(value)
+    return 0
+
+
+def is_expert_parallel_capable(hf_config) -> bool:
+    """True if the model can be expert-parallel sharded.
+
+    Transformers exposes ``base_model_ep_plan`` on MoE configs whose experts support
+    expert parallelism (e.g. gpt_oss, llama4; qwen3_moe in transformers>=5.7). Dense models
+    and MoE configs without the plan return False -> use plain FSDP2. For VLMs the plan
+    lives on ``text_config`` (the nested LM submodel), so check both.
+    """
+    return (
+        bool(getattr(_ep_config(hf_config), "base_model_ep_plan", None))
+        and num_experts_from_config(hf_config) > 0
+    )
+
+
+def choose_ep_dp_degrees(world_size: int, num_experts: int) -> tuple[int, int]:
+    """Pick ``(expert_parallel, expert_data_parallel)`` degrees with ``EP * EDP == world_size``.
+
+    EP must divide both ``num_experts`` (whole experts per rank) and ``world_size`` (a valid mesh
+    factorization); we take the largest such EP, and each expert shard is then replicated over the
+    remaining ``EDP = world_size // EP`` ranks. Falls back to ``(1, world_size)`` when the model is not
+    MoE or no valid EP degree exists. Heuristic -- callers may override.
+    """
+    if num_experts <= 0 or world_size <= 1:
+        return 1, world_size
+    for ep in range(world_size, 0, -1):
+        if world_size % ep == 0 and num_experts % ep == 0:
+            return ep, world_size // ep
+    return 1, world_size
+
+
+def _estimate_dense_param_bytes(hf_config) -> "int | None":
+    """Best-effort estimate of the DENSE (non-expert) parameter bytes from an HF config.
+
+    embeddings + lm_head + per-layer attention/norms (+ the dense FFN when the model is NOT MoE, since a MoE's FFN is
+    the expert weights). Returns ``None`` when the config lacks the fields to estimate (the caller then treats the dense
+    part as 'fits'). A deliberately CONSERVATIVE upper bound (ignores GQA and tied embeddings) -- it exists to catch a
+    dense part that clearly cannot be replicated, not to be exact.
+    """
+    cfg = getattr(hf_config, "text_config", None) or hf_config
+    h = getattr(cfg, "hidden_size", None)
+    n_layers = getattr(cfg, "num_hidden_layers", None)
+    if not h or not n_layers:
+        return None
+    vocab = getattr(cfg, "vocab_size", None) or getattr(hf_config, "vocab_size", None) or 0
+    inter = getattr(cfg, "intermediate_size", None) or 4 * h
+    attn = 4 * h * h  # q,k,v,o (upper bound; ignores GQA head reduction)
+    norms = 4 * h  # ~2 layernorms per layer
+    ffn = (
+        0 if is_expert_parallel_capable(hf_config) else 3 * h * inter
+    )  # gated FFN; a MoE's FFN is expert
+    per_layer = attn + norms + ffn
+    embed = 2 * vocab * h  # embed + lm_head (untied upper bound)
+    bytes_per = 2  # bf16/fp16 default (the pre-quant load dtype)
+    return (embed + n_layers * per_layer) * bytes_per
+
+
+def _dense_fits_per_gpu(
+    hf_config, device_mem_bytes: "int | None" = None, headroom: float = 0.6
+) -> bool:
+    """Fit predicate for the dense knob.
+
+    does the DENSE footprint fit ONE GPU (leaving ``headroom`` for activations + comm buffers + any EP-local experts)?
+    CONSERVATIVE -- returns ``True`` ('fits' -> replicate / data-parallel) whenever it cannot estimate the dense size or
+    the device memory, so the data-parallel default is preserved. This is the pluggable decision point of the auto-
+    resolver; a heuristic or a recipe hint may substitute for it without changing callers.
+    """
+    dense_bytes = _estimate_dense_param_bytes(hf_config)
+    if not dense_bytes:
+        return True
+    if device_mem_bytes is None:
+        try:
+            import torch
+
+            device_mem_bytes = torch.cuda.get_device_properties(
+                torch.cuda.current_device()
+            ).total_memory
+        except Exception:
+            return True
+    return dense_bytes < headroom * device_mem_bytes
+
+
+def _resolve_ep_size(hf_config, world_size: int, ep_size) -> int:
+    """Resolve the expert-parallel degree.
+
+    An explicit ``ep_size`` wins; ``auto`` -> the largest EP degree dividing both the expert count and
+    ``world_size`` for an EP-capable model, else ``1``.
+    """
+    if ep_size not in (None, "auto", "0", 0):
+        return int(ep_size)
+    if world_size > 1 and is_expert_parallel_capable(hf_config):
+        ep, _ = choose_ep_dp_degrees(world_size, num_experts_from_config(hf_config))
+        return ep
+    return 1
+
+
+def decide_distribution_plan(
+    hf_config,
+    world_size: int,
+    *,
+    ep_size="auto",
+    fsdp_size="auto",
+    device_mem_bytes: "int | None" = None,
+):
+    """Decide how a HuggingFace MoE model is distributed for PTQ: a :class:`DistributionPlan`.
+
+    Two knobs, each an explicit value or ``"auto"`` (explicit always wins):
+
+    * ``ep_size`` -- ``auto`` -> the largest EP degree dividing both the expert count and
+      ``world_size`` (``1`` for a model that cannot be expert-parallel).
+    * ``fsdp_size`` -- ``auto`` -> ``1`` (replicate the dense part) when it fits one GPU, else
+      ``world_size`` (FSDP2-shard it); ``"world"`` -> ``world_size``.
+
+    On a distributed run the plan is decided ONCE, on rank 0, and broadcast; every other rank adopts rank 0's
+    decision instead of making its own. The inputs are all read locally -- the env knobs, this GPU's memory,
+    the config on this rank's filesystem -- so ranks deciding independently could disagree, and a
+    disagreement surfaces as a hang at mesh creation rather than as an error. Collective: every rank must
+    call it. If rank 0 fails to decide, every rank raises.
+    """
+    from .distributed import is_distributed_run
+
+    decide_kwargs = {
+        "ep_size": ep_size,
+        "fsdp_size": fsdp_size,
+        "device_mem_bytes": device_mem_bytes,
+    }
+    if not is_distributed_run():
+        return _decide_distribution_plan_on_rank0(hf_config, world_size, **decide_kwargs)
+
+    plan, error = None, None
+    if dist.get_rank() == 0:
+        try:
+            plan = _decide_distribution_plan_on_rank0(hf_config, world_size, **decide_kwargs)
+        except Exception as e:
+            error = e
+    # Ship the failure as text, not the exception object: the other ranks only need to stop, and an
+    # exception is not guaranteed to pickle.
+    decision = [(plan, None if error is None else f"{type(error).__name__}: {error}")]
+    dist.broadcast_object_list(decision, src=0)
+    if error is not None:
+        raise error
+    plan, rank0_error = decision[0]
+    if rank0_error is not None:
+        raise RuntimeError(f"rank 0 failed to decide the distribution plan: {rank0_error}")
+    return plan
+
+
+def _decide_distribution_plan_on_rank0(
+    hf_config, world_size: int, *, ep_size, fsdp_size, device_mem_bytes: "int | None"
+):
+    """The decision itself (see :func:`decide_distribution_plan`); run on rank 0 alone when distributed."""
+    from .distributed import DistributionPlan
+
+    ep = _resolve_ep_size(hf_config, world_size, ep_size)
+    if fsdp_size in (None, "auto"):
+        fsdp = 1 if _dense_fits_per_gpu(hf_config, device_mem_bytes) else world_size
+    elif str(fsdp_size).lower() in ("world", "full"):
+        fsdp = world_size
+    else:
+        fsdp = int(fsdp_size)
+    fsdp = max(1, min(fsdp, world_size))
+    return DistributionPlan(world_size=world_size, ep_size=ep, fsdp_size=fsdp)
+
+
 # ---------------------------------------------------------------------------------------------------
 # DP=EP all-to-all expert transport (`torch_a2a`)
 #
diff --git a/tests/unit/torch/utils/test_ep_distribution.py b/tests/unit/torch/utils/test_ep_distribution.py
new file mode 100644
index 00000000000..38ac661db80
--- /dev/null
+++ b/tests/unit/torch/utils/test_ep_distribution.py
@@ -0,0 +1,196 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for EP degree selection and the decided distribution plan."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch.distributed as dist
+
+import modelopt.torch.utils.distributed as mto_dist
+import modelopt.torch.utils.moe as moe
+from modelopt.torch.utils.moe import (
+    _resolve_ep_size,
+    choose_ep_dp_degrees,
+    decide_distribution_plan,
+    is_expert_parallel_capable,
+    num_experts_from_config,
+)
+
+
+def cfg(**kw):
+    return SimpleNamespace(**kw)
+
+
+def moe_cfg(experts=64, **kw):
+    """An EP-capable MoE config: transformers advertises EP support via ``base_model_ep_plan``."""
+    return cfg(num_experts=experts, base_model_ep_plan={"experts": "grouped"}, **kw)
+
+
+# ------------------------------------------------------------------------------- expert counting
+
+
+def test_dense_config_reports_no_experts():
+    assert num_experts_from_config(cfg(hidden_size=128)) == 0
+
+
+@pytest.mark.parametrize(
+    "attr", ["num_experts", "num_local_experts", "n_routed_experts", "num_routed_experts"]
+)
+def test_every_spelling_of_the_expert_count_is_understood(attr):
+    """Architectures disagree on the field name; missing one would silently disable EP."""
+    assert num_experts_from_config(cfg(**{attr: 16})) == 16
+
+
+@pytest.mark.parametrize("nest", ["text_config", "llm_config"])
+def test_expert_count_is_found_on_a_vlms_nested_text_model(nest):
+    """For a VLM the MoE is the nested LM, so the root config has no expert count of its own."""
+    assert num_experts_from_config(cfg(**{nest: cfg(n_routed_experts=32)})) == 32
+
+
+# ----------------------------------------------------------------------------------- capability
+
+
+def test_moe_with_an_ep_plan_is_capable():
+    assert is_expert_parallel_capable(moe_cfg())
+
+
+def test_moe_without_an_ep_plan_is_not_capable():
+    """No ``base_model_ep_plan`` means transformers cannot shard these experts -> use plain FSDP2."""
+    assert not is_expert_parallel_capable(cfg(num_experts=64))
+
+
+def test_ep_plan_without_experts_is_not_capable():
+    assert not is_expert_parallel_capable(cfg(base_model_ep_plan={"experts": "grouped"}))
+
+
+def test_capability_is_detected_through_a_vlm_nesting():
+    assert is_expert_parallel_capable(cfg(text_config=moe_cfg(experts=8)))
+
+
+# ------------------------------------------------------------------------------- degree selection
+
+
+def test_largest_valid_ep_degree_is_chosen():
+    assert choose_ep_dp_degrees(8, 64) == (8, 1)
+
+
+def test_ep_must_divide_the_expert_count_so_ranks_hold_whole_experts():
+    assert choose_ep_dp_degrees(8, 6) == (2, 4)
+
+
+def test_ep_falls_back_to_one_when_nothing_divides():
+    assert choose_ep_dp_degrees(8, 5) == (1, 8)
+
+
+@pytest.mark.parametrize(("world", "experts"), [(8, 0), (1, 64)])
+def test_no_expert_parallelism_for_dense_or_single_rank(world, experts):
+    assert choose_ep_dp_degrees(world, experts) == (1, world)
+
+
+@pytest.mark.parametrize("world", [2, 4, 8, 16])
+def test_degrees_always_factor_the_world(world):
+    ep, dp = choose_ep_dp_degrees(world, 64)
+    assert ep * dp == world
+
+
+def test_an_explicit_ep_size_overrides_auto_resolution():
+    assert _resolve_ep_size(moe_cfg(), world_size=8, ep_size=2) == 2
+
+
+def test_auto_ep_is_one_for_a_model_that_cannot_shard_experts():
+    assert _resolve_ep_size(cfg(num_experts=64), world_size=8, ep_size="auto") == 1
+
+
+# ------------------------------------------------------------------------------- decided plan
+
+
+def test_auto_plan_takes_the_largest_ep_degree_dividing_the_experts():
+    plan = decide_distribution_plan(moe_cfg(experts=6), world_size=8)
+    assert plan.ep_size == 2
+    assert plan.edp == 4
+
+
+def test_an_ep_degree_that_does_not_divide_the_world_is_rejected():
+    with pytest.raises(ValueError, match="must divide world_size"):
+        decide_distribution_plan(moe_cfg(), world_size=8, ep_size=3)
+
+
+@pytest.mark.parametrize(("knob", "expected"), [("world", 8), ("full", 8), (4, 4), (99, 8), (0, 1)])
+def test_dense_fsdp_knob_is_honoured_and_clamped_to_the_world(knob, expected):
+    plan = decide_distribution_plan(moe_cfg(), world_size=8, fsdp_size=knob)
+    assert plan.fsdp_size == expected
+
+
+def test_unknown_dense_size_defaults_to_replication():
+    """With no way to estimate the dense footprint, stay on the cheaper replicate path."""
+    assert decide_distribution_plan(moe_cfg(), world_size=8).fsdp_size == 1
+
+
+# ------------------------------------------------------------------------------- rank-0 decision
+
+
+def _as_rank(monkeypatch, rank, rank0_decision=None):
+    """Pretend to be ``rank`` of a multi-rank run; returns what this rank handed the broadcast.
+
+    The fake broadcast delivers ``rank0_decision`` to a non-zero rank, as rank 0 would have sent it.
+    """
+    sent = []
+
+    def fake_broadcast(objs, src):
+        assert src == 0
+        sent.append(objs[0])
+        if rank != 0:
+            objs[0] = rank0_decision
+
+    monkeypatch.setattr(mto_dist, "is_distributed_run", lambda: True)
+    monkeypatch.setattr(dist, "get_rank", lambda group=None: rank)
+    monkeypatch.setattr(dist, "broadcast_object_list", fake_broadcast)
+    return sent
+
+
+def test_rank0_decides_and_broadcasts_its_plan(monkeypatch):
+    sent = _as_rank(monkeypatch, 0)
+    plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=4)
+    assert sent == [(plan, None)]
+    assert plan.ep_size == 4
+
+
+def test_other_ranks_adopt_rank0s_plan_without_deciding(monkeypatch):
+    """A non-zero rank's own inputs (here an EP knob of 2) must not leak into the plan."""
+    rank0_plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=4)
+    _as_rank(monkeypatch, 1, rank0_decision=(rank0_plan, None))
+
+    def must_not_decide(*args, **kwargs):
+        raise AssertionError("a non-zero rank decided its own plan")
+
+    monkeypatch.setattr(moe, "_decide_distribution_plan_on_rank0", must_not_decide)
+    assert decide_distribution_plan(moe_cfg(), world_size=8, ep_size=2) is rank0_plan
+
+
+def test_rank0_failure_is_broadcast_before_it_raises(monkeypatch):
+    """Rank 0 must still reach the broadcast on failure, or every other rank hangs in it."""
+    sent = _as_rank(monkeypatch, 0)
+    with pytest.raises(ValueError, match="must divide world_size"):
+        decide_distribution_plan(moe_cfg(), world_size=8, ep_size=3)
+    assert sent[0][0] is None
+    assert "ValueError" in sent[0][1]
+
+
+def test_other_ranks_raise_when_rank0_failed(monkeypatch):
+    _as_rank(monkeypatch, 1, rank0_decision=(None, "ValueError: bad degrees"))
+    with pytest.raises(RuntimeError, match=r"rank 0 failed to decide.*bad degrees"):
+        decide_distribution_plan(moe_cfg(), world_size=8)

From e2037f88eb7181a488c24368f028ae7704b35960 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 13 Oct 2026 17:58:48 +0000
Subject: [PATCH 02/44] Decide the plan through the session

decide_distribution_plan hand-rolled the rank-0 decision: broadcast the
plan, or rank 0's error so the other ranks raise instead of hanging.
Inside a ModelOptSession it now goes through session.decide, which does
the same and keeps the result as session.plan, the same on every rank.
Without a session the process decides alone, as a single process always
did.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/moe.py                   | 53 ++++++--------
 .../unit/torch/utils/test_ep_distribution.py  | 73 ++++++++++++++-----
 2 files changed, 76 insertions(+), 50 deletions(-)

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index d922ee54603..5c7c17dfa15 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -198,38 +198,29 @@ def decide_distribution_plan(
     * ``fsdp_size`` -- ``auto`` -> ``1`` (replicate the dense part) when it fits one GPU, else
       ``world_size`` (FSDP2-shard it); ``"world"`` -> ``world_size``.
 
-    On a distributed run the plan is decided ONCE, on rank 0, and broadcast; every other rank adopts rank 0's
-    decision instead of making its own. The inputs are all read locally -- the env knobs, this GPU's memory,
-    the config on this rank's filesystem -- so ranks deciding independently could disagree, and a
-    disagreement surfaces as a hang at mesh creation rather than as an error. Collective: every rank must
-    call it. If rank 0 fails to decide, every rank raises.
+    Inside a :class:`~modelopt.torch.utils.distributed.ModelOptSession` the session decides it and
+    keeps it as ``session.plan``: on a distributed run once, on rank 0, and every other rank adopts
+    rank 0's decision instead of making its own (:meth:`~modelopt.torch.utils.distributed.ModelOptSession.decide`).
+    The inputs are all read locally -- the env knobs, this GPU's memory, the config on this rank's
+    filesystem -- so ranks deciding independently could disagree, and a disagreement surfaces as a
+    hang at mesh creation rather than as an error. Collective: every rank must call it. If rank 0
+    fails to decide, every rank raises. Without a session, this process decides alone.
     """
-    from .distributed import is_distributed_run
-
-    decide_kwargs = {
-        "ep_size": ep_size,
-        "fsdp_size": fsdp_size,
-        "device_mem_bytes": device_mem_bytes,
-    }
-    if not is_distributed_run():
-        return _decide_distribution_plan_on_rank0(hf_config, world_size, **decide_kwargs)
-
-    plan, error = None, None
-    if dist.get_rank() == 0:
-        try:
-            plan = _decide_distribution_plan_on_rank0(hf_config, world_size, **decide_kwargs)
-        except Exception as e:
-            error = e
-    # Ship the failure as text, not the exception object: the other ranks only need to stop, and an
-    # exception is not guaranteed to pickle.
-    decision = [(plan, None if error is None else f"{type(error).__name__}: {error}")]
-    dist.broadcast_object_list(decision, src=0)
-    if error is not None:
-        raise error
-    plan, rank0_error = decision[0]
-    if rank0_error is not None:
-        raise RuntimeError(f"rank 0 failed to decide the distribution plan: {rank0_error}")
-    return plan
+    from .distributed import get_active_mo_session
+
+    def decide():
+        return _decide_distribution_plan_on_rank0(
+            hf_config,
+            world_size,
+            ep_size=ep_size,
+            fsdp_size=fsdp_size,
+            device_mem_bytes=device_mem_bytes,
+        )
+
+    session = get_active_mo_session(required=False)
+    if not session.initialized:
+        return decide()
+    return session.decide("plan", decide)
 
 
 def _decide_distribution_plan_on_rank0(
diff --git a/tests/unit/torch/utils/test_ep_distribution.py b/tests/unit/torch/utils/test_ep_distribution.py
index 38ac661db80..87adc3122f8 100644
--- a/tests/unit/torch/utils/test_ep_distribution.py
+++ b/tests/unit/torch/utils/test_ep_distribution.py
@@ -20,8 +20,8 @@
 import pytest
 import torch.distributed as dist
 
-import modelopt.torch.utils.distributed as mto_dist
 import modelopt.torch.utils.moe as moe
+from modelopt.torch.utils.distributed import get_active_mo_session
 from modelopt.torch.utils.moe import (
     _resolve_ep_size,
     choose_ep_dp_degrees,
@@ -143,54 +143,89 @@ def test_unknown_dense_size_defaults_to_replication():
 # ------------------------------------------------------------------------------- rank-0 decision
 
 
-def _as_rank(monkeypatch, rank, rank0_decision=None):
-    """Pretend to be ``rank`` of a multi-rank run; returns what this rank handed the broadcast.
+@pytest.fixture
+def session():
+    """Enter the active session as ``rank`` of an 8-rank run."""
+    active = get_active_mo_session(required=False)
 
-    The fake broadcast delivers ``rank0_decision`` to a non-zero rank, as rank 0 would have sent it.
+    def enter(rank):
+        active.initialize(
+            distributed=True,
+            rank=rank,
+            world_size=8,
+            local_rank=rank,
+            device=None,
+            is_main=not rank,
+        )
+        return active
+
+    yield enter
+    active._reset()
+
+
+def _as_rank(monkeypatch, rank, rank0_message=None):
+    """Answer the session's broadcast as ``rank``; returns what this rank handed it.
+
+    A non-zero rank receives ``rank0_message``, as rank 0 would have sent it.
     """
     sent = []
 
-    def fake_broadcast(objs, src):
-        assert src == 0
-        sent.append(objs[0])
+    def broadcast_object_list(message, group=None, group_src=None):
+        sent.append(message[0])
         if rank != 0:
-            objs[0] = rank0_decision
+            message[0] = rank0_message
 
-    monkeypatch.setattr(mto_dist, "is_distributed_run", lambda: True)
     monkeypatch.setattr(dist, "get_rank", lambda group=None: rank)
-    monkeypatch.setattr(dist, "broadcast_object_list", fake_broadcast)
+    monkeypatch.setattr(dist, "broadcast_object_list", broadcast_object_list)
     return sent
 
 
-def test_rank0_decides_and_broadcasts_its_plan(monkeypatch):
+def test_rank0_decides_and_keeps_its_plan_on_the_session(monkeypatch, session):
+    active = session(0)
     sent = _as_rank(monkeypatch, 0)
+
     plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=4)
-    assert sent == [(plan, None)]
+
+    ((name, sent_plan, error),) = sent
+    assert (name, sent_plan.ep_size, error) == ("plan", 4, None)
     assert plan.ep_size == 4
+    assert active.plan is plan
 
 
-def test_other_ranks_adopt_rank0s_plan_without_deciding(monkeypatch):
+def test_other_ranks_adopt_rank0s_plan_without_deciding(monkeypatch, session):
     """A non-zero rank's own inputs (here an EP knob of 2) must not leak into the plan."""
     rank0_plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=4)
-    _as_rank(monkeypatch, 1, rank0_decision=(rank0_plan, None))
+    active = session(1)
+    _as_rank(monkeypatch, 1, rank0_message=("plan", rank0_plan, None))
 
     def must_not_decide(*args, **kwargs):
         raise AssertionError("a non-zero rank decided its own plan")
 
     monkeypatch.setattr(moe, "_decide_distribution_plan_on_rank0", must_not_decide)
     assert decide_distribution_plan(moe_cfg(), world_size=8, ep_size=2) is rank0_plan
+    assert active.plan is rank0_plan
 
 
-def test_rank0_failure_is_broadcast_before_it_raises(monkeypatch):
+def test_rank0_failure_is_broadcast_before_it_raises(monkeypatch, session):
     """Rank 0 must still reach the broadcast on failure, or every other rank hangs in it."""
+    session(0)
     sent = _as_rank(monkeypatch, 0)
     with pytest.raises(ValueError, match="must divide world_size"):
         decide_distribution_plan(moe_cfg(), world_size=8, ep_size=3)
-    assert sent[0][0] is None
-    assert "ValueError" in sent[0][1]
+    ((_, plan, error),) = sent
+    assert plan is None
+    assert "ValueError" in error
 
 
-def test_other_ranks_raise_when_rank0_failed(monkeypatch):
-    _as_rank(monkeypatch, 1, rank0_decision=(None, "ValueError: bad degrees"))
+def test_other_ranks_raise_when_rank0_failed(monkeypatch, session):
+    session(1)
+    _as_rank(monkeypatch, 1, rank0_message=("plan", None, "ValueError: bad degrees"))
     with pytest.raises(RuntimeError, match=r"rank 0 failed to decide.*bad degrees"):
         decide_distribution_plan(moe_cfg(), world_size=8)
+
+
+def test_without_a_session_the_process_decides_alone(monkeypatch):
+    monkeypatch.setattr(
+        dist, "broadcast_object_list", lambda *a, **k: pytest.fail("a collective ran")
+    )
+    assert decide_distribution_plan(moe_cfg(), world_size=8, ep_size=4).ep_size == 4

From 4cde41f8a6d2b6db0aa46a4c701cbbc62d249526 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 14 Oct 2026 02:37:41 +0000
Subject: [PATCH 03/44] Decide the plan through the session's plan handle

decide_distribution_plan sets session.plan with session.plan.decide(fn): rank 0 decides, and every
rank holds and returns rank 0's plan, read back as session.plan.value.

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

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index 5c7c17dfa15..632f8b66398 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -220,7 +220,7 @@ def decide():
     session = get_active_mo_session(required=False)
     if not session.initialized:
         return decide()
-    return session.decide("plan", decide)
+    return session.plan.decide(decide)
 
 
 def _decide_distribution_plan_on_rank0(
diff --git a/tests/unit/torch/utils/test_ep_distribution.py b/tests/unit/torch/utils/test_ep_distribution.py
index 87adc3122f8..0cf70dcfa31 100644
--- a/tests/unit/torch/utils/test_ep_distribution.py
+++ b/tests/unit/torch/utils/test_ep_distribution.py
@@ -189,7 +189,7 @@ def test_rank0_decides_and_keeps_its_plan_on_the_session(monkeypatch, session):
     ((name, sent_plan, error),) = sent
     assert (name, sent_plan.ep_size, error) == ("plan", 4, None)
     assert plan.ep_size == 4
-    assert active.plan is plan
+    assert active.plan.value is plan
 
 
 def test_other_ranks_adopt_rank0s_plan_without_deciding(monkeypatch, session):
@@ -202,8 +202,9 @@ def must_not_decide(*args, **kwargs):
         raise AssertionError("a non-zero rank decided its own plan")
 
     monkeypatch.setattr(moe, "_decide_distribution_plan_on_rank0", must_not_decide)
-    assert decide_distribution_plan(moe_cfg(), world_size=8, ep_size=2) is rank0_plan
-    assert active.plan is rank0_plan
+    plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=2)
+    assert repr(plan) == repr(rank0_plan)
+    assert active.plan.value is plan
 
 
 def test_rank0_failure_is_broadcast_before_it_raises(monkeypatch, session):

From 896ac7ce32de104812185ab55bcc827f5abfe184 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 14 Oct 2026 03:08:01 +0000
Subject: [PATCH 04/44] Decide the plan through the renamed broadcast_from

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/moe.py                    | 13 +++++++------
 tests/unit/torch/utils/test_ep_distribution.py |  2 +-
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index 632f8b66398..5b1b402f496 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -200,11 +200,12 @@ def decide_distribution_plan(
 
     Inside a :class:`~modelopt.torch.utils.distributed.ModelOptSession` the session decides it and
     keeps it as ``session.plan``: on a distributed run once, on rank 0, and every other rank adopts
-    rank 0's decision instead of making its own (:meth:`~modelopt.torch.utils.distributed.ModelOptSession.decide`).
-    The inputs are all read locally -- the env knobs, this GPU's memory, the config on this rank's
-    filesystem -- so ranks deciding independently could disagree, and a disagreement surfaces as a
-    hang at mesh creation rather than as an error. Collective: every rank must call it. If rank 0
-    fails to decide, every rank raises. Without a session, this process decides alone.
+    rank 0's decision instead of making its own
+    (:meth:`~modelopt.torch.utils.distributed.ModelOptSession.broadcast_from`). The inputs are all
+    read locally -- the env knobs, this GPU's memory, the config on this rank's filesystem -- so
+    ranks deciding independently could disagree, and a disagreement surfaces as a hang at mesh
+    creation rather than as an error. Collective: every rank must call it. If rank 0 fails to decide,
+    every rank raises. Without a session, this process decides alone.
     """
     from .distributed import get_active_mo_session
 
@@ -220,7 +221,7 @@ def decide():
     session = get_active_mo_session(required=False)
     if not session.initialized:
         return decide()
-    return session.plan.decide(decide)
+    return session.plan.broadcast_from(decide)
 
 
 def _decide_distribution_plan_on_rank0(
diff --git a/tests/unit/torch/utils/test_ep_distribution.py b/tests/unit/torch/utils/test_ep_distribution.py
index 0cf70dcfa31..0a52d521bd6 100644
--- a/tests/unit/torch/utils/test_ep_distribution.py
+++ b/tests/unit/torch/utils/test_ep_distribution.py
@@ -221,7 +221,7 @@ def test_rank0_failure_is_broadcast_before_it_raises(monkeypatch, session):
 def test_other_ranks_raise_when_rank0_failed(monkeypatch, session):
     session(1)
     _as_rank(monkeypatch, 1, rank0_message=("plan", None, "ValueError: bad degrees"))
-    with pytest.raises(RuntimeError, match=r"rank 0 failed to decide.*bad degrees"):
+    with pytest.raises(RuntimeError, match=r"rank 0 failed to compute.*bad degrees"):
         decide_distribution_plan(moe_cfg(), world_size=8)
 
 

From 7212b50d1c643f536d7efd50d3bc2b32cdcdd375 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 14 Oct 2026 17:16:39 +0000
Subject: [PATCH 05/44] Decide the plan on rank 0 and broadcast it, without a
 session to keep it on

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/moe.py                   | 29 ++++++-------
 .../unit/torch/utils/test_ep_distribution.py  | 41 ++++---------------
 2 files changed, 20 insertions(+), 50 deletions(-)

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index 5b1b402f496..7cc524b7135 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -198,30 +198,25 @@ def decide_distribution_plan(
     * ``fsdp_size`` -- ``auto`` -> ``1`` (replicate the dense part) when it fits one GPU, else
       ``world_size`` (FSDP2-shard it); ``"world"`` -> ``world_size``.
 
-    Inside a :class:`~modelopt.torch.utils.distributed.ModelOptSession` the session decides it and
-    keeps it as ``session.plan``: on a distributed run once, on rank 0, and every other rank adopts
-    rank 0's decision instead of making its own
-    (:meth:`~modelopt.torch.utils.distributed.ModelOptSession.broadcast_from`). The inputs are all
-    read locally -- the env knobs, this GPU's memory, the config on this rank's filesystem -- so
-    ranks deciding independently could disagree, and a disagreement surfaces as a hang at mesh
-    creation rather than as an error. Collective: every rank must call it. If rank 0 fails to decide,
-    every rank raises. Without a session, this process decides alone.
+    On a distributed run it is decided once, on rank 0, and every other rank adopts rank 0's
+    decision instead of making its own (:func:`~modelopt.torch.utils.distributed.broadcast_from`).
+    The inputs are all read locally -- the env knobs, this GPU's memory, the config on this rank's
+    filesystem -- so ranks deciding independently could disagree, and a disagreement surfaces as a
+    hang at mesh creation rather than as an error. Collective on a distributed run: every rank must
+    call it. If rank 0 fails to decide, every rank raises. A single process decides alone.
     """
-    from .distributed import get_active_mo_session
+    from .distributed import broadcast_from
 
-    def decide():
-        return _decide_distribution_plan_on_rank0(
+    return broadcast_from(
+        "plan",
+        lambda: _decide_distribution_plan_on_rank0(
             hf_config,
             world_size,
             ep_size=ep_size,
             fsdp_size=fsdp_size,
             device_mem_bytes=device_mem_bytes,
-        )
-
-    session = get_active_mo_session(required=False)
-    if not session.initialized:
-        return decide()
-    return session.plan.broadcast_from(decide)
+        ),
+    )
 
 
 def _decide_distribution_plan_on_rank0(
diff --git a/tests/unit/torch/utils/test_ep_distribution.py b/tests/unit/torch/utils/test_ep_distribution.py
index 0a52d521bd6..3f795311cb7 100644
--- a/tests/unit/torch/utils/test_ep_distribution.py
+++ b/tests/unit/torch/utils/test_ep_distribution.py
@@ -21,7 +21,6 @@
 import torch.distributed as dist
 
 import modelopt.torch.utils.moe as moe
-from modelopt.torch.utils.distributed import get_active_mo_session
 from modelopt.torch.utils.moe import (
     _resolve_ep_size,
     choose_ep_dp_degrees,
@@ -143,28 +142,8 @@ def test_unknown_dense_size_defaults_to_replication():
 # ------------------------------------------------------------------------------- rank-0 decision
 
 
-@pytest.fixture
-def session():
-    """Enter the active session as ``rank`` of an 8-rank run."""
-    active = get_active_mo_session(required=False)
-
-    def enter(rank):
-        active.initialize(
-            distributed=True,
-            rank=rank,
-            world_size=8,
-            local_rank=rank,
-            device=None,
-            is_main=not rank,
-        )
-        return active
-
-    yield enter
-    active._reset()
-
-
 def _as_rank(monkeypatch, rank, rank0_message=None):
-    """Answer the session's broadcast as ``rank``; returns what this rank handed it.
+    """Run as ``rank`` of an 8-rank run, answering its broadcast; returns what this rank handed it.
 
     A non-zero rank receives ``rank0_message``, as rank 0 would have sent it.
     """
@@ -175,13 +154,14 @@ def broadcast_object_list(message, group=None, group_src=None):
         if rank != 0:
             message[0] = rank0_message
 
+    monkeypatch.setattr(dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(dist, "get_world_size", lambda group=None: 8)
     monkeypatch.setattr(dist, "get_rank", lambda group=None: rank)
     monkeypatch.setattr(dist, "broadcast_object_list", broadcast_object_list)
     return sent
 
 
-def test_rank0_decides_and_keeps_its_plan_on_the_session(monkeypatch, session):
-    active = session(0)
+def test_rank0_decides_and_sends_its_plan(monkeypatch):
     sent = _as_rank(monkeypatch, 0)
 
     plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=4)
@@ -189,13 +169,11 @@ def test_rank0_decides_and_keeps_its_plan_on_the_session(monkeypatch, session):
     ((name, sent_plan, error),) = sent
     assert (name, sent_plan.ep_size, error) == ("plan", 4, None)
     assert plan.ep_size == 4
-    assert active.plan.value is plan
 
 
-def test_other_ranks_adopt_rank0s_plan_without_deciding(monkeypatch, session):
+def test_other_ranks_adopt_rank0s_plan_without_deciding(monkeypatch):
     """A non-zero rank's own inputs (here an EP knob of 2) must not leak into the plan."""
     rank0_plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=4)
-    active = session(1)
     _as_rank(monkeypatch, 1, rank0_message=("plan", rank0_plan, None))
 
     def must_not_decide(*args, **kwargs):
@@ -204,12 +182,10 @@ def must_not_decide(*args, **kwargs):
     monkeypatch.setattr(moe, "_decide_distribution_plan_on_rank0", must_not_decide)
     plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=2)
     assert repr(plan) == repr(rank0_plan)
-    assert active.plan.value is plan
 
 
-def test_rank0_failure_is_broadcast_before_it_raises(monkeypatch, session):
+def test_rank0_failure_is_broadcast_before_it_raises(monkeypatch):
     """Rank 0 must still reach the broadcast on failure, or every other rank hangs in it."""
-    session(0)
     sent = _as_rank(monkeypatch, 0)
     with pytest.raises(ValueError, match="must divide world_size"):
         decide_distribution_plan(moe_cfg(), world_size=8, ep_size=3)
@@ -218,14 +194,13 @@ def test_rank0_failure_is_broadcast_before_it_raises(monkeypatch, session):
     assert "ValueError" in error
 
 
-def test_other_ranks_raise_when_rank0_failed(monkeypatch, session):
-    session(1)
+def test_other_ranks_raise_when_rank0_failed(monkeypatch):
     _as_rank(monkeypatch, 1, rank0_message=("plan", None, "ValueError: bad degrees"))
     with pytest.raises(RuntimeError, match=r"rank 0 failed to compute.*bad degrees"):
         decide_distribution_plan(moe_cfg(), world_size=8)
 
 
-def test_without_a_session_the_process_decides_alone(monkeypatch):
+def test_a_single_process_decides_alone(monkeypatch):
     monkeypatch.setattr(
         dist, "broadcast_object_list", lambda *a, **k: pytest.fail("a collective ran")
     )

From 7102844828941c867429c9a7a17b71622d97ba4a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 14 Oct 2026 19:43:22 +0000
Subject: [PATCH 06/44] Decide a model's distribution in
 resolve_distribution_plan, from its meta model

resolve_distribution_plan is now the one place a run decides how a model is spread over its ranks,
for every model rather than only an MoE. It moves here from load_pretrained's routing, with the
MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE knobs, and decide_distribution_plan, a second public way
to decide, is gone.

Rank 0 decides from the model built on the meta device, by the class that builds it whole: a
causal LM, or a VLM through the class that carries its LM head. The routed experts are the ones its
base_model_ep_plan names, found the way the expert-parallel load will find them. Whether the dense
part fits one GPU is judged from the parameter bytes that every rank would hold whole, rather than
from a config estimate that ignored GQA and tied embeddings and gave up on configs without the
fields it read. An explicit EP degree must divide the routed-expert count, and MODELOPT_FSDP_SIZE
either replicates the dense part (1) or shards it over the whole world (world): a part of the world
is refused.

A single process decides without building anything: it holds the whole model, unless dense='fsdp'
asks for FSDP2 on one rank.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/env.py                   |  43 ++
 modelopt/torch/utils/moe.py                   | 163 ++---
 .../torch/utils/plugins/model_load_utils.py   | 182 +++++-
 .../unit/torch/utils/test_ep_distribution.py  | 207 -------
 .../utils/test_resolve_distribution_plan.py   | 573 ++++++++++++++++++
 5 files changed, 824 insertions(+), 344 deletions(-)
 create mode 100644 modelopt/torch/utils/env.py
 delete mode 100644 tests/unit/torch/utils/test_ep_distribution.py
 create mode 100644 tests/unit/torch/utils/test_resolve_distribution_plan.py

diff --git a/modelopt/torch/utils/env.py b/modelopt/torch/utils/env.py
new file mode 100644
index 00000000000..e8e7827ef07
--- /dev/null
+++ b/modelopt/torch/utils/env.py
@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The ``MODELOPT_*`` environment-variable knobs of distributed HuggingFace model loading.
+
+Each knob is a small accessor that reads ``os.environ`` *at call time*, not at import: a launcher sets
+them before it forks the distributed workers, and tests may override them, so the live value is what
+matters. Callers do ``from modelopt.torch.utils import env`` then ``env.ep_degree()``.
+
+Distribution plan (see ``resolve_distribution_plan``):
+  MODELOPT_EP_DEGREE        expert-parallel degree; ``""``/unset -> ``"auto"``.
+  MODELOPT_FSDP_SIZE        whether to FSDP2-shard the dense part: ``"1"`` replicates it, ``"world"``
+                            shards it over every rank; ``""``/unset -> ``"auto"``.
+"""
+
+import os
+
+__all__ = [
+    "ep_degree",
+    "fsdp_size",
+]
+
+
+def ep_degree() -> str:
+    """Expert-parallel degree knob for ``resolve_distribution_plan``; ``""``/unset -> ``"auto"``."""
+    return os.environ.get("MODELOPT_EP_DEGREE") or "auto"
+
+
+def fsdp_size() -> str:
+    """Dense-sharding knob for ``resolve_distribution_plan``; ``""``/unset -> ``"auto"``."""
+    return os.environ.get("MODELOPT_FSDP_SIZE") or "auto"
diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index 7cc524b7135..e487e3fceea 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -15,8 +15,8 @@
 
 """HuggingFace-coupled MoE expert-parallel (EP) helpers.
 
-Capability detection and EP degree selection: deciding how a model distributes, from the model,
-this GPU's memory and the knobs.
+Finding a model's routed experts, choosing the EP degree they are sharded over, and the DP=EP
+all-to-all expert transport (``torch_a2a``).
 
 The all-to-all transport is plain ``torch.distributed.all_to_all_single`` (NCCL) wrapped around
 transformers' ``_grouped_linear``. ``register_torch_a2a()`` registers it into
@@ -25,11 +25,15 @@
 """
 
 import warnings
+from typing import TYPE_CHECKING
 
 import torch
 import torch.distributed as dist
 from torch.distributed.nn.functional import all_to_all_single as autograd_all_to_all_single
 
+if TYPE_CHECKING:
+    from torch import nn
+
 
 def require_transformers(min_version: str, feature: str) -> None:
     """Fail loudly if the installed transformers is older than ``min_version`` (the floor for ``feature``).
@@ -86,20 +90,6 @@ def num_experts_from_config(hf_config) -> int:
     return 0
 
 
-def is_expert_parallel_capable(hf_config) -> bool:
-    """True if the model can be expert-parallel sharded.
-
-    Transformers exposes ``base_model_ep_plan`` on MoE configs whose experts support
-    expert parallelism (e.g. gpt_oss, llama4; qwen3_moe in transformers>=5.7). Dense models
-    and MoE configs without the plan return False -> use plain FSDP2. For VLMs the plan
-    lives on ``text_config`` (the nested LM submodel), so check both.
-    """
-    return (
-        bool(getattr(_ep_config(hf_config), "base_model_ep_plan", None))
-        and num_experts_from_config(hf_config) > 0
-    )
-
-
 def choose_ep_dp_degrees(world_size: int, num_experts: int) -> tuple[int, int]:
     """Pick ``(expert_parallel, expert_data_parallel)`` degrees with ``EP * EDP == world_size``.
 
@@ -116,124 +106,39 @@ def choose_ep_dp_degrees(world_size: int, num_experts: int) -> tuple[int, int]:
     return 1, world_size
 
 
-def _estimate_dense_param_bytes(hf_config) -> "int | None":
-    """Best-effort estimate of the DENSE (non-expert) parameter bytes from an HF config.
+# The styles a base_model_ep_plan gives a layer's routed-experts container. A plan that styles only
+# the experts' grouped-GEMM parameters names the containers through their parents.
+_EXPERTS_STYLES = ("moe_tp_experts", "nemotron_experts_ep")
 
-    embeddings + lm_head + per-layer attention/norms (+ the dense FFN when the model is NOT MoE, since a MoE's FFN is
-    the expert weights). Returns ``None`` when the config lacks the fields to estimate (the caller then treats the dense
-    part as 'fits'). A deliberately CONSERVATIVE upper bound (ignores GQA and tied embeddings) -- it exists to catch a
-    dense part that clearly cannot be replicated, not to be exact.
-    """
-    cfg = getattr(hf_config, "text_config", None) or hf_config
-    h = getattr(cfg, "hidden_size", None)
-    n_layers = getattr(cfg, "num_hidden_layers", None)
-    if not h or not n_layers:
-        return None
-    vocab = getattr(cfg, "vocab_size", None) or getattr(hf_config, "vocab_size", None) or 0
-    inter = getattr(cfg, "intermediate_size", None) or 4 * h
-    attn = 4 * h * h  # q,k,v,o (upper bound; ignores GQA head reduction)
-    norms = 4 * h  # ~2 layernorms per layer
-    ffn = (
-        0 if is_expert_parallel_capable(hf_config) else 3 * h * inter
-    )  # gated FFN; a MoE's FFN is expert
-    per_layer = attn + norms + ffn
-    embed = 2 * vocab * h  # embed + lm_head (untied upper bound)
-    bytes_per = 2  # bf16/fp16 default (the pre-quant load dtype)
-    return (embed + n_layers * per_layer) * bytes_per
-
-
-def _dense_fits_per_gpu(
-    hf_config, device_mem_bytes: "int | None" = None, headroom: float = 0.6
-) -> bool:
-    """Fit predicate for the dense knob.
-
-    does the DENSE footprint fit ONE GPU (leaving ``headroom`` for activations + comm buffers + any EP-local experts)?
-    CONSERVATIVE -- returns ``True`` ('fits' -> replicate / data-parallel) whenever it cannot estimate the dense size or
-    the device memory, so the data-parallel default is preserved. This is the pluggable decision point of the auto-
-    resolver; a heuristic or a recipe hint may substitute for it without changing callers.
-    """
-    dense_bytes = _estimate_dense_param_bytes(hf_config)
-    if not dense_bytes:
-        return True
-    if device_mem_bytes is None:
-        try:
-            import torch
-
-            device_mem_bytes = torch.cuda.get_device_properties(
-                torch.cuda.current_device()
-            ).total_memory
-        except Exception:
-            return True
-    return dense_bytes < headroom * device_mem_bytes
 
+def _routed_expert_keys(ep_plan: dict[str, str]) -> set[str]:
+    """The keys of ``ep_plan``, a ``base_model_ep_plan``, that name a routed-experts container."""
+    keys = {key for key, style in ep_plan.items() if style in _EXPERTS_STYLES}
+    return keys or {
+        key.rsplit(".", 1)[0] for key, style in ep_plan.items() if style == "grouped_gemm"
+    }
 
-def _resolve_ep_size(hf_config, world_size: int, ep_size) -> int:
-    """Resolve the expert-parallel degree.
 
-    An explicit ``ep_size`` wins; ``auto`` -> the largest EP degree dividing both the expert count and
-    ``world_size`` for an EP-capable model, else ``1``.
-    """
-    if ep_size not in (None, "auto", "0", 0):
-        return int(ep_size)
-    if world_size > 1 and is_expert_parallel_capable(hf_config):
-        ep, _ = choose_ep_dp_degrees(world_size, num_experts_from_config(hf_config))
-        return ep
-    return 1
-
-
-def decide_distribution_plan(
-    hf_config,
-    world_size: int,
-    *,
-    ep_size="auto",
-    fsdp_size="auto",
-    device_mem_bytes: "int | None" = None,
-):
-    """Decide how a HuggingFace MoE model is distributed for PTQ: a :class:`DistributionPlan`.
-
-    Two knobs, each an explicit value or ``"auto"`` (explicit always wins):
-
-    * ``ep_size`` -- ``auto`` -> the largest EP degree dividing both the expert count and
-      ``world_size`` (``1`` for a model that cannot be expert-parallel).
-    * ``fsdp_size`` -- ``auto`` -> ``1`` (replicate the dense part) when it fits one GPU, else
-      ``world_size`` (FSDP2-shard it); ``"world"`` -> ``world_size``.
-
-    On a distributed run it is decided once, on rank 0, and every other rank adopts rank 0's
-    decision instead of making its own (:func:`~modelopt.torch.utils.distributed.broadcast_from`).
-    The inputs are all read locally -- the env knobs, this GPU's memory, the config on this rank's
-    filesystem -- so ranks deciding independently could disagree, and a disagreement surfaces as a
-    hang at mesh creation rather than as an error. Collective on a distributed run: every rank must
-    call it. If rank 0 fails to decide, every rank raises. A single process decides alone.
+def _routed_expert_modules(model: "nn.Module", ep_plan: dict[str, str]) -> "dict[str, nn.Module]":
+    """``model``'s routed-expert containers, ``{name: module}``, as its ``base_model_ep_plan`` names them.
+
+    A module is one when its name, with layer indices wildcarded as transformers does, is a container
+    key or ends with one: the keys are relative to the base model, so they omit the ``model.`` or
+    ``language_model.model.`` a causal LM or a VLM puts in front. Deciding a model's distribution and
+    loading it expert-parallel both find the experts this way, so the two cannot disagree.
     """
-    from .distributed import broadcast_from
-
-    return broadcast_from(
-        "plan",
-        lambda: _decide_distribution_plan_on_rank0(
-            hf_config,
-            world_size,
-            ep_size=ep_size,
-            fsdp_size=fsdp_size,
-            device_mem_bytes=device_mem_bytes,
-        ),
-    )
-
-
-def _decide_distribution_plan_on_rank0(
-    hf_config, world_size: int, *, ep_size, fsdp_size, device_mem_bytes: "int | None"
-):
-    """The decision itself (see :func:`decide_distribution_plan`); run on rank 0 alone when distributed."""
-    from .distributed import DistributionPlan
-
-    ep = _resolve_ep_size(hf_config, world_size, ep_size)
-    if fsdp_size in (None, "auto"):
-        fsdp = 1 if _dense_fits_per_gpu(hf_config, device_mem_bytes) else world_size
-    elif str(fsdp_size).lower() in ("world", "full"):
-        fsdp = world_size
-    else:
-        fsdp = int(fsdp_size)
-    fsdp = max(1, min(fsdp, world_size))
-    return DistributionPlan(world_size=world_size, ep_size=ep, fsdp_size=fsdp)
+    from transformers.integrations.tensor_parallel import replace_layer_number_by_wildcard
+
+    keys = _routed_expert_keys(ep_plan)
+    return {
+        name: module
+        for name, module in model.named_modules()
+        if any(
+            (generic := replace_layer_number_by_wildcard(name)) == key
+            or generic.endswith("." + key)
+            for key in keys
+        )
+    }
 
 
 # ---------------------------------------------------------------------------------------------------
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 45a8e3ecff5..f8418e9d888 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""HuggingFace model loading: :func:`load_pretrained`, and parallel FSDP2 loading."""
+"""HuggingFace model loading: :func:`load_pretrained`, and the plan to distribute a model by."""
 
 import logging
 import os
@@ -33,11 +33,13 @@
 from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
 from torch.distributed.fsdp import FSDPModule
 from torch.distributed.tensor import DTensor
-from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, PretrainedConfig
 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.utils import env
 from modelopt.torch.utils.distributed import (
+    DistributionPlan,
     barrier,
     broadcast_state_dict,
     fsdp2_wrap,
@@ -194,8 +196,9 @@ def build_meta_causal_lm(
     trust_remote_code: bool,
     attn_implementation: str | None,
     hf_config=None,
+    auto_class: Any = None,
 ):
-    """Build a meta-init causal LM (no real storage allocated)."""
+    """Build a meta-init causal LM (no real storage allocated), or ``auto_class``'s model."""
     if hf_config is None:
         config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code}
         if attn_implementation is not None:
@@ -207,12 +210,14 @@ def build_meta_causal_lm(
     dtype = getattr(hf_config, "dtype", None) or torch.bfloat16
     from accelerate import init_empty_weights  # only real callers of this function need it
 
-    model_factory = AutoModelForCausalLM
-    if is_fsdp2_text_only_vlm(hf_config):
-        # Only newer Transformers versions supporting these models expose this factory.
-        from transformers import AutoModelForImageTextToText
+    model_factory = auto_class
+    if model_factory is None:
+        model_factory = AutoModelForCausalLM
+        if is_fsdp2_text_only_vlm(hf_config):
+            # Only newer Transformers versions supporting these models expose this factory.
+            from transformers import AutoModelForImageTextToText
 
-        model_factory = AutoModelForImageTextToText
+            model_factory = AutoModelForImageTextToText
     with init_empty_weights(include_buffers=False):
         model = model_factory.from_config(
             hf_config, dtype=dtype, trust_remote_code=trust_remote_code
@@ -882,3 +887,164 @@ def parallel_load_and_prepare_fsdp2(
     if hasattr(model, "tie_weights"):
         model.tie_weights()
     return model
+
+
+def _auto_class(hf_config, require_lm_head: bool = True) -> Any:
+    """The auto class that builds ``hf_config``'s whole model.
+
+    A causal LM builds through ``AutoModelForCausalLM``. A VLM -- a nested text config, a vision
+    config, or a ``*ForConditionalGeneration`` architecture -- builds through ``AutoModel``, which
+    leaves its LM head out. ``require_lm_head`` builds it with the head: through the remote
+    ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` there
+    and no image-text class, else through ``AutoModelForImageTextToText`` if that recognizes the
+    config.
+    """
+    architectures = getattr(hf_config, "architectures", None) or []
+    is_vlm = (
+        getattr(hf_config, "text_config", None) is not None
+        or getattr(hf_config, "vision_config", None) is not None
+        or any("conditionalgeneration" in str(arch).lower() for arch in architectures)
+    )
+    if not is_vlm:
+        return AutoModelForCausalLM
+    if not require_lm_head:
+        return AutoModel
+    auto_map = getattr(hf_config, "auto_map", None) or {}
+    remote_with_head = (
+        "conditionalgeneration" in str(auto_map.get("AutoModelForCausalLM", "")).lower()
+    )
+    if remote_with_head and "AutoModelForImageTextToText" not in auto_map:
+        return AutoModelForCausalLM
+    try:
+        from transformers import AutoModelForImageTextToText
+    except ImportError:
+        return AutoModel
+    if type(hf_config) in AutoModelForImageTextToText._model_mapping:
+        return AutoModelForImageTextToText
+    return AutoModel
+
+
+# The share of one GPU's memory a replicated dense part may take. The rest is left to calibration's
+# activations and communication buffers, and to the rank's share of the experts.
+_REPLICATED_SHARE = 0.6
+
+
+def _gpu_memory() -> int | None:
+    """This rank's GPU memory in bytes, or ``None`` without a GPU."""
+    if not torch.cuda.is_available():
+        return None
+    return torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory
+
+
+def _dense_from_env(world_size: int) -> str:
+    """``MODELOPT_FSDP_SIZE`` as a dense layout, or ``"auto"``."""
+    value = env.fsdp_size().lower()
+    if value == "auto":
+        return "auto"
+    if value == "1":
+        return "replicate"
+    if value in ("world", "full", str(world_size)):
+        return "fsdp"
+    raise ValueError(
+        f"MODELOPT_FSDP_SIZE={value!r}: expected auto, 1 (replicate the dense part) or world "
+        f"(FSDP2-shard it over all {world_size} ranks)."
+    )
+
+
+def _decide_distribution(
+    ckpt_path: str, hf_config, world_size: int, ep_size, dense, trust_remote_code: bool
+) -> DistributionPlan:
+    """The plan :func:`resolve_distribution_plan` returns, decided by this rank alone."""
+    from modelopt.torch.utils.moe import (
+        _ep_config,
+        _routed_expert_modules,
+        choose_ep_dp_degrees,
+        num_experts_from_config,
+    )
+
+    ep_size = env.ep_degree() if ep_size is None else ep_size
+    ep_size = "auto" if str(ep_size) in ("auto", "0") else int(ep_size)
+    dense = _dense_from_env(world_size) if dense is None else dense
+    if world_size == 1:
+        # Nothing to shard over, so nothing about the model matters.
+        return DistributionPlan(
+            ep_size=1 if ep_size == "auto" else ep_size,
+            dense="replicate" if dense == "auto" else dense,
+        )
+    try:
+        model = build_meta_causal_lm(
+            ckpt_path, trust_remote_code, None, hf_config, auto_class=_auto_class(hf_config)
+        )
+    except Exception as e:
+        raise RuntimeError(
+            f"Deciding how to distribute {hf_config.model_type} needs its model built on the meta "
+            f"device, which failed: {type(e).__name__}: {e}"
+        ) from e
+    ep_plan = getattr(_ep_config(hf_config), "base_model_ep_plan", None) or {}
+    experts = _routed_expert_modules(model, ep_plan)
+    num_experts = num_experts_from_config(hf_config) if experts else 0
+    if ep_size == "auto":
+        ep_size = choose_ep_dp_degrees(world_size, num_experts)[0]
+    elif ep_size > 1 and (num_experts == 0 or num_experts % ep_size):
+        raise ValueError(
+            f"ep_size={ep_size} must divide the routed experts {hf_config.model_type} shards "
+            f"({num_experts})."
+        )
+    if dense == "auto":
+        # Every rank holds a whole copy of what expert parallelism does not shard.
+        sharded = (
+            {id(p) for m in experts.values() for p in m.parameters()} if ep_size > 1 else set()
+        )
+        replicated = sum(
+            p.numel() * p.element_size() for p in model.parameters() if id(p) not in sharded
+        )
+        memory = _gpu_memory()
+        too_big = memory is not None and replicated > _REPLICATED_SHARE * memory
+        dense = "fsdp" if too_big else "replicate"
+    return DistributionPlan(world_size, ep_size, dense)
+
+
+def resolve_distribution_plan(
+    ckpt_path: str,
+    *,
+    trust_remote_code: bool = False,
+    ep_size: int | str | None = None,
+    dense: str | None = None,
+) -> tuple[DistributionPlan, PretrainedConfig]:
+    """Decide how this run distributes ``ckpt_path``'s model over its ranks.
+
+    The one place a run decides it. :func:`load_pretrained` loads the model by the plan this returns
+    and records it on the model, and every later step reads it back from there
+    (:func:`~modelopt.torch.utils.distributed.distribution_plan`). Returns ``(plan, hf_config)``,
+    ``hf_config`` being the config a sharded load takes back as
+    ``load_pretrained(ckpt_path, distribution=plan, hf_config=hf_config)``.
+
+    Rank 0 decides, from the model built on the meta device: from its parameters rather than
+    estimates of them, and from its routed experts as an expert-parallel load finds them. Every rank
+    takes rank 0's plan (:func:`~modelopt.torch.utils.distributed.broadcast_from`): the inputs are
+    read locally, ranks deciding for themselves could disagree, and ranks loading different plans
+    hang. Collective: every rank must call it.
+
+    * ``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``.
+    * ``dense``: ``"replicate"`` or ``"fsdp"`` for the rest of the model, or ``"auto"``, which
+      FSDP2-shards it when a whole copy would take more than 60% of one GPU's memory.
+
+    ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` (``1`` replicates, ``world``
+    shards) from rank 0's environment, both defaulting to ``"auto"``. A single process holds the
+    whole model unless ``dense="fsdp"`` asks for FSDP2 on one rank.
+    """
+    from modelopt.torch.utils import distributed as dist
+
+    _, ckpt_path = ensure_local_checkpoint(ckpt_path)
+    hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code)
+    world_size = dist.size()
+    plan = dist.broadcast_from(
+        "plan",
+        lambda: _decide_distribution(
+            ckpt_path, hf_config, world_size, ep_size, dense, trust_remote_code
+        ),
+    )
+    if plan != DistributionPlan() and dist.is_master():
+        print(f"Distribution: {plan!r}", flush=True)
+    return plan, hf_config
diff --git a/tests/unit/torch/utils/test_ep_distribution.py b/tests/unit/torch/utils/test_ep_distribution.py
deleted file mode 100644
index 3f795311cb7..00000000000
--- a/tests/unit/torch/utils/test_ep_distribution.py
+++ /dev/null
@@ -1,207 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Unit tests for EP degree selection and the decided distribution plan."""
-
-from types import SimpleNamespace
-
-import pytest
-import torch.distributed as dist
-
-import modelopt.torch.utils.moe as moe
-from modelopt.torch.utils.moe import (
-    _resolve_ep_size,
-    choose_ep_dp_degrees,
-    decide_distribution_plan,
-    is_expert_parallel_capable,
-    num_experts_from_config,
-)
-
-
-def cfg(**kw):
-    return SimpleNamespace(**kw)
-
-
-def moe_cfg(experts=64, **kw):
-    """An EP-capable MoE config: transformers advertises EP support via ``base_model_ep_plan``."""
-    return cfg(num_experts=experts, base_model_ep_plan={"experts": "grouped"}, **kw)
-
-
-# ------------------------------------------------------------------------------- expert counting
-
-
-def test_dense_config_reports_no_experts():
-    assert num_experts_from_config(cfg(hidden_size=128)) == 0
-
-
-@pytest.mark.parametrize(
-    "attr", ["num_experts", "num_local_experts", "n_routed_experts", "num_routed_experts"]
-)
-def test_every_spelling_of_the_expert_count_is_understood(attr):
-    """Architectures disagree on the field name; missing one would silently disable EP."""
-    assert num_experts_from_config(cfg(**{attr: 16})) == 16
-
-
-@pytest.mark.parametrize("nest", ["text_config", "llm_config"])
-def test_expert_count_is_found_on_a_vlms_nested_text_model(nest):
-    """For a VLM the MoE is the nested LM, so the root config has no expert count of its own."""
-    assert num_experts_from_config(cfg(**{nest: cfg(n_routed_experts=32)})) == 32
-
-
-# ----------------------------------------------------------------------------------- capability
-
-
-def test_moe_with_an_ep_plan_is_capable():
-    assert is_expert_parallel_capable(moe_cfg())
-
-
-def test_moe_without_an_ep_plan_is_not_capable():
-    """No ``base_model_ep_plan`` means transformers cannot shard these experts -> use plain FSDP2."""
-    assert not is_expert_parallel_capable(cfg(num_experts=64))
-
-
-def test_ep_plan_without_experts_is_not_capable():
-    assert not is_expert_parallel_capable(cfg(base_model_ep_plan={"experts": "grouped"}))
-
-
-def test_capability_is_detected_through_a_vlm_nesting():
-    assert is_expert_parallel_capable(cfg(text_config=moe_cfg(experts=8)))
-
-
-# ------------------------------------------------------------------------------- degree selection
-
-
-def test_largest_valid_ep_degree_is_chosen():
-    assert choose_ep_dp_degrees(8, 64) == (8, 1)
-
-
-def test_ep_must_divide_the_expert_count_so_ranks_hold_whole_experts():
-    assert choose_ep_dp_degrees(8, 6) == (2, 4)
-
-
-def test_ep_falls_back_to_one_when_nothing_divides():
-    assert choose_ep_dp_degrees(8, 5) == (1, 8)
-
-
-@pytest.mark.parametrize(("world", "experts"), [(8, 0), (1, 64)])
-def test_no_expert_parallelism_for_dense_or_single_rank(world, experts):
-    assert choose_ep_dp_degrees(world, experts) == (1, world)
-
-
-@pytest.mark.parametrize("world", [2, 4, 8, 16])
-def test_degrees_always_factor_the_world(world):
-    ep, dp = choose_ep_dp_degrees(world, 64)
-    assert ep * dp == world
-
-
-def test_an_explicit_ep_size_overrides_auto_resolution():
-    assert _resolve_ep_size(moe_cfg(), world_size=8, ep_size=2) == 2
-
-
-def test_auto_ep_is_one_for_a_model_that_cannot_shard_experts():
-    assert _resolve_ep_size(cfg(num_experts=64), world_size=8, ep_size="auto") == 1
-
-
-# ------------------------------------------------------------------------------- decided plan
-
-
-def test_auto_plan_takes_the_largest_ep_degree_dividing_the_experts():
-    plan = decide_distribution_plan(moe_cfg(experts=6), world_size=8)
-    assert plan.ep_size == 2
-    assert plan.edp == 4
-
-
-def test_an_ep_degree_that_does_not_divide_the_world_is_rejected():
-    with pytest.raises(ValueError, match="must divide world_size"):
-        decide_distribution_plan(moe_cfg(), world_size=8, ep_size=3)
-
-
-@pytest.mark.parametrize(("knob", "expected"), [("world", 8), ("full", 8), (4, 4), (99, 8), (0, 1)])
-def test_dense_fsdp_knob_is_honoured_and_clamped_to_the_world(knob, expected):
-    plan = decide_distribution_plan(moe_cfg(), world_size=8, fsdp_size=knob)
-    assert plan.fsdp_size == expected
-
-
-def test_unknown_dense_size_defaults_to_replication():
-    """With no way to estimate the dense footprint, stay on the cheaper replicate path."""
-    assert decide_distribution_plan(moe_cfg(), world_size=8).fsdp_size == 1
-
-
-# ------------------------------------------------------------------------------- rank-0 decision
-
-
-def _as_rank(monkeypatch, rank, rank0_message=None):
-    """Run as ``rank`` of an 8-rank run, answering its broadcast; returns what this rank handed it.
-
-    A non-zero rank receives ``rank0_message``, as rank 0 would have sent it.
-    """
-    sent = []
-
-    def broadcast_object_list(message, group=None, group_src=None):
-        sent.append(message[0])
-        if rank != 0:
-            message[0] = rank0_message
-
-    monkeypatch.setattr(dist, "is_initialized", lambda: True)
-    monkeypatch.setattr(dist, "get_world_size", lambda group=None: 8)
-    monkeypatch.setattr(dist, "get_rank", lambda group=None: rank)
-    monkeypatch.setattr(dist, "broadcast_object_list", broadcast_object_list)
-    return sent
-
-
-def test_rank0_decides_and_sends_its_plan(monkeypatch):
-    sent = _as_rank(monkeypatch, 0)
-
-    plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=4)
-
-    ((name, sent_plan, error),) = sent
-    assert (name, sent_plan.ep_size, error) == ("plan", 4, None)
-    assert plan.ep_size == 4
-
-
-def test_other_ranks_adopt_rank0s_plan_without_deciding(monkeypatch):
-    """A non-zero rank's own inputs (here an EP knob of 2) must not leak into the plan."""
-    rank0_plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=4)
-    _as_rank(monkeypatch, 1, rank0_message=("plan", rank0_plan, None))
-
-    def must_not_decide(*args, **kwargs):
-        raise AssertionError("a non-zero rank decided its own plan")
-
-    monkeypatch.setattr(moe, "_decide_distribution_plan_on_rank0", must_not_decide)
-    plan = decide_distribution_plan(moe_cfg(), world_size=8, ep_size=2)
-    assert repr(plan) == repr(rank0_plan)
-
-
-def test_rank0_failure_is_broadcast_before_it_raises(monkeypatch):
-    """Rank 0 must still reach the broadcast on failure, or every other rank hangs in it."""
-    sent = _as_rank(monkeypatch, 0)
-    with pytest.raises(ValueError, match="must divide world_size"):
-        decide_distribution_plan(moe_cfg(), world_size=8, ep_size=3)
-    ((_, plan, error),) = sent
-    assert plan is None
-    assert "ValueError" in error
-
-
-def test_other_ranks_raise_when_rank0_failed(monkeypatch):
-    _as_rank(monkeypatch, 1, rank0_message=("plan", None, "ValueError: bad degrees"))
-    with pytest.raises(RuntimeError, match=r"rank 0 failed to compute.*bad degrees"):
-        decide_distribution_plan(moe_cfg(), world_size=8)
-
-
-def test_a_single_process_decides_alone(monkeypatch):
-    monkeypatch.setattr(
-        dist, "broadcast_object_list", lambda *a, **k: pytest.fail("a collective ran")
-    )
-    assert decide_distribution_plan(moe_cfg(), world_size=8, ep_size=4).ep_size == 4
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
new file mode 100644
index 00000000000..7a63aa599f8
--- /dev/null
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -0,0 +1,573 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""How ``resolve_distribution_plan`` decides a model's distribution, and the MoE helpers it uses.
+
+Nothing here needs a process group or a GPU: the process group's rank and size are answered the way
+a torchrun launch would, and the configs are tiny real transformers configs served from memory, so
+the decision builds their real models on the meta device.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+transformers = pytest.importorskip("transformers")
+
+from packaging.version import Version
+from transformers import (
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    LlamaConfig,
+    Qwen2_5_VLConfig,
+    Qwen3MoeConfig,
+)
+
+import modelopt.torch.utils.distributed as dist
+from modelopt.torch.utils import env
+from modelopt.torch.utils.distributed import DistributionPlan
+from modelopt.torch.utils.moe import (
+    _routed_expert_modules,
+    choose_ep_dp_degrees,
+    num_experts_from_config,
+)
+from modelopt.torch.utils.plugins import model_load_utils
+from modelopt.torch.utils.plugins.model_load_utils import _auto_class, resolve_distribution_plan
+
+_TINY = {
+    "hidden_size": 16,
+    "intermediate_size": 32,
+    "num_hidden_layers": 1,
+    "num_attention_heads": 2,
+    "num_key_value_heads": 1,
+    "vocab_size": 32,
+}
+
+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",
+)
+
+
+def _qwen3_moe(num_experts=8):
+    """Qwen3-MoE, whose base_model_ep_plan ships with transformers itself."""
+    return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY)
+
+
+def _bytes(module):
+    return sum(p.numel() * p.element_size() for p in module.parameters())
+
+
+def _meta(config):
+    """``config``'s model as the decision builds it, in the same dtype."""
+    return model_load_utils.build_meta_causal_lm(None, False, None, config)
+
+
+@pytest.fixture(autouse=True)
+def _no_distribution_knobs(monkeypatch):
+    """A shell that exports the knobs would otherwise steer every decision below."""
+    monkeypatch.delenv("MODELOPT_EP_DEGREE", raising=False)
+    monkeypatch.delenv("MODELOPT_FSDP_SIZE", raising=False)
+
+
+@pytest.fixture
+def run_as(monkeypatch):
+    """Run as rank ``rank`` of a ``world_size``-rank launch, deciding for itself.
+
+    How the ranks agree on one plan is tested at the end of this file.
+    """
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0)
+    monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None)
+
+    def enter(world_size, rank=0):
+        monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1)
+        monkeypatch.setattr(dist, "size", lambda group=None: world_size)
+        monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0)
+
+    return enter
+
+
+@pytest.fixture
+def serve(monkeypatch):
+    """Answer every config read with ``config`` rather than a checkpoint on disk."""
+
+    def serve(config):
+        monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **k: config)
+        return config
+
+    return serve
+
+
+@pytest.fixture
+def gpu_memory(monkeypatch):
+    """Give this rank a GPU of ``size`` bytes."""
+
+    def give(size):
+        monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: size)
+
+    return give
+
+
+# ----------------------------------------------------------------------------------- env knobs
+
+
+@pytest.mark.parametrize(
+    ("knob", "name"),
+    [(env.ep_degree, "MODELOPT_EP_DEGREE"), (env.fsdp_size, "MODELOPT_FSDP_SIZE")],
+)
+def test_distribution_knobs_default_to_auto(monkeypatch, knob, name):
+    """Unset and exported-but-empty both mean "let the plan decide"."""
+    assert knob() == "auto"
+    monkeypatch.setenv(name, "")
+    assert knob() == "auto"
+
+
+@pytest.mark.parametrize(
+    ("knob", "name", "other"),
+    [
+        (env.ep_degree, "MODELOPT_EP_DEGREE", env.fsdp_size),
+        (env.fsdp_size, "MODELOPT_FSDP_SIZE", env.ep_degree),
+    ],
+)
+def test_distribution_knobs_are_read_when_called(monkeypatch, knob, name, other):
+    """A value snapshotted at import would miss what the launcher sets later; and each knob reads
+    only its own variable."""
+    monkeypatch.setenv(name, "2")
+    assert knob() == "2"
+    monkeypatch.setenv(name, "8")
+    assert knob() == "8"
+    assert other() == "auto"
+
+
+# ------------------------------------------------------------------------------- expert counting
+
+
+def test_dense_config_reports_no_experts():
+    assert num_experts_from_config(SimpleNamespace(hidden_size=128)) == 0
+
+
+@pytest.mark.parametrize(
+    "attr", ["num_experts", "num_local_experts", "n_routed_experts", "num_routed_experts"]
+)
+def test_every_spelling_of_the_expert_count_is_understood(attr):
+    """Architectures disagree on the field name; missing one would silently disable EP."""
+    assert num_experts_from_config(SimpleNamespace(**{attr: 16})) == 16
+
+
+@pytest.mark.parametrize("nest", ["text_config", "llm_config"])
+def test_expert_count_is_found_on_a_vlms_nested_text_model(nest):
+    """For a VLM the MoE is the nested LM, so the root config has no expert count of its own."""
+    nested = SimpleNamespace(n_routed_experts=32)
+    assert num_experts_from_config(SimpleNamespace(**{nest: nested})) == 32
+
+
+# ------------------------------------------------------------------------------- degree selection
+
+
+def test_largest_valid_ep_degree_is_chosen():
+    assert choose_ep_dp_degrees(8, 64) == (8, 1)
+
+
+def test_ep_must_divide_the_expert_count_so_ranks_hold_whole_experts():
+    assert choose_ep_dp_degrees(8, 6) == (2, 4)
+
+
+def test_ep_falls_back_to_one_when_nothing_divides():
+    assert choose_ep_dp_degrees(8, 5) == (1, 8)
+
+
+@pytest.mark.parametrize(("world", "experts"), [(8, 0), (1, 64)])
+def test_no_expert_parallelism_for_dense_or_single_rank(world, experts):
+    assert choose_ep_dp_degrees(world, experts) == (1, world)
+
+
+@pytest.mark.parametrize("world", [2, 4, 8, 16])
+def test_degrees_always_factor_the_world(world):
+    ep, dp = choose_ep_dp_degrees(world, 64)
+    assert ep * dp == world
+
+
+# ---------------------------------------------------------------------------- routed experts
+
+
+class _MoeLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.self_attn = nn.Linear(2, 2)
+        self.mlp = nn.Module()
+        self.mlp.gate = nn.Linear(2, 4, bias=False)  # the router
+        self.mlp.experts = nn.Module()
+        self.mlp.experts.gate_up_proj = nn.Parameter(torch.zeros(4, 4, 2))
+        self.mlp.experts.down_proj = nn.Parameter(torch.zeros(4, 2, 2))
+        self.mlp.shared_experts = nn.Linear(2, 4, bias=False)
+
+
+def _model_with_decoder_at(prefix):
+    """A two-layer MoE whose decoder (``layers.<i>``) sits under ``prefix`` (``""``: the decoder)."""
+    decoder = nn.Module()
+    decoder.embed_tokens = nn.Embedding(8, 2)
+    decoder.layers = nn.ModuleList(_MoeLayer() for _ in range(2))
+    if not prefix:
+        return decoder
+    root = parent = nn.Module()
+    *path, leaf = prefix.rstrip(".").split(".")
+    for name in path:
+        parent.add_module(name, nn.Module())
+        parent = getattr(parent, name)
+    parent.add_module(leaf, decoder)
+    return root
+
+
+_STOCK_PLAN = {
+    "layers.*.mlp.gate": "ep_router",
+    "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+    "layers.*.mlp.experts.down_proj": "grouped_gemm",
+    "layers.*.mlp.experts": "moe_tp_experts",
+}
+
+
+@pytest.mark.parametrize(
+    "plan",
+    [
+        _STOCK_PLAN,
+        {"layers.*.mlp.experts": "moe_tp_experts"},
+        {"layers.*.mlp.experts": "nemotron_experts_ep"},
+        {k: v for k, v in _STOCK_PLAN.items() if v == "grouped_gemm"},
+    ],
+    ids=["stock", "experts-style-only", "nemotron-style-only", "grouped-gemm-only"],
+)
+@pytest.mark.parametrize(
+    "prefix",
+    ["", "model.", "model.language_model.", "language_model.model."],
+    ids=["base-model", "causal-lm", "vlm", "legacy-vlm"],
+)
+def test_routed_experts_are_the_plans_containers_under_any_prefix(plan, prefix):
+    """The base-relative plan keys match the decoder wherever the model nests it, at every layer
+    index, the container named by its style or as its grouped-GEMM parameters' parent. The router,
+    the shared experts, and a module whose name merely ends in a key's text are not routed
+    experts."""
+    model = _model_with_decoder_at(prefix)
+    model.add_module("sublayers", nn.ModuleList([_MoeLayer()]))  # "sublayers.0.mlp.experts"
+
+    found = _routed_expert_modules(model, plan)
+
+    modules = dict(model.named_modules())
+    assert found == {
+        name: modules[name] for name in [f"{prefix}layers.{i}.mlp.experts" for i in (0, 1)]
+    }
+
+
+def test_a_model_without_an_ep_plan_has_no_routed_experts():
+    assert _routed_expert_modules(_model_with_decoder_at("model."), {}) == {}
+
+
+# ------------------------------------------------------------------------------ model class
+
+
+def _vlm_config(**kwargs):
+    return SimpleNamespace(text_config=SimpleNamespace(), architectures=[], **kwargs)
+
+
+def _tiny_qwen2_5_vl():
+    return Qwen2_5_VLConfig(
+        text_config={**_TINY, "max_position_embeddings": 16},
+        vision_config={
+            "depth": 1,
+            "hidden_size": 16,
+            "intermediate_size": 16,
+            "num_heads": 2,
+            "out_hidden_size": 16,
+            "patch_size": 2,
+            "spatial_merge_size": 1,
+            "temporal_patch_size": 1,
+            "fullatt_block_indexes": [0],
+        },
+    )
+
+
+@pytest.mark.parametrize(
+    ("config", "require_lm_head", "expected"),
+    [
+        (SimpleNamespace(architectures=["LlamaForCausalLM"]), True, AutoModelForCausalLM),
+        (_vlm_config(), False, AutoModel),
+        (
+            _vlm_config(auto_map={"AutoModelForCausalLM": "m.KimiK25ForConditionalGeneration"}),
+            True,
+            AutoModelForCausalLM,
+        ),
+        (_tiny_qwen2_5_vl(), True, AutoModelForImageTextToText),
+        (_vlm_config(), True, AutoModel),
+    ],
+    ids=[
+        "causal-lm",
+        "vlm-backbone",
+        "remote-vlm-with-head",
+        "native-vlm-with-head",
+        "unknown-vlm",
+    ],
+)
+def test_the_class_that_builds_the_whole_model(config, require_lm_head, expected):
+    """A VLM's AutoModel leaves the LM head out: with the head required it builds through the class
+    that carries it, keeping the headless backbone only when no class recognizes the config."""
+    assert _auto_class(config, require_lm_head) is expected
+
+
+# -------------------------------------------------------------------------------- the decision
+
+
+@pytest.mark.parametrize("world_size", [None, 1], ids=["no_process_group", "one_rank"])
+def test_a_single_process_holds_the_whole_model(run_as, serve, monkeypatch, tmp_path, world_size):
+    """With no ranks to shard over, even an MoE that knobs ask to shard loads whole, and no model is
+    built to decide it."""
+    if world_size is not None:
+        run_as(world_size)
+    served = serve(_qwen3_moe())
+    monkeypatch.setattr(
+        model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model")
+    )
+
+    assert resolve_distribution_plan(str(tmp_path)) == (DistributionPlan(), served)
+    assert resolve_distribution_plan(str(tmp_path), dense="fsdp")[0] == DistributionPlan(
+        dense="fsdp"
+    )
+
+
+@needs_qwen3_moe_ep_plan
+@pytest.mark.parametrize("rank", [0, 1])
+def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, capsys, rank):
+    """Auto picks the largest EP degree dividing both the world and the expert count, hands back
+    the config it decided from, and only rank 0 reports the decision."""
+    run_as(4, rank)
+    served = serve(_qwen3_moe(num_experts=8))
+
+    plan, hf_config = resolve_distribution_plan(str(tmp_path))
+
+    assert plan == DistributionPlan(4, ep_size=4, dense="replicate")
+    assert hf_config is served
+    assert ("Distribution:" in capsys.readouterr().out) == (rank == 0)
+
+
+@needs_qwen3_moe_ep_plan
+def test_no_ep_degree_dividing_the_experts_leaves_them_whole(run_as, serve, tmp_path):
+    run_as(2)
+    serve(_qwen3_moe(num_experts=3))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2)
+
+
+def test_an_moe_without_an_ep_plan_is_not_expert_parallel(run_as, serve, tmp_path):
+    """The config counts 8 experts, but with no base_model_ep_plan nothing could shard them."""
+    run_as(4)
+    config = _qwen3_moe(num_experts=8)
+    config.base_model_ep_plan = None
+    serve(config)
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
+
+
+@needs_qwen3_moe_ep_plan
+def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path):
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", "0")
+    run_as(4)
+    serve(_qwen3_moe(num_experts=8))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)
+
+
+def test_a_dense_model_that_fits_one_gpu_is_replicated(run_as, serve, gpu_memory, tmp_path):
+    run_as(4)
+    config = serve(LlamaConfig(**_TINY))
+    gpu_memory(2 * _bytes(_meta(config)))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
+
+
+def test_a_dense_model_too_big_for_one_gpu_is_fsdp_sharded(run_as, serve, gpu_memory, tmp_path):
+    """Whole, it would leave less than 40% of the GPU to calibration."""
+    run_as(4)
+    config = serve(LlamaConfig(**_TINY))
+    gpu_memory(_bytes(_meta(config)))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, dense="fsdp")
+
+
+@needs_qwen3_moe_ep_plan
+def test_only_what_expert_parallelism_leaves_whole_must_fit(run_as, serve, gpu_memory, tmp_path):
+    """The routed experts are sharded at EP 4, so only the rest needs to fit one GPU; with no EP the
+    whole MoE does, and here it does not."""
+    run_as(4)
+    config = serve(_qwen3_moe(num_experts=8))
+    model = _meta(config)
+    experts = sum(_bytes(layer.mlp.experts) for layer in model.model.layers)
+    gpu_memory(int((_bytes(model) - experts / 2) / 0.6))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)
+    assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan(
+        4, dense="fsdp"
+    )
+
+
+@needs_qwen3_moe_ep_plan
+def test_the_env_knobs_override_the_automatic_plan(run_as, serve, monkeypatch, tmp_path):
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", "2")
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+    run_as(4)
+    serve(_qwen3_moe(num_experts=8))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 2, "fsdp")
+
+
+@needs_qwen3_moe_ep_plan
+def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path):
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", "4")
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+    run_as(4)
+    serve(_qwen3_moe(num_experts=8))
+
+    plan, _ = resolve_distribution_plan(str(tmp_path), ep_size=2, dense="replicate")
+
+    assert plan == DistributionPlan(4, 2, "replicate")
+
+
+@pytest.mark.parametrize(
+    ("value", "dense"), [("1", "replicate"), ("4", "fsdp"), ("full", "fsdp"), ("WORLD", "fsdp")]
+)
+def test_the_dense_knob_replicates_or_shards_over_the_whole_world(
+    run_as, serve, monkeypatch, tmp_path, value, dense
+):
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", value)
+    run_as(4)
+    serve(LlamaConfig(**_TINY))
+
+    assert resolve_distribution_plan(str(tmp_path))[0].dense == dense
+
+
+def test_a_dense_knob_sharding_over_part_of_the_world_is_refused(
+    run_as, serve, monkeypatch, tmp_path
+):
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "2")
+    run_as(4)
+    serve(LlamaConfig(**_TINY))
+
+    with pytest.raises(ValueError, match=r"expected auto, 1 .* or world"):
+        resolve_distribution_plan(str(tmp_path))
+
+
+@pytest.mark.parametrize(
+    ("config", "ep_size", "experts"),
+    [(LlamaConfig(**_TINY), 2, 0), (_qwen3_moe(num_experts=6), 4, 6)],
+    ids=["dense-model", "experts-not-divisible"],
+)
+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")
+    run_as(4)
+    serve(config)
+
+    with pytest.raises(ValueError, match=rf"ep_size={ep_size} must divide .* \({experts}\)"):
+        resolve_distribution_plan(str(tmp_path), ep_size=ep_size)
+
+
+def test_a_vlm_is_decided_from_its_whole_model(run_as, serve, gpu_memory, tmp_path):
+    """AutoModelForCausalLM does not build every VLM: the decision builds what the load builds."""
+    run_as(2)
+    serve(_tiny_qwen2_5_vl())
+    gpu_memory(1)
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp")
+
+
+def test_a_model_that_cannot_be_built_on_meta_names_itself(run_as, serve, monkeypatch, tmp_path):
+    def unbuildable(*args, **kwargs):
+        raise NotImplementedError("custom init touches the weights")
+
+    run_as(2)
+    serve(LlamaConfig(**_TINY))
+    monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", unbuildable)
+
+    with pytest.raises(RuntimeError, match=r"distribute llama .* custom init touches the weights"):
+        resolve_distribution_plan(str(tmp_path))
+
+
+# ------------------------------------------------------------------------------ rank 0 decides
+
+
+def _as_rank(monkeypatch, rank, rank0_message=None):
+    """Run as ``rank`` of a 4-rank run, answering its broadcast; returns what this rank handed it.
+
+    A non-zero rank receives ``rank0_message``, as rank 0 would have sent it.
+    """
+    sent = []
+
+    def broadcast_object_list(message, group=None, group_src=None):
+        sent.append(message[0])
+        if rank != 0:
+            message[0] = rank0_message
+
+    monkeypatch.setattr(dist, "is_distributed_run", lambda: True)
+    monkeypatch.setattr(dist, "size", lambda group=None: 4)
+    monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0)
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: rank)
+    monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list)
+    return sent
+
+
+def test_rank0_decides_and_sends_its_plan(monkeypatch, serve, tmp_path):
+    sent = _as_rank(monkeypatch, 0)
+    serve(LlamaConfig(**_TINY))
+
+    plan, _ = resolve_distribution_plan(str(tmp_path), dense="fsdp")
+
+    assert sent == [("plan", DistributionPlan(4, dense="fsdp"), None)]
+    assert plan == DistributionPlan(4, dense="fsdp")
+
+
+def test_other_ranks_adopt_rank0s_plan_without_deciding(monkeypatch, serve, tmp_path):
+    """A non-zero rank's own knobs, here a dense layout of its own, must not leak into the plan."""
+    rank0_plan = DistributionPlan(4, dense="fsdp")
+    _as_rank(monkeypatch, 1, rank0_message=("plan", rank0_plan, None))
+    serve(LlamaConfig(**_TINY))
+    monkeypatch.setattr(
+        model_load_utils, "_decide_distribution", lambda *a: pytest.fail("rank 1 decided")
+    )
+
+    assert resolve_distribution_plan(str(tmp_path), dense="replicate")[0] == rank0_plan
+
+
+def test_rank0s_failure_reaches_the_broadcast_before_it_raises(monkeypatch, serve, tmp_path):
+    """Rank 0 must still reach the broadcast on failure, or every other rank hangs in it."""
+    sent = _as_rank(monkeypatch, 0)
+    serve(LlamaConfig(**_TINY))
+
+    with pytest.raises(ValueError, match="must divide"):
+        resolve_distribution_plan(str(tmp_path), ep_size=2)
+    ((_, plan, error),) = sent
+    assert plan is None
+    assert "ValueError" in error
+
+
+def test_other_ranks_raise_when_rank0_failed(monkeypatch, serve, tmp_path):
+    _as_rank(monkeypatch, 1, rank0_message=("plan", None, "ValueError: bad degrees"))
+    serve(LlamaConfig(**_TINY))
+
+    with pytest.raises(RuntimeError, match=r"rank 0 failed to compute plan.*bad degrees"):
+        resolve_distribution_plan(str(tmp_path))

From 5f3a86cc9c285183602bc88525179b4c81b9b3d6 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 14 Oct 2026 21:29:23 +0000
Subject: [PATCH 07/44] Keep each env knob on one docstring line, as Sphinx
 parses it

The module docstring lists its knobs as an indented block under a
paragraph line, which reStructuredText reads as a definition list. A
knob wrapped onto a deeper-indented second line is unexpected
indentation there, and the docs build fails on it.

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

diff --git a/modelopt/torch/utils/env.py b/modelopt/torch/utils/env.py
index e8e7827ef07..2ca20e97b32 100644
--- a/modelopt/torch/utils/env.py
+++ b/modelopt/torch/utils/env.py
@@ -21,8 +21,7 @@
 
 Distribution plan (see ``resolve_distribution_plan``):
   MODELOPT_EP_DEGREE        expert-parallel degree; ``""``/unset -> ``"auto"``.
-  MODELOPT_FSDP_SIZE        whether to FSDP2-shard the dense part: ``"1"`` replicates it, ``"world"``
-                            shards it over every rank; ``""``/unset -> ``"auto"``.
+  MODELOPT_FSDP_SIZE        dense layout: ``"1"`` replicates, ``"world"`` FSDP2-shards; ``""``/unset -> ``"auto"``.
 """
 
 import os

From db4cfbbb5eaefa55f0ac2cf255cf021cf1b5e9d8 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 14 Oct 2026 21:42:10 +0000
Subject: [PATCH 08/44] Name the world size as a value MODELOPT_FSDP_SIZE
 accepts

MODELOPT_FSDP_SIZE FSDP2-shards the dense part for "world" or for the
world size as a number, which is what launchers set (64 on 64 ranks),
but the docs named only "world". Say so in the env module and in
resolve_distribution_plan. The module lists its knobs as a bullet list
now, so an entry can wrap without breaking the Sphinx build.

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

diff --git a/modelopt/torch/utils/env.py b/modelopt/torch/utils/env.py
index 2ca20e97b32..090fa8050cd 100644
--- a/modelopt/torch/utils/env.py
+++ b/modelopt/torch/utils/env.py
@@ -20,8 +20,10 @@
 matters. Callers do ``from modelopt.torch.utils import env`` then ``env.ep_degree()``.
 
 Distribution plan (see ``resolve_distribution_plan``):
-  MODELOPT_EP_DEGREE        expert-parallel degree; ``""``/unset -> ``"auto"``.
-  MODELOPT_FSDP_SIZE        dense layout: ``"1"`` replicates, ``"world"`` FSDP2-shards; ``""``/unset -> ``"auto"``.
+
+- ``MODELOPT_EP_DEGREE``: expert-parallel degree; ``""``/unset -> ``"auto"``.
+- ``MODELOPT_FSDP_SIZE``: the dense layout. ``"1"`` replicates it; ``"world"``, or the world size as a
+  number, FSDP2-shards it over every rank; ``""``/unset -> ``"auto"``.
 """
 
 import os
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index f8418e9d888..e4cd9cc8707 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1030,9 +1030,10 @@ def resolve_distribution_plan(
     * ``dense``: ``"replicate"`` or ``"fsdp"`` for the rest of the model, or ``"auto"``, which
       FSDP2-shards it when a whole copy would take more than 60% of one GPU's memory.
 
-    ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` (``1`` replicates, ``world``
-    shards) from rank 0's environment, both defaulting to ``"auto"``. A single process holds the
-    whole model unless ``dense="fsdp"`` asks for FSDP2 on one rank.
+    ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` (``1`` replicates; ``world``,
+    or the world size as a number, shards) from rank 0's environment, both defaulting to
+    ``"auto"``. A single process holds the whole model unless ``dense="fsdp"`` asks for FSDP2 on
+    one rank.
     """
     from modelopt.torch.utils import distributed as dist
 

From 5c34be95bc9924c2822669800814e18633e8d148 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 15 Oct 2026 21:05:00 +0000
Subject: [PATCH 09/44] Say which GPU the automatic dense placement is sized
 against

The plan decision sizes dense="auto" against rank 0's GPU, by its total
memory, and every rank takes that decision; the docstring said only "one
GPU's memory", leaving the assumption that the run's GPUs are alike
unstated.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index e4cd9cc8707..a0009d2ee59 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1028,7 +1028,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``.
     * ``dense``: ``"replicate"`` or ``"fsdp"`` for the rest of the model, or ``"auto"``, which
-      FSDP2-shards it when a whole copy would take more than 60% of one GPU's memory.
+      FSDP2-shards it when a whole copy would take more than 60% of one GPU's memory. That GPU is
+      rank 0's, by its total memory, so ``"auto"`` assumes the run's GPUs are alike.
 
     ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` (``1`` replicates; ``world``,
     or the world size as a number, shards) from rank 0's environment, both defaulting to

From 9b157ac6e97dc5645265f74d67784a98e4502fee Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 02:17:46 +0000
Subject: [PATCH 10/44] Fail every rank when one cannot read the checkpoint,
 before the plan's broadcast

resolve_distribution_plan made the checkpoint local and read its config
on each rank before the broadcasts rank 0 sends the plan through. A rank
that could not see the checkpoint raised there alone, and rank 0 waited in
the broadcast for it. The read now goes through the failure-carrying
all_reduce, so every rank raises with the failing rank's error first. The
tests that answer a launch's broadcasts answer its gather too.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   | 14 ++++-
 .../utils/test_resolve_distribution_plan.py   | 58 +++++++++++++++++--
 2 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index a0009d2ee59..0436167917e 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1038,8 +1038,18 @@ def resolve_distribution_plan(
     """
     from modelopt.torch.utils import distributed as dist
 
-    _, ckpt_path = ensure_local_checkpoint(ckpt_path)
-    hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code)
+    # 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]] = []
+
+    def checkpoint_read() -> bool:
+        _, local_path = ensure_local_checkpoint(ckpt_path)
+        config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code)
+        read.append((local_path, config))
+        return True
+
+    dist.all_reduce("checkpoint read", checkpoint_read, all)
+    ckpt_path, hf_config = read[0]
     world_size = dist.size()
     plan = dist.broadcast_from(
         "plan",
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 7a63aa599f8..e7db51dc14c 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -15,12 +15,16 @@
 
 """How ``resolve_distribution_plan`` decides a model's distribution, and the MoE helpers it uses.
 
-Nothing here needs a process group or a GPU: the process group's rank and size are answered the way
-a torchrun launch would, and the configs are tiny real transformers configs served from memory, so
-the decision builds their real models on the meta device.
+Nothing here needs a GPU, and only the last test a process group: elsewhere the process group's
+rank and size are answered the way a torchrun launch would, and the configs are tiny real
+transformers configs served from memory, so the decision builds their real models on the meta device.
 """
 
+import os
+from datetime import timedelta
+from functools import partial
 from types import SimpleNamespace
+from unittest import mock
 
 import pytest
 import torch
@@ -85,19 +89,27 @@ def _no_distribution_knobs(monkeypatch):
     monkeypatch.delenv("MODELOPT_FSDP_SIZE", raising=False)
 
 
+def _gathered_alike(gathered, mine, group=None):
+    """Answer a gather as if every rank handed it what this rank did."""
+    gathered[:] = [mine] * len(gathered)
+
+
 @pytest.fixture
 def run_as(monkeypatch):
     """Run as rank ``rank`` of a ``world_size``-rank launch, deciding for itself.
 
-    How the ranks agree on one plan is tested at the end of this file.
+    Every rank reads the checkpoint as this one does. How the ranks agree on one plan is tested at
+    the end of this file.
     """
     monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0)
     monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None)
+    monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike)
 
     def enter(world_size, rank=0):
         monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1)
         monkeypatch.setattr(dist, "size", lambda group=None: world_size)
         monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0)
+        monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size)
 
     return enter
 
@@ -514,7 +526,8 @@ def unbuildable(*args, **kwargs):
 def _as_rank(monkeypatch, rank, rank0_message=None):
     """Run as ``rank`` of a 4-rank run, answering its broadcast; returns what this rank handed it.
 
-    A non-zero rank receives ``rank0_message``, as rank 0 would have sent it.
+    A non-zero rank receives ``rank0_message``, as rank 0 would have sent it. Every rank reads the
+    checkpoint as this one does.
     """
     sent = []
 
@@ -527,7 +540,9 @@ def broadcast_object_list(message, group=None, group_src=None):
     monkeypatch.setattr(dist, "size", lambda group=None: 4)
     monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0)
     monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: rank)
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 4)
     monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list)
+    monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike)
     return sent
 
 
@@ -571,3 +586,36 @@ def test_other_ranks_raise_when_rank0_failed(monkeypatch, serve, tmp_path):
 
     with pytest.raises(RuntimeError, match=r"rank 0 failed to compute plan.*bad degrees"):
         resolve_distribution_plan(str(tmp_path))
+
+
+def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt):
+    # A rank left waiting in a collective its peer never joins gives up in seconds on this group.
+    group = torch.distributed.new_group(timeout=timedelta(seconds=10))
+    isdir = os.path.isdir
+    with (
+        mock.patch.object(
+            torch.distributed,
+            "broadcast_object_list",
+            partial(torch.distributed.broadcast_object_list, group=group),
+        ),
+        mock.patch.object(
+            torch.distributed,
+            "all_gather_object",
+            partial(torch.distributed.all_gather_object, group=group),
+        ),
+        mock.patch("os.path.isdir", lambda path: (rank == 0 or path != ckpt) and isdir(path)),
+    ):
+        try:
+            resolve_distribution_plan(ckpt)
+        except RuntimeError as error:
+            raised = str(error)
+        else:
+            raised = "no error"
+    # Rank 1's own error, carried to rank 0, rather than rank 0 timing out in the plan's broadcast.
+    assert "which this rank cannot see" in raised, raised
+
+
+def test_a_rank_that_cannot_see_the_checkpoint_fails_every_rank(cpu_dist_workers, tmp_path):
+    """Before the plan's broadcast, which rank 0 would otherwise wait in for a rank that raised."""
+    LlamaConfig(**_TINY).save_pretrained(tmp_path)
+    cpu_dist_workers.run(_one_rank_cannot_see_the_checkpoint_job, ckpt=str(tmp_path))

From 17adc73989cf89187d923bfd8750d0d72f69406e Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 01:18:29 +0000
Subject: [PATCH 11/44] Read no distribution knob in a single process

A single process read MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE before
returning its whole-model plan, so a shell still exporting a multi-GPU
run's knobs reshaped a one-process load: an EP degree of 4 made a plan
the world of one cannot hold and raised, MODELOPT_FSDP_SIZE=world asked a
plain load for FSDP2, and a size of 2 raised. With nothing to shard over,
the knobs are now not read; the caller's own arguments still count, so
dense="fsdp" is still FSDP2 on one rank.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   | 18 +++++++++-------
 .../utils/test_resolve_distribution_plan.py   | 21 ++++++++++++++++---
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 0436167917e..142f348f203 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -962,15 +962,17 @@ def _decide_distribution(
         num_experts_from_config,
     )
 
-    ep_size = env.ep_degree() if ep_size is None else ep_size
-    ep_size = "auto" if str(ep_size) in ("auto", "0") else int(ep_size)
-    dense = _dense_from_env(world_size) if dense is None else dense
     if world_size == 1:
-        # Nothing to shard over, so nothing about the model matters.
+        # 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.
         return DistributionPlan(
-            ep_size=1 if ep_size == "auto" else ep_size,
-            dense="replicate" if dense == "auto" else dense,
+            ep_size=1 if ep_size is None or str(ep_size) in ("auto", "0") else int(ep_size),
+            dense="replicate" if dense in (None, "auto") else dense,
         )
+    ep_size = env.ep_degree() if ep_size is None else ep_size
+    ep_size = "auto" if str(ep_size) in ("auto", "0") else int(ep_size)
+    dense = _dense_from_env(world_size) if dense is None else dense
     try:
         model = build_meta_causal_lm(
             ckpt_path, trust_remote_code, None, hf_config, auto_class=_auto_class(hf_config)
@@ -1033,8 +1035,8 @@ def resolve_distribution_plan(
 
     ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` (``1`` replicates; ``world``,
     or the world size as a number, shards) from rank 0's environment, both defaulting to
-    ``"auto"``. A single process holds the whole model unless ``dense="fsdp"`` asks for FSDP2 on
-    one rank.
+    ``"auto"``. A single process reads neither and holds the whole model, unless ``dense="fsdp"``
+    asks for FSDP2 on one rank.
     """
     from modelopt.torch.utils import distributed as dist
 
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index e7db51dc14c..5f8ac52afa3 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -342,12 +342,27 @@ def test_the_class_that_builds_the_whole_model(config, require_lm_head, expected
 # -------------------------------------------------------------------------------- the decision
 
 
+@pytest.mark.parametrize(
+    "knob",
+    [
+        None,
+        ("MODELOPT_EP_DEGREE", "4"),
+        ("MODELOPT_FSDP_SIZE", "world"),
+        ("MODELOPT_FSDP_SIZE", "2"),
+    ],
+    ids=["no_knob", "ep_degree", "fsdp_world", "fsdp_part"],
+)
 @pytest.mark.parametrize("world_size", [None, 1], ids=["no_process_group", "one_rank"])
-def test_a_single_process_holds_the_whole_model(run_as, serve, monkeypatch, tmp_path, world_size):
-    """With no ranks to shard over, even an MoE that knobs ask to shard loads whole, and no model is
-    built to decide it."""
+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."""
     if world_size is not None:
         run_as(world_size)
+    if knob is not None:
+        monkeypatch.setenv(*knob)
     served = serve(_qwen3_moe())
     monkeypatch.setattr(
         model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model")

From 0673d4d9dba31a4b658f35071bb29a386a233dd2 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 07:09:53 +0000
Subject: [PATCH 12/44] Test on two real ranks that every rank takes the plan
 rank 0 decided

Two tests played rank 0 and rank 1 against a stand-in broadcast and
asserted the message it was handed, and two more repeated
broadcast_from's own failure tests. One pooled job replaces them: both
gloo ranks resolve the plan of a real checkpoint, rank 1 with a knob of
its own that would decide FSDP2, and both must load the plan rank 0
decided. The failure tests now live with broadcast_from's, on real
ranks too, and the checkpoint-read job takes its bounded group from the
shared helper.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../utils/test_resolve_distribution_plan.py   | 103 +++++-------------
 1 file changed, 30 insertions(+), 73 deletions(-)

diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 5f8ac52afa3..53f18079aa6 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -15,19 +15,20 @@
 
 """How ``resolve_distribution_plan`` decides a model's distribution, and the MoE helpers it uses.
 
-Nothing here needs a GPU, and only the last test a process group: elsewhere the process group's
-rank and size are answered the way a torchrun launch would, and the configs are tiny real
-transformers configs served from memory, so the decision builds their real models on the meta device.
+Nothing here needs a GPU, and only the tests at the end a process group, two gloo ranks: elsewhere
+the process group's rank and size are answered the way a torchrun launch would, and the configs are
+tiny real transformers configs served from memory, so the decision builds their real models on the
+meta device.
 """
 
 import os
-from datetime import timedelta
-from functools import partial
+from contextlib import contextmanager
 from types import SimpleNamespace
 from unittest import mock
 
 import pytest
 import torch
+from _test_utils.torch.distributed.utils import bounded_collectives
 from torch import nn
 
 transformers = pytest.importorskip("transformers")
@@ -538,86 +539,42 @@ def unbuildable(*args, **kwargs):
 # ------------------------------------------------------------------------------ rank 0 decides
 
 
-def _as_rank(monkeypatch, rank, rank0_message=None):
-    """Run as ``rank`` of a 4-rank run, answering its broadcast; returns what this rank handed it.
+@contextmanager
+def _knobs(**values):
+    """This rank's distribution knobs, exactly ``values``, within the block.
 
-    A non-zero rank receives ``rank0_message``, as rank 0 would have sent it. Every rank reads the
-    checkpoint as this one does.
+    A pooled worker keeps the environment it was spawned with, so the knobs are cleared first. The
+    environment is restored on the way out.
     """
-    sent = []
-
-    def broadcast_object_list(message, group=None, group_src=None):
-        sent.append(message[0])
-        if rank != 0:
-            message[0] = rank0_message
-
-    monkeypatch.setattr(dist, "is_distributed_run", lambda: True)
-    monkeypatch.setattr(dist, "size", lambda group=None: 4)
-    monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0)
-    monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: rank)
-    monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 4)
-    monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list)
-    monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike)
-    return sent
-
+    with mock.patch.dict(os.environ):
+        for name in ("MODELOPT_EP_DEGREE", "MODELOPT_FSDP_SIZE"):
+            os.environ.pop(name, None)
+        os.environ.update(values)
+        yield
 
-def test_rank0_decides_and_sends_its_plan(monkeypatch, serve, tmp_path):
-    sent = _as_rank(monkeypatch, 0)
-    serve(LlamaConfig(**_TINY))
 
-    plan, _ = resolve_distribution_plan(str(tmp_path), dense="fsdp")
+def _every_rank_takes_rank0s_plan_job(rank, size, ckpt, knobs, expected):
+    with bounded_collectives(), _knobs(**knobs[rank]):
+        plan, _ = resolve_distribution_plan(ckpt)
+    assert plan == expected, (rank, plan)
 
-    assert sent == [("plan", DistributionPlan(4, dense="fsdp"), None)]
-    assert plan == DistributionPlan(4, dense="fsdp")
 
-
-def test_other_ranks_adopt_rank0s_plan_without_deciding(monkeypatch, serve, tmp_path):
-    """A non-zero rank's own knobs, here a dense layout of its own, must not leak into the plan."""
-    rank0_plan = DistributionPlan(4, dense="fsdp")
-    _as_rank(monkeypatch, 1, rank0_message=("plan", rank0_plan, None))
-    serve(LlamaConfig(**_TINY))
-    monkeypatch.setattr(
-        model_load_utils, "_decide_distribution", lambda *a: pytest.fail("rank 1 decided")
+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)
+    cpu_dist_workers.run(
+        _every_rank_takes_rank0s_plan_job,
+        ckpt=str(tmp_path),
+        knobs=[{}, {"MODELOPT_FSDP_SIZE": "world"}],
+        expected=DistributionPlan(2),
     )
 
-    assert resolve_distribution_plan(str(tmp_path), dense="replicate")[0] == rank0_plan
-
-
-def test_rank0s_failure_reaches_the_broadcast_before_it_raises(monkeypatch, serve, tmp_path):
-    """Rank 0 must still reach the broadcast on failure, or every other rank hangs in it."""
-    sent = _as_rank(monkeypatch, 0)
-    serve(LlamaConfig(**_TINY))
-
-    with pytest.raises(ValueError, match="must divide"):
-        resolve_distribution_plan(str(tmp_path), ep_size=2)
-    ((_, plan, error),) = sent
-    assert plan is None
-    assert "ValueError" in error
-
-
-def test_other_ranks_raise_when_rank0_failed(monkeypatch, serve, tmp_path):
-    _as_rank(monkeypatch, 1, rank0_message=("plan", None, "ValueError: bad degrees"))
-    serve(LlamaConfig(**_TINY))
-
-    with pytest.raises(RuntimeError, match=r"rank 0 failed to compute plan.*bad degrees"):
-        resolve_distribution_plan(str(tmp_path))
-
 
 def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt):
-    # A rank left waiting in a collective its peer never joins gives up in seconds on this group.
-    group = torch.distributed.new_group(timeout=timedelta(seconds=10))
     isdir = os.path.isdir
     with (
-        mock.patch.object(
-            torch.distributed,
-            "broadcast_object_list",
-            partial(torch.distributed.broadcast_object_list, group=group),
-        ),
-        mock.patch.object(
-            torch.distributed,
-            "all_gather_object",
-            partial(torch.distributed.all_gather_object, group=group),
-        ),
+        bounded_collectives(),
         mock.patch("os.path.isdir", lambda path: (rank == 0 or path != ckpt) and isdir(path)),
     ):
         try:

From ec3b0b410e3aeee2e355e3279a73da0f12e2e64c Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 07:11:06 +0000
Subject: [PATCH 13/44] Leave require_transformers and the headless VLM build
 to their first caller

Nothing in the distribution decision calls moe.require_transformers or
asks _auto_class for a VLM without its LM head: the expert-parallel
loader is the first to, so both, with the test case for the headless
backbone, arrive with it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/moe.py                   | 18 -------------
 .../torch/utils/plugins/model_load_utils.py   | 11 +++-----
 .../utils/test_resolve_distribution_plan.py   | 26 +++++++------------
 3 files changed, 13 insertions(+), 42 deletions(-)

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index e487e3fceea..4a4335ba111 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -35,24 +35,6 @@
     from torch import nn
 
 
-def require_transformers(min_version: str, feature: str) -> None:
-    """Fail loudly if the installed transformers is older than ``min_version`` (the floor for ``feature``).
-
-    Callers enforce the version where the needed APIs first exist, so a too-old transformers errors
-    clearly instead of silently mis-behaving (e.g. transformers 5.5 silently shards nothing for EP).
-    See each call site for the empirically-determined floor.
-    """
-    import transformers
-    from packaging.version import parse
-
-    have = transformers.__version__
-    if parse(have) < parse(min_version):
-        raise RuntimeError(
-            f"{feature} requires transformers >= {min_version}, but {have} is installed. "
-            f"Upgrade with: pip install -U 'transformers>={min_version}'."
-        )
-
-
 def _ep_config(hf_config):
     """Return the config that carries the EP plan / expert count.
 
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 142f348f203..823cec3883d 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -889,15 +889,14 @@ def parallel_load_and_prepare_fsdp2(
     return model
 
 
-def _auto_class(hf_config, require_lm_head: bool = True) -> Any:
-    """The auto class that builds ``hf_config``'s whole model.
+def _auto_class(hf_config) -> Any:
+    """The auto class that builds ``hf_config``'s whole model, LM head included.
 
     A causal LM builds through ``AutoModelForCausalLM``. A VLM -- a nested text config, a vision
-    config, or a ``*ForConditionalGeneration`` architecture -- builds through ``AutoModel``, which
-    leaves its LM head out. ``require_lm_head`` builds it with the head: through the remote
+    config, or a ``*ForConditionalGeneration`` architecture -- builds through the remote
     ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` there
     and no image-text class, else through ``AutoModelForImageTextToText`` if that recognizes the
-    config.
+    config, else through ``AutoModel``, which leaves the head out.
     """
     architectures = getattr(hf_config, "architectures", None) or []
     is_vlm = (
@@ -907,8 +906,6 @@ def _auto_class(hf_config, require_lm_head: bool = True) -> Any:
     )
     if not is_vlm:
         return AutoModelForCausalLM
-    if not require_lm_head:
-        return AutoModel
     auto_map = getattr(hf_config, "auto_map", None) or {}
     remote_with_head = (
         "conditionalgeneration" in str(auto_map.get("AutoModelForCausalLM", "")).lower()
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 53f18079aa6..3e48b2fe279 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -314,30 +314,22 @@ def _tiny_qwen2_5_vl():
 
 
 @pytest.mark.parametrize(
-    ("config", "require_lm_head", "expected"),
+    ("config", "expected"),
     [
-        (SimpleNamespace(architectures=["LlamaForCausalLM"]), True, AutoModelForCausalLM),
-        (_vlm_config(), False, AutoModel),
+        (SimpleNamespace(architectures=["LlamaForCausalLM"]), AutoModelForCausalLM),
         (
             _vlm_config(auto_map={"AutoModelForCausalLM": "m.KimiK25ForConditionalGeneration"}),
-            True,
             AutoModelForCausalLM,
         ),
-        (_tiny_qwen2_5_vl(), True, AutoModelForImageTextToText),
-        (_vlm_config(), True, AutoModel),
-    ],
-    ids=[
-        "causal-lm",
-        "vlm-backbone",
-        "remote-vlm-with-head",
-        "native-vlm-with-head",
-        "unknown-vlm",
+        (_tiny_qwen2_5_vl(), AutoModelForImageTextToText),
+        (_vlm_config(), AutoModel),
     ],
+    ids=["causal-lm", "remote-vlm-with-head", "native-vlm-with-head", "unknown-vlm"],
 )
-def test_the_class_that_builds_the_whole_model(config, require_lm_head, expected):
-    """A VLM's AutoModel leaves the LM head out: with the head required it builds through the class
-    that carries it, keeping the headless backbone only when no class recognizes the config."""
-    assert _auto_class(config, require_lm_head) is expected
+def test_the_class_that_builds_the_whole_model(config, expected):
+    """A VLM's AutoModel leaves the LM head out: it builds through the class that carries the head,
+    keeping the headless backbone only when no class recognizes the config."""
+    assert _auto_class(config) is expected
 
 
 # -------------------------------------------------------------------------------- the decision

From c1f2b1c52b330e7ca2e295805d65dab2318f1352 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 17:55:34 +0000
Subject: [PATCH 14/44] Trim the plan decision's tests

Fold near-duplicates into one test per behavior, and drop tests that
other tests already cover:
- test_the_expert_count_is_read_under_every_spelling takes the dense,
  four-spelling and two VLM-nesting cases.
- test_the_largest_ep_degree_dividing_the_world_and_the_experts_is_chosen
  takes the five degree cases. The ep * dp == world test goes: every
  case checks both degrees.
- test_a_dense_model_is_fsdp_sharded_only_when_too_big_for_one_gpu takes
  the fits and too-big cases.
- test_the_dense_knob_replicates_or_shards_over_the_whole_world takes the
  part-of-the-world refusal and an exported-empty case, and
  test_an_ep_degree_of_0_is_auto an exported-empty one. Both env-knob unit
  tests go: the decision tests fail if an empty or unset knob stops
  meaning auto, if a knob is read at import, or if one knob reads the
  other's variable.
- test_a_model_without_an_ep_plan_has_no_routed_experts goes:
  test_an_moe_without_an_ep_plan_is_not_expert_parallel fails if an
  empty plan finds experts.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../utils/test_resolve_distribution_plan.py   | 168 ++++++++----------
 1 file changed, 70 insertions(+), 98 deletions(-)

diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 3e48b2fe279..2e406178dff 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -44,7 +44,6 @@
 )
 
 import modelopt.torch.utils.distributed as dist
-from modelopt.torch.utils import env
 from modelopt.torch.utils.distributed import DistributionPlan
 from modelopt.torch.utils.moe import (
     _routed_expert_modules,
@@ -136,83 +135,56 @@ def give(size):
     return give
 
 
-# ----------------------------------------------------------------------------------- env knobs
-
-
-@pytest.mark.parametrize(
-    ("knob", "name"),
-    [(env.ep_degree, "MODELOPT_EP_DEGREE"), (env.fsdp_size, "MODELOPT_FSDP_SIZE")],
-)
-def test_distribution_knobs_default_to_auto(monkeypatch, knob, name):
-    """Unset and exported-but-empty both mean "let the plan decide"."""
-    assert knob() == "auto"
-    monkeypatch.setenv(name, "")
-    assert knob() == "auto"
+# ------------------------------------------------------------------------------- expert counting
 
 
 @pytest.mark.parametrize(
-    ("knob", "name", "other"),
+    ("config", "experts"),
     [
-        (env.ep_degree, "MODELOPT_EP_DEGREE", env.fsdp_size),
-        (env.fsdp_size, "MODELOPT_FSDP_SIZE", env.ep_degree),
+        (SimpleNamespace(hidden_size=128), 0),
+        (SimpleNamespace(num_experts=16), 16),
+        (SimpleNamespace(num_local_experts=16), 16),
+        (SimpleNamespace(n_routed_experts=16), 16),
+        (SimpleNamespace(num_routed_experts=16), 16),
+        (SimpleNamespace(text_config=SimpleNamespace(n_routed_experts=32)), 32),
+        (SimpleNamespace(llm_config=SimpleNamespace(n_routed_experts=32)), 32),
+    ],
+    ids=[
+        "dense",
+        "num_experts",
+        "num_local_experts",
+        "n_routed_experts",
+        "num_routed_experts",
+        "vlm-text_config",
+        "vlm-llm_config",
     ],
 )
-def test_distribution_knobs_are_read_when_called(monkeypatch, knob, name, other):
-    """A value snapshotted at import would miss what the launcher sets later; and each knob reads
-    only its own variable."""
-    monkeypatch.setenv(name, "2")
-    assert knob() == "2"
-    monkeypatch.setenv(name, "8")
-    assert knob() == "8"
-    assert other() == "auto"
+def test_the_expert_count_is_read_under_every_spelling(config, experts):
+    """Architectures disagree on the field name, and a VLM's MoE is its nested LM, whose root
+    config has no expert count of its own: missing either would silently disable EP."""
+    assert num_experts_from_config(config) == experts
 
 
-# ------------------------------------------------------------------------------- expert counting
-
-
-def test_dense_config_reports_no_experts():
-    assert num_experts_from_config(SimpleNamespace(hidden_size=128)) == 0
+# ------------------------------------------------------------------------------- degree selection
 
 
 @pytest.mark.parametrize(
-    "attr", ["num_experts", "num_local_experts", "n_routed_experts", "num_routed_experts"]
+    ("world", "experts", "degrees"),
+    [
+        (8, 64, (8, 1)),
+        (8, 6, (2, 4)),
+        (8, 5, (1, 8)),
+        (8, 0, (1, 8)),
+        (1, 64, (1, 1)),
+    ],
+    ids=["largest-valid", "divides-the-experts", "nothing-divides", "dense", "single-rank"],
 )
-def test_every_spelling_of_the_expert_count_is_understood(attr):
-    """Architectures disagree on the field name; missing one would silently disable EP."""
-    assert num_experts_from_config(SimpleNamespace(**{attr: 16})) == 16
-
-
-@pytest.mark.parametrize("nest", ["text_config", "llm_config"])
-def test_expert_count_is_found_on_a_vlms_nested_text_model(nest):
-    """For a VLM the MoE is the nested LM, so the root config has no expert count of its own."""
-    nested = SimpleNamespace(n_routed_experts=32)
-    assert num_experts_from_config(SimpleNamespace(**{nest: nested})) == 32
-
-
-# ------------------------------------------------------------------------------- degree selection
-
-
-def test_largest_valid_ep_degree_is_chosen():
-    assert choose_ep_dp_degrees(8, 64) == (8, 1)
-
-
-def test_ep_must_divide_the_expert_count_so_ranks_hold_whole_experts():
-    assert choose_ep_dp_degrees(8, 6) == (2, 4)
-
-
-def test_ep_falls_back_to_one_when_nothing_divides():
-    assert choose_ep_dp_degrees(8, 5) == (1, 8)
-
-
-@pytest.mark.parametrize(("world", "experts"), [(8, 0), (1, 64)])
-def test_no_expert_parallelism_for_dense_or_single_rank(world, experts):
-    assert choose_ep_dp_degrees(world, experts) == (1, world)
-
-
-@pytest.mark.parametrize("world", [2, 4, 8, 16])
-def test_degrees_always_factor_the_world(world):
-    ep, dp = choose_ep_dp_degrees(world, 64)
-    assert ep * dp == world
+def test_the_largest_ep_degree_dividing_the_world_and_the_experts_is_chosen(
+    world, experts, degrees
+):
+    """EP must divide the expert count, so every rank holds whole experts; with nothing dividing, a
+    dense model, or a single rank, there is no expert parallelism."""
+    assert choose_ep_dp_degrees(world, experts) == degrees
 
 
 # ---------------------------------------------------------------------------- routed experts
@@ -285,10 +257,6 @@ def test_routed_experts_are_the_plans_containers_under_any_prefix(plan, prefix):
     }
 
 
-def test_a_model_without_an_ep_plan_has_no_routed_experts():
-    assert _routed_expert_modules(_model_with_decoder_at("model."), {}) == {}
-
-
 # ------------------------------------------------------------------------------ model class
 
 
@@ -401,29 +369,28 @@ def test_an_moe_without_an_ep_plan_is_not_expert_parallel(run_as, serve, tmp_pat
 
 
 @needs_qwen3_moe_ep_plan
-def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path):
-    monkeypatch.setenv("MODELOPT_EP_DEGREE", "0")
+@pytest.mark.parametrize("value", ["0", ""], ids=["zero", "exported-empty"])
+def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value):
+    """And so is one exported empty: unset, the plan decides too."""
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", value)
     run_as(4)
     serve(_qwen3_moe(num_experts=8))
 
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)
 
 
-def test_a_dense_model_that_fits_one_gpu_is_replicated(run_as, serve, gpu_memory, tmp_path):
-    run_as(4)
-    config = serve(LlamaConfig(**_TINY))
-    gpu_memory(2 * _bytes(_meta(config)))
-
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
-
-
-def test_a_dense_model_too_big_for_one_gpu_is_fsdp_sharded(run_as, serve, gpu_memory, tmp_path):
-    """Whole, it would leave less than 40% of the GPU to calibration."""
+@pytest.mark.parametrize(
+    ("gpus_worth", "dense"), [(2, "replicate"), (1, "fsdp")], ids=["fits", "too-big"]
+)
+def test_a_dense_model_is_fsdp_sharded_only_when_too_big_for_one_gpu(
+    run_as, serve, gpu_memory, tmp_path, gpus_worth, dense
+):
+    """On a GPU only its own size, the whole model would leave less than 40% of it to calibration."""
     run_as(4)
     config = serve(LlamaConfig(**_TINY))
-    gpu_memory(_bytes(_meta(config)))
+    gpu_memory(gpus_worth * _bytes(_meta(config)))
 
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, dense="fsdp")
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, dense=dense)
 
 
 @needs_qwen3_moe_ep_plan
@@ -465,27 +432,32 @@ def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path):
 
 
 @pytest.mark.parametrize(
-    ("value", "dense"), [("1", "replicate"), ("4", "fsdp"), ("full", "fsdp"), ("WORLD", "fsdp")]
+    ("value", "dense"),
+    [
+        ("1", "replicate"),
+        ("4", "fsdp"),
+        ("full", "fsdp"),
+        ("WORLD", "fsdp"),
+        ("", "fsdp"),
+        ("2", None),
+    ],
+    ids=["1", "4", "full", "WORLD", "exported-empty", "part-of-the-world"],
 )
 def test_the_dense_knob_replicates_or_shards_over_the_whole_world(
-    run_as, serve, monkeypatch, tmp_path, value, dense
+    run_as, serve, gpu_memory, monkeypatch, tmp_path, value, dense
 ):
+    """Exported empty, the knob is auto, which shards a model too big for one GPU. Sharding over
+    part of the world is refused: no layout does that."""
     monkeypatch.setenv("MODELOPT_FSDP_SIZE", value)
     run_as(4)
     serve(LlamaConfig(**_TINY))
+    gpu_memory(1)  # too small for the whole model, so only auto tells "" from "1"
 
-    assert resolve_distribution_plan(str(tmp_path))[0].dense == dense
-
-
-def test_a_dense_knob_sharding_over_part_of_the_world_is_refused(
-    run_as, serve, monkeypatch, tmp_path
-):
-    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "2")
-    run_as(4)
-    serve(LlamaConfig(**_TINY))
-
-    with pytest.raises(ValueError, match=r"expected auto, 1 .* or world"):
-        resolve_distribution_plan(str(tmp_path))
+    if dense is None:
+        with pytest.raises(ValueError, match=r"expected auto, 1 .* or world"):
+            resolve_distribution_plan(str(tmp_path))
+    else:
+        assert resolve_distribution_plan(str(tmp_path))[0].dense == dense
 
 
 @pytest.mark.parametrize(

From eafb18f97c938bebf1b0a76915de104b1ad91b96 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 22 Oct 2026 19:05:25 +0000
Subject: [PATCH 15/44] Date env and test_resolve_distribution_plan 2026, as
 LICENSE_HEADER does

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

diff --git a/modelopt/torch/utils/env.py b/modelopt/torch/utils/env.py
index 090fa8050cd..7e27e7857c3 100644
--- a/modelopt/torch/utils/env.py
+++ b/modelopt/torch/utils/env.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 2e406178dff..dc7c7a001db 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

From 75a6a6f54e3548c2b79a80a943b28c4ad257bc89 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 24 Oct 2026 19:59:44 +0000
Subject: [PATCH 16/44] Trust a checkpoint's remote class only under
 trust_remote_code, and name a bad EP knob

_auto_class took a VLM checkpoint's remote AutoModelForCausalLM whenever its auto_map named one,
whether or not the load trusts the checkpoint's code, so an untrusted load built through a class it
could not import instead of the native image-text class. It now takes the remote class only when
trust_remote_code is set. A MODELOPT_EP_DEGREE that is not a number failed with int()'s bare
ValueError, which the other ranks get as "invalid literal"; the error now names the knob.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 823cec3883d..617e1a61998 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -889,14 +889,15 @@ def parallel_load_and_prepare_fsdp2(
     return model
 
 
-def _auto_class(hf_config) -> Any:
+def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
     """The auto class that builds ``hf_config``'s whole model, LM head included.
 
     A causal LM builds through ``AutoModelForCausalLM``. A VLM -- a nested text config, a vision
     config, or a ``*ForConditionalGeneration`` architecture -- builds through the remote
-    ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` there
-    and no image-text class, else through ``AutoModelForImageTextToText`` if that recognizes the
-    config, else through ``AutoModel``, which leaves the head out.
+    ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` there,
+    its code is trusted and it registers no image-text class, else through
+    ``AutoModelForImageTextToText`` if that recognizes the config, else through ``AutoModel``, which
+    leaves the head out.
     """
     architectures = getattr(hf_config, "architectures", None) or []
     is_vlm = (
@@ -908,7 +909,8 @@ def _auto_class(hf_config) -> Any:
         return AutoModelForCausalLM
     auto_map = getattr(hf_config, "auto_map", None) or {}
     remote_with_head = (
-        "conditionalgeneration" in str(auto_map.get("AutoModelForCausalLM", "")).lower()
+        trust_remote_code
+        and "conditionalgeneration" in str(auto_map.get("AutoModelForCausalLM", "")).lower()
     )
     if remote_with_head and "AutoModelForImageTextToText" not in auto_map:
         return AutoModelForCausalLM
@@ -948,6 +950,18 @@ 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``."""
+    value = env.ep_degree() if ep_size is None else ep_size
+    if str(value) in ("auto", "0"):
+        return "auto"
+    try:
+        return int(value)
+    except ValueError:
+        name = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size"
+        raise ValueError(f"{name}={value!r}: expected auto, 0 (auto) or an EP degree.") from None
+
+
 def _decide_distribution(
     ckpt_path: str, hf_config, world_size: int, ep_size, dense, trust_remote_code: bool
 ) -> DistributionPlan:
@@ -967,12 +981,15 @@ def _decide_distribution(
             ep_size=1 if ep_size is None or str(ep_size) in ("auto", "0") else int(ep_size),
             dense="replicate" if dense in (None, "auto") else dense,
         )
-    ep_size = env.ep_degree() if ep_size is None else ep_size
-    ep_size = "auto" if str(ep_size) in ("auto", "0") else int(ep_size)
+    ep_size = _ep_degree(ep_size)
     dense = _dense_from_env(world_size) if dense is None else dense
     try:
         model = build_meta_causal_lm(
-            ckpt_path, trust_remote_code, None, hf_config, auto_class=_auto_class(hf_config)
+            ckpt_path,
+            trust_remote_code,
+            None,
+            hf_config,
+            auto_class=_auto_class(hf_config, trust_remote_code),
         )
     except Exception as e:
         raise RuntimeError(
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index dc7c7a001db..41a6261678f 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -281,23 +281,31 @@ def _tiny_qwen2_5_vl():
     )
 
 
+_REMOTE_VLM = {"AutoModelForCausalLM": "m.KimiK25ForConditionalGeneration"}
+
+
 @pytest.mark.parametrize(
-    ("config", "expected"),
+    ("config", "trust_remote_code", "expected"),
     [
-        (SimpleNamespace(architectures=["LlamaForCausalLM"]), AutoModelForCausalLM),
-        (
-            _vlm_config(auto_map={"AutoModelForCausalLM": "m.KimiK25ForConditionalGeneration"}),
-            AutoModelForCausalLM,
-        ),
-        (_tiny_qwen2_5_vl(), AutoModelForImageTextToText),
-        (_vlm_config(), AutoModel),
+        (SimpleNamespace(architectures=["LlamaForCausalLM"]), False, AutoModelForCausalLM),
+        (_vlm_config(auto_map=_REMOTE_VLM), True, AutoModelForCausalLM),
+        (_vlm_config(auto_map=_REMOTE_VLM), False, AutoModel),
+        (_tiny_qwen2_5_vl(), False, AutoModelForImageTextToText),
+        (_vlm_config(), False, AutoModel),
+    ],
+    ids=[
+        "causal-lm",
+        "remote-vlm-with-head",
+        "remote-vlm-untrusted",
+        "native-vlm-with-head",
+        "unknown-vlm",
     ],
-    ids=["causal-lm", "remote-vlm-with-head", "native-vlm-with-head", "unknown-vlm"],
 )
-def test_the_class_that_builds_the_whole_model(config, expected):
+def test_the_class_that_builds_the_whole_model(config, trust_remote_code, expected):
     """A VLM's AutoModel leaves the LM head out: it builds through the class that carries the head,
-    keeping the headless backbone only when no class recognizes the config."""
-    assert _auto_class(config) is expected
+    keeping the headless backbone only when no class recognizes the config. A remote class counts
+    only when the checkpoint's code is trusted."""
+    assert _auto_class(config, trust_remote_code) is expected
 
 
 # -------------------------------------------------------------------------------- the decision
@@ -368,6 +376,15 @@ 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)
 
 
+def test_an_ep_degree_that_is_not_a_number_names_the_knob(run_as, serve, monkeypatch, tmp_path):
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", "world")
+    run_as(4)
+    serve(LlamaConfig(**_TINY))
+
+    with pytest.raises(ValueError, match="MODELOPT_EP_DEGREE='world'"):
+        resolve_distribution_plan(str(tmp_path))
+
+
 @needs_qwen3_moe_ep_plan
 @pytest.mark.parametrize("value", ["0", ""], ids=["zero", "exported-empty"])
 def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value):

From 2b8563c3c4b0c5fa8b2dfa2bf990b14cde8aeede Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 24 Oct 2026 22:32:00 +0000
Subject: [PATCH 17/44] Count a rank's expert share in the dense fit, and take
 the Hub options

dense="auto" compared only what every rank holds whole with 60% of a GPU, leaving out the rank's
1/ep_size share of the routed experts, so it could replicate the dense part of an expert-dominated
MoE that fits only with that part sharded: Qwen3-235B-A22B at EP 8 holds about 73 GB a rank on
80 GB GPUs with the dense part replicated, 59 GB with it sharded. The fit now counts that share.

resolve_distribution_plan takes revision, token, cache_dir, local_files_only and subfolder, as
load_pretrained does, and reads the checkpoint and its config with them, so a plan is decided for
the checkpoint the load reads rather than the Hub default branch's, and a gated repo, another cache
or a checkpoint in a subfolder can be planned. Its docstring no longer quotes a load_pretrained
call this branch does not take.

A single process parses its own ep_size through _ep_degree, so a bad one gets the error a
distributed run gives it. The EP-degree helper returns the degree alone, the data-parallel degree
being DistributionPlan.edp, and it and the expert-count helper are private like their neighbours.
MODELOPT_FSDP_SIZE no longer accepts the undocumented "full".

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/moe.py                   | 25 +++---
 .../torch/utils/plugins/model_load_utils.py   | 69 ++++++++++-----
 .../utils/test_resolve_distribution_plan.py   | 87 +++++++++++++------
 3 files changed, 121 insertions(+), 60 deletions(-)

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index 4a4335ba111..35a2390450f 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -52,7 +52,7 @@ def _ep_config(hf_config):
     return hf_config
 
 
-def num_experts_from_config(hf_config) -> int:
+def _num_experts_from_config(hf_config) -> int:
     """Return the MoE expert count from an HF config (0 if the model is not MoE).
 
     Recurses into ``.text_config`` / ``.llm_config`` so VLMs whose nested text submodel is the MoE
@@ -72,20 +72,19 @@ def num_experts_from_config(hf_config) -> int:
     return 0
 
 
-def choose_ep_dp_degrees(world_size: int, num_experts: int) -> tuple[int, int]:
-    """Pick ``(expert_parallel, expert_data_parallel)`` degrees with ``EP * EDP == world_size``.
+def _choose_ep_degree(world_size: int, num_experts: int) -> int:
+    """The largest EP degree dividing both ``world_size`` and ``num_experts``; ``1`` if none does.
 
-    EP must divide both ``num_experts`` (whole experts per rank) and ``world_size`` (a valid mesh
-    factorization); we take the largest such EP, and each expert shard is then replicated over the
-    remaining ``EDP = world_size // EP`` ranks. Falls back to ``(1, world_size)`` when the model is not
-    MoE or no valid EP degree exists. Heuristic -- callers may override.
+    EP must divide ``num_experts`` (whole experts per rank) and ``world_size`` (a valid mesh
+    factorization). Each expert shard is then replicated over the remaining ranks
+    (``DistributionPlan.edp``). Heuristic -- callers may override.
     """
-    if num_experts <= 0 or world_size <= 1:
-        return 1, world_size
-    for ep in range(world_size, 0, -1):
-        if world_size % ep == 0 and num_experts % ep == 0:
-            return ep, world_size // ep
-    return 1, world_size
+    if num_experts <= 0:
+        return 1
+    return next(
+        (ep for ep in range(world_size, 0, -1) if world_size % ep == 0 and num_experts % ep == 0),
+        1,
+    )
 
 
 # The styles a base_model_ep_plan gives a layer's routed-experts container. A plan that styles only
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 617e1a61998..745aeb29390 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -923,8 +923,9 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
     return AutoModel
 
 
-# The share of one GPU's memory a replicated dense part may take. The rest is left to calibration's
-# activations and communication buffers, and to the rank's share of the experts.
+# The share of one GPU's memory a rank may take with the dense part replicated: that whole copy and
+# the rank's share of the routed experts. The rest is left to calibration's activations and
+# communication buffers.
 _REPLICATED_SHARE = 0.6
 
 
@@ -942,7 +943,7 @@ def _dense_from_env(world_size: int) -> str:
         return "auto"
     if value == "1":
         return "replicate"
-    if value in ("world", "full", str(world_size)):
+    if value in ("world", str(world_size)):
         return "fsdp"
     raise ValueError(
         f"MODELOPT_FSDP_SIZE={value!r}: expected auto, 1 (replicate the dense part) or world "
@@ -967,18 +968,19 @@ def _decide_distribution(
 ) -> DistributionPlan:
     """The plan :func:`resolve_distribution_plan` returns, decided by this rank alone."""
     from modelopt.torch.utils.moe import (
+        _choose_ep_degree,
         _ep_config,
+        _num_experts_from_config,
         _routed_expert_modules,
-        choose_ep_dp_degrees,
-        num_experts_from_config,
     )
 
     if world_size == 1:
         # 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)
         return DistributionPlan(
-            ep_size=1 if ep_size is None or str(ep_size) in ("auto", "0") else int(ep_size),
+            ep_size=single if isinstance(single, int) else 1,
             dense="replicate" if dense in (None, "auto") else dense,
         )
     ep_size = _ep_degree(ep_size)
@@ -998,24 +1000,26 @@ 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)
-    num_experts = num_experts_from_config(hf_config) if experts else 0
+    num_experts = _num_experts_from_config(hf_config) if experts else 0
     if ep_size == "auto":
-        ep_size = choose_ep_dp_degrees(world_size, num_experts)[0]
+        ep_size = _choose_ep_degree(world_size, num_experts)
     elif ep_size > 1 and (num_experts == 0 or num_experts % ep_size):
         raise ValueError(
             f"ep_size={ep_size} must divide the routed experts {hf_config.model_type} shards "
             f"({num_experts})."
         )
     if dense == "auto":
-        # Every rank holds a whole copy of what expert parallelism does not shard.
+        # A rank holds a whole copy of what expert parallelism does not shard, and its 1/ep_size
+        # share of the routed experts it does.
         sharded = (
             {id(p) for m in experts.values() for p in m.parameters()} if ep_size > 1 else set()
         )
-        replicated = sum(
-            p.numel() * p.element_size() for p in model.parameters() if id(p) not in sharded
+        held = sum(
+            p.numel() * p.element_size() // (ep_size if id(p) in sharded else 1)
+            for p in model.parameters()
         )
         memory = _gpu_memory()
-        too_big = memory is not None and replicated > _REPLICATED_SHARE * memory
+        too_big = memory is not None and held > _REPLICATED_SHARE * memory
         dense = "fsdp" if too_big else "replicate"
     return DistributionPlan(world_size, ep_size, dense)
 
@@ -1026,14 +1030,19 @@ def resolve_distribution_plan(
     trust_remote_code: bool = False,
     ep_size: int | str | None = None,
     dense: str | None = None,
+    revision: str | None = None,
+    token: str | bool | None = None,
+    cache_dir: str | os.PathLike | None = None,
+    local_files_only: bool = False,
+    subfolder: str | None = None,
 ) -> tuple[DistributionPlan, PretrainedConfig]:
     """Decide how this run distributes ``ckpt_path``'s model over its ranks.
 
-    The one place a run decides it. :func:`load_pretrained` loads the model by the plan this returns
-    and records it on the model, and every later step reads it back from there
+    The one place a run decides it: a distributed load loads the model by the plan this returns and
+    records it on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), and
+    every later step reads it back from there
     (:func:`~modelopt.torch.utils.distributed.distribution_plan`). Returns ``(plan, hf_config)``,
-    ``hf_config`` being the config a sharded load takes back as
-    ``load_pretrained(ckpt_path, distribution=plan, hf_config=hf_config)``.
+    ``hf_config`` being the config the plan was decided from, for a sharded load to build from.
 
     Rank 0 decides, from the model built on the meta device: from its parameters rather than
     estimates of them, and from its routed experts as an expert-parallel load finds them. Every rank
@@ -1044,13 +1053,19 @@ 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``.
     * ``dense``: ``"replicate"`` or ``"fsdp"`` for the rest of the model, or ``"auto"``, which
-      FSDP2-shards it when a whole copy would take more than 60% of one GPU's memory. That GPU is
-      rank 0's, by its total memory, so ``"auto"`` assumes the run's GPUs are alike.
+      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"``
+      assumes the run's GPUs are alike.
 
     ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` (``1`` replicates; ``world``,
     or the world size as a number, shards) from rank 0's environment, both defaulting to
     ``"auto"``. A single process reads neither and holds the whole model, unless ``dense="fsdp"``
     asks for FSDP2 on one rank.
+
+    ``revision``, ``token``, ``cache_dir``, ``local_files_only`` and ``subfolder`` locate the
+    checkpoint as :func:`load_pretrained` takes them, so the plan is decided for what the load reads:
+    all but ``subfolder`` resolve a Hub model ID (:func:`ensure_local_checkpoint`), and the config
+    read takes all but ``revision``, for custom code its ``auto_map`` names in another Hub repo.
     """
     from modelopt.torch.utils import distributed as dist
 
@@ -1059,8 +1074,22 @@ def resolve_distribution_plan(
     read: list[tuple[str, PretrainedConfig]] = []
 
     def checkpoint_read() -> bool:
-        _, local_path = ensure_local_checkpoint(ckpt_path)
-        config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code)
+        _, local_path = ensure_local_checkpoint(
+            ckpt_path,
+            revision=revision,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        if subfolder:
+            local_path = os.path.join(local_path, subfolder)
+        config = AutoConfig.from_pretrained(
+            local_path,
+            trust_remote_code=trust_remote_code,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
         read.append((local_path, config))
         return True
 
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 41a6261678f..9bb8a886a6d 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -46,9 +46,9 @@
 import modelopt.torch.utils.distributed as dist
 from modelopt.torch.utils.distributed import DistributionPlan
 from modelopt.torch.utils.moe import (
+    _choose_ep_degree,
+    _num_experts_from_config,
     _routed_expert_modules,
-    choose_ep_dp_degrees,
-    num_experts_from_config,
 )
 from modelopt.torch.utils.plugins import model_load_utils
 from modelopt.torch.utils.plugins.model_load_utils import _auto_class, resolve_distribution_plan
@@ -162,29 +162,21 @@ def give(size):
 def test_the_expert_count_is_read_under_every_spelling(config, experts):
     """Architectures disagree on the field name, and a VLM's MoE is its nested LM, whose root
     config has no expert count of its own: missing either would silently disable EP."""
-    assert num_experts_from_config(config) == experts
+    assert _num_experts_from_config(config) == experts
 
 
 # ------------------------------------------------------------------------------- degree selection
 
 
 @pytest.mark.parametrize(
-    ("world", "experts", "degrees"),
-    [
-        (8, 64, (8, 1)),
-        (8, 6, (2, 4)),
-        (8, 5, (1, 8)),
-        (8, 0, (1, 8)),
-        (1, 64, (1, 1)),
-    ],
+    ("world", "experts", "degree"),
+    [(8, 64, 8), (8, 6, 2), (8, 5, 1), (8, 0, 1), (1, 64, 1)],
     ids=["largest-valid", "divides-the-experts", "nothing-divides", "dense", "single-rank"],
 )
-def test_the_largest_ep_degree_dividing_the_world_and_the_experts_is_chosen(
-    world, experts, degrees
-):
+def test_the_largest_ep_degree_dividing_the_world_and_the_experts_is_chosen(world, experts, degree):
     """EP must divide the expert count, so every rank holds whole experts; with nothing dividing, a
     dense model, or a single rank, there is no expert parallelism."""
-    assert choose_ep_dp_degrees(world, experts) == degrees
+    assert _choose_ep_degree(world, experts) == degree
 
 
 # ---------------------------------------------------------------------------- routed experts
@@ -376,13 +368,25 @@ 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)
 
 
-def test_an_ep_degree_that_is_not_a_number_names_the_knob(run_as, serve, monkeypatch, tmp_path):
-    monkeypatch.setenv("MODELOPT_EP_DEGREE", "world")
-    run_as(4)
+@pytest.mark.parametrize(
+    ("world_size", "source"),
+    [(4, "MODELOPT_EP_DEGREE"), (1, "ep_size")],
+    ids=["knob", "argument-in-a-single-process"],
+)
+def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from(
+    run_as, serve, monkeypatch, tmp_path, world_size, source
+):
+    """A single process takes its own argument as given, but parses it as the decision does."""
+    run_as(world_size)
     serve(LlamaConfig(**_TINY))
+    arguments = {}
+    if source == "ep_size":
+        arguments["ep_size"] = "world"
+    else:
+        monkeypatch.setenv(source, "world")
 
-    with pytest.raises(ValueError, match="MODELOPT_EP_DEGREE='world'"):
-        resolve_distribution_plan(str(tmp_path))
+    with pytest.raises(ValueError, match=f"{source}='world'"):
+        resolve_distribution_plan(str(tmp_path), **arguments)
 
 
 @needs_qwen3_moe_ep_plan
@@ -411,16 +415,25 @@ def test_a_dense_model_is_fsdp_sharded_only_when_too_big_for_one_gpu(
 
 
 @needs_qwen3_moe_ep_plan
-def test_only_what_expert_parallelism_leaves_whole_must_fit(run_as, serve, gpu_memory, tmp_path):
-    """The routed experts are sharded at EP 4, so only the rest needs to fit one GPU; with no EP the
-    whole MoE does, and here it does not."""
+@pytest.mark.parametrize(
+    ("experts_fitting", "dense"),
+    [(1 / 2, "replicate"), (1 / 8, "fsdp")],
+    ids=["fits", "the-experts-share-does-not"],
+)
+def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit(
+    run_as, serve, gpu_memory, tmp_path, experts_fitting, dense
+):
+    """At EP 4 a rank holds the rest of the model whole and a quarter of the routed experts, so the
+    dense part is sharded when that does not fit one GPU, however small the rest; with no EP the
+    whole MoE must fit, and here it does not."""
     run_as(4)
     config = serve(_qwen3_moe(num_experts=8))
     model = _meta(config)
     experts = sum(_bytes(layer.mlp.experts) for layer in model.model.layers)
-    gpu_memory(int((_bytes(model) - experts / 2) / 0.6))
+    # A GPU 60% of which holds the rest of the model and that share of the experts.
+    gpu_memory(int((_bytes(model) - experts + experts_fitting * experts) / 0.6))
 
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4, dense)
     assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan(
         4, dense="fsdp"
     )
@@ -453,12 +466,11 @@ def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path):
     [
         ("1", "replicate"),
         ("4", "fsdp"),
-        ("full", "fsdp"),
         ("WORLD", "fsdp"),
         ("", "fsdp"),
         ("2", None),
     ],
-    ids=["1", "4", "full", "WORLD", "exported-empty", "part-of-the-world"],
+    ids=["1", "4", "WORLD", "exported-empty", "part-of-the-world"],
 )
 def test_the_dense_knob_replicates_or_shards_over_the_whole_world(
     run_as, serve, gpu_memory, monkeypatch, tmp_path, value, dense
@@ -517,6 +529,27 @@ def unbuildable(*args, **kwargs):
         resolve_distribution_plan(str(tmp_path))
 
 
+def test_the_hub_options_locate_the_checkpoint_the_load_reads(tmp_path):
+    """``revision``, ``cache_dir``, ``local_files_only`` and ``subfolder`` find the checkpoint as
+    ``load_pretrained`` takes them, so the plan and its config are those of what the load reads."""
+    repo = tmp_path / "models--org--model"
+    (repo / "refs").mkdir(parents=True)
+    for revision, commit, hidden_size in (("main", "a" * 40, 16), ("v2", "b" * 40, 32)):
+        (repo / "refs" / revision).write_text(commit)
+        config = LlamaConfig(**{**_TINY, "hidden_size": hidden_size})
+        config.save_pretrained(repo / "snapshots" / commit / "nested")
+
+    _, config = resolve_distribution_plan(
+        "org/model",
+        revision="v2",
+        cache_dir=str(tmp_path),
+        local_files_only=True,
+        subfolder="nested",
+    )
+
+    assert config.hidden_size == 32
+
+
 # ------------------------------------------------------------------------------ rank 0 decides
 
 

From ec865ee64c297cee3cc7a6386e707490b3c2063b Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 24 Oct 2026 22:40:12 +0000
Subject: [PATCH 18/44] Build a trusted remote VLM through the image-text class
 it registers

_auto_class took AutoModelForImageTextToText only for a config transformers maps to it natively, so
a trusted remote VLM whose auto_map registers an image-text class, and whose config class is its
own, fell through to AutoModel: the headless backbone, which leaves the LM head out of the dense
fit. It now takes the registered class when the checkpoint's code is trusted, the class the remote
causal-LM branch already defers to for such a checkpoint.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 745aeb29390..5a847e1412d 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -896,8 +896,8 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
     config, or a ``*ForConditionalGeneration`` architecture -- builds through the remote
     ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` there,
     its code is trusted and it registers no image-text class, else through
-    ``AutoModelForImageTextToText`` if that recognizes the config, else through ``AutoModel``, which
-    leaves the head out.
+    ``AutoModelForImageTextToText`` if that recognizes the config or, its code trusted, the
+    checkpoint registers one, else through ``AutoModel``, which leaves the head out.
     """
     architectures = getattr(hf_config, "architectures", None) or []
     is_vlm = (
@@ -918,7 +918,9 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
         from transformers import AutoModelForImageTextToText
     except ImportError:
         return AutoModel
-    if type(hf_config) in AutoModelForImageTextToText._model_mapping:
+    if type(hf_config) in AutoModelForImageTextToText._model_mapping or (
+        trust_remote_code and "AutoModelForImageTextToText" in auto_map
+    ):
         return AutoModelForImageTextToText
     return AutoModel
 
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 9bb8a886a6d..58d10151d88 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -274,6 +274,10 @@ def _tiny_qwen2_5_vl():
 
 
 _REMOTE_VLM = {"AutoModelForCausalLM": "m.KimiK25ForConditionalGeneration"}
+_REMOTE_IMAGE_TEXT_VLM = {
+    **_REMOTE_VLM,
+    "AutoModelForImageTextToText": "m.KimiK25ForConditionalGeneration",
+}
 
 
 @pytest.mark.parametrize(
@@ -282,6 +286,7 @@ def _tiny_qwen2_5_vl():
         (SimpleNamespace(architectures=["LlamaForCausalLM"]), False, AutoModelForCausalLM),
         (_vlm_config(auto_map=_REMOTE_VLM), True, AutoModelForCausalLM),
         (_vlm_config(auto_map=_REMOTE_VLM), False, AutoModel),
+        (_vlm_config(auto_map=_REMOTE_IMAGE_TEXT_VLM), True, AutoModelForImageTextToText),
         (_tiny_qwen2_5_vl(), False, AutoModelForImageTextToText),
         (_vlm_config(), False, AutoModel),
     ],
@@ -289,6 +294,7 @@ def _tiny_qwen2_5_vl():
         "causal-lm",
         "remote-vlm-with-head",
         "remote-vlm-untrusted",
+        "remote-image-text-vlm",
         "native-vlm-with-head",
         "unknown-vlm",
     ],

From a889e441cd5ddec21dd0bb279bebbbe50f7085e7 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sun, 25 Oct 2026 08:56:15 +0000
Subject: [PATCH 19/44] Leave the auto class as it was after building a model
 on the meta device

Building a remote-code model through an auto class registers its class there, as it is, and a later
load through that auto class then builds the registered class, without the GenerationMixin a first
load adds. Rank 0 builds one more meta model than its peers, to decide a distribution, so its model
alone lacked generate. Under expert parallelism Kimi-K3's rank 0 then failed or skipped hf_ptq's
preview while the other ranks decoded in lockstep, and the run deadlocked. build_meta_causal_lm now
drops the registrations its build adds, so every rank's load builds the same class.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   | 24 +++++++++++++----
 .../utils/test_resolve_distribution_plan.py   | 27 +++++++++++++++++++
 2 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 5a847e1412d..3354148ae99 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -198,7 +198,14 @@ def build_meta_causal_lm(
     hf_config=None,
     auto_class: Any = None,
 ):
-    """Build a meta-init causal LM (no real storage allocated), or ``auto_class``'s model."""
+    """Build a meta-init causal LM (no real storage allocated), or ``auto_class``'s model.
+
+    It leaves ``auto_class`` as it found it. Building a remote-code model through an auto class
+    registers its class there, as it is, and a later load through that auto class builds the
+    registered class, without the ``GenerationMixin`` a first load adds: a rank that built one more
+    meta model than its peers, as rank 0 deciding a distribution does, would load its model without
+    the ``generate`` theirs have.
+    """
     if hf_config is None:
         config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code}
         if attn_implementation is not None:
@@ -218,10 +225,17 @@ def build_meta_causal_lm(
             from transformers import AutoModelForImageTextToText
 
             model_factory = AutoModelForImageTextToText
-    with init_empty_weights(include_buffers=False):
-        model = model_factory.from_config(
-            hf_config, dtype=dtype, trust_remote_code=trust_remote_code
-        )
+    registered = getattr(getattr(model_factory, "_model_mapping", None), "_extra_content", None)
+    kept = dict(registered) if registered is not None else {}
+    try:
+        with init_empty_weights(include_buffers=False):
+            model = model_factory.from_config(
+                hf_config, dtype=dtype, trust_remote_code=trust_remote_code
+            )
+    finally:
+        if registered is not None:
+            registered.clear()
+            registered.update(kept)
     model.eval()
     return model
 
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 58d10151d88..aed9cc285a3 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -535,6 +535,33 @@ def unbuildable(*args, **kwargs):
         resolve_distribution_plan(str(tmp_path))
 
 
+def test_a_meta_build_leaves_the_auto_class_as_it_found_it(monkeypatch):
+    """Building a remote-code model through an auto class registers its class there, as it is, and a
+    later load through that auto class builds the registered class, without the ``generate`` a first
+    load adds. Rank 0 alone builds one to decide a distribution: its load alone would lack it."""
+
+    class ProbeConfig(transformers.PretrainedConfig):
+        model_type = "meta_build_probe"
+
+    class ProbeModel(transformers.PreTrainedModel):
+        config_class = ProbeConfig
+
+    mapping = AutoModelForCausalLM._model_mapping
+    monkeypatch.setattr(mapping, "_extra_content", dict(mapping._extra_content))
+    from_config = AutoModelForCausalLM.from_config
+
+    def remote_from_config(config, **kwargs):
+        # What transformers' from_config does for a remote-code config before it builds.
+        AutoModelForCausalLM.register(ProbeConfig, ProbeModel, exist_ok=True)
+        return from_config(config, **kwargs)
+
+    monkeypatch.setattr(AutoModelForCausalLM, "from_config", remote_from_config)
+
+    model_load_utils.build_meta_causal_lm(None, False, None, LlamaConfig(**_TINY))
+
+    assert ProbeConfig not in mapping
+
+
 def test_the_hub_options_locate_the_checkpoint_the_load_reads(tmp_path):
     """``revision``, ``cache_dir``, ``local_files_only`` and ``subfolder`` find the checkpoint as
     ``load_pretrained`` takes them, so the plan and its config are those of what the load reads."""

From 432b131f260e38320af22f789ae4840947e89983 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sun, 25 Oct 2026 18:21:33 +0000
Subject: [PATCH 20/44] Read the auto class's registrations without a fallback

build_meta_causal_lm read transformers' private _model_mapping._extra_content through getattr
defaults, so a release that moved it would have turned the restore into a silent no-op, and with it
brought back rank 0's model without generate and the deadlock it caused. It now reads the attribute
directly: a move fails the build at once.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 3354148ae99..a7574d8b78b 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -225,17 +225,18 @@ def build_meta_causal_lm(
             from transformers import AutoModelForImageTextToText
 
             model_factory = AutoModelForImageTextToText
-    registered = getattr(getattr(model_factory, "_model_mapping", None), "_extra_content", None)
-    kept = dict(registered) if registered is not None else {}
+    # Where transformers records the class a remote-code build registers. Read without a fallback:
+    # if a release moves it, this fails here rather than leave one rank's model without generate.
+    registered = model_factory._model_mapping._extra_content
+    kept = dict(registered)
     try:
         with init_empty_weights(include_buffers=False):
             model = model_factory.from_config(
                 hf_config, dtype=dtype, trust_remote_code=trust_remote_code
             )
     finally:
-        if registered is not None:
-            registered.clear()
-            registered.update(kept)
+        registered.clear()
+        registered.update(kept)
     model.eval()
     return model
 

From bf7557b943fb0220561b11d9fbccf579b60ba8a4 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 01:27:20 +0000
Subject: [PATCH 21/44] Name a plan's routed experts by both styles, not the
 first one found

The grouped-GEMM parameters' parents counted only when no container style appeared in the plan, so a
plan styling some layers' expert containers and only the grouped-GEMM parameters of others lost the
second kind, both in the decision and in the expert-parallel load. Every transformers 5.18 plan with
both styles names the same containers both ways, so the union changes nothing for them.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/moe.py                          |  8 ++++----
 .../torch/utils/test_resolve_distribution_plan.py    | 12 ++++++++----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index 35a2390450f..3e59099c45e 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -87,15 +87,15 @@ def _choose_ep_degree(world_size: int, num_experts: int) -> int:
     )
 
 
-# The styles a base_model_ep_plan gives a layer's routed-experts container. A plan that styles only
-# the experts' grouped-GEMM parameters names the containers through their parents.
+# The styles a base_model_ep_plan gives a layer's routed-experts container. A plan that styles the
+# experts' grouped-GEMM parameters names the containers through their parents too, so a plan that
+# styles some containers one way and others the other names them all.
 _EXPERTS_STYLES = ("moe_tp_experts", "nemotron_experts_ep")
 
 
 def _routed_expert_keys(ep_plan: dict[str, str]) -> set[str]:
     """The keys of ``ep_plan``, a ``base_model_ep_plan``, that name a routed-experts container."""
-    keys = {key for key, style in ep_plan.items() if style in _EXPERTS_STYLES}
-    return keys or {
+    return {key for key, style in ep_plan.items() if style in _EXPERTS_STYLES} | {
         key.rsplit(".", 1)[0] for key, style in ep_plan.items() if style == "grouped_gemm"
     }
 
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index aed9cc285a3..512903f7f95 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -225,8 +225,12 @@ def _model_with_decoder_at(prefix):
         {"layers.*.mlp.experts": "moe_tp_experts"},
         {"layers.*.mlp.experts": "nemotron_experts_ep"},
         {k: v for k, v in _STOCK_PLAN.items() if v == "grouped_gemm"},
+        {
+            "layers.*.block_sparse_moe.experts": "moe_tp_experts",
+            "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+        },
     ],
-    ids=["stock", "experts-style-only", "nemotron-style-only", "grouped-gemm-only"],
+    ids=["stock", "experts-style-only", "nemotron-style-only", "grouped-gemm-only", "hybrid"],
 )
 @pytest.mark.parametrize(
     "prefix",
@@ -235,9 +239,9 @@ def _model_with_decoder_at(prefix):
 )
 def test_routed_experts_are_the_plans_containers_under_any_prefix(plan, prefix):
     """The base-relative plan keys match the decoder wherever the model nests it, at every layer
-    index, the container named by its style or as its grouped-GEMM parameters' parent. The router,
-    the shared experts, and a module whose name merely ends in a key's text are not routed
-    experts."""
+    index, the container named by its style or as its grouped-GEMM parameters' parent, also in a
+    plan that names some containers one way and some the other. The router, the shared experts, and
+    a module whose name merely ends in a key's text are not routed experts."""
     model = _model_with_decoder_at(prefix)
     model.add_module("sublayers", nn.ModuleList([_MoeLayer()]))  # "sublayers.0.mlp.experts"
 

From 4a9f7df16d21408c14fcb3e8ad1b0611dc6fa7d3 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 01:27:20 +0000
Subject: [PATCH 22/44] Build an encoder-decoder model through its seq2seq
 class, head included

Its *ForConditionalGeneration architecture reads like a VLM's, so T5, BART, Whisper and Qwen2-Audio
fell through to AutoModel, which leaves the LM head out of the parameters the dense fit counts. The
seq2seq auto classes now claim their configs first.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index a7574d8b78b..1a0d18b1874 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -912,7 +912,8 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
     ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` there,
     its code is trusted and it registers no image-text class, else through
     ``AutoModelForImageTextToText`` if that recognizes the config or, its code trusted, the
-    checkpoint registers one, else through ``AutoModel``, which leaves the head out.
+    checkpoint registers one, else through ``AutoModel``, which leaves the head out. An
+    encoder-decoder model, whose architecture reads the same, builds through its seq2seq class.
     """
     architectures = getattr(hf_config, "architectures", None) or []
     is_vlm = (
@@ -929,6 +930,11 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
     )
     if remote_with_head and "AutoModelForImageTextToText" not in auto_map:
         return AutoModelForCausalLM
+    from transformers import AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq
+
+    for seq2seq in (AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq):
+        if type(hf_config) in seq2seq._model_mapping:
+            return seq2seq
     try:
         from transformers import AutoModelForImageTextToText
     except ImportError:
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 512903f7f95..4b66f57f248 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -38,9 +38,13 @@
     AutoModel,
     AutoModelForCausalLM,
     AutoModelForImageTextToText,
+    AutoModelForSeq2SeqLM,
+    AutoModelForSpeechSeq2Seq,
     LlamaConfig,
     Qwen2_5_VLConfig,
     Qwen3MoeConfig,
+    T5Config,
+    WhisperConfig,
 )
 
 import modelopt.torch.utils.distributed as dist
@@ -293,6 +297,12 @@ def _tiny_qwen2_5_vl():
         (_vlm_config(auto_map=_REMOTE_IMAGE_TEXT_VLM), True, AutoModelForImageTextToText),
         (_tiny_qwen2_5_vl(), False, AutoModelForImageTextToText),
         (_vlm_config(), False, AutoModel),
+        (T5Config(architectures=["T5ForConditionalGeneration"]), False, AutoModelForSeq2SeqLM),
+        (
+            WhisperConfig(architectures=["WhisperForConditionalGeneration"]),
+            False,
+            AutoModelForSpeechSeq2Seq,
+        ),
     ],
     ids=[
         "causal-lm",
@@ -301,12 +311,15 @@ def _tiny_qwen2_5_vl():
         "remote-image-text-vlm",
         "native-vlm-with-head",
         "unknown-vlm",
+        "encoder-decoder",
+        "speech-encoder-decoder",
     ],
 )
 def test_the_class_that_builds_the_whole_model(config, trust_remote_code, expected):
     """A VLM's AutoModel leaves the LM head out: it builds through the class that carries the head,
     keeping the headless backbone only when no class recognizes the config. A remote class counts
-    only when the checkpoint's code is trusted."""
+    only when the checkpoint's code is trusted. An encoder-decoder model's architecture reads like a
+    VLM's, but its seq2seq class carries its head."""
     assert _auto_class(config, trust_remote_code) is expected
 
 

From ab81f87a3b9e2f790b84530bab1cf2d74bef2e4a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 01:27:20 +0000
Subject: [PATCH 23/44] Read the EP degree's auto in any case, as the FSDP knob
 reads its values

MODELOPT_FSDP_SIZE=Auto or WORLD worked, but MODELOPT_EP_DEGREE=AUTO was refused as neither auto nor
a degree.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 1a0d18b1874..6f402f791e5 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -977,7 +977,7 @@ 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``."""
     value = env.ep_degree() if ep_size is None else ep_size
-    if str(value) in ("auto", "0"):
+    if str(value).lower() in ("auto", "0"):
         return "auto"
     try:
         return int(value)
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 4b66f57f248..5072e0fe426 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -413,9 +413,9 @@ def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from(
 
 
 @needs_qwen3_moe_ep_plan
-@pytest.mark.parametrize("value", ["0", ""], ids=["zero", "exported-empty"])
+@pytest.mark.parametrize("value", ["0", "", "AUTO"], ids=["zero", "exported-empty", "capitals"])
 def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value):
-    """And so is one exported empty: unset, the plan decides too."""
+    """And so is one exported empty, and auto in any case: the plan decides."""
     monkeypatch.setenv("MODELOPT_EP_DEGREE", value)
     run_as(4)
     serve(_qwen3_moe(num_experts=8))

From 5d21a79c1f1fe4b5385019795bd9bc117aeb4650 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 01:27:20 +0000
Subject: [PATCH 24/44] Count the routed experts in the config that holds their
 plan

For a VLM whose nested text config holds the MoE, an expert-count field on the root config was read
first, so an automatic ep_size could be chosen against a count the plan does not describe.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py     |  2 +-
 .../utils/test_resolve_distribution_plan.py     | 17 +++++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 6f402f791e5..d312d8cdeac 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1023,7 +1023,7 @@ 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)
-    num_experts = _num_experts_from_config(hf_config) if experts else 0
+    num_experts = _num_experts_from_config(_ep_config(hf_config)) if experts else 0
     if ep_size == "auto":
         ep_size = _choose_ep_degree(world_size, num_experts)
     elif ep_size > 1 and (num_experts == 0 or num_experts % ep_size):
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 5072e0fe426..cb209577a2a 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -540,6 +540,23 @@ def test_a_vlm_is_decided_from_its_whole_model(run_as, serve, gpu_memory, tmp_pa
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp")
 
 
+def test_the_expert_count_is_read_from_the_config_holding_the_plan(
+    run_as, serve, monkeypatch, tmp_path
+):
+    """A VLM whose nested config holds the MoE and its plan counts the experts there, not by an
+    expert-ish field on its root config."""
+    text = SimpleNamespace(
+        base_model_ep_plan={"layers.*.mlp.experts": "moe_tp_experts"}, num_experts=8
+    )
+    serve(SimpleNamespace(model_type="vlm", architectures=[], text_config=text, num_experts=3))
+    monkeypatch.setattr(
+        model_load_utils, "build_meta_causal_lm", lambda *a, **k: _model_with_decoder_at("model.")
+    )
+    run_as(4)
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)
+
+
 def test_a_model_that_cannot_be_built_on_meta_names_itself(run_as, serve, monkeypatch, tmp_path):
     def unbuildable(*args, **kwargs):
         raise NotImplementedError("custom init touches the weights")

From b16d580278fc466d8b4b2f15e5f5be3e7535f5b3 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 01:27:20 +0000
Subject: [PATCH 25/44] Say that the automatic dense layout replicates on a
 rank 0 without a GPU

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 d312d8cdeac..f2296894886 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1078,7 +1078,7 @@ def resolve_distribution_plan(
     * ``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"``
-      assumes the run's GPUs are alike.
+      assumes the run's GPUs are alike; a rank 0 without a GPU replicates it.
 
     ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` (``1`` replicates; ``world``,
     or the world size as a number, shards) from rank 0's environment, both defaulting to

From 2092b0cf1521a15fdc871af9b2d1f07b521ffd58 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 18:54:41 +0000
Subject: [PATCH 26/44] Name the knob an EP degree came from when it cannot
 shard the experts

The refusal always said ep_size=, also for a degree read from
MODELOPT_EP_DEGREE, so a user who set the variable was told to change an
argument they never passed. It now names what the degree came from.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index f2296894886..43941f5dd06 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -997,6 +997,8 @@ def _decide_distribution(
         _routed_expert_modules,
     )
 
+    # What the caller set the EP degree with, for a refusal to name it.
+    knob = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size"
     if world_size == 1:
         # 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
@@ -1028,7 +1030,7 @@ def _decide_distribution(
         ep_size = _choose_ep_degree(world_size, num_experts)
     elif ep_size > 1 and (num_experts == 0 or num_experts % ep_size):
         raise ValueError(
-            f"ep_size={ep_size} must divide the routed experts {hf_config.model_type} shards "
+            f"{knob}={ep_size} must divide the routed experts {hf_config.model_type} shards "
             f"({num_experts})."
         )
     if dense == "auto":

From 9af25255f4a399896d311ddad490192dfeab2d5c Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 18:55:51 +0000
Subject: [PATCH 27/44] Print the decided distribution with print_rank_0

It is the module's rank-0 printer, already used for the load's other
notice, and flushes as the explicit print did.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 43941f5dd06..aae3ac19897 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1127,6 +1127,6 @@ def checkpoint_read() -> bool:
             ckpt_path, hf_config, world_size, ep_size, dense, trust_remote_code
         ),
     )
-    if plan != DistributionPlan() and dist.is_master():
-        print(f"Distribution: {plan!r}", flush=True)
+    if plan != DistributionPlan():
+        print_rank_0(f"Distribution: {plan!r}")
     return plan, hf_config

From a048b74fd820401ac3c98ef613affc236d6285a3 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 19:49:20 +0000
Subject: [PATCH 28/44] Leave the auto classes a remote model's parts register
 in as found

The meta build kept only the auto class it built through. A remote
VLM's __init__ can build its text model through AutoModelForCausalLM
while the whole model builds through another auto class, and that
registration stayed: rank 0's later load through AutoModelForCausalLM
would then build the registered class, without the generate its peers'
load adds. The build now also keeps AutoModel's and
AutoModelForCausalLM's registrations as it found them.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   | 27 +++++++++-------
 .../utils/test_resolve_distribution_plan.py   | 31 +++++++++++++++++++
 2 files changed, 47 insertions(+), 11 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index aae3ac19897..3b7950da1b4 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -200,11 +200,12 @@ def build_meta_causal_lm(
 ):
     """Build a meta-init causal LM (no real storage allocated), or ``auto_class``'s model.
 
-    It leaves ``auto_class`` as it found it. Building a remote-code model through an auto class
-    registers its class there, as it is, and a later load through that auto class builds the
-    registered class, without the ``GenerationMixin`` a first load adds: a rank that built one more
-    meta model than its peers, as rank 0 deciding a distribution does, would load its model without
-    the ``generate`` theirs have.
+    It leaves ``auto_class`` as it found it, and the auto classes a remote model's ``__init__``
+    builds its parts through (``AutoModel``, ``AutoModelForCausalLM``). Building a remote-code
+    model through an auto class registers its class there, as it is, and a later load through that
+    auto class builds the registered class, without the ``GenerationMixin`` a first load adds: a
+    rank that built one more meta model than its peers, as rank 0 deciding a distribution does,
+    would load its model without the ``generate`` theirs have.
     """
     if hf_config is None:
         config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code}
@@ -225,18 +226,22 @@ def build_meta_causal_lm(
             from transformers import AutoModelForImageTextToText
 
             model_factory = AutoModelForImageTextToText
-    # Where transformers records the class a remote-code build registers. Read without a fallback:
-    # if a release moves it, this fails here rather than leave one rank's model without generate.
-    registered = model_factory._model_mapping._extra_content
-    kept = dict(registered)
+    # Where transformers records the class a remote-code build registers: in the auto class built
+    # through, and in those a remote VLM builds its towers through. Read without a fallback: if a
+    # release moves it, this fails here rather than leave one rank's model without generate.
+    kept = [
+        (cls._model_mapping._extra_content, dict(cls._model_mapping._extra_content))
+        for cls in (model_factory, AutoModel, AutoModelForCausalLM)
+    ]
     try:
         with init_empty_weights(include_buffers=False):
             model = model_factory.from_config(
                 hf_config, dtype=dtype, trust_remote_code=trust_remote_code
             )
     finally:
-        registered.clear()
-        registered.update(kept)
+        for registered, saved in kept:
+            registered.clear()
+            registered.update(saved)
     model.eval()
     return model
 
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index cb209577a2a..1ec85fb3f13 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -569,6 +569,37 @@ def unbuildable(*args, **kwargs):
         resolve_distribution_plan(str(tmp_path))
 
 
+def test_a_meta_build_leaves_the_auto_class_a_part_registers_in_as_it_found_it(monkeypatch):
+    """A remote VLM's ``__init__`` can build its text model through ``AutoModelForCausalLM`` while
+    the whole model builds through another auto class. That registers the text model's class in
+    ``AutoModelForCausalLM`` too, which would leave rank 0's later load of it without the
+    ``generate`` its peers' have."""
+
+    class ProbeConfig(transformers.PretrainedConfig):
+        model_type = "meta_build_part_probe"
+
+    class ProbeModel(transformers.PreTrainedModel):
+        config_class = ProbeConfig
+
+    for auto in (AutoModel, AutoModelForCausalLM):
+        mapping = auto._model_mapping
+        monkeypatch.setattr(mapping, "_extra_content", dict(mapping._extra_content))
+    from_config = AutoModel.from_config
+
+    def remote_vlm_from_config(config, **kwargs):
+        # What a remote VLM's __init__ does for its text model before the whole model is built.
+        AutoModelForCausalLM.register(ProbeConfig, ProbeModel, exist_ok=True)
+        return from_config(config, **kwargs)
+
+    monkeypatch.setattr(AutoModel, "from_config", remote_vlm_from_config)
+
+    model_load_utils.build_meta_causal_lm(
+        None, False, None, LlamaConfig(**_TINY), auto_class=AutoModel
+    )
+
+    assert ProbeConfig not in AutoModelForCausalLM._model_mapping
+
+
 def test_a_meta_build_leaves_the_auto_class_as_it_found_it(monkeypatch):
     """Building a remote-code model through an auto class registers its class there, as it is, and a
     later load through that auto class builds the registered class, without the ``generate`` a first

From a334adc1102831df4abe24b9971ba9cb402c7cfd Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 19:54:49 +0000
Subject: [PATCH 29/44] Refuse an EP plan whose expert containers match no
 module

A config with routed experts whose base_model_ep_plan names containers
the built model has none of (a renamed submodule, a hand-written plan)
left the decision with no experts to count, so the automatic degree
became 1 and every expert was replicated on every rank without a word;
the expert-parallel load finds the experts the same way, so it could
not have sharded them either. That is refused now, naming the
containers, unless the degree is set to 1.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   | 13 ++++++++++++
 .../utils/test_resolve_distribution_plan.py   | 20 +++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 3b7950da1b4..1225db42b60 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1031,6 +1031,19 @@ def _decide_distribution(
     ep_plan = getattr(_ep_config(hf_config), "base_model_ep_plan", None) or {}
     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
+        # without them would replicate every expert on every rank, and the expert-parallel load
+        # finds them the same way.
+        from modelopt.torch.utils.moe import _routed_expert_keys
+
+        containers = ", ".join(sorted(_routed_expert_keys(ep_plan))) or "none"
+        raise ValueError(
+            f"{hf_config.model_type} has routed experts and a base_model_ep_plan, but none of the "
+            f"plan's expert containers ({containers}) names a module of the built model, so they "
+            f"can be neither counted nor sharded. Set {knob}=1 to load it without expert "
+            "parallelism."
+        )
     if ep_size == "auto":
         ep_size = _choose_ep_degree(world_size, num_experts)
     elif ep_size > 1 and (num_experts == 0 or num_experts % ep_size):
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 1ec85fb3f13..b9d21a37d4f 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -391,6 +391,26 @@ 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"])
+def test_an_ep_plan_naming_no_module_of_the_model_is_refused(
+    run_as, serve, monkeypatch, tmp_path, degree
+):
+    """A plan whose expert containers match no module of the built model (a renamed submodule, a
+    hand-written plan) leaves no expert to count or shard: the automatic degree would replicate
+    every expert on every rank without a word. A degree set to 1 asks for exactly that."""
+    run_as(4)
+    config = _qwen3_moe(num_experts=8)
+    config.base_model_ep_plan = {"layers.*.mlp.renamed_experts": "moe_tp_experts"}
+    serve(config)
+
+    if degree is None:
+        with pytest.raises(ValueError, match=r"renamed_experts\) names .* MODELOPT_EP_DEGREE=1"):
+            resolve_distribution_plan(str(tmp_path))
+    else:
+        monkeypatch.setenv("MODELOPT_EP_DEGREE", degree)
+        assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
+
+
 @pytest.mark.parametrize(
     ("world_size", "source"),
     [(4, "MODELOPT_EP_DEGREE"), (1, "ep_size")],

From 6f85cf7771317bf5b72432e9abd38348b1a6f43e Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 19:58:01 +0000
Subject: [PATCH 30/44] Name the knob a degree came from when it cannot divide
 the world

An explicit EP degree that does not divide the world size was refused
by DistributionPlan, as ep_size=..., also when MODELOPT_EP_DEGREE set
it. The decision now refuses it first and names the knob or argument
the degree came from, as its refusal for the experts does.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 1225db42b60..5411ea1ba3e 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1046,6 +1046,8 @@ def _decide_distribution(
         )
     if ep_size == "auto":
         ep_size = _choose_ep_degree(world_size, num_experts)
+    elif ep_size > 1 and world_size % ep_size:
+        raise ValueError(f"{knob}={ep_size} must divide the world size ({world_size}).")
     elif ep_size > 1 and (num_experts == 0 or num_experts % ep_size):
         raise ValueError(
             f"{knob}={ep_size} must divide the routed experts {hf_config.model_type} shards "
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index b9d21a37d4f..14c60b7b359 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -411,6 +411,24 @@ 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("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
+):
+    """Every rank belongs to one EP group, so the degree must divide the world; the refusal names
+    the knob or argument the degree came from, as the one for the experts does."""
+    run_as(4)
+    serve(LlamaConfig(**_TINY))
+    arguments = {}
+    if source == "ep_size":
+        arguments["ep_size"] = 3
+    else:
+        monkeypatch.setenv(source, "3")
+
+    with pytest.raises(ValueError, match=rf"{source}=3 must divide the world size \(4\)"):
+        resolve_distribution_plan(str(tmp_path), **arguments)
+
+
 @pytest.mark.parametrize(
     ("world_size", "source"),
     [(4, "MODELOPT_EP_DEGREE"), (1, "ep_size")],

From 2b1792d8a35ebcb6af7a86c014ddc8ae7cf3be02 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 21:04:12 +0000
Subject: [PATCH 31/44] Wildcard each module name once when finding the routed
 experts

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

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index 3e59099c45e..adcaaf02b14 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -111,15 +111,12 @@ def _routed_expert_modules(model: "nn.Module", ep_plan: dict[str, str]) -> "dict
     from transformers.integrations.tensor_parallel import replace_layer_number_by_wildcard
 
     keys = _routed_expert_keys(ep_plan)
-    return {
-        name: module
-        for name, module in model.named_modules()
-        if any(
-            (generic := replace_layer_number_by_wildcard(name)) == key
-            or generic.endswith("." + key)
-            for key in keys
-        )
-    }
+    found = {}
+    for name, module in model.named_modules():
+        generic = replace_layer_number_by_wildcard(name)
+        if any(generic == key or generic.endswith("." + key) for key in keys):
+            found[name] = module
+    return found
 
 
 # ---------------------------------------------------------------------------------------------------

From c1e5ed40f3ff43c393dd1fca9ce2bf7c4d8e69d7 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 21:04:29 +0000
Subject: [PATCH 32/44] Say which auto-class registrations the meta build's
 restore assumes

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 5411ea1ba3e..3d91d7db0ce 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -205,7 +205,9 @@ def build_meta_causal_lm(
     model through an auto class registers its class there, as it is, and a later load through that
     auto class builds the registered class, without the ``GenerationMixin`` a first load adds: a
     rank that built one more meta model than its peers, as rank 0 deciding a distribution does,
-    would load its model without the ``generate`` theirs have.
+    would load its model without the ``generate`` theirs have. That assumes a remote model registers
+    its classes while it is built: one its module registered when imported is undone as well, and a
+    later load in the process, which finds the module imported, does not register it again.
     """
     if hf_config is None:
         config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code}

From c16c587e344d7b7fab905fc543a39d8430c9d09c Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 21:04:53 +0000
Subject: [PATCH 33/44] Say a caller resolves the checkpoint before deciding
 its plan

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/plugins/model_load_utils.py | 5 ++++-
 1 file 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 3d91d7db0ce..5bec81ae271 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1112,7 +1112,10 @@ def resolve_distribution_plan(
     ``revision``, ``token``, ``cache_dir``, ``local_files_only`` and ``subfolder`` locate the
     checkpoint as :func:`load_pretrained` takes them, so the plan is decided for what the load reads:
     all but ``subfolder`` resolve a Hub model ID (:func:`ensure_local_checkpoint`), and the config
-    read takes all but ``revision``, for custom code its ``auto_map`` names in another Hub repo.
+    read takes all but ``revision``, for custom code its ``auto_map`` names in another Hub repo. It
+    returns no path, so a caller that loads the checkpoint resolves it first, with
+    :func:`ensure_local_checkpoint`, and passes the local directory: the plan is then decided for
+    the snapshot the load reads.
     """
     from modelopt.torch.utils import distributed as dist
 

From 76216dfbcce892adcd996e930e115357f6fadba6 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 21:51:37 +0000
Subject: [PATCH 34/44] Build a multimodal causal LM through
 AutoModelForCausalLM for its plan

_auto_class fell back to the headless AutoModel for any multimodal config that neither the
image-text nor a seq2seq class recognizes, also when AutoModelForCausalLM builds it whole, as it
does Phi-4-multimodal (the one such config in transformers 5.5.4 and 5.18). The decision then
weighed the model without its LM head. AutoModelForImageTextToText exists in every supported
release, so its import guard goes.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 5bec81ae271..8037aeed2e1 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -919,7 +919,8 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
     ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` there,
     its code is trusted and it registers no image-text class, else through
     ``AutoModelForImageTextToText`` if that recognizes the config or, its code trusted, the
-    checkpoint registers one, else through ``AutoModel``, which leaves the head out. An
+    checkpoint registers one, else through ``AutoModelForCausalLM`` if that recognizes it (as it
+    does Phi-4-multimodal), else through ``AutoModel``, which leaves the head out. An
     encoder-decoder model, whose architecture reads the same, builds through its seq2seq class.
     """
     architectures = getattr(hf_config, "architectures", None) or []
@@ -937,19 +938,21 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
     )
     if remote_with_head and "AutoModelForImageTextToText" not in auto_map:
         return AutoModelForCausalLM
-    from transformers import AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq
+    from transformers import (
+        AutoModelForImageTextToText,
+        AutoModelForSeq2SeqLM,
+        AutoModelForSpeechSeq2Seq,
+    )
 
     for seq2seq in (AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq):
         if type(hf_config) in seq2seq._model_mapping:
             return seq2seq
-    try:
-        from transformers import AutoModelForImageTextToText
-    except ImportError:
-        return AutoModel
     if type(hf_config) in AutoModelForImageTextToText._model_mapping or (
         trust_remote_code and "AutoModelForImageTextToText" in auto_map
     ):
         return AutoModelForImageTextToText
+    if type(hf_config) in AutoModelForCausalLM._model_mapping:
+        return AutoModelForCausalLM
     return AutoModel
 
 
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 14c60b7b359..c154411357c 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -41,6 +41,7 @@
     AutoModelForSeq2SeqLM,
     AutoModelForSpeechSeq2Seq,
     LlamaConfig,
+    Phi4MultimodalConfig,
     Qwen2_5_VLConfig,
     Qwen3MoeConfig,
     T5Config,
@@ -296,6 +297,7 @@ def _tiny_qwen2_5_vl():
         (_vlm_config(auto_map=_REMOTE_VLM), False, AutoModel),
         (_vlm_config(auto_map=_REMOTE_IMAGE_TEXT_VLM), True, AutoModelForImageTextToText),
         (_tiny_qwen2_5_vl(), False, AutoModelForImageTextToText),
+        (Phi4MultimodalConfig(), False, AutoModelForCausalLM),
         (_vlm_config(), False, AutoModel),
         (T5Config(architectures=["T5ForConditionalGeneration"]), False, AutoModelForSeq2SeqLM),
         (
@@ -310,6 +312,7 @@ def _tiny_qwen2_5_vl():
         "remote-vlm-untrusted",
         "remote-image-text-vlm",
         "native-vlm-with-head",
+        "native-multimodal-causal-lm",
         "unknown-vlm",
         "encoder-decoder",
         "speech-encoder-decoder",

From d61e9474073e0f1dea026879b12173c22141b866 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 23:02:07 +0000
Subject: [PATCH 35/44] Say each rank returns its own copy of the config, rank
 0's being what decides the plan

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 8037aeed2e1..d0a046b72f6 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1092,7 +1092,8 @@ def resolve_distribution_plan(
     records it on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), and
     every later step reads it back from there
     (:func:`~modelopt.torch.utils.distributed.distribution_plan`). Returns ``(plan, hf_config)``,
-    ``hf_config`` being the config the plan was decided from, for a sharded load to build from.
+    ``hf_config`` being this rank's own copy of the config (rank 0 decides the plan from its copy),
+    for a sharded load to build from.
 
     Rank 0 decides, from the model built on the meta device: from its parameters rather than
     estimates of them, and from its routed experts as an expert-parallel load finds them. Every rank

From ffb5707be2cd134dff663bdc0c871c8d7163e222 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 00:30:30 +0000
Subject: [PATCH 36/44] Say why the automatic EP degree is the largest one

The largest degree holds the fewest experts per rank, leaving the most memory
for activations, which calibration batches fill; any degree above 1 pays the
same all-to-alls per MoE layer.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/moe.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index adcaaf02b14..09ce803599d 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -77,7 +77,9 @@ def _choose_ep_degree(world_size: int, num_experts: int) -> int:
 
     EP must divide ``num_experts`` (whole experts per rank) and ``world_size`` (a valid mesh
     factorization). Each expert shard is then replicated over the remaining ranks
-    (``DistributionPlan.edp``). Heuristic -- callers may override.
+    (``DistributionPlan.edp``). The largest degree holds the fewest experts per rank, which leaves
+    the most memory for activations, i.e. for calibration batches; any degree above 1 pays the same
+    two all-to-alls per MoE layer. Heuristic -- callers may override.
     """
     if num_experts <= 0:
         return 1

From 5c0307386ccfcbd34895d46f82135ad0b09dc8be Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 01:12:42 +0000
Subject: [PATCH 37/44] Import the plan decision's MoE helpers and auto classes
 at module top

Nothing deferred them: modelopt.torch.utils.moe does not import this module, and
transformers is imported at the top already. The refusal's second import of
_routed_expert_keys goes with them.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index d0a046b72f6..08d54a8f313 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -33,7 +33,15 @@
 from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
 from torch.distributed.fsdp import FSDPModule
 from torch.distributed.tensor import DTensor
-from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, PretrainedConfig
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    AutoModelForSeq2SeqLM,
+    AutoModelForSpeechSeq2Seq,
+    PretrainedConfig,
+)
 from transformers.conversion_mapping import get_model_conversion_mapping
 from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key
 
@@ -46,6 +54,13 @@
     is_initialized,
 )
 from modelopt.torch.utils.logging import print_rank_0, warn_rank_0
+from modelopt.torch.utils.moe import (
+    _choose_ep_degree,
+    _ep_config,
+    _num_experts_from_config,
+    _routed_expert_keys,
+    _routed_expert_modules,
+)
 from modelopt.torch.utils.plugins.hf_checkpoint_utils import (
     checkpoint_tensor_files,
     ensure_local_checkpoint,
@@ -938,12 +953,6 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
     )
     if remote_with_head and "AutoModelForImageTextToText" not in auto_map:
         return AutoModelForCausalLM
-    from transformers import (
-        AutoModelForImageTextToText,
-        AutoModelForSeq2SeqLM,
-        AutoModelForSpeechSeq2Seq,
-    )
-
     for seq2seq in (AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq):
         if type(hf_config) in seq2seq._model_mapping:
             return seq2seq
@@ -1000,13 +1009,6 @@ def _decide_distribution(
     ckpt_path: str, hf_config, world_size: int, ep_size, dense, trust_remote_code: bool
 ) -> DistributionPlan:
     """The plan :func:`resolve_distribution_plan` returns, decided by this rank alone."""
-    from modelopt.torch.utils.moe import (
-        _choose_ep_degree,
-        _ep_config,
-        _num_experts_from_config,
-        _routed_expert_modules,
-    )
-
     # What the caller set the EP degree with, for a refusal to name it.
     knob = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size"
     if world_size == 1:
@@ -1040,8 +1042,6 @@ def _decide_distribution(
         # The plan names routed experts the built model has none of under those names: deciding
         # without them would replicate every expert on every rank, and the expert-parallel load
         # finds them the same way.
-        from modelopt.torch.utils.moe import _routed_expert_keys
-
         containers = ", ".join(sorted(_routed_expert_keys(ep_plan))) or "none"
         raise ValueError(
             f"{hf_config.model_type} has routed experts and a base_model_ep_plan, but none of the "

From 539b57d4a8137b9cb0ea5bea9d216f07f797a465 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 08:00:27 +0000
Subject: [PATCH 38/44] Import the distributed helpers' module at the top too

The plan decision imported modelopt.torch.utils.distributed inside the
function, with nothing left to defer it: the module already imports names
from it at the top. It is bound as a module alias, which the tests patch.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 08d54a8f313..55b8dbfc7f5 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -45,6 +45,7 @@
 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.utils import distributed as dist
 from modelopt.torch.utils import env
 from modelopt.torch.utils.distributed import (
     DistributionPlan,
@@ -1121,8 +1122,6 @@ def resolve_distribution_plan(
     :func:`ensure_local_checkpoint`, and passes the local directory: the plan is then decided for
     the snapshot the load reads.
     """
-    from modelopt.torch.utils import distributed as dist
-
     # 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]] = []

From 71d090d0f3b27c929ffa663a712ecef956c089eb Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 00:36:56 +0000
Subject: [PATCH 39/44] Count a tied embedding and LM head once when sizing the
 model

init_empty_weights re-wraps every parameter it registers, which unties what
from_config tied, so the plan decision counted a tied embedding and LM head
twice. It now re-ties the meta model before summing its bytes, as the loaded
model holds them. The dense-fit test is sized from a normally built model, tied
and untied, at either side of the 60% share.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   |  3 +++
 .../utils/test_resolve_distribution_plan.py   | 21 ++++++++++++-------
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 55b8dbfc7f5..7f6057db8a5 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1031,6 +1031,9 @@ def _decide_distribution(
             hf_config,
             auto_class=_auto_class(hf_config, trust_remote_code),
         )
+        # init_empty_weights re-wraps every parameter it registers, which unties what from_config
+        # tied: re-tie, so a tied embedding and LM head count once, as the loaded model holds them.
+        model.tie_weights()
     except Exception as e:
         raise RuntimeError(
             f"Deciding how to distribute {hf_config.model_type} needs its model built on the meta "
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index c154411357c..4cfc6690787 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -21,6 +21,7 @@
 meta device.
 """
 
+import math
 import os
 from contextlib import contextmanager
 from types import SimpleNamespace
@@ -464,16 +465,20 @@ def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value):
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)
 
 
+@pytest.mark.parametrize("tied", [False, True], ids=["untied", "tied-embeddings"])
 @pytest.mark.parametrize(
-    ("gpus_worth", "dense"), [(2, "replicate"), (1, "fsdp")], ids=["fits", "too-big"]
+    ("share", "dense"), [(0.59, "replicate"), (0.61, "fsdp")], ids=["fits", "too-big"]
 )
 def test_a_dense_model_is_fsdp_sharded_only_when_too_big_for_one_gpu(
-    run_as, serve, gpu_memory, tmp_path, gpus_worth, dense
+    run_as, serve, gpu_memory, tmp_path, share, dense, tied
 ):
-    """On a GPU only its own size, the whole model would leave less than 40% of it to calibration."""
+    """Taking more than 60% of a GPU, the whole model would leave too little to calibration. It is
+    sized as it loads, a tied embedding and LM head counting once, though its meta build unties
+    them."""
     run_as(4)
-    config = serve(LlamaConfig(**_TINY))
-    gpu_memory(gpus_worth * _bytes(_meta(config)))
+    config = serve(LlamaConfig(**_TINY, tie_word_embeddings=tied))
+    loaded = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16)
+    gpu_memory(math.ceil(_bytes(loaded) / share))
 
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, dense=dense)
 
@@ -590,9 +595,9 @@ def test_the_expert_count_is_read_from_the_config_holding_the_plan(
         base_model_ep_plan={"layers.*.mlp.experts": "moe_tp_experts"}, num_experts=8
     )
     serve(SimpleNamespace(model_type="vlm", architectures=[], text_config=text, num_experts=3))
-    monkeypatch.setattr(
-        model_load_utils, "build_meta_causal_lm", lambda *a, **k: _model_with_decoder_at("model.")
-    )
+    built = _model_with_decoder_at("model.")
+    built.tie_weights = lambda: None  # a PreTrainedModel's; nothing here is tied
+    monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: built)
     run_as(4)
 
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)

From d42fc09d58929de257f11ae35ed443241f8d0bb8 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 00:37:05 +0000
Subject: [PATCH 40/44] Build an encoder-decoder config that names no
 architecture through its seq2seq class

The plan decision took a config with no nested text or vision config and no
*ForConditionalGeneration architecture for a causal LM, so a bare T5Config or
WhisperConfig failed its meta build. is_encoder_decoder now sends it to the
seq2seq dispatch.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/utils/plugins/model_load_utils.py   |  5 ++--
 .../utils/test_resolve_distribution_plan.py   | 23 +++++++++++--------
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 7f6057db8a5..32009a9cb47 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -937,7 +937,8 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
     ``AutoModelForImageTextToText`` if that recognizes the config or, its code trusted, the
     checkpoint registers one, else through ``AutoModelForCausalLM`` if that recognizes it (as it
     does Phi-4-multimodal), else through ``AutoModel``, which leaves the head out. An
-    encoder-decoder model, whose architecture reads the same, builds through its seq2seq class.
+    encoder-decoder model (``is_encoder_decoder``, whether or not it names its architecture, which
+    reads like a VLM's) builds through its seq2seq class.
     """
     architectures = getattr(hf_config, "architectures", None) or []
     is_vlm = (
@@ -945,7 +946,7 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any:
         or getattr(hf_config, "vision_config", None) is not None
         or any("conditionalgeneration" in str(arch).lower() for arch in architectures)
     )
-    if not is_vlm:
+    if not is_vlm and not getattr(hf_config, "is_encoder_decoder", False):
         return AutoModelForCausalLM
     auto_map = getattr(hf_config, "auto_map", None) or {}
     remote_with_head = (
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
index 4cfc6690787..1d0ee32a32b 100644
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py
@@ -301,11 +301,7 @@ def _tiny_qwen2_5_vl():
         (Phi4MultimodalConfig(), False, AutoModelForCausalLM),
         (_vlm_config(), False, AutoModel),
         (T5Config(architectures=["T5ForConditionalGeneration"]), False, AutoModelForSeq2SeqLM),
-        (
-            WhisperConfig(architectures=["WhisperForConditionalGeneration"]),
-            False,
-            AutoModelForSpeechSeq2Seq,
-        ),
+        (WhisperConfig(), False, AutoModelForSpeechSeq2Seq),
     ],
     ids=[
         "causal-lm",
@@ -323,7 +319,7 @@ def test_the_class_that_builds_the_whole_model(config, trust_remote_code, expect
     """A VLM's AutoModel leaves the LM head out: it builds through the class that carries the head,
     keeping the headless backbone only when no class recognizes the config. A remote class counts
     only when the checkpoint's code is trusted. An encoder-decoder model's architecture reads like a
-    VLM's, but its seq2seq class carries its head."""
+    VLM's, but its seq2seq class carries its head, also for a config that names no architecture."""
     assert _auto_class(config, trust_remote_code) is expected
 
 
@@ -577,10 +573,19 @@ def test_an_ep_degree_that_cannot_shard_the_experts_is_refused(
         resolve_distribution_plan(str(tmp_path), ep_size=ep_size)
 
 
-def test_a_vlm_is_decided_from_its_whole_model(run_as, serve, gpu_memory, tmp_path):
-    """AutoModelForCausalLM does not build every VLM: the decision builds what the load builds."""
+@pytest.mark.parametrize(
+    "config",
+    [
+        _tiny_qwen2_5_vl(),
+        T5Config(d_model=16, d_ff=32, d_kv=8, num_layers=1, num_heads=2, vocab_size=32),
+    ],
+    ids=["vlm", "encoder-decoder-naming-no-architecture"],
+)
+def test_a_model_is_decided_from_its_whole_model(run_as, serve, gpu_memory, tmp_path, config):
+    """AutoModelForCausalLM builds neither every VLM nor an encoder-decoder model, whose config need
+    not name its architecture: the decision builds what the load builds."""
     run_as(2)
-    serve(_tiny_qwen2_5_vl())
+    serve(config)
     gpu_memory(1)
 
     assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp")

From 0ac9e27eb97d42c042c033b06475a21e5874d11b Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 00:37:15 +0000
Subject: [PATCH 41/44] Say why the routed-expert match imports transformers
 locally

transformers is the optional [hf] extra, and modelopt.torch.utils.moe must
import without it (test_moe_a2a imports it in the torch-only install).

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

diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py
index 09ce803599d..98c4bc2453c 100644
--- a/modelopt/torch/utils/moe.py
+++ b/modelopt/torch/utils/moe.py
@@ -110,6 +110,7 @@ def _routed_expert_modules(model: "nn.Module", ep_plan: dict[str, str]) -> "dict
     ``language_model.model.`` a causal LM or a VLM puts in front. Deciding a model's distribution and
     loading it expert-parallel both find the experts this way, so the two cannot disagree.
     """
+    # Local: transformers is the optional [hf] extra, and this module must import without it.
     from transformers.integrations.tensor_parallel import replace_layer_number_by_wildcard
 
     keys = _routed_expert_keys(ep_plan)

From cff385b79298aed7803a6ea6180af9b99cf00e92 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 01:04:45 +0000
Subject: [PATCH 42/44] Move the plan decision's tests into the files of the
 modules they cover

CONTRIBUTING.md names a test file after the production module it covers. The
tests of resolve_distribution_plan, _auto_class and build_meta_causal_lm move
into test_model_load_utils.py, and those of moe.py's expert counting, EP
degree and routed-expert match into test_moe_a2a.py, that module's test file:
tests/unit already has a test_moe.py (modelopt/torch/models/moe.py's), and
pytest refuses two test modules of one basename in a run. The plain-module MoE
the two files share is in _test_utils/torch/moe_models.py. Test bodies are
unchanged apart from the helpers' names, and the routed-expert test skips
without transformers, as test_moe_a2a.py imports without it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/_test_utils/torch/moe_models.py         |  51 ++
 .../unit/torch/utils/test_model_load_utils.py | 621 +++++++++++++-
 tests/unit/torch/utils/test_moe_a2a.py        | 104 ++-
 .../utils/test_resolve_distribution_plan.py   | 756 ------------------
 4 files changed, 772 insertions(+), 760 deletions(-)
 create mode 100644 tests/_test_utils/torch/moe_models.py
 delete mode 100644 tests/unit/torch/utils/test_resolve_distribution_plan.py

diff --git a/tests/_test_utils/torch/moe_models.py b/tests/_test_utils/torch/moe_models.py
new file mode 100644
index 00000000000..709944f7b34
--- /dev/null
+++ b/tests/_test_utils/torch/moe_models.py
@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Plain-module stand-ins for an MoE model: a decoder whose layers hold routed experts, under any
+prefix a causal LM or a VLM puts it.
+"""
+
+import torch
+from torch import nn
+
+
+class MoeLayer(nn.Module):
+    """A decoder layer: attention, and an MLP of a router, routed experts and shared experts."""
+
+    def __init__(self):
+        super().__init__()
+        self.self_attn = nn.Linear(2, 2)
+        self.mlp = nn.Module()
+        self.mlp.gate = nn.Linear(2, 4, bias=False)  # the router
+        self.mlp.experts = nn.Module()
+        self.mlp.experts.gate_up_proj = nn.Parameter(torch.zeros(4, 4, 2))
+        self.mlp.experts.down_proj = nn.Parameter(torch.zeros(4, 2, 2))
+        self.mlp.shared_experts = nn.Linear(2, 4, bias=False)
+
+
+def model_with_decoder_at(prefix):
+    """A two-layer MoE whose decoder (``layers.<i>``) sits under ``prefix`` (``""``: the decoder)."""
+    decoder = nn.Module()
+    decoder.embed_tokens = nn.Embedding(8, 2)
+    decoder.layers = nn.ModuleList(MoeLayer() for _ in range(2))
+    if not prefix:
+        return decoder
+    root = parent = nn.Module()
+    *path, leaf = prefix.rstrip(".").split(".")
+    for name in path:
+        parent.add_module(name, nn.Module())
+        parent = getattr(parent, name)
+    parent.add_module(leaf, decoder)
+    return root
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index c2a90c33fb0..16fc0e79e92 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -13,28 +13,61 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Pure-function tests for ``modelopt.torch.utils.plugins.model_load_utils``."""
+"""Tests for ``modelopt.torch.utils.plugins.model_load_utils``.
+
+Nothing here needs a GPU. The tests of ``resolve_distribution_plan`` answer the process group's rank
+and size as a torchrun launch would, and serve tiny real transformers configs from memory, so the
+decision builds their real models on the meta device; only the last ones use a process group, two
+gloo ranks.
+"""
+
+import math
+import os
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest import mock
 
 import pytest
 import torch
+from _test_utils.torch.distributed.utils import bounded_collectives
+from _test_utils.torch.moe_models import model_with_decoder_at
 from safetensors.torch import save_file
 
 pytest.importorskip("accelerate")
-pytest.importorskip("transformers")
+transformers = pytest.importorskip("transformers")
 from _test_utils.torch.transformers_models import (
     _TinyQwen35Tokenizer,
     get_tiny_qwen3_5_moe_vl,
     get_tiny_qwen3_5_vl,
 )
-from transformers import AutoConfig
+from packaging.version import Version
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    AutoModelForSeq2SeqLM,
+    AutoModelForSpeechSeq2Seq,
+    LlamaConfig,
+    Phi4MultimodalConfig,
+    Qwen2_5_VLConfig,
+    Qwen3MoeConfig,
+    T5Config,
+    WhisperConfig,
+)
 
+import modelopt.torch.utils.distributed as dist
 from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
+from modelopt.torch.utils.distributed import DistributionPlan
+from modelopt.torch.utils.plugins import model_load_utils
 from modelopt.torch.utils.plugins.model_load_utils import (
+    _auto_class,
     _conversion_plan,
     _convert_keys,
     _resolve_target,
     build_meta_causal_lm,
     record_unplaced_source_keys,
+    resolve_distribution_plan,
     unplaced_source_keys,
 )
 
@@ -153,3 +186,585 @@ def test_meta_qwen_vlm_preserves_architecture_and_vision(model_type):
         == config.text_config.num_hidden_layers
     )
     assert all(parameter.is_meta for parameter in model.parameters())
+
+
+# -------------------------------------------------------------- deciding how a model is distributed
+
+
+_TINY = {
+    "hidden_size": 16,
+    "intermediate_size": 32,
+    "num_hidden_layers": 1,
+    "num_attention_heads": 2,
+    "num_key_value_heads": 1,
+    "vocab_size": 32,
+}
+
+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",
+)
+
+
+def _qwen3_moe(num_experts=8):
+    """Qwen3-MoE, whose base_model_ep_plan ships with transformers itself."""
+    return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY)
+
+
+def _bytes(module):
+    return sum(p.numel() * p.element_size() for p in module.parameters())
+
+
+def _meta(config):
+    """``config``'s model as the decision builds it, in the same dtype."""
+    return model_load_utils.build_meta_causal_lm(None, False, None, config)
+
+
+@pytest.fixture(autouse=True)
+def _no_distribution_knobs(monkeypatch):
+    """A shell that exports the knobs would otherwise steer every decision below."""
+    monkeypatch.delenv("MODELOPT_EP_DEGREE", raising=False)
+    monkeypatch.delenv("MODELOPT_FSDP_SIZE", raising=False)
+
+
+def _gathered_alike(gathered, mine, group=None):
+    """Answer a gather as if every rank handed it what this rank did."""
+    gathered[:] = [mine] * len(gathered)
+
+
+@pytest.fixture
+def run_as(monkeypatch):
+    """Run as rank ``rank`` of a ``world_size``-rank launch, deciding for itself.
+
+    Every rank reads the checkpoint as this one does. How the ranks agree on one plan is tested at
+    the end of this file.
+    """
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0)
+    monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None)
+    monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike)
+
+    def enter(world_size, rank=0):
+        monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1)
+        monkeypatch.setattr(dist, "size", lambda group=None: world_size)
+        monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0)
+        monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size)
+
+    return enter
+
+
+@pytest.fixture
+def serve(monkeypatch):
+    """Answer every config read with ``config`` rather than a checkpoint on disk."""
+
+    def serve(config):
+        monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **k: config)
+        return config
+
+    return serve
+
+
+@pytest.fixture
+def gpu_memory(monkeypatch):
+    """Give this rank a GPU of ``size`` bytes."""
+
+    def give(size):
+        monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: size)
+
+    return give
+
+
+# ------------------------------------------------------------------------------ model class
+
+
+def _vlm_config(**kwargs):
+    return SimpleNamespace(text_config=SimpleNamespace(), architectures=[], **kwargs)
+
+
+def _tiny_qwen2_5_vl():
+    return Qwen2_5_VLConfig(
+        text_config={**_TINY, "max_position_embeddings": 16},
+        vision_config={
+            "depth": 1,
+            "hidden_size": 16,
+            "intermediate_size": 16,
+            "num_heads": 2,
+            "out_hidden_size": 16,
+            "patch_size": 2,
+            "spatial_merge_size": 1,
+            "temporal_patch_size": 1,
+            "fullatt_block_indexes": [0],
+        },
+    )
+
+
+_REMOTE_VLM = {"AutoModelForCausalLM": "m.KimiK25ForConditionalGeneration"}
+_REMOTE_IMAGE_TEXT_VLM = {
+    **_REMOTE_VLM,
+    "AutoModelForImageTextToText": "m.KimiK25ForConditionalGeneration",
+}
+
+
+@pytest.mark.parametrize(
+    ("config", "trust_remote_code", "expected"),
+    [
+        (SimpleNamespace(architectures=["LlamaForCausalLM"]), False, AutoModelForCausalLM),
+        (_vlm_config(auto_map=_REMOTE_VLM), True, AutoModelForCausalLM),
+        (_vlm_config(auto_map=_REMOTE_VLM), False, AutoModel),
+        (_vlm_config(auto_map=_REMOTE_IMAGE_TEXT_VLM), True, AutoModelForImageTextToText),
+        (_tiny_qwen2_5_vl(), False, AutoModelForImageTextToText),
+        (Phi4MultimodalConfig(), False, AutoModelForCausalLM),
+        (_vlm_config(), False, AutoModel),
+        (T5Config(architectures=["T5ForConditionalGeneration"]), False, AutoModelForSeq2SeqLM),
+        (WhisperConfig(), False, AutoModelForSpeechSeq2Seq),
+    ],
+    ids=[
+        "causal-lm",
+        "remote-vlm-with-head",
+        "remote-vlm-untrusted",
+        "remote-image-text-vlm",
+        "native-vlm-with-head",
+        "native-multimodal-causal-lm",
+        "unknown-vlm",
+        "encoder-decoder",
+        "speech-encoder-decoder",
+    ],
+)
+def test_the_class_that_builds_the_whole_model(config, trust_remote_code, expected):
+    """A VLM's AutoModel leaves the LM head out: it builds through the class that carries the head,
+    keeping the headless backbone only when no class recognizes the config. A remote class counts
+    only when the checkpoint's code is trusted. An encoder-decoder model's architecture reads like a
+    VLM's, but its seq2seq class carries its head, also for a config that names no architecture."""
+    assert _auto_class(config, trust_remote_code) is expected
+
+
+# -------------------------------------------------------------------------------- the decision
+
+
+@pytest.mark.parametrize(
+    "knob",
+    [
+        None,
+        ("MODELOPT_EP_DEGREE", "4"),
+        ("MODELOPT_FSDP_SIZE", "world"),
+        ("MODELOPT_FSDP_SIZE", "2"),
+    ],
+    ids=["no_knob", "ep_degree", "fsdp_world", "fsdp_part"],
+)
+@pytest.mark.parametrize("world_size", [None, 1], ids=["no_process_group", "one_rank"])
+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."""
+    if world_size is not None:
+        run_as(world_size)
+    if knob is not None:
+        monkeypatch.setenv(*knob)
+    served = serve(_qwen3_moe())
+    monkeypatch.setattr(
+        model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model")
+    )
+
+    assert resolve_distribution_plan(str(tmp_path)) == (DistributionPlan(), served)
+    assert resolve_distribution_plan(str(tmp_path), dense="fsdp")[0] == DistributionPlan(
+        dense="fsdp"
+    )
+
+
+@needs_qwen3_moe_ep_plan
+@pytest.mark.parametrize("rank", [0, 1])
+def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, capsys, rank):
+    """Auto picks the largest EP degree dividing both the world and the expert count, hands back
+    the config it decided from, and only rank 0 reports the decision."""
+    run_as(4, rank)
+    served = serve(_qwen3_moe(num_experts=8))
+
+    plan, hf_config = resolve_distribution_plan(str(tmp_path))
+
+    assert plan == DistributionPlan(4, ep_size=4, dense="replicate")
+    assert hf_config is served
+    assert ("Distribution:" in capsys.readouterr().out) == (rank == 0)
+
+
+@needs_qwen3_moe_ep_plan
+def test_no_ep_degree_dividing_the_experts_leaves_them_whole(run_as, serve, tmp_path):
+    run_as(2)
+    serve(_qwen3_moe(num_experts=3))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2)
+
+
+def test_an_moe_without_an_ep_plan_is_not_expert_parallel(run_as, serve, tmp_path):
+    """The config counts 8 experts, but with no base_model_ep_plan nothing could shard them."""
+    run_as(4)
+    config = _qwen3_moe(num_experts=8)
+    config.base_model_ep_plan = None
+    serve(config)
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
+
+
+@pytest.mark.parametrize("degree", [None, "1"], ids=["auto", "set-to-1"])
+def test_an_ep_plan_naming_no_module_of_the_model_is_refused(
+    run_as, serve, monkeypatch, tmp_path, degree
+):
+    """A plan whose expert containers match no module of the built model (a renamed submodule, a
+    hand-written plan) leaves no expert to count or shard: the automatic degree would replicate
+    every expert on every rank without a word. A degree set to 1 asks for exactly that."""
+    run_as(4)
+    config = _qwen3_moe(num_experts=8)
+    config.base_model_ep_plan = {"layers.*.mlp.renamed_experts": "moe_tp_experts"}
+    serve(config)
+
+    if degree is None:
+        with pytest.raises(ValueError, match=r"renamed_experts\) names .* MODELOPT_EP_DEGREE=1"):
+            resolve_distribution_plan(str(tmp_path))
+    else:
+        monkeypatch.setenv("MODELOPT_EP_DEGREE", degree)
+        assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
+
+
+@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
+):
+    """Every rank belongs to one EP group, so the degree must divide the world; the refusal names
+    the knob or argument the degree came from, as the one for the experts does."""
+    run_as(4)
+    serve(LlamaConfig(**_TINY))
+    arguments = {}
+    if source == "ep_size":
+        arguments["ep_size"] = 3
+    else:
+        monkeypatch.setenv(source, "3")
+
+    with pytest.raises(ValueError, match=rf"{source}=3 must divide the world size \(4\)"):
+        resolve_distribution_plan(str(tmp_path), **arguments)
+
+
+@pytest.mark.parametrize(
+    ("world_size", "source"),
+    [(4, "MODELOPT_EP_DEGREE"), (1, "ep_size")],
+    ids=["knob", "argument-in-a-single-process"],
+)
+def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from(
+    run_as, serve, monkeypatch, tmp_path, world_size, source
+):
+    """A single process takes its own argument as given, but parses it as the decision does."""
+    run_as(world_size)
+    serve(LlamaConfig(**_TINY))
+    arguments = {}
+    if source == "ep_size":
+        arguments["ep_size"] = "world"
+    else:
+        monkeypatch.setenv(source, "world")
+
+    with pytest.raises(ValueError, match=f"{source}='world'"):
+        resolve_distribution_plan(str(tmp_path), **arguments)
+
+
+@needs_qwen3_moe_ep_plan
+@pytest.mark.parametrize("value", ["0", "", "AUTO"], ids=["zero", "exported-empty", "capitals"])
+def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value):
+    """And so is one exported empty, and auto in any case: the plan decides."""
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", value)
+    run_as(4)
+    serve(_qwen3_moe(num_experts=8))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)
+
+
+@pytest.mark.parametrize("tied", [False, True], ids=["untied", "tied-embeddings"])
+@pytest.mark.parametrize(
+    ("share", "dense"), [(0.59, "replicate"), (0.61, "fsdp")], ids=["fits", "too-big"]
+)
+def test_a_dense_model_is_fsdp_sharded_only_when_too_big_for_one_gpu(
+    run_as, serve, gpu_memory, tmp_path, share, dense, tied
+):
+    """Taking more than 60% of a GPU, the whole model would leave too little to calibration. It is
+    sized as it loads, a tied embedding and LM head counting once, though its meta build unties
+    them."""
+    run_as(4)
+    config = serve(LlamaConfig(**_TINY, tie_word_embeddings=tied))
+    loaded = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16)
+    gpu_memory(math.ceil(_bytes(loaded) / share))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, dense=dense)
+
+
+@needs_qwen3_moe_ep_plan
+@pytest.mark.parametrize(
+    ("experts_fitting", "dense"),
+    [(1 / 2, "replicate"), (1 / 8, "fsdp")],
+    ids=["fits", "the-experts-share-does-not"],
+)
+def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit(
+    run_as, serve, gpu_memory, tmp_path, experts_fitting, dense
+):
+    """At EP 4 a rank holds the rest of the model whole and a quarter of the routed experts, so the
+    dense part is sharded when that does not fit one GPU, however small the rest; with no EP the
+    whole MoE must fit, and here it does not."""
+    run_as(4)
+    config = serve(_qwen3_moe(num_experts=8))
+    model = _meta(config)
+    experts = sum(_bytes(layer.mlp.experts) for layer in model.model.layers)
+    # A GPU 60% of which holds the rest of the model and that share of the experts.
+    gpu_memory(int((_bytes(model) - experts + experts_fitting * experts) / 0.6))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4, dense)
+    assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan(
+        4, dense="fsdp"
+    )
+
+
+@needs_qwen3_moe_ep_plan
+def test_the_env_knobs_override_the_automatic_plan(run_as, serve, monkeypatch, tmp_path):
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", "2")
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+    run_as(4)
+    serve(_qwen3_moe(num_experts=8))
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 2, "fsdp")
+
+
+@needs_qwen3_moe_ep_plan
+def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path):
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", "4")
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+    run_as(4)
+    serve(_qwen3_moe(num_experts=8))
+
+    plan, _ = resolve_distribution_plan(str(tmp_path), ep_size=2, dense="replicate")
+
+    assert plan == DistributionPlan(4, 2, "replicate")
+
+
+@pytest.mark.parametrize(
+    ("value", "dense"),
+    [
+        ("1", "replicate"),
+        ("4", "fsdp"),
+        ("WORLD", "fsdp"),
+        ("", "fsdp"),
+        ("2", None),
+    ],
+    ids=["1", "4", "WORLD", "exported-empty", "part-of-the-world"],
+)
+def test_the_dense_knob_replicates_or_shards_over_the_whole_world(
+    run_as, serve, gpu_memory, monkeypatch, tmp_path, value, dense
+):
+    """Exported empty, the knob is auto, which shards a model too big for one GPU. Sharding over
+    part of the world is refused: no layout does that."""
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", value)
+    run_as(4)
+    serve(LlamaConfig(**_TINY))
+    gpu_memory(1)  # too small for the whole model, so only auto tells "" from "1"
+
+    if dense is None:
+        with pytest.raises(ValueError, match=r"expected auto, 1 .* or world"):
+            resolve_distribution_plan(str(tmp_path))
+    else:
+        assert resolve_distribution_plan(str(tmp_path))[0].dense == dense
+
+
+@pytest.mark.parametrize(
+    ("config", "ep_size", "experts"),
+    [(LlamaConfig(**_TINY), 2, 0), (_qwen3_moe(num_experts=6), 4, 6)],
+    ids=["dense-model", "experts-not-divisible"],
+)
+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")
+    run_as(4)
+    serve(config)
+
+    with pytest.raises(ValueError, match=rf"ep_size={ep_size} must divide .* \({experts}\)"):
+        resolve_distribution_plan(str(tmp_path), ep_size=ep_size)
+
+
+@pytest.mark.parametrize(
+    "config",
+    [
+        _tiny_qwen2_5_vl(),
+        T5Config(d_model=16, d_ff=32, d_kv=8, num_layers=1, num_heads=2, vocab_size=32),
+    ],
+    ids=["vlm", "encoder-decoder-naming-no-architecture"],
+)
+def test_a_model_is_decided_from_its_whole_model(run_as, serve, gpu_memory, tmp_path, config):
+    """AutoModelForCausalLM builds neither every VLM nor an encoder-decoder model, whose config need
+    not name its architecture: the decision builds what the load builds."""
+    run_as(2)
+    serve(config)
+    gpu_memory(1)
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp")
+
+
+def test_the_expert_count_is_read_from_the_config_holding_the_plan(
+    run_as, serve, monkeypatch, tmp_path
+):
+    """A VLM whose nested config holds the MoE and its plan counts the experts there, not by an
+    expert-ish field on its root config."""
+    text = SimpleNamespace(
+        base_model_ep_plan={"layers.*.mlp.experts": "moe_tp_experts"}, num_experts=8
+    )
+    serve(SimpleNamespace(model_type="vlm", architectures=[], text_config=text, num_experts=3))
+    built = model_with_decoder_at("model.")
+    built.tie_weights = lambda: None  # a PreTrainedModel's; nothing here is tied
+    monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: built)
+    run_as(4)
+
+    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)
+
+
+def test_a_model_that_cannot_be_built_on_meta_names_itself(run_as, serve, monkeypatch, tmp_path):
+    def unbuildable(*args, **kwargs):
+        raise NotImplementedError("custom init touches the weights")
+
+    run_as(2)
+    serve(LlamaConfig(**_TINY))
+    monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", unbuildable)
+
+    with pytest.raises(RuntimeError, match=r"distribute llama .* custom init touches the weights"):
+        resolve_distribution_plan(str(tmp_path))
+
+
+def test_a_meta_build_leaves_the_auto_class_a_part_registers_in_as_it_found_it(monkeypatch):
+    """A remote VLM's ``__init__`` can build its text model through ``AutoModelForCausalLM`` while
+    the whole model builds through another auto class. That registers the text model's class in
+    ``AutoModelForCausalLM`` too, which would leave rank 0's later load of it without the
+    ``generate`` its peers' have."""
+
+    class ProbeConfig(transformers.PretrainedConfig):
+        model_type = "meta_build_part_probe"
+
+    class ProbeModel(transformers.PreTrainedModel):
+        config_class = ProbeConfig
+
+    for auto in (AutoModel, AutoModelForCausalLM):
+        mapping = auto._model_mapping
+        monkeypatch.setattr(mapping, "_extra_content", dict(mapping._extra_content))
+    from_config = AutoModel.from_config
+
+    def remote_vlm_from_config(config, **kwargs):
+        # What a remote VLM's __init__ does for its text model before the whole model is built.
+        AutoModelForCausalLM.register(ProbeConfig, ProbeModel, exist_ok=True)
+        return from_config(config, **kwargs)
+
+    monkeypatch.setattr(AutoModel, "from_config", remote_vlm_from_config)
+
+    model_load_utils.build_meta_causal_lm(
+        None, False, None, LlamaConfig(**_TINY), auto_class=AutoModel
+    )
+
+    assert ProbeConfig not in AutoModelForCausalLM._model_mapping
+
+
+def test_a_meta_build_leaves_the_auto_class_as_it_found_it(monkeypatch):
+    """Building a remote-code model through an auto class registers its class there, as it is, and a
+    later load through that auto class builds the registered class, without the ``generate`` a first
+    load adds. Rank 0 alone builds one to decide a distribution: its load alone would lack it."""
+
+    class ProbeConfig(transformers.PretrainedConfig):
+        model_type = "meta_build_probe"
+
+    class ProbeModel(transformers.PreTrainedModel):
+        config_class = ProbeConfig
+
+    mapping = AutoModelForCausalLM._model_mapping
+    monkeypatch.setattr(mapping, "_extra_content", dict(mapping._extra_content))
+    from_config = AutoModelForCausalLM.from_config
+
+    def remote_from_config(config, **kwargs):
+        # What transformers' from_config does for a remote-code config before it builds.
+        AutoModelForCausalLM.register(ProbeConfig, ProbeModel, exist_ok=True)
+        return from_config(config, **kwargs)
+
+    monkeypatch.setattr(AutoModelForCausalLM, "from_config", remote_from_config)
+
+    model_load_utils.build_meta_causal_lm(None, False, None, LlamaConfig(**_TINY))
+
+    assert ProbeConfig not in mapping
+
+
+def test_the_hub_options_locate_the_checkpoint_the_load_reads(tmp_path):
+    """``revision``, ``cache_dir``, ``local_files_only`` and ``subfolder`` find the checkpoint as
+    ``load_pretrained`` takes them, so the plan and its config are those of what the load reads."""
+    repo = tmp_path / "models--org--model"
+    (repo / "refs").mkdir(parents=True)
+    for revision, commit, hidden_size in (("main", "a" * 40, 16), ("v2", "b" * 40, 32)):
+        (repo / "refs" / revision).write_text(commit)
+        config = LlamaConfig(**{**_TINY, "hidden_size": hidden_size})
+        config.save_pretrained(repo / "snapshots" / commit / "nested")
+
+    _, config = resolve_distribution_plan(
+        "org/model",
+        revision="v2",
+        cache_dir=str(tmp_path),
+        local_files_only=True,
+        subfolder="nested",
+    )
+
+    assert config.hidden_size == 32
+
+
+# ------------------------------------------------------------------------------ rank 0 decides
+
+
+@contextmanager
+def _knobs(**values):
+    """This rank's distribution knobs, exactly ``values``, within the block.
+
+    A pooled worker keeps the environment it was spawned with, so the knobs are cleared first. The
+    environment is restored on the way out.
+    """
+    with mock.patch.dict(os.environ):
+        for name in ("MODELOPT_EP_DEGREE", "MODELOPT_FSDP_SIZE"):
+            os.environ.pop(name, None)
+        os.environ.update(values)
+        yield
+
+
+def _every_rank_takes_rank0s_plan_job(rank, size, ckpt, knobs, expected):
+    with bounded_collectives(), _knobs(**knobs[rank]):
+        plan, _ = resolve_distribution_plan(ckpt)
+    assert plan == expected, (rank, plan)
+
+
+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)
+    cpu_dist_workers.run(
+        _every_rank_takes_rank0s_plan_job,
+        ckpt=str(tmp_path),
+        knobs=[{}, {"MODELOPT_FSDP_SIZE": "world"}],
+        expected=DistributionPlan(2),
+    )
+
+
+def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt):
+    isdir = os.path.isdir
+    with (
+        bounded_collectives(),
+        mock.patch("os.path.isdir", lambda path: (rank == 0 or path != ckpt) and isdir(path)),
+    ):
+        try:
+            resolve_distribution_plan(ckpt)
+        except RuntimeError as error:
+            raised = str(error)
+        else:
+            raised = "no error"
+    # Rank 1's own error, carried to rank 0, rather than rank 0 timing out in the plan's broadcast.
+    assert "which this rank cannot see" in raised, raised
+
+
+def test_a_rank_that_cannot_see_the_checkpoint_fails_every_rank(cpu_dist_workers, tmp_path):
+    """Before the plan's broadcast, which rank 0 would otherwise wait in for a rank that raised."""
+    LlamaConfig(**_TINY).save_pretrained(tmp_path)
+    cpu_dist_workers.run(_one_rank_cannot_see_the_checkpoint_job, ckpt=str(tmp_path))
diff --git a/tests/unit/torch/utils/test_moe_a2a.py b/tests/unit/torch/utils/test_moe_a2a.py
index a53c334ec51..7a872183040 100644
--- a/tests/unit/torch/utils/test_moe_a2a.py
+++ b/tests/unit/torch/utils/test_moe_a2a.py
@@ -23,19 +23,30 @@
 
 The routing-weight gradient check guards that the routing weights keep their gradient: they never
 travel, and weight their tokens' expert results on the rank that holds them.
+
+The tests at the end cover the module's helpers that count a model's routed experts, choose the EP
+degree they are sharded over, and find them by the model's ``base_model_ep_plan``.
 """
 
 import copy
 import warnings
 from datetime import timedelta
+from types import SimpleNamespace
 
 import pytest
 import torch
 import torch.distributed as dist
+from _test_utils.torch.moe_models import MoeLayer, model_with_decoder_at
 from packaging.version import Version
+from torch import nn
 
 import modelopt.torch.utils.moe as moe
-from modelopt.torch.utils.moe import a2a_moe
+from modelopt.torch.utils.moe import (
+    _choose_ep_degree,
+    _num_experts_from_config,
+    _routed_expert_modules,
+    a2a_moe,
+)
 
 # Small, fully-deterministic problem. E must be divisible by every world size under test.
 E, TOP_K, H, T = 4, 2, 8, 6
@@ -401,3 +412,94 @@ def test_fp32_routing_weights_give_grouped_mms_result_at_ep1(cpu_dist_workers):
     if Version(transformers.__version__) < Version("5.6"):
         pytest.skip("expert parallelism needs transformers >= 5.6")
     cpu_dist_workers.run(_fp32_routing_weights_job)
+
+
+# ------------------------------------------------------------------------------- expert counting
+
+
+@pytest.mark.parametrize(
+    ("config", "experts"),
+    [
+        (SimpleNamespace(hidden_size=128), 0),
+        (SimpleNamespace(num_experts=16), 16),
+        (SimpleNamespace(num_local_experts=16), 16),
+        (SimpleNamespace(n_routed_experts=16), 16),
+        (SimpleNamespace(num_routed_experts=16), 16),
+        (SimpleNamespace(text_config=SimpleNamespace(n_routed_experts=32)), 32),
+        (SimpleNamespace(llm_config=SimpleNamespace(n_routed_experts=32)), 32),
+    ],
+    ids=[
+        "dense",
+        "num_experts",
+        "num_local_experts",
+        "n_routed_experts",
+        "num_routed_experts",
+        "vlm-text_config",
+        "vlm-llm_config",
+    ],
+)
+def test_the_expert_count_is_read_under_every_spelling(config, experts):
+    """Architectures disagree on the field name, and a VLM's MoE is its nested LM, whose root
+    config has no expert count of its own: missing either would silently disable EP."""
+    assert _num_experts_from_config(config) == experts
+
+
+# ------------------------------------------------------------------------------- degree selection
+
+
+@pytest.mark.parametrize(
+    ("world", "experts", "degree"),
+    [(8, 64, 8), (8, 6, 2), (8, 5, 1), (8, 0, 1), (1, 64, 1)],
+    ids=["largest-valid", "divides-the-experts", "nothing-divides", "dense", "single-rank"],
+)
+def test_the_largest_ep_degree_dividing_the_world_and_the_experts_is_chosen(world, experts, degree):
+    """EP must divide the expert count, so every rank holds whole experts; with nothing dividing, a
+    dense model, or a single rank, there is no expert parallelism."""
+    assert _choose_ep_degree(world, experts) == degree
+
+
+# ---------------------------------------------------------------------------- routed experts
+
+
+_STOCK_PLAN = {
+    "layers.*.mlp.gate": "ep_router",
+    "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+    "layers.*.mlp.experts.down_proj": "grouped_gemm",
+    "layers.*.mlp.experts": "moe_tp_experts",
+}
+
+
+@pytest.mark.parametrize(
+    "plan",
+    [
+        _STOCK_PLAN,
+        {"layers.*.mlp.experts": "moe_tp_experts"},
+        {"layers.*.mlp.experts": "nemotron_experts_ep"},
+        {k: v for k, v in _STOCK_PLAN.items() if v == "grouped_gemm"},
+        {
+            "layers.*.block_sparse_moe.experts": "moe_tp_experts",
+            "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+        },
+    ],
+    ids=["stock", "experts-style-only", "nemotron-style-only", "grouped-gemm-only", "hybrid"],
+)
+@pytest.mark.parametrize(
+    "prefix",
+    ["", "model.", "model.language_model.", "language_model.model."],
+    ids=["base-model", "causal-lm", "vlm", "legacy-vlm"],
+)
+def test_routed_experts_are_the_plans_containers_under_any_prefix(plan, prefix):
+    """The base-relative plan keys match the decoder wherever the model nests it, at every layer
+    index, the container named by its style or as its grouped-GEMM parameters' parent, also in a
+    plan that names some containers one way and some the other. The router, the shared experts, and
+    a module whose name merely ends in a key's text are not routed experts."""
+    pytest.importorskip("transformers.integrations.tensor_parallel")
+    model = model_with_decoder_at(prefix)
+    model.add_module("sublayers", nn.ModuleList([MoeLayer()]))  # "sublayers.0.mlp.experts"
+
+    found = _routed_expert_modules(model, plan)
+
+    modules = dict(model.named_modules())
+    assert found == {
+        name: modules[name] for name in [f"{prefix}layers.{i}.mlp.experts" for i in (0, 1)]
+    }
diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py
deleted file mode 100644
index 1d0ee32a32b..00000000000
--- a/tests/unit/torch/utils/test_resolve_distribution_plan.py
+++ /dev/null
@@ -1,756 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""How ``resolve_distribution_plan`` decides a model's distribution, and the MoE helpers it uses.
-
-Nothing here needs a GPU, and only the tests at the end a process group, two gloo ranks: elsewhere
-the process group's rank and size are answered the way a torchrun launch would, and the configs are
-tiny real transformers configs served from memory, so the decision builds their real models on the
-meta device.
-"""
-
-import math
-import os
-from contextlib import contextmanager
-from types import SimpleNamespace
-from unittest import mock
-
-import pytest
-import torch
-from _test_utils.torch.distributed.utils import bounded_collectives
-from torch import nn
-
-transformers = pytest.importorskip("transformers")
-
-from packaging.version import Version
-from transformers import (
-    AutoModel,
-    AutoModelForCausalLM,
-    AutoModelForImageTextToText,
-    AutoModelForSeq2SeqLM,
-    AutoModelForSpeechSeq2Seq,
-    LlamaConfig,
-    Phi4MultimodalConfig,
-    Qwen2_5_VLConfig,
-    Qwen3MoeConfig,
-    T5Config,
-    WhisperConfig,
-)
-
-import modelopt.torch.utils.distributed as dist
-from modelopt.torch.utils.distributed import DistributionPlan
-from modelopt.torch.utils.moe import (
-    _choose_ep_degree,
-    _num_experts_from_config,
-    _routed_expert_modules,
-)
-from modelopt.torch.utils.plugins import model_load_utils
-from modelopt.torch.utils.plugins.model_load_utils import _auto_class, resolve_distribution_plan
-
-_TINY = {
-    "hidden_size": 16,
-    "intermediate_size": 32,
-    "num_hidden_layers": 1,
-    "num_attention_heads": 2,
-    "num_key_value_heads": 1,
-    "vocab_size": 32,
-}
-
-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",
-)
-
-
-def _qwen3_moe(num_experts=8):
-    """Qwen3-MoE, whose base_model_ep_plan ships with transformers itself."""
-    return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY)
-
-
-def _bytes(module):
-    return sum(p.numel() * p.element_size() for p in module.parameters())
-
-
-def _meta(config):
-    """``config``'s model as the decision builds it, in the same dtype."""
-    return model_load_utils.build_meta_causal_lm(None, False, None, config)
-
-
-@pytest.fixture(autouse=True)
-def _no_distribution_knobs(monkeypatch):
-    """A shell that exports the knobs would otherwise steer every decision below."""
-    monkeypatch.delenv("MODELOPT_EP_DEGREE", raising=False)
-    monkeypatch.delenv("MODELOPT_FSDP_SIZE", raising=False)
-
-
-def _gathered_alike(gathered, mine, group=None):
-    """Answer a gather as if every rank handed it what this rank did."""
-    gathered[:] = [mine] * len(gathered)
-
-
-@pytest.fixture
-def run_as(monkeypatch):
-    """Run as rank ``rank`` of a ``world_size``-rank launch, deciding for itself.
-
-    Every rank reads the checkpoint as this one does. How the ranks agree on one plan is tested at
-    the end of this file.
-    """
-    monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0)
-    monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None)
-    monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike)
-
-    def enter(world_size, rank=0):
-        monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1)
-        monkeypatch.setattr(dist, "size", lambda group=None: world_size)
-        monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0)
-        monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size)
-
-    return enter
-
-
-@pytest.fixture
-def serve(monkeypatch):
-    """Answer every config read with ``config`` rather than a checkpoint on disk."""
-
-    def serve(config):
-        monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **k: config)
-        return config
-
-    return serve
-
-
-@pytest.fixture
-def gpu_memory(monkeypatch):
-    """Give this rank a GPU of ``size`` bytes."""
-
-    def give(size):
-        monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: size)
-
-    return give
-
-
-# ------------------------------------------------------------------------------- expert counting
-
-
-@pytest.mark.parametrize(
-    ("config", "experts"),
-    [
-        (SimpleNamespace(hidden_size=128), 0),
-        (SimpleNamespace(num_experts=16), 16),
-        (SimpleNamespace(num_local_experts=16), 16),
-        (SimpleNamespace(n_routed_experts=16), 16),
-        (SimpleNamespace(num_routed_experts=16), 16),
-        (SimpleNamespace(text_config=SimpleNamespace(n_routed_experts=32)), 32),
-        (SimpleNamespace(llm_config=SimpleNamespace(n_routed_experts=32)), 32),
-    ],
-    ids=[
-        "dense",
-        "num_experts",
-        "num_local_experts",
-        "n_routed_experts",
-        "num_routed_experts",
-        "vlm-text_config",
-        "vlm-llm_config",
-    ],
-)
-def test_the_expert_count_is_read_under_every_spelling(config, experts):
-    """Architectures disagree on the field name, and a VLM's MoE is its nested LM, whose root
-    config has no expert count of its own: missing either would silently disable EP."""
-    assert _num_experts_from_config(config) == experts
-
-
-# ------------------------------------------------------------------------------- degree selection
-
-
-@pytest.mark.parametrize(
-    ("world", "experts", "degree"),
-    [(8, 64, 8), (8, 6, 2), (8, 5, 1), (8, 0, 1), (1, 64, 1)],
-    ids=["largest-valid", "divides-the-experts", "nothing-divides", "dense", "single-rank"],
-)
-def test_the_largest_ep_degree_dividing_the_world_and_the_experts_is_chosen(world, experts, degree):
-    """EP must divide the expert count, so every rank holds whole experts; with nothing dividing, a
-    dense model, or a single rank, there is no expert parallelism."""
-    assert _choose_ep_degree(world, experts) == degree
-
-
-# ---------------------------------------------------------------------------- routed experts
-
-
-class _MoeLayer(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.self_attn = nn.Linear(2, 2)
-        self.mlp = nn.Module()
-        self.mlp.gate = nn.Linear(2, 4, bias=False)  # the router
-        self.mlp.experts = nn.Module()
-        self.mlp.experts.gate_up_proj = nn.Parameter(torch.zeros(4, 4, 2))
-        self.mlp.experts.down_proj = nn.Parameter(torch.zeros(4, 2, 2))
-        self.mlp.shared_experts = nn.Linear(2, 4, bias=False)
-
-
-def _model_with_decoder_at(prefix):
-    """A two-layer MoE whose decoder (``layers.<i>``) sits under ``prefix`` (``""``: the decoder)."""
-    decoder = nn.Module()
-    decoder.embed_tokens = nn.Embedding(8, 2)
-    decoder.layers = nn.ModuleList(_MoeLayer() for _ in range(2))
-    if not prefix:
-        return decoder
-    root = parent = nn.Module()
-    *path, leaf = prefix.rstrip(".").split(".")
-    for name in path:
-        parent.add_module(name, nn.Module())
-        parent = getattr(parent, name)
-    parent.add_module(leaf, decoder)
-    return root
-
-
-_STOCK_PLAN = {
-    "layers.*.mlp.gate": "ep_router",
-    "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
-    "layers.*.mlp.experts.down_proj": "grouped_gemm",
-    "layers.*.mlp.experts": "moe_tp_experts",
-}
-
-
-@pytest.mark.parametrize(
-    "plan",
-    [
-        _STOCK_PLAN,
-        {"layers.*.mlp.experts": "moe_tp_experts"},
-        {"layers.*.mlp.experts": "nemotron_experts_ep"},
-        {k: v for k, v in _STOCK_PLAN.items() if v == "grouped_gemm"},
-        {
-            "layers.*.block_sparse_moe.experts": "moe_tp_experts",
-            "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
-        },
-    ],
-    ids=["stock", "experts-style-only", "nemotron-style-only", "grouped-gemm-only", "hybrid"],
-)
-@pytest.mark.parametrize(
-    "prefix",
-    ["", "model.", "model.language_model.", "language_model.model."],
-    ids=["base-model", "causal-lm", "vlm", "legacy-vlm"],
-)
-def test_routed_experts_are_the_plans_containers_under_any_prefix(plan, prefix):
-    """The base-relative plan keys match the decoder wherever the model nests it, at every layer
-    index, the container named by its style or as its grouped-GEMM parameters' parent, also in a
-    plan that names some containers one way and some the other. The router, the shared experts, and
-    a module whose name merely ends in a key's text are not routed experts."""
-    model = _model_with_decoder_at(prefix)
-    model.add_module("sublayers", nn.ModuleList([_MoeLayer()]))  # "sublayers.0.mlp.experts"
-
-    found = _routed_expert_modules(model, plan)
-
-    modules = dict(model.named_modules())
-    assert found == {
-        name: modules[name] for name in [f"{prefix}layers.{i}.mlp.experts" for i in (0, 1)]
-    }
-
-
-# ------------------------------------------------------------------------------ model class
-
-
-def _vlm_config(**kwargs):
-    return SimpleNamespace(text_config=SimpleNamespace(), architectures=[], **kwargs)
-
-
-def _tiny_qwen2_5_vl():
-    return Qwen2_5_VLConfig(
-        text_config={**_TINY, "max_position_embeddings": 16},
-        vision_config={
-            "depth": 1,
-            "hidden_size": 16,
-            "intermediate_size": 16,
-            "num_heads": 2,
-            "out_hidden_size": 16,
-            "patch_size": 2,
-            "spatial_merge_size": 1,
-            "temporal_patch_size": 1,
-            "fullatt_block_indexes": [0],
-        },
-    )
-
-
-_REMOTE_VLM = {"AutoModelForCausalLM": "m.KimiK25ForConditionalGeneration"}
-_REMOTE_IMAGE_TEXT_VLM = {
-    **_REMOTE_VLM,
-    "AutoModelForImageTextToText": "m.KimiK25ForConditionalGeneration",
-}
-
-
-@pytest.mark.parametrize(
-    ("config", "trust_remote_code", "expected"),
-    [
-        (SimpleNamespace(architectures=["LlamaForCausalLM"]), False, AutoModelForCausalLM),
-        (_vlm_config(auto_map=_REMOTE_VLM), True, AutoModelForCausalLM),
-        (_vlm_config(auto_map=_REMOTE_VLM), False, AutoModel),
-        (_vlm_config(auto_map=_REMOTE_IMAGE_TEXT_VLM), True, AutoModelForImageTextToText),
-        (_tiny_qwen2_5_vl(), False, AutoModelForImageTextToText),
-        (Phi4MultimodalConfig(), False, AutoModelForCausalLM),
-        (_vlm_config(), False, AutoModel),
-        (T5Config(architectures=["T5ForConditionalGeneration"]), False, AutoModelForSeq2SeqLM),
-        (WhisperConfig(), False, AutoModelForSpeechSeq2Seq),
-    ],
-    ids=[
-        "causal-lm",
-        "remote-vlm-with-head",
-        "remote-vlm-untrusted",
-        "remote-image-text-vlm",
-        "native-vlm-with-head",
-        "native-multimodal-causal-lm",
-        "unknown-vlm",
-        "encoder-decoder",
-        "speech-encoder-decoder",
-    ],
-)
-def test_the_class_that_builds_the_whole_model(config, trust_remote_code, expected):
-    """A VLM's AutoModel leaves the LM head out: it builds through the class that carries the head,
-    keeping the headless backbone only when no class recognizes the config. A remote class counts
-    only when the checkpoint's code is trusted. An encoder-decoder model's architecture reads like a
-    VLM's, but its seq2seq class carries its head, also for a config that names no architecture."""
-    assert _auto_class(config, trust_remote_code) is expected
-
-
-# -------------------------------------------------------------------------------- the decision
-
-
-@pytest.mark.parametrize(
-    "knob",
-    [
-        None,
-        ("MODELOPT_EP_DEGREE", "4"),
-        ("MODELOPT_FSDP_SIZE", "world"),
-        ("MODELOPT_FSDP_SIZE", "2"),
-    ],
-    ids=["no_knob", "ep_degree", "fsdp_world", "fsdp_part"],
-)
-@pytest.mark.parametrize("world_size", [None, 1], ids=["no_process_group", "one_rank"])
-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."""
-    if world_size is not None:
-        run_as(world_size)
-    if knob is not None:
-        monkeypatch.setenv(*knob)
-    served = serve(_qwen3_moe())
-    monkeypatch.setattr(
-        model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model")
-    )
-
-    assert resolve_distribution_plan(str(tmp_path)) == (DistributionPlan(), served)
-    assert resolve_distribution_plan(str(tmp_path), dense="fsdp")[0] == DistributionPlan(
-        dense="fsdp"
-    )
-
-
-@needs_qwen3_moe_ep_plan
-@pytest.mark.parametrize("rank", [0, 1])
-def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, capsys, rank):
-    """Auto picks the largest EP degree dividing both the world and the expert count, hands back
-    the config it decided from, and only rank 0 reports the decision."""
-    run_as(4, rank)
-    served = serve(_qwen3_moe(num_experts=8))
-
-    plan, hf_config = resolve_distribution_plan(str(tmp_path))
-
-    assert plan == DistributionPlan(4, ep_size=4, dense="replicate")
-    assert hf_config is served
-    assert ("Distribution:" in capsys.readouterr().out) == (rank == 0)
-
-
-@needs_qwen3_moe_ep_plan
-def test_no_ep_degree_dividing_the_experts_leaves_them_whole(run_as, serve, tmp_path):
-    run_as(2)
-    serve(_qwen3_moe(num_experts=3))
-
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2)
-
-
-def test_an_moe_without_an_ep_plan_is_not_expert_parallel(run_as, serve, tmp_path):
-    """The config counts 8 experts, but with no base_model_ep_plan nothing could shard them."""
-    run_as(4)
-    config = _qwen3_moe(num_experts=8)
-    config.base_model_ep_plan = None
-    serve(config)
-
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
-
-
-@pytest.mark.parametrize("degree", [None, "1"], ids=["auto", "set-to-1"])
-def test_an_ep_plan_naming_no_module_of_the_model_is_refused(
-    run_as, serve, monkeypatch, tmp_path, degree
-):
-    """A plan whose expert containers match no module of the built model (a renamed submodule, a
-    hand-written plan) leaves no expert to count or shard: the automatic degree would replicate
-    every expert on every rank without a word. A degree set to 1 asks for exactly that."""
-    run_as(4)
-    config = _qwen3_moe(num_experts=8)
-    config.base_model_ep_plan = {"layers.*.mlp.renamed_experts": "moe_tp_experts"}
-    serve(config)
-
-    if degree is None:
-        with pytest.raises(ValueError, match=r"renamed_experts\) names .* MODELOPT_EP_DEGREE=1"):
-            resolve_distribution_plan(str(tmp_path))
-    else:
-        monkeypatch.setenv("MODELOPT_EP_DEGREE", degree)
-        assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
-
-
-@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
-):
-    """Every rank belongs to one EP group, so the degree must divide the world; the refusal names
-    the knob or argument the degree came from, as the one for the experts does."""
-    run_as(4)
-    serve(LlamaConfig(**_TINY))
-    arguments = {}
-    if source == "ep_size":
-        arguments["ep_size"] = 3
-    else:
-        monkeypatch.setenv(source, "3")
-
-    with pytest.raises(ValueError, match=rf"{source}=3 must divide the world size \(4\)"):
-        resolve_distribution_plan(str(tmp_path), **arguments)
-
-
-@pytest.mark.parametrize(
-    ("world_size", "source"),
-    [(4, "MODELOPT_EP_DEGREE"), (1, "ep_size")],
-    ids=["knob", "argument-in-a-single-process"],
-)
-def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from(
-    run_as, serve, monkeypatch, tmp_path, world_size, source
-):
-    """A single process takes its own argument as given, but parses it as the decision does."""
-    run_as(world_size)
-    serve(LlamaConfig(**_TINY))
-    arguments = {}
-    if source == "ep_size":
-        arguments["ep_size"] = "world"
-    else:
-        monkeypatch.setenv(source, "world")
-
-    with pytest.raises(ValueError, match=f"{source}='world'"):
-        resolve_distribution_plan(str(tmp_path), **arguments)
-
-
-@needs_qwen3_moe_ep_plan
-@pytest.mark.parametrize("value", ["0", "", "AUTO"], ids=["zero", "exported-empty", "capitals"])
-def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value):
-    """And so is one exported empty, and auto in any case: the plan decides."""
-    monkeypatch.setenv("MODELOPT_EP_DEGREE", value)
-    run_as(4)
-    serve(_qwen3_moe(num_experts=8))
-
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)
-
-
-@pytest.mark.parametrize("tied", [False, True], ids=["untied", "tied-embeddings"])
-@pytest.mark.parametrize(
-    ("share", "dense"), [(0.59, "replicate"), (0.61, "fsdp")], ids=["fits", "too-big"]
-)
-def test_a_dense_model_is_fsdp_sharded_only_when_too_big_for_one_gpu(
-    run_as, serve, gpu_memory, tmp_path, share, dense, tied
-):
-    """Taking more than 60% of a GPU, the whole model would leave too little to calibration. It is
-    sized as it loads, a tied embedding and LM head counting once, though its meta build unties
-    them."""
-    run_as(4)
-    config = serve(LlamaConfig(**_TINY, tie_word_embeddings=tied))
-    loaded = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16)
-    gpu_memory(math.ceil(_bytes(loaded) / share))
-
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, dense=dense)
-
-
-@needs_qwen3_moe_ep_plan
-@pytest.mark.parametrize(
-    ("experts_fitting", "dense"),
-    [(1 / 2, "replicate"), (1 / 8, "fsdp")],
-    ids=["fits", "the-experts-share-does-not"],
-)
-def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit(
-    run_as, serve, gpu_memory, tmp_path, experts_fitting, dense
-):
-    """At EP 4 a rank holds the rest of the model whole and a quarter of the routed experts, so the
-    dense part is sharded when that does not fit one GPU, however small the rest; with no EP the
-    whole MoE must fit, and here it does not."""
-    run_as(4)
-    config = serve(_qwen3_moe(num_experts=8))
-    model = _meta(config)
-    experts = sum(_bytes(layer.mlp.experts) for layer in model.model.layers)
-    # A GPU 60% of which holds the rest of the model and that share of the experts.
-    gpu_memory(int((_bytes(model) - experts + experts_fitting * experts) / 0.6))
-
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4, dense)
-    assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan(
-        4, dense="fsdp"
-    )
-
-
-@needs_qwen3_moe_ep_plan
-def test_the_env_knobs_override_the_automatic_plan(run_as, serve, monkeypatch, tmp_path):
-    monkeypatch.setenv("MODELOPT_EP_DEGREE", "2")
-    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
-    run_as(4)
-    serve(_qwen3_moe(num_experts=8))
-
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 2, "fsdp")
-
-
-@needs_qwen3_moe_ep_plan
-def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path):
-    monkeypatch.setenv("MODELOPT_EP_DEGREE", "4")
-    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
-    run_as(4)
-    serve(_qwen3_moe(num_experts=8))
-
-    plan, _ = resolve_distribution_plan(str(tmp_path), ep_size=2, dense="replicate")
-
-    assert plan == DistributionPlan(4, 2, "replicate")
-
-
-@pytest.mark.parametrize(
-    ("value", "dense"),
-    [
-        ("1", "replicate"),
-        ("4", "fsdp"),
-        ("WORLD", "fsdp"),
-        ("", "fsdp"),
-        ("2", None),
-    ],
-    ids=["1", "4", "WORLD", "exported-empty", "part-of-the-world"],
-)
-def test_the_dense_knob_replicates_or_shards_over_the_whole_world(
-    run_as, serve, gpu_memory, monkeypatch, tmp_path, value, dense
-):
-    """Exported empty, the knob is auto, which shards a model too big for one GPU. Sharding over
-    part of the world is refused: no layout does that."""
-    monkeypatch.setenv("MODELOPT_FSDP_SIZE", value)
-    run_as(4)
-    serve(LlamaConfig(**_TINY))
-    gpu_memory(1)  # too small for the whole model, so only auto tells "" from "1"
-
-    if dense is None:
-        with pytest.raises(ValueError, match=r"expected auto, 1 .* or world"):
-            resolve_distribution_plan(str(tmp_path))
-    else:
-        assert resolve_distribution_plan(str(tmp_path))[0].dense == dense
-
-
-@pytest.mark.parametrize(
-    ("config", "ep_size", "experts"),
-    [(LlamaConfig(**_TINY), 2, 0), (_qwen3_moe(num_experts=6), 4, 6)],
-    ids=["dense-model", "experts-not-divisible"],
-)
-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")
-    run_as(4)
-    serve(config)
-
-    with pytest.raises(ValueError, match=rf"ep_size={ep_size} must divide .* \({experts}\)"):
-        resolve_distribution_plan(str(tmp_path), ep_size=ep_size)
-
-
-@pytest.mark.parametrize(
-    "config",
-    [
-        _tiny_qwen2_5_vl(),
-        T5Config(d_model=16, d_ff=32, d_kv=8, num_layers=1, num_heads=2, vocab_size=32),
-    ],
-    ids=["vlm", "encoder-decoder-naming-no-architecture"],
-)
-def test_a_model_is_decided_from_its_whole_model(run_as, serve, gpu_memory, tmp_path, config):
-    """AutoModelForCausalLM builds neither every VLM nor an encoder-decoder model, whose config need
-    not name its architecture: the decision builds what the load builds."""
-    run_as(2)
-    serve(config)
-    gpu_memory(1)
-
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp")
-
-
-def test_the_expert_count_is_read_from_the_config_holding_the_plan(
-    run_as, serve, monkeypatch, tmp_path
-):
-    """A VLM whose nested config holds the MoE and its plan counts the experts there, not by an
-    expert-ish field on its root config."""
-    text = SimpleNamespace(
-        base_model_ep_plan={"layers.*.mlp.experts": "moe_tp_experts"}, num_experts=8
-    )
-    serve(SimpleNamespace(model_type="vlm", architectures=[], text_config=text, num_experts=3))
-    built = _model_with_decoder_at("model.")
-    built.tie_weights = lambda: None  # a PreTrainedModel's; nothing here is tied
-    monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: built)
-    run_as(4)
-
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4)
-
-
-def test_a_model_that_cannot_be_built_on_meta_names_itself(run_as, serve, monkeypatch, tmp_path):
-    def unbuildable(*args, **kwargs):
-        raise NotImplementedError("custom init touches the weights")
-
-    run_as(2)
-    serve(LlamaConfig(**_TINY))
-    monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", unbuildable)
-
-    with pytest.raises(RuntimeError, match=r"distribute llama .* custom init touches the weights"):
-        resolve_distribution_plan(str(tmp_path))
-
-
-def test_a_meta_build_leaves_the_auto_class_a_part_registers_in_as_it_found_it(monkeypatch):
-    """A remote VLM's ``__init__`` can build its text model through ``AutoModelForCausalLM`` while
-    the whole model builds through another auto class. That registers the text model's class in
-    ``AutoModelForCausalLM`` too, which would leave rank 0's later load of it without the
-    ``generate`` its peers' have."""
-
-    class ProbeConfig(transformers.PretrainedConfig):
-        model_type = "meta_build_part_probe"
-
-    class ProbeModel(transformers.PreTrainedModel):
-        config_class = ProbeConfig
-
-    for auto in (AutoModel, AutoModelForCausalLM):
-        mapping = auto._model_mapping
-        monkeypatch.setattr(mapping, "_extra_content", dict(mapping._extra_content))
-    from_config = AutoModel.from_config
-
-    def remote_vlm_from_config(config, **kwargs):
-        # What a remote VLM's __init__ does for its text model before the whole model is built.
-        AutoModelForCausalLM.register(ProbeConfig, ProbeModel, exist_ok=True)
-        return from_config(config, **kwargs)
-
-    monkeypatch.setattr(AutoModel, "from_config", remote_vlm_from_config)
-
-    model_load_utils.build_meta_causal_lm(
-        None, False, None, LlamaConfig(**_TINY), auto_class=AutoModel
-    )
-
-    assert ProbeConfig not in AutoModelForCausalLM._model_mapping
-
-
-def test_a_meta_build_leaves_the_auto_class_as_it_found_it(monkeypatch):
-    """Building a remote-code model through an auto class registers its class there, as it is, and a
-    later load through that auto class builds the registered class, without the ``generate`` a first
-    load adds. Rank 0 alone builds one to decide a distribution: its load alone would lack it."""
-
-    class ProbeConfig(transformers.PretrainedConfig):
-        model_type = "meta_build_probe"
-
-    class ProbeModel(transformers.PreTrainedModel):
-        config_class = ProbeConfig
-
-    mapping = AutoModelForCausalLM._model_mapping
-    monkeypatch.setattr(mapping, "_extra_content", dict(mapping._extra_content))
-    from_config = AutoModelForCausalLM.from_config
-
-    def remote_from_config(config, **kwargs):
-        # What transformers' from_config does for a remote-code config before it builds.
-        AutoModelForCausalLM.register(ProbeConfig, ProbeModel, exist_ok=True)
-        return from_config(config, **kwargs)
-
-    monkeypatch.setattr(AutoModelForCausalLM, "from_config", remote_from_config)
-
-    model_load_utils.build_meta_causal_lm(None, False, None, LlamaConfig(**_TINY))
-
-    assert ProbeConfig not in mapping
-
-
-def test_the_hub_options_locate_the_checkpoint_the_load_reads(tmp_path):
-    """``revision``, ``cache_dir``, ``local_files_only`` and ``subfolder`` find the checkpoint as
-    ``load_pretrained`` takes them, so the plan and its config are those of what the load reads."""
-    repo = tmp_path / "models--org--model"
-    (repo / "refs").mkdir(parents=True)
-    for revision, commit, hidden_size in (("main", "a" * 40, 16), ("v2", "b" * 40, 32)):
-        (repo / "refs" / revision).write_text(commit)
-        config = LlamaConfig(**{**_TINY, "hidden_size": hidden_size})
-        config.save_pretrained(repo / "snapshots" / commit / "nested")
-
-    _, config = resolve_distribution_plan(
-        "org/model",
-        revision="v2",
-        cache_dir=str(tmp_path),
-        local_files_only=True,
-        subfolder="nested",
-    )
-
-    assert config.hidden_size == 32
-
-
-# ------------------------------------------------------------------------------ rank 0 decides
-
-
-@contextmanager
-def _knobs(**values):
-    """This rank's distribution knobs, exactly ``values``, within the block.
-
-    A pooled worker keeps the environment it was spawned with, so the knobs are cleared first. The
-    environment is restored on the way out.
-    """
-    with mock.patch.dict(os.environ):
-        for name in ("MODELOPT_EP_DEGREE", "MODELOPT_FSDP_SIZE"):
-            os.environ.pop(name, None)
-        os.environ.update(values)
-        yield
-
-
-def _every_rank_takes_rank0s_plan_job(rank, size, ckpt, knobs, expected):
-    with bounded_collectives(), _knobs(**knobs[rank]):
-        plan, _ = resolve_distribution_plan(ckpt)
-    assert plan == expected, (rank, plan)
-
-
-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)
-    cpu_dist_workers.run(
-        _every_rank_takes_rank0s_plan_job,
-        ckpt=str(tmp_path),
-        knobs=[{}, {"MODELOPT_FSDP_SIZE": "world"}],
-        expected=DistributionPlan(2),
-    )
-
-
-def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt):
-    isdir = os.path.isdir
-    with (
-        bounded_collectives(),
-        mock.patch("os.path.isdir", lambda path: (rank == 0 or path != ckpt) and isdir(path)),
-    ):
-        try:
-            resolve_distribution_plan(ckpt)
-        except RuntimeError as error:
-            raised = str(error)
-        else:
-            raised = "no error"
-    # Rank 1's own error, carried to rank 0, rather than rank 0 timing out in the plan's broadcast.
-    assert "which this rank cannot see" in raised, raised
-
-
-def test_a_rank_that_cannot_see_the_checkpoint_fails_every_rank(cpu_dist_workers, tmp_path):
-    """Before the plan's broadcast, which rank 0 would otherwise wait in for a rank that raised."""
-    LlamaConfig(**_TINY).save_pretrained(tmp_path)
-    cpu_dist_workers.run(_one_rank_cannot_see_the_checkpoint_job, ckpt=str(tmp_path))

From 09329ac9322b3f6158049f8bb0c5ee66df88fa86 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 01:35:35 +0000
Subject: [PATCH 43/44] Drop build_meta_causal_lm's local
 AutoModelForImageTextToText import

The module imports it at its top since this PR, and the repo's transformers floor
(5.5) has it, so the local import and its comment were dead.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 32009a9cb47..95f4fa2b1a0 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -240,9 +240,6 @@ def build_meta_causal_lm(
     if model_factory is None:
         model_factory = AutoModelForCausalLM
         if is_fsdp2_text_only_vlm(hf_config):
-            # Only newer Transformers versions supporting these models expose this factory.
-            from transformers import AutoModelForImageTextToText
-
             model_factory = AutoModelForImageTextToText
     # Where transformers records the class a remote-code build registers: in the auto class built
     # through, and in those a remote VLM builds its towers through. Read without a fallback: if a

From 2a6290b470a661fdd24802d443a506247029e2f0 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 01:35:44 +0000
Subject: [PATCH 44/44] Refuse routed experts whose config gives no expert
 count, unless EP is 1

When the plan's expert containers are in the built model but the config spells
its expert count under none of the fields the decision reads, the automatic
degree came out 1 and every expert stayed on every rank without a word. It is
now refused, naming the fields and the knob or argument to set to 1, as the
refusal of a plan naming no module of the model is.

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

diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 95f4fa2b1a0..babf7f11054 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -1051,6 +1051,15 @@ def _decide_distribution(
             f"can be neither counted nor sharded. Set {knob}=1 to load it without expert "
             "parallelism."
         )
+    if experts and not num_experts and ep_size != 1:
+        # The mirror case: the containers are there, but no count to choose a degree by, which
+        # would leave every expert replicated on every rank.
+        raise ValueError(
+            f"{hf_config.model_type}'s base_model_ep_plan names routed experts the built model has, "
+            "but its config gives no expert count (num_experts, num_local_experts, "
+            f"n_routed_experts or num_routed_experts) to shard them by. Set {knob}=1 to load it "
+            "without expert parallelism."
+        )
     if ep_size == "auto":
         ep_size = _choose_ep_degree(world_size, num_experts)
     elif ep_size > 1 and world_size % ep_size:
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 16fc0e79e92..2c225dafc14 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -425,6 +425,28 @@ 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"])
+def test_routed_experts_the_config_does_not_count_are_refused(
+    run_as, serve, monkeypatch, tmp_path, degree
+):
+    """The mirror case: the plan's containers are in the built model, but the config counts its
+    experts under no field the decision reads, which would leave every expert replicated on every
+    rank without a word. A degree set to 1 asks for exactly that."""
+    plan = {"layers.*.mlp.experts": "moe_tp_experts"}
+    serve(SimpleNamespace(model_type="moe", architectures=[], base_model_ep_plan=plan))
+    built = model_with_decoder_at("model.")
+    built.tie_weights = lambda: None  # a PreTrainedModel's; nothing here is tied
+    monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: built)
+    run_as(4)
+
+    if degree is None:
+        with pytest.raises(ValueError, match=r"no expert count .* MODELOPT_EP_DEGREE=1"):
+            resolve_distribution_plan(str(tmp_path))
+    else:
+        monkeypatch.setenv("MODELOPT_EP_DEGREE", degree)
+        assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4)
+
+
 @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