[OMNIML-6060] Load an expert-parallel plan through load_pretrained

CHANGELOG.rst

@@ -33,7 +33,7 @@ Changelog
 *Misc*
 
 - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs.
-- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it FSDP2-shards the model or gives every rank a whole copy, by the plan the new ``resolve_distribution_plan`` decides from the model and the ``MODELOPT_FSDP_SIZE`` variable, or by a plan passed as its new ``distribution`` argument. The loaded model records the plan (``modelopt.torch.utils.distributed.distribution_plan``). Until the expert-parallel load lands, an MoE planned expert-parallel, as the default ``MODELOPT_EP_DEGREE=auto`` plans one whose config has a ``base_model_ep_plan``, is refused before it loads: set ``MODELOPT_EP_DEGREE=1`` (with ``MODELOPT_FSDP_SIZE=world`` to FSDP2-shard it) or pass a plan with ``ep_size=1``.
+- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it shards an MoE's routed experts across the ranks, FSDP2-shards the model, or gives every rank a whole copy, by the plan the new ``resolve_distribution_plan`` decides from the model and the ``MODELOPT_EP_DEGREE`` / ``MODELOPT_FSDP_SIZE`` variables, or by a plan passed as its new ``distribution`` argument. The loaded model records the plan (``modelopt.torch.utils.distributed.distribution_plan``).
 
 **Backward Breaking Changes**
 

modelopt/torch/utils/plugins/model_load_utils.py

@@ -20,6 +20,7 @@
 import inspect
 import json
 import logging
+import math
 import os
 import re
 from collections import defaultdict
@@ -413,10 +414,20 @@ def record_unplaced_source_keys(
     return keys
 
 
+_EXPERT_PARALLEL = "expert-parallel"
 _FSDP2 = "FSDP2"
 # What each sharded load takes besides the checkpoint. These loads build the model class and place
 # it themselves, so from_pretrained's own arguments (device_map, max_memory, ...) do not reach them.
 _SHARDED_LOAD_OPTIONS = {
+    _EXPERT_PARALLEL: frozenset(
+        {
+            "trust_remote_code",
+            "attn_implementation",
+            "hf_config",
+            "freeze",
+            "require_lm_head",
+        }
+    ),
     _FSDP2: frozenset(
         {
             "trust_remote_code",
@@ -439,15 +450,11 @@ def _route(plan: DistributionPlan) -> str | None:
             f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}."
         )
     if plan.is_expert_parallel:
-        raise NotImplementedError(
-            f"{plan!r} shards the routed experts, which load_pretrained has no load for. Plan "
-            "ep_size=1 (MODELOPT_EP_DEGREE=1) to give every rank the whole model, or shard it "
-            "with FSDP2."
-        )
+        return _EXPERT_PARALLEL
     return _FSDP2 if plan.is_fsdp_dense else None
 
 
-def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) -> nn.Module:
+def _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_options) -> nn.Module:
     """The sharded load ``route`` names, refusing what it cannot honor before it starts."""
     from modelopt.torch.models.hf import checkpoint_has_mtp
     from modelopt.torch.utils import distributed as dist
@@ -487,6 +494,15 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options)
         )
     trust_remote_code = kwargs.pop("trust_remote_code", False)
     hf_config = kwargs.pop("hf_config", None)
+    if route == _EXPERT_PARALLEL:
+        return load_and_prepare_ep(
+            ckpt_path,
+            plan,
+            trust_remote_code=trust_remote_code,
+            hf_config=hf_config,
+            hub_options=hub_options,
+            **kwargs,
+        )
     if hf_config is None:
         hf_config = AutoConfig.from_pretrained(
             ckpt_path, trust_remote_code=trust_remote_code, **hub_options
@@ -539,20 +555,22 @@ def load_pretrained(
     the model over the run's ranks by. Every rank takes rank 0's, or, when rank 0 passes none,
     :func:`resolve_distribution_plan`'s. By the plan:
 
-    * ``dense="fsdp"``: a causal LM FSDP2-sharded over the whole process group, each rank reading its
-      own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``, ``mp_policy`` and
-      ``broadcast_chunk_size`` ``kwargs`` may set). A single process launched with
+    * ``ep_size > 1``: the routed experts sharded as they load, the dense part replicated or
+      FSDP2-sharded as the plan says (:func:`load_and_prepare_ep`, whose
+      ``freeze`` and ``require_lm_head`` ``kwargs`` may set);
+    * otherwise ``dense="fsdp"``: a causal LM FSDP2-sharded over the whole process group, each rank
+      reading its own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``,
+      ``mp_policy`` and ``broadcast_chunk_size`` ``kwargs`` may set). A single process launched with
       ``torchrun --nproc_per_node=1`` has no group to shard through, so the load sets up one of one
       rank, which :func:`~modelopt.torch.utils.distributed.distributed_run` destroys with the run; a
       process started without a launcher cannot set one up and is refused. It reads the
-      checkpoint's tensors as stored, so it refuses a quantized checkpoint.
+      checkpoint's tensors as stored, so it refuses a quantized checkpoint;
     * otherwise ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to
       ``AutoModelForCausalLM``), placed as ``kwargs`` say. Every rank of a distributed run gets the
       whole model, ``device_map`` defaulting to the rank's GPU.
 
-    A plan with ``ep_size > 1`` raises: no load here shards the routed experts. On a distributed run
-    every rank loads the plan rank 0 was handed. A rank of a multi-rank launch whose process group
-    is not set up is refused on every route: set it up first, with
+    On a distributed run every rank loads the plan rank 0 was handed. A rank of a multi-rank launch
+    whose process group is not set up is refused on every route: set it up first, with
     :func:`~modelopt.torch.utils.distributed.distributed_run`. The plan is recorded on the model
     (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), where calibration and export
     read it back.
@@ -569,11 +587,12 @@ def load_pretrained(
     * A Hub model ID is made local first (:func:`ensure_local_checkpoint`, which takes the
       ``revision``, ``token``, ``cache_dir`` and ``local_files_only`` ``kwargs``): the export reads
       the source checkpoint from disk. Under ``torch.distributed`` every rank must call this. The
-      config read and the whole-model load still take all but ``revision``: an ``auto_map`` entry
-      can name custom code in another Hub repo, which Transformers fetches with them. A sharded
-      load builds its model from the config with ``from_config``, which takes none, so such code
-      comes with the environment's Hub settings. A ``subfolder`` is part of where the checkpoint
-      is: the config, the loading hook and the record below read it, as the load does.
+      config read, the whole-model load and the expert-parallel load still take all but
+      ``revision``: an ``auto_map`` entry can name custom code in another Hub repo, which
+      Transformers fetches with them. The FSDP2 load builds its model from the config with
+      ``from_config``, which takes none, so such code comes with the environment's Hub settings.
+      A ``subfolder`` is part of where the checkpoint is: the config, the loading hook and the
+      record below read it, as the load does.
     * The model type's loading hook runs
       (:func:`~modelopt.torch.models.hf.prepare_model_for_loading`), e.g. to keep the MTP head a
       Nemotron-H checkpoint carries. ``model_type`` selects it; when omitted, the type of the
@@ -674,7 +693,7 @@ def handed() -> DistributionPlan | None:
             ).model_type
         )
     if route is not None:
-        model = _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options)
+        model = _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_options)
         # A sharded load finds the unplaced keys while it places the weights.
         keys = model._modelopt_unplaced_source_keys
     else:
@@ -1320,6 +1339,9 @@ def _decide_distribution(
         )
     ep_size = _ep_degree(ep_size)
     dense = _dense_from_env(world_size) if dense is None else dense
+    # A replication both knobs ask for is the caller's choice, and loads; one reached from "auto"
+    # would leave an expert-parallel MoE whole on every rank unasked, and is refused below.
+    asked_to_replicate = ep_size == 1 and dense == "replicate"
     if (
         auto_class not in (None, AutoModelForCausalLM)
         and ep_size in ("auto", 1)
@@ -1417,6 +1439,39 @@ def _decide_distribution(
         raise NotImplementedError(
             f"The dense part of {hf_config.model_type} cannot be sharded {unshardable}.{remedy}"
         )
+    if experts and asked_to_replicate:
+        print(
+            f"{hf_config.model_type}: ep_size=1 and dense='replicate' as asked, so every rank "
+            "holds the whole MoE.",
+            flush=True,
+        )
+    elif experts and ep_size == 1 and dense == "replicate":
+        # Name a way to shard the experts only where one exists: a count the config spells none of
+        # the ways ModelOpt reads leaves no EP degree to name, as does one sharing no factor above 1
+        # with the world size.
+        if not num_experts:
+            shard = (
+                "Its config names its routed-expert count none of the ways ModelOpt reads, so "
+                "no EP degree can shard its experts. Shard"
+            )
+        elif math.gcd(num_experts, world_size) == 1:
+            shard = (
+                f"No EP degree above 1 divides both its {num_experts} experts and the world size, "
+                f"so none can shard its experts on {world_size} ranks. Shard"
+            )
+        else:
+            shard = (
+                f"Shard its experts (MODELOPT_EP_DEGREE above 1, dividing both its {num_experts} "
+                "experts and the world size), shard"
+            )
+        raise RuntimeError(
+            f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 at "
+            f"world_size={world_size}, so every rank would hold the whole MoE. {shard} the whole "
+            "model with FSDP2 (MODELOPT_FSDP_SIZE=world, for an unquantized checkpoint), load it in "
+            "a single process, or replicate it on every rank outright: MODELOPT_EP_DEGREE=1 with "
+            "MODELOPT_FSDP_SIZE=1, or "
+            f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))."
+        )
     return DistributionPlan(world_size, ep_size, dense)
 
 
@@ -1474,6 +1529,10 @@ def resolve_distribution_plan(
 
     ``config``, the config the load will use when the caller has one, is decided from instead of
     the checkpoint's, which is then not read: each rank takes a copy, as building writes to it.
+
+    A plan that leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for
+    exactly that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and
+    ``MODELOPT_FSDP_SIZE=1``) replicates it, noting that every rank holds the whole MoE.
     """
     from modelopt.torch.utils import distributed as dist
 
@@ -2240,6 +2299,7 @@ def load_and_prepare_ep(
     freeze: bool = True,
     hf_config=None,
     require_lm_head: bool = True,
+    hub_options: dict | None = None,
 ) -> nn.Module:
     """Load + shard a MoE causal LM for expert-parallel PTQ via the DP=EP all-to-all transport.
 
@@ -2248,7 +2308,9 @@ def load_and_prepare_ep(
     does (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), as
     :func:`load_pretrained` does for every load. The model builds with its LM head where an auto
     class builds one; ``require_lm_head=False`` builds a VLM without it, through ``AutoModel`` (a
-    causal LM always keeps its head).
+    causal LM always keeps its head). ``hub_options`` (``token``, ``cache_dir``,
+    ``local_files_only``, as :func:`load_pretrained` takes them) go to its config read and
+    ``from_pretrained``, which fetch the custom code an ``auto_map`` names in another Hub repo.
 
     Builds an ``(edp, ep)`` device mesh and loads experts EP-sharded over the ``ep`` submesh
     (``enable_expert_parallel=True`` -- each rank owns ``num_experts / ep`` experts as grouped-GEMM
@@ -2333,8 +2395,9 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A
     from transformers.distributed import DistributedConfig
 
     world_size, ep_degree, edp = plan.world_size, plan.ep_size, plan.edp
+    hub_options = dict(hub_options or {})
     if hf_config is None:
-        config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code}
+        config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code, **hub_options}
         if attn_implementation is not None:
             config_kwargs["attn_implementation"] = attn_implementation
         hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs)
@@ -2436,6 +2499,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A
         "device_mesh": ep_submesh,
         "distributed_config": distributed_config,
         "trust_remote_code": trust_remote_code,
+        **hub_options,
     }
     if not _dc_has_tp_plan:
         from_pretrained_kwargs["tp_plan"] = "auto"

tests/gpu/torch/utils/test_load_pretrained_distribution.py

@@ -15,24 +15,26 @@
 
 """``load_pretrained`` in a distributed run, on real GPUs.
 
-Left to decide, it FSDP2-shards a dense model when asked to, and otherwise gives every rank the whole
-model on its own GPU. Whichever load runs, every rank ends up with the checkpoint's weights and the
-load's record of what it could not place.
+Left to decide, it shards an expert-parallel MoE's experts, FSDP2-shards a dense model when asked to,
+and otherwise gives every rank the whole model on its own GPU. Whichever load runs, every rank ends
+up with the checkpoint's weights and the load's record of what it could not place.
 """
 
 import os
 
 import pytest
 import torch
+from packaging.version import Version
 
 transformers = pytest.importorskip("transformers")
 
 from torch.distributed.tensor import DTensor
-from transformers import AutoModelForCausalLM, LlamaConfig
+from transformers import AutoModelForCausalLM, LlamaConfig, Qwen3MoeConfig
 
 from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan
 from modelopt.torch.utils.plugins.model_load_utils import load_pretrained
 
+NUM_EXPERTS = 8
 _TINY = {
     "hidden_size": 64,
     "intermediate_size": 128,
@@ -68,14 +70,31 @@ def _load_and_check(rank, size, ckpt_dir, knobs, expected):
         if expected == "replicated":
             assert not sharded
             assert all(param.device == device for param in params.values())
-        else:
+        elif expected == "fsdp2":
             assert {name for name in params if ".layers." in name} <= sharded
+        else:
+            # Each rank holds its contiguous share of the routed experts, and the rest whole.
+            experts = {name for name in params if ".mlp.experts." in name}
+            assert experts and not sharded - experts
+            per_rank = NUM_EXPERTS // size
+            for name in experts:
+                local = params.pop(name)
+                local = local.to_local() if isinstance(local, DTensor) else local
+                assert local.shape[0] == per_rank
+                owned = whole_params[name][rank * per_rank : (rank + 1) * per_rank]
+                torch.testing.assert_close(local.cpu(), owned)
         for name, param in params.items():  # the same order on every rank: full_tensor gathers
             torch.testing.assert_close(_full(param).cpu(), whole_params[name])
         assert model._modelopt_unplaced_source_keys == []
         assert model._modelopt_source_checkpoint == ckpt_dir
-        dense = "replicate" if expected == "replicated" else "fsdp"
-        assert distribution_plan(model) == DistributionPlan(size, dense=dense)
+        assert (
+            distribution_plan(model)
+            == {
+                "replicated": DistributionPlan(size),
+                "fsdp2": DistributionPlan(size, dense="fsdp"),
+                "expert-parallel": DistributionPlan(size, ep_size=size),
+            }[expected]
+        )
     finally:
         for knob in knobs:
             del os.environ[knob]
@@ -89,3 +108,16 @@ def test_a_dense_model_loads_whole_on_every_rank(need_2_gpus, dist_workers, tmp_
 def test_a_dense_model_asked_to_shard_loads_fsdp2_sharded(need_2_gpus, dist_workers, tmp_path):
     _save(LlamaConfig(**_TINY), tmp_path)
     dist_workers.run(_load_and_check, str(tmp_path), {"MODELOPT_FSDP_SIZE": "world"}, "fsdp2")
+
+
+@pytest.mark.skipif(
+    Version(transformers.__version__) < Version("5.7"),
+    reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7",
+)
+def test_an_expert_parallel_moe_loads_with_its_experts_sharded(need_2_gpus, dist_workers, tmp_path):
+    if NUM_EXPERTS % dist_workers.world_size:
+        pytest.skip(f"{NUM_EXPERTS} experts do not shard over {dist_workers.world_size} GPUs")
+    config = Qwen3MoeConfig(moe_intermediate_size=32, num_experts=NUM_EXPERTS, **_TINY)
+    config.num_experts_per_tok = 2
+    _save(config, tmp_path)
+    dist_workers.run(_load_and_check, str(tmp_path), {}, "expert-parallel")

tests/unit/torch/utils/test_ep_loader.py

@@ -518,7 +518,8 @@ def next_check(experts):
 
 def test_the_weights_load_in_the_checkpoints_dtype_through_its_class(ep_world_of_two, monkeypatch):
     """The checkpoint's ``dtype`` reaches ``from_pretrained``, and the class that builds it is chosen
-    with the trust the caller gave its code."""
+    with the trust the caller gave its code. The Hub options reach the config read and
+    ``from_pretrained``, which fetch the custom code an ``auto_map`` names in another repo."""
 
     class _LoadedError(Exception):
         """Stops the load at from_pretrained."""
@@ -540,12 +541,21 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False):
     )
     monkeypatch.setattr(model_load_utils, "_auto_class", auto_class)
     config = SimpleNamespace(dtype=torch.float16, base_model_ep_plan=_EP_PLAN)
+    reads = []
+    monkeypatch.setattr(
+        "transformers.AutoConfig.from_pretrained", lambda *a, **k: reads.append(k) or config
+    )
 
     with pytest.raises(_LoadedError):
         load_and_prepare_ep(
-            "org/model", DistributionPlan(2, ep_size=2), trust_remote_code=True, hf_config=config
+            "org/model",
+            DistributionPlan(2, ep_size=2),
+            trust_remote_code=True,
+            hub_options={"local_files_only": True},
         )
     assert loaded["dtype"] is torch.float16
+    assert [read["local_files_only"] for read in reads] == [True]
+    assert loaded["local_files_only"] is True
     assert classes == [True]