From 707499f834ee3705825dda2265d6b66025072a98 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:19:38 +0000 Subject: [PATCH 01/32] Load an expert-parallel plan through load_pretrained An expert-parallel DistributionPlan now loads through load_and_prepare_ep, handed the plan itself, the load's options and its Hub options. The expert-parallel load replicates the dense part, so a plan that also FSDP2-shards it is refused before anything loads, naming the ways out. The route's checks run with the FSDP2 load's in _sharded_load_kwargs, inside the collective load checks, and the load itself in _load_sharded. load_pretrained no longer takes an EP degree left to the decision as 1, as loads expert-parallel, as one MODELOPT_EP_DEGREE asks for does. The plan decision refuses to leave an expert-parallel MoE whole on every rank unasked, naming the expert count an EP degree must divide and every way out. A replication both knobs ask for loads, with a note, and the decision's refusals of experts it can neither count nor shard by advise both knobs. One commit for the PR's former history, re-parented onto the expert-parallel loader's core when #2741 was split, and adapted to #2779's collective load checks and its test-file move. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- .../torch/utils/plugins/model_load_utils.py | 152 +++++++++++--- .../gpu/torch/utils/test_model_load_utils.py | 48 ++++- .../unit/torch/utils/test_model_load_utils.py | 187 ++++++++++++++---- 4 files changed, 311 insertions(+), 78 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0b48641edec..1f947089d83 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +43,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 loads a model distributed under ``torchrun``, FSDP2-sharded or whole on every rank, by the plan ``resolve_distribution_plan`` decides or one passed as its new ``distribution`` argument. +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now loads a model distributed under ``torchrun``, with an MoE's routed experts sharded across the ranks, FSDP2-sharded or whole on every rank, by the plan ``resolve_distribution_plan`` decides or one passed as its new ``distribution`` argument. **Backward Breaking Changes** diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 705e21d5a93..705ffccf906 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -19,6 +19,7 @@ import copy import dataclasses import logging +import math import os import re import warnings @@ -451,12 +452,22 @@ 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. _FSDP2_ONLY = frozenset({"cpu_offload", "mp_policy", "broadcast_chunk_size"}) _SHARDED_LOAD_OPTIONS = { - _FSDP2: _FSDP2_ONLY | {"trust_remote_code", "attn_implementation", "hf_config"} + _EXPERT_PARALLEL: frozenset( + { + "trust_remote_code", + "attn_implementation", + "hf_config", + "freeze", + "require_lm_head", + } + ), + _FSDP2: _FSDP2_ONLY | {"trust_remote_code", "attn_implementation", "hf_config"}, } # The FSDP2-only values that ask nothing of a whole load: the FSDP2 load's defaults, and None where # that leaves the option unset (a None broadcast_chunk_size asks for all of a source's layers). @@ -470,11 +481,20 @@ 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." - ) + if plan.is_fsdp_dense: + # Refused before anything loads, where the remedies can be named: the expert-parallel + # load replicates the dense part, so it cannot honor this plan. + raise NotImplementedError( + f"{plan!r} FSDP2-shards the dense part, but the expert-parallel load replicates it " + "on every rank. Sharding the dense part under expert parallelism is not supported " + "yet. A plan left to decide shards it when a whole copy of it and a rank's share " + "of the routed experts would take more than " + f"{_REPLICATED_SHARE:.0%} of a GPU. Replicate it anyway (MODELOPT_FSDP_SIZE=1, or " + "a plan with dense='replicate'), or shard the whole model with FSDP2 " + "(MODELOPT_EP_DEGREE=1 with MODELOPT_FSDP_SIZE=world) if the checkpoint is an " + "unquantized causal LM." + ) + return _EXPERT_PARALLEL return _FSDP2 if plan.is_fsdp_dense else None @@ -560,7 +580,15 @@ def _refuse_sharded_load_options( def _sharded_load_kwargs( - route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided: bool, measured: bool + route, + plan, + ckpt_path, + auto_class, + model_type, + kwargs, + hub_options, + decided: bool, + measured: bool, ) -> dict: """What the sharded load ``route`` names takes, refusing what it cannot honor.""" from modelopt.torch.models.hf import checkpoint_has_mtp @@ -576,6 +604,14 @@ def _sharded_load_kwargs( ) trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) + if route == _EXPERT_PARALLEL: + return { + "plan": 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 @@ -632,8 +668,10 @@ def _has_own_gpu() -> bool: return _own_gpu_index() is not None -def _load_sharded(ckpt_path, load_kwargs) -> nn.Module: - """The FSDP2 load, by the arguments :func:`_sharded_load_kwargs` checked.""" +def _load_sharded(route, ckpt_path, load_kwargs) -> nn.Module: + """The sharded load ``route`` names, by the arguments :func:`_sharded_load_kwargs` checked.""" + if route == _EXPERT_PARALLEL: + return load_and_prepare_ep(ckpt_path, **load_kwargs) if not dist.is_initialized(): # A single process's group of one, which the sharded model needs while it lives. dist.setup() @@ -659,6 +697,10 @@ def load_pretrained( ``distribution`` is the :class:`~modelopt.torch.utils.distributed.DistributionPlan` to load by: every rank takes rank 0's or, when rank 0 passes none, :func:`resolve_distribution_plan`'s. + * ``ep_size > 1``: :func:`load_and_prepare_ep` shards an MoE's routed experts over the process + group as they load and replicates the dense part, so a plan with ``dense="fsdp"`` as well is + refused. It builds and places the model itself, so it takes its own options but no + ``auto_class``, placement or ``dtype``, and refuses a quantized checkpoint. * ``dense="fsdp"``: :func:`parallel_load_and_prepare_fsdp2` shards a causal LM, or a text-only VLM it supports, over the process group (one process needs ``torchrun --nproc_per_node=1``). It builds and places the model itself, so it takes its own options but no ``auto_class``, @@ -667,8 +709,7 @@ def load_pretrained( ``AutoModelForCausalLM``); on a distributed run every rank loads the whole model, onto its own GPU unless ``device_map`` is passed or ``init_quantized_weights`` places it. - No load here shards the routed experts: a plan with ``ep_size > 1`` is refused. A refusal on any - rank raises on every rank, before the load. The plan is recorded on the model + A refusal on any rank raises on every rank, before the load. The plan is recorded on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) for calibration and export, so :func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap` refuses it: to shard a model yourself, load it with ``from_pretrained``. :doc:`/guides/_distributed_loading` has the routing in full. @@ -677,10 +718,11 @@ 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 a 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 the model with ``from_config``, which takes none. 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, a 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 the model with ``from_config``, which + takes none. 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 @@ -749,9 +791,9 @@ def handed() -> DistributionPlan | None: decided = plan is None and dist.size() > 1 # Each rank checks its own arguments; a refusal on one fails every rank before the fetch. with _refused_together("load arguments", alone): - # No load here shards the routed experts yet, so an EP degree left to the decision is 1. - # Parsed in this check, a knob one rank cannot parse fails every rank. - ep_size = 1 if decided and _ep_degree(None) == "auto" else None + if decided: + # Parsed in this check, a knob one rank cannot parse fails every rank. + _ep_degree(None) if kwargs.pop("output_loading_info", False): raise ValueError( "load_pretrained returns the model alone and does not take " @@ -803,7 +845,6 @@ def handed() -> DistributionPlan | None: config=config, auto_class=auto_class, trust_remote_code=trust_remote_code, - ep_size=ep_size, **hub_options, ) if _route(plan) is not None: @@ -846,11 +887,19 @@ def handed() -> DistributionPlan | None: {} if route is None else _sharded_load_kwargs( - route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided, measured + route, + plan, + ckpt_path, + auto_class, + model_type, + kwargs, + hub_options, + decided, + measured, ) ) if route is not None: - model = _load_sharded(ckpt_path, load_kwargs) + model = _load_sharded(route, ckpt_path, load_kwargs) # A sharded load finds the unplaced keys while it places the weights. keys = model._modelopt_unplaced_source_keys else: @@ -1412,8 +1461,9 @@ def _decide_distribution( ckpt_path: str, hf_config, world_size: int, ep_size, dense, trust_remote_code: bool, auto_class ) -> DistributionPlan: """The plan :func:`resolve_distribution_plan` returns, decided by this rank alone.""" - # What the caller set the EP degree with, for a refusal to name it. + # What the caller set the EP degree and the dense layout with, for a refusal to name them. knob = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size" + replicate = "MODELOPT_FSDP_SIZE=1" if dense is None else "dense='replicate'" 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 @@ -1425,6 +1475,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) or _init_quantized_weights()) and ep_size in ("auto", 1) @@ -1471,8 +1524,8 @@ def _decide_distribution( 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." + f"can be neither counted nor sharded. Set {knob}=1 and {replicate} to load it without " + "expert parallelism, whole on every rank." ) 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 @@ -1480,8 +1533,8 @@ def _decide_distribution( 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." + f"n_routed_experts or num_routed_experts) to shard them by. Set {knob}=1 and " + f"{replicate} to load it without expert parallelism, whole on every rank." ) if ep_size == "auto": ep_size = _choose_ep_degree(world_size, num_experts) @@ -1505,6 +1558,38 @@ def _decide_distribution( memory = _gpu_memory() too_big = memory is not None and held > _REPLICATED_SHARE * memory dense = "fsdp" if too_big else "replicate" + if experts and asked_to_replicate: + print_rank_0( + f"{hf_config.model_type}: ep_size=1 and dense='replicate' as asked, so every rank " + "holds the whole MoE." + ) + 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) @@ -1561,6 +1646,10 @@ def resolve_distribution_plan( is the class the load builds with: with one the sharded loads refuse, or inside ``init_quantized_weights``, and nothing asking to shard, the plan loads the model whole, decided without a build. + + A plan that leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for + exactly that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and + ``MODELOPT_FSDP_SIZE=1``) replicates it, noting that every rank holds the whole MoE. """ # 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. @@ -1683,6 +1772,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 a MoE causal LM expert-parallel for PTQ: each rank holds its share of the routed experts. @@ -1700,7 +1790,9 @@ def load_and_prepare_ep( The model builds with its LM head where an auto class builds one; ``require_lm_head=False`` builds a VLM without it. Without ``hf_config`` every rank reads the checkpoint's config itself; - on many ranks, pass it. The caller records the plan on the model + on many ranks, pass it. ``hub_options`` (``token``, ``cache_dir``, ``local_files_only``, as + :func:`load_pretrained` takes them) go to the checkpoint's download, its config read and + ``from_pretrained``. The caller records the plan on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), as :func:`load_pretrained` does; :func:`~modelopt.torch.utils.moe.expert_parallel_mesh` then finds its mesh. """ @@ -1719,6 +1811,7 @@ def load_and_prepare_ep( f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}." ) world_size, ep_degree, edp = plan.world_size, plan.ep_size, plan.edp + hub_options = dict(hub_options or {}) # The ep_plan styles this load keeps: grouped_gemm, and those it swaps for its torch_a2a styles. # The checks and the swapped plan both read it. swap = { @@ -1735,7 +1828,7 @@ def load_and_prepare_ep( "5.7", "ModelOpt DP=EP expert-parallel PTQ (transformers MoE EP framework)" ) 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) @@ -1844,7 +1937,7 @@ def gpu() -> tuple[bool, bool, str | None]: # The load and the export both read the local snapshot, made only once the checks pass, so a # refused load downloads nothing. with _refused_together("the local checkpoint", alone=False): - _, ckpt_path = ensure_local_checkpoint(ckpt_path) + _, ckpt_path = ensure_local_checkpoint(ckpt_path, **hub_options) # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over # the whole world (``dp``). @@ -1871,6 +1964,7 @@ def gpu() -> tuple[bool, bool, str | None]: "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" diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py index 54b96461485..eb2685c68b2 100644 --- a/tests/gpu/torch/utils/test_model_load_utils.py +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -22,10 +22,12 @@ import pytest import torch +import transformers from _test_utils.torch.transformers_models import create_tiny_llama_dir +from packaging.version import Version from safetensors.torch import load_file from torch.distributed.tensor import DTensor -from transformers import AutoModelForCausalLM, LlamaConfig +from transformers import AutoModelForCausalLM, LlamaConfig, Qwen3MoeConfig from modelopt.torch.export.unified_export_hf import export_hf_checkpoint from modelopt.torch.utils.distributed import ( @@ -187,6 +189,7 @@ def test_carry_over_unplaced_weights(dist_workers, tmp_path): # ------------------------------------------------------------------ load_pretrained, left to decide +NUM_EXPERTS = 8 _TINY = { "hidden_size": 64, "intermediate_size": 128, @@ -210,9 +213,10 @@ def _full(tensor): def _load_and_check(rank, size, ckpt_dir, knobs, expected): """Load on every rank of the run with the plan left to decide, and check what each rank holds. - Left to decide, ``load_pretrained`` 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, ``load_pretrained`` shards an expert-parallel MoE's experts, FSDP2-shards a dense + model when asked to, and otherwise gives every rank the whole model on its own GPU. Whichever + load runs, every rank ends up with the checkpoint's weights and the load's record of what it + could not place. """ device = torch.device(f"cuda:{rank}") os.environ.update(knobs) @@ -227,14 +231,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] @@ -248,3 +269,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") diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 5f07787ddee..7414c1192b6 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -54,6 +54,7 @@ AutoModelForImageTextToText, AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq, + DeepseekV3Config, FineGrainedFP8Config, LlamaConfig, LlamaForCausalLM, @@ -233,6 +234,24 @@ def _qwen3_moe(num_experts=8): return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY) +def _deepseek_v3(n_routed_experts=8): + """DeepSeek-V3 without a plan of its own, so only its model spec can make it EP-capable. + + Its one layer is an MoE layer: DeepSeek-V3's first layers are dense by default. + """ + config = DeepseekV3Config( + n_routed_experts=n_routed_experts, + first_k_dense_replace=0, + moe_intermediate_size=8, + n_group=1, + topk_group=1, + num_experts_per_tok=1, + **_TINY, + ) + config.base_model_ep_plan = None # holds even if a later transformers ships a plan + return config + + def _bytes(module): return sum(p.numel() * p.element_size() for p in module.parameters()) @@ -440,14 +459,6 @@ def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, c 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) @@ -464,17 +475,22 @@ def test_an_ep_plan_naming_no_module_of_the_model_is_refused( ): """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.""" + every expert on every rank without a word. A degree of 1 with the dense part replicated 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"): + with pytest.raises( + ValueError, + match=r"renamed_experts\) names .* MODELOPT_EP_DEGREE=1 and MODELOPT_FSDP_SIZE=1", + ): resolve_distribution_plan(str(tmp_path)) else: monkeypatch.setenv("MODELOPT_EP_DEGREE", degree) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "1") assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) @@ -484,7 +500,7 @@ def test_routed_experts_the_config_does_not_count_are_refused( ): """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.""" + rank without a word. A degree of 1 with the dense part replicated 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.") @@ -493,10 +509,13 @@ def test_routed_experts_the_config_does_not_count_are_refused( run_as(4) if degree is None: - with pytest.raises(ValueError, match=r"no expert count .* MODELOPT_EP_DEGREE=1"): + with pytest.raises( + ValueError, match=r"no expert count .* MODELOPT_EP_DEGREE=1 and MODELOPT_FSDP_SIZE=1" + ): resolve_distribution_plan(str(tmp_path)) else: monkeypatch.setenv("MODELOPT_EP_DEGREE", degree) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "1") assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) @@ -551,6 +570,54 @@ def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from( resolve_distribution_plan(str(tmp_path), **arguments) +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize( + ("config", "ep_degree"), + [ + pytest.param(_qwen3_moe(num_experts=3), None, id="auto_finds_no_degree"), + pytest.param(_qwen3_moe(num_experts=8), "1", id="explicit_degree_one"), + ], +) +def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( + run_as, serve, monkeypatch, tmp_path, config, ep_degree +): + """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none. With + nothing sharded either, every rank would hold the whole MoE: that must fail, not run.""" + if ep_degree is not None: + monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) + run_as(2) + serve(config) + + with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): + resolve_distribution_plan(str(tmp_path)) + + +@needs_qwen3_moe_ep_plan +def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(run_as, serve, monkeypatch, tmp_path): + """Asked for no expert parallelism but for dense sharding, the MoE is sharded whole by FSDP2, + which is not the waste of a replica on every rank.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + run_as(2) + serve(_qwen3_moe(num_experts=8)) + + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp") + + +@needs_qwen3_moe_ep_plan +def test_an_ep_capable_moe_is_replicated_when_both_knobs_ask_for_it( + run_as, serve, monkeypatch, tmp_path +): + """Both knobs asking to shard nothing is a replication the caller chose, not one the plan + reached on its own: every rank holds the whole MoE.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "1") + run_as(2) + serve(_qwen3_moe(num_experts=8)) + + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2) + + @needs_qwen3_moe_ep_plan @pytest.mark.parametrize( "value", @@ -903,6 +970,11 @@ def stub(*args, **kwargs): return stub + monkeypatch.setattr( + model_load_utils, + "load_and_prepare_ep", + recording("ep", model_load_utils.load_and_prepare_ep), + ) monkeypatch.setattr( model_load_utils, "parallel_load_and_prepare_fsdp2", @@ -1224,36 +1296,52 @@ def test_a_model_class_given_by_name_is_refused_by_name( @needs_qwen3_moe_ep_plan @pytest.mark.parametrize("knob", [None, "4"], ids=["left-to-decide", "asked-for"]) -def test_a_load_it_decides_keeps_the_experts_whole( - run_as, serve, loaders, monkeypatch, tmp_path, knob +def test_a_load_it_decides_shards_the_experts( + run_as, serve, loaders, a_gpu, monkeypatch, tmp_path, knob ): - """No load here shards the routed experts yet, so an MoE whose EP degree is left to the decision - is decided without expert parallelism and loads, where an automatic degree would be one the - load refuses; a degree MODELOPT_EP_DEGREE asks for is still refused, not ignored.""" + """An MoE whose EP degree is left to the decision is decided expert-parallel and loads through + the expert-parallel load, as one MODELOPT_EP_DEGREE asks for does.""" run_as(4) serve(_qwen3_moe()) if knob is not None: monkeypatch.setenv("MODELOPT_EP_DEGREE", knob) - monkeypatch.setattr( - transformers.AutoModelForCausalLM, - "from_pretrained", - classmethod(lambda cls, *a, **k: torch.nn.Linear(1, 1)), - ) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) - if knob is not None: - with pytest.raises(NotImplementedError, match=r"ep_size=4.*MODELOPT_EP_DEGREE=1"): - load_pretrained(str(tmp_path)) - return - assert distribution_plan(load_pretrained(str(tmp_path))) == DistributionPlan(4) + model = load_pretrained(str(tmp_path)) + assert loaders["ep"]["plan"] == DistributionPlan(4, ep_size=4) + assert distribution_plan(model) == DistributionPlan(4, ep_size=4) -def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, tmp_path): - """No load here shards the routed experts: an EP plan fails before anything loads.""" - run_as(2) - serve(Qwen3MoeConfig(num_experts=8, moe_intermediate_size=8, **_TINY)) +@pytest.mark.parametrize("dense", ["replicate", "fsdp"]) +def test_an_expert_parallel_plan_loads_through_the_ep_loader( + run_as, serve, loaders, auto_class, tmp_path, dense +): + """An EP plan loads through load_and_prepare_ep, with the plan and the Hub options the load was + given, and through no other load. That load replicates the dense part, so a plan that also + FSDP2-shards it is refused before anything loads, naming the ways out.""" + run_as(4) + serve(_deepseek_v3(n_routed_experts=8)) + plan = DistributionPlan(4, ep_size=4, dense=dense) + options = { + "trust_remote_code": True, + "attn_implementation": "eager", + "freeze": False, + "local_files_only": True, + } - with pytest.raises(NotImplementedError, match=r"ep_size=2.*MODELOPT_EP_DEGREE=1"): - load_pretrained(str(tmp_path), auto_class, distribution=DistributionPlan(2, ep_size=2)) + if dense == "fsdp": + with pytest.raises( + NotImplementedError, match=r"load replicates it on every rank\..*\(MODELOPT_FSDP_SIZE=1" + ): + load_pretrained(str(tmp_path), distribution=plan, **options) + else: + model = load_pretrained(str(tmp_path), distribution=plan, **options) + call = loaders.pop("ep") + assert (call["ckpt_path"], call["plan"]) == (str(tmp_path), plan) + assert (call["trust_remote_code"], call["attn_implementation"]) == (True, "eager") + assert call["freeze"] is False + assert call["hub_options"] == {"local_files_only": True} + assert distribution_plan(model) == plan assert loaders == {} assert auto_class.calls == [] @@ -1443,6 +1531,7 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, cpu_dist_workers.run(_one_rank_fails_to_load_job, ckpt=str(tmp_path)) +@pytest.mark.parametrize("sharding", ["expert-parallel", "fsdp2"]) @pytest.mark.parametrize( ("change", "error", "match"), [ @@ -1468,22 +1557,21 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, ids=["placement", "dtype", "torch_dtype", "model_class", "mtp_head", "two_configs"], ) def test_a_sharded_load_refuses_what_it_cannot_honor( - run_as, serve, loaders, process_group, monkeypatch, tmp_path, change, error, match + run_as, serve, loaders, process_group, monkeypatch, tmp_path, sharding, change, error, match ): """Placement, a model class, a head only a whole model builds, two configs to load with: none reaches a sharded load, which fails before it starts rather than load something other than what was asked.""" run_as(2) - serve(LlamaConfig(**_TINY)) + expert_parallel = sharding == "expert-parallel" + serve(_deepseek_v3() if expert_parallel else LlamaConfig(**_TINY)) if change.get("mtp"): monkeypatch.setattr(models_hf, "checkpoint_has_mtp", lambda *a: True) + plan = DistributionPlan(2, ep_size=2) if expert_parallel else DistributionPlan(2, dense="fsdp") with pytest.raises(error, match=match): load_pretrained( - str(tmp_path), - change.get("auto_class"), - distribution=DistributionPlan(2, dense="fsdp"), - **change.get("kwargs", {}), + str(tmp_path), change.get("auto_class"), distribution=plan, **change.get("kwargs", {}) ) assert loaders == {} @@ -2206,7 +2294,9 @@ def next_check(*args, **kwargs): 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 checkpoint's download, 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.""" @@ -2228,12 +2318,27 @@ 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, downloads = [], [] + monkeypatch.setattr( + "transformers.AutoConfig.from_pretrained", lambda *a, **k: reads.append(k) or config + ) + monkeypatch.setattr( + model_load_utils, + "ensure_local_checkpoint", + lambda p, **k: downloads.append(k) or (None, p), + ) 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 downloads == [{"local_files_only": True}] + assert loaded["local_files_only"] is True assert classes == [True] From 119c2943c653d9e87ebe11a6a2c094978862f3a2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:19:55 +0000 Subject: [PATCH 02/32] Say the FSDP2 way out of a whole MoE takes an unquantized causal LM The refusal of an expert-parallel MoE left whole on every rank offered FSDP2 "for an unquantized checkpoint", but the FSDP2 load builds a causal LM and refuses an MoE VLM as well, as the refusal of an expert-parallel plan that shards the dense part already says. Signed-off-by: Shengliang Xu --- 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 705ffccf906..6ca16a288dc 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1585,7 +1585,7 @@ def _decide_distribution( 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 " + "model with FSDP2 (MODELOPT_FSDP_SIZE=world, for an unquantized causal LM), 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'))." From 2a836208ddc350d6177f24aafb3c6c0bd6fd284f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:37:52 +0000 Subject: [PATCH 03/32] Name the way to a whole model when the expert-parallel load refuses a placement The refusal of an option a sharded load does not take named a way out only for the FSDP2 load. An expert-parallel MoE needs both knobs to load whole on every rank, since the plan decision refuses to replicate one unasked, so the expert-parallel load's refusal now names MODELOPT_EP_DEGREE=1 with MODELOPT_FSDP_SIZE=1 for a decided plan, and a plan with an EP degree of 1 for one handed in, through _whole_instead, which now takes the route. The FSDP2 load's text is unchanged. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 21 +++++++++++++------ .../unit/torch/utils/test_model_load_utils.py | 7 ++++++- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 6ca16a288dc..2cb7628c371 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -529,10 +529,19 @@ def _init_quantized_weights() -> bool: return getattr(AutoModelForCausalLM.from_pretrained, "_modelopt_init_quantized_weights", False) -def _whole_instead(decided: bool) -> str: - """How to have every rank load the model whole: the knob counts only for a decided plan.""" - plan = f"distribution=DistributionPlan({dist.size()}, dense='replicate')" - return f"MODELOPT_FSDP_SIZE=1 or {plan}" if decided else plan +def _whole_instead(decided: bool, route: str = _FSDP2) -> str: + """How to have every rank load the model whole: the knob counts only for a decided plan. + + Instead of the expert-parallel load it takes both knobs, as the plan decision refuses to + replicate an expert-parallel MoE unasked. + """ + if route == _EXPERT_PARALLEL: + plan = f"distribution=DistributionPlan({dist.size()}, 1, 'replicate')" + knobs = "MODELOPT_EP_DEGREE=1 with MODELOPT_FSDP_SIZE=1" + else: + plan = f"distribution=DistributionPlan({dist.size()}, dense='replicate')" + knobs = "MODELOPT_FSDP_SIZE=1" + return f"{knobs} or {plan}" if decided else plan def _refuse_sharded_load_options( @@ -555,7 +564,7 @@ def _refuse_sharded_load_options( raise NotImplementedError( f"The {route} load builds the model without from_pretrained, so it would skip " "init_quantized_weights' quantization; load the model whole on every rank: " - f"{_whole_instead(decided)}." + f"{_whole_instead(decided, route)}." ) if auto_class not in (None, AutoModelForCausalLM): name = getattr(auto_class, "__name__", repr(auto_class)) @@ -575,7 +584,7 @@ def _refuse_sharded_load_options( raise TypeError( f"{unknown} do not apply to the {route} load, which places the model itself; it " f"takes {sorted(_SHARDED_LOAD_OPTIONS[route])}. To place the model yourself, have " - f"every rank load it whole: {_whole_instead(decided)}." + f"every rank load it whole: {_whole_instead(decided, route)}." ) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 7414c1192b6..6fbe41d32d3 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1569,10 +1569,15 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( monkeypatch.setattr(models_hf, "checkpoint_has_mtp", lambda *a: True) plan = DistributionPlan(2, ep_size=2) if expert_parallel else DistributionPlan(2, dense="fsdp") - with pytest.raises(error, match=match): + with pytest.raises(error, match=match) as refused: load_pretrained( str(tmp_path), change.get("auto_class"), distribution=plan, **change.get("kwargs", {}) ) + if "device_map" in change.get("kwargs", {}): + # The way to a whole model it names, for the plan handed in: an expert-parallel MoE's is a + # plan with an EP degree of 1. + whole_moe = "DistributionPlan(2, 1, 'replicate')" + assert (whole_moe in str(refused.value)) == expert_parallel, refused.value assert loaders == {} From d07b1a1cca6d53ddf524ed7c55a001b20ae731f1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:39:58 +0000 Subject: [PATCH 04/32] Name FSDP2's mp_policy only in the FSDP2 load's dtype refusal The dtype a sharded load refuses is refused for both sharded loads now, but the refusal ended by pointing at FSDP2's mp_policy, which the expert-parallel load does not take. Name it for the FSDP2 load alone. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 5 +++-- tests/unit/torch/utils/test_model_load_utils.py | 3 +++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 2cb7628c371..abf47322264 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -575,10 +575,11 @@ def _refuse_sharded_load_options( ) if dtypes: # A sharded load keeps the checkpoint's dtype. name = dtypes[0] + # Only the FSDP2 load computes in another dtype, through its mp_policy. + computes = "; FSDP2's mp_policy sets the dtype it computes in" if route == _FSDP2 else "" raise TypeError( f"The {route} load reads the checkpoint's tensors as stored and keeps their dtype, so " - f"{name}={kwargs[name]!r} does not apply to it; FSDP2's mp_policy sets the dtype it " - "computes in." + f"{name}={kwargs[name]!r} does not apply to it{computes}." ) if unknown: raise TypeError( diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 6fbe41d32d3..5bf458c5055 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1578,6 +1578,9 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( # plan with an EP degree of 1. whole_moe = "DistributionPlan(2, 1, 'replicate')" assert (whole_moe in str(refused.value)) == expert_parallel, refused.value + if {"dtype", "torch_dtype"} & set(change.get("kwargs", {})): + # Only the FSDP2 load takes an mp_policy to compute in another dtype. + assert ("mp_policy" in str(refused.value)) != expert_parallel, refused.value assert loaders == {} From a60df78ca2aec47e48cf7b00f9035ea50f530f85 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:40:23 +0000 Subject: [PATCH 05/32] Refuse a rank without its own GPU before the expert-parallel load too The sharded loads refuse a rank whose local rank numbers no GPU of its own before they start, with the other ranks, rather than leave them in the load's collectives. The expert-parallel route returned before that check, so such a rank reached the load's meshes, which put it on another rank's GPU. Both routes now run the check, and the refusal names the way to a whole model for its route, as the placement refusal does. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 16 ++++++++++++---- tests/unit/torch/utils/test_model_load_utils.py | 15 +++++++++------ 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index abf47322264..a10a174ac64 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -610,11 +610,12 @@ def _sharded_load_kwargs( # The head is built on a whole model only; carried over, it would export unquantized. raise NotImplementedError( f"The {route} load cannot build the MTP head this {model_type} checkpoint carries. To " - f"load it whole on every rank: {_whole_instead(decided)}." + f"load it whole on every rank: {_whole_instead(decided, route)}." ) trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) if route == _EXPERT_PARALLEL: + _refuse_a_rank_without_its_gpu(route, decided) return { "plan": plan, "trust_remote_code": trust_remote_code, @@ -650,14 +651,21 @@ def _sharded_load_kwargs( f"The {route} load shards through a process group, even on one rank, which a process " "started without a launcher cannot set up; launch it with `torchrun --nproc_per_node=1`." ) + _refuse_a_rank_without_its_gpu(route, decided) + return {"trust_remote_code": trust_remote_code, "hf_config": hf_config, **kwargs} + + +def _refuse_a_rank_without_its_gpu(route, decided: bool) -> None: + """Refuse a rank without a GPU of its own before the ``route`` load, with the other ranks. + + Inside the load, the others would wait for this rank in its collectives. + """ if not _has_own_gpu(): - # Refused with the other ranks, not inside the load, where they would wait for this one. raise RuntimeError( f"The {route} load shards onto each rank's own GPU, and this rank has none: it sees " f"{torch.cuda.device_count()} GPU(s) as local rank {dist.local_rank()}. To load it " - f"whole on every rank: {_whole_instead(decided)}." + f"whole on every rank: {_whole_instead(decided, route)}." ) - return {"trust_remote_code": trust_remote_code, "hf_config": hf_config, **kwargs} def _own_gpu_index() -> int | None: diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 5bf458c5055..82b0be7c060 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1314,7 +1314,7 @@ def test_a_load_it_decides_shards_the_experts( @pytest.mark.parametrize("dense", ["replicate", "fsdp"]) def test_an_expert_parallel_plan_loads_through_the_ep_loader( - run_as, serve, loaders, auto_class, tmp_path, dense + run_as, serve, loaders, auto_class, a_gpu, tmp_path, dense ): """An EP plan loads through load_and_prepare_ep, with the plan and the Hub options the load was given, and through no other load. That load replicates the dense part, so a plan that also @@ -1553,27 +1553,30 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, TypeError, "config and hf_config", ), + ({"gpus": 0}, RuntimeError, "this rank has none"), ], - ids=["placement", "dtype", "torch_dtype", "model_class", "mtp_head", "two_configs"], + ids=["placement", "dtype", "torch_dtype", "model_class", "mtp_head", "two_configs", "no_gpu"], ) def test_a_sharded_load_refuses_what_it_cannot_honor( run_as, serve, loaders, process_group, monkeypatch, tmp_path, sharding, change, error, match ): - """Placement, a model class, a head only a whole model builds, two configs to load with: none - reaches a sharded load, which fails before it starts rather than load something other than - what was asked.""" + """Placement, a model class, a head only a whole model builds, two configs to load with, a rank + without a GPU of its own: none reaches a sharded load, which fails before it starts rather than + load something other than what was asked.""" run_as(2) expert_parallel = sharding == "expert-parallel" serve(_deepseek_v3() if expert_parallel else LlamaConfig(**_TINY)) if change.get("mtp"): monkeypatch.setattr(models_hf, "checkpoint_has_mtp", lambda *a: True) + if "gpus" in change: # a host that shows this rank no GPU + monkeypatch.setattr(torch.cuda, "device_count", lambda: change["gpus"]) plan = DistributionPlan(2, ep_size=2) if expert_parallel else DistributionPlan(2, dense="fsdp") with pytest.raises(error, match=match) as refused: load_pretrained( str(tmp_path), change.get("auto_class"), distribution=plan, **change.get("kwargs", {}) ) - if "device_map" in change.get("kwargs", {}): + if "device_map" in change.get("kwargs", {}) or "gpus" in change: # The way to a whole model it names, for the plan handed in: an expert-parallel MoE's is a # plan with an EP degree of 1. whole_moe = "DistributionPlan(2, 1, 'replicate')" From ccd48747e62ccd6f5c99e57cba955c9d6666b397 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:31:54 +0000 Subject: [PATCH 06/32] Refuse the expert-parallel load's own options on a whole-model load freeze and require_lm_head are options only the expert-parallel load takes. A plan that loads the model whole on every rank, which the left-whole refusal recommends, handed them on to from_pretrained, which failed on an unexpected keyword argument from deep inside transformers. Refuse them before anything loads, naming the load that takes them, as the FSDP2 load's own options already are. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 11 +++++++++++ tests/unit/torch/utils/test_model_load_utils.py | 12 ++++++++++-- 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 a10a174ac64..b1ea33e69cf 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -474,6 +474,11 @@ def record_unplaced_source_keys( _ASKS_NOTHING = {"cpu_offload": (False, None), "mp_policy": (None,), "broadcast_chunk_size": (8,)} +# The options only the expert-parallel load takes, which a whole-model load would hand on to +# from_pretrained. +_EP_ONLY_LOAD_OPTIONS = _SHARDED_LOAD_OPTIONS[_EXPERT_PARALLEL] - _SHARDED_LOAD_OPTIONS[_FSDP2] + + def _route(plan: DistributionPlan) -> str | None: """The sharded load ``plan`` takes, or ``None`` when every rank loads the whole model.""" if plan.world_size != dist.size(): @@ -883,6 +888,12 @@ def handed() -> DistributionPlan | None: f"{sharded} only apply to the FSDP2 load, but {plan!r} gives every rank the whole " "model; pass a plan with dense='fsdp' to shard it." ) + # As do the options only the expert-parallel load takes. + if route is None and (sharded := sorted(set(kwargs) & _EP_ONLY_LOAD_OPTIONS)): + raise TypeError( + f"{sharded} only apply to the expert-parallel load, but {plan!r} gives every rank " + "the whole model; pass a plan with ep_size > 1 to shard an MoE's routed experts." + ) measured = decided and route == _FSDP2 and _dense_from_env(plan.world_size) == "auto" if route is not None and config is not None: # The sharded load builds from hf_config, a copy, as building writes to the config. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 82b0be7c060..dc824632c1b 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1633,6 +1633,13 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( TypeError, r"\['cpu_offload'\] only apply", ), + ( + DistributionPlan(2), + "causal-lm", + {"freeze": False}, + TypeError, + r"\['freeze'\] only apply to the expert-parallel load", + ), ], ids=[ "another-world", @@ -1642,6 +1649,7 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( "not-a-plan", "hf-config-on-a-whole-load", "fsdp2-option-on-a-whole-load", + "ep-option-on-a-whole-load", ], ) def test_a_load_refuses_what_it_cannot_load( @@ -1661,8 +1669,8 @@ def test_a_load_refuses_what_it_cannot_load( read as stored, even when only a VLM's language model is quantized and its remote code builds it as a causal LM; a VLM but a text-only one it supports, which FSDP2 would build through AutoModelForCausalLM, refused or as its language model alone; anything but a plan; and the - sharded loads' config, or an option only the FSDP2 load takes, which from_pretrained would take - for a model argument and fail far from here, or drop.""" + sharded loads' config, or an option only one sharded load takes, which from_pretrained would + take for a model argument and fail far from here, or drop.""" run_as(2) config = serve(transformers.LlavaConfig() if checkpoint == "vlm" else LlamaConfig(**_TINY)) if checkpoint == "quantized": From 822d7adf618cdb8f6975f64c0d7d61f9c5098aee Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:50:24 +0000 Subject: [PATCH 07/32] Say the expert-parallel load freezes the model unless told not to load_and_prepare_ep freezes the model it loads by default, which the FSDP2 and whole-model loads do not, so the same load_pretrained call returns a frozen model only when its plan is expert-parallel. Say so where load_pretrained lists its loads. Signed-off-by: Shengliang Xu --- 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 b1ea33e69cf..fbacb90392e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -723,7 +723,8 @@ def load_pretrained( * ``ep_size > 1``: :func:`load_and_prepare_ep` shards an MoE's routed experts over the process group as they load and replicates the dense part, so a plan with ``dense="fsdp"`` as well is refused. It builds and places the model itself, so it takes its own options but no - ``auto_class``, placement or ``dtype``, and refuses a quantized checkpoint. + ``auto_class``, placement or ``dtype``, and refuses a quantized checkpoint. Unlike the other + loads, it freezes the model unless ``freeze=False``. * ``dense="fsdp"``: :func:`parallel_load_and_prepare_fsdp2` shards a causal LM, or a text-only VLM it supports, over the process group (one process needs ``torchrun --nproc_per_node=1``). It builds and places the model itself, so it takes its own options but no ``auto_class``, From dd0bc322e8c95f758228230769b9bef8daa1f57a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:50:33 +0000 Subject: [PATCH 08/32] Name the caller's own knobs in the left-whole refusal The refusal of an MoE left whole on every rank named the environment knobs for every way out, while its sibling refusals name what the caller set. A resolve_distribution_plan caller who passed ep_size or dense as arguments, which override the knobs, would set the named knob, run again and meet the same refusal. Name the arguments there instead. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 10 +++++----- tests/unit/torch/utils/test_model_load_utils.py | 16 ++++++++++++---- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index fbacb90392e..3546d8310d6 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1494,6 +1494,7 @@ def _decide_distribution( # What the caller set the EP degree and the dense layout with, for a refusal to name them. knob = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size" replicate = "MODELOPT_FSDP_SIZE=1" if dense is None else "dense='replicate'" + fsdp = "MODELOPT_FSDP_SIZE=world" if dense is None else "dense='fsdp'" 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 @@ -1609,15 +1610,14 @@ def _decide_distribution( ) else: shard = ( - f"Shard its experts (MODELOPT_EP_DEGREE above 1, dividing both its {num_experts} " - "experts and the world size), shard" + f"Shard its experts ({knob} 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 causal LM), load it in " - "a single process, or replicate it on every rank outright: MODELOPT_EP_DEGREE=1 with " - "MODELOPT_FSDP_SIZE=1, or " + f"model with FSDP2 ({fsdp}, for an unquantized causal LM), load it in a single " + f"process, or replicate it on every rank outright: {knob}=1 with {replicate}, or " f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))." ) return DistributionPlan(world_size, ep_size, dense) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index dc824632c1b..34761c2b96d 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -576,20 +576,28 @@ def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from( [ pytest.param(_qwen3_moe(num_experts=3), None, id="auto_finds_no_degree"), pytest.param(_qwen3_moe(num_experts=8), "1", id="explicit_degree_one"), + pytest.param(_qwen3_moe(num_experts=8), 1, id="explicit_degree_one_as_an_argument"), ], ) def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( run_as, serve, monkeypatch, tmp_path, config, ep_degree ): """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none. With - nothing sharded either, every rank would hold the whole MoE: that must fail, not run.""" - if ep_degree is not None: + nothing sharded either, every rank would hold the whole MoE: that must fail, not run. The ways + out are named as the caller set the plan, so arguments, which override the knobs, are named.""" + if isinstance(ep_degree, str): monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) + arguments = {"ep_size": ep_degree, "dense": "auto"} if isinstance(ep_degree, int) else {} run_as(2) serve(config) - with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): - resolve_distribution_plan(str(tmp_path)) + with pytest.raises( + RuntimeError, match="can be expert-parallel, but the plan resolved" + ) as refused: + resolve_distribution_plan(str(tmp_path), **arguments) + if arguments: # a knob the arguments override would lead straight back here + assert "ep_size=1 with dense='replicate'" in str(refused.value), refused.value + assert "MODELOPT_" not in str(refused.value), refused.value @needs_qwen3_moe_ep_plan From c29014ee47a3516b2097b9c99dc510c791a6f9ca Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:17:29 +0000 Subject: [PATCH 09/32] Refuse an MoE left whole on every rank with a ValueError The plan decision refuses what its inputs resolve to with ValueError: a plan naming no module of the built model, a config with no expert count, an EP degree dividing neither the world nor the experts. The refusal of an MoE left whole on every rank says the same kind of thing but raised RuntimeError, so a caller catching ValueError around resolve_distribution_plan missed it alone. Raise ValueError there too. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 3 ++- tests/unit/torch/utils/test_model_load_utils.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 3546d8310d6..30f04ee52ae 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1613,7 +1613,8 @@ def _decide_distribution( f"Shard its experts ({knob} above 1, dividing both its {num_experts} experts and " "the world size), shard" ) - raise RuntimeError( + # A ValueError, as the decision's other refusals of what its inputs resolve to are. + raise ValueError( f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 at " f"world_size={world_size}, so every rank would hold the whole MoE. {shard} the whole " f"model with FSDP2 ({fsdp}, for an unquantized causal LM), load it in a single " diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 34761c2b96d..d68c4e3b731 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -592,7 +592,7 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( serve(config) with pytest.raises( - RuntimeError, match="can be expert-parallel, but the plan resolved" + ValueError, match="can be expert-parallel, but the plan resolved" ) as refused: resolve_distribution_plan(str(tmp_path), **arguments) if arguments: # a knob the arguments override would lead straight back here From 713019d737d7515dbf29eacc62855ca20ad2af7b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:18:59 +0000 Subject: [PATCH 10/32] Read the checkpoint's config once for an expert-parallel load too A load handed a sharded plan and no config reads the checkpoint's config for its model type. The FSDP2 load then builds from that read, but the expert-parallel load read the config again on every rank. Hand it the same read, as the FSDP2 load gets. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 4 ++-- tests/unit/torch/utils/test_model_load_utils.py | 8 +++++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 30f04ee52ae..f868c7d0a61 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -903,8 +903,8 @@ def handed() -> DistributionPlan | None: kwargs["hf_config"] = copy.deepcopy(config) if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. - if route == _FSDP2 and hf_config is None and config is None: - # The FSDP2 load builds from the checkpoint's config too: read it once, for both. + if route is not None and hf_config is None and config is None: + # A sharded load builds from the checkpoint's config too: read it once, for both. hf_config = AutoConfig.from_pretrained( ckpt_path, trust_remote_code=trust_remote_code, **hub_options ) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index d68c4e3b731..9b27824ec01 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1325,10 +1325,11 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader( run_as, serve, loaders, auto_class, a_gpu, tmp_path, dense ): """An EP plan loads through load_and_prepare_ep, with the plan and the Hub options the load was - given, and through no other load. That load replicates the dense part, so a plan that also - FSDP2-shards it is refused before anything loads, naming the ways out.""" + given, and the checkpoint's config, read once for both, and through no other load. That load + replicates the dense part, so a plan that also FSDP2-shards it is refused before anything + loads, naming the ways out.""" run_as(4) - serve(_deepseek_v3(n_routed_experts=8)) + config = serve(_deepseek_v3(n_routed_experts=8)) plan = DistributionPlan(4, ep_size=4, dense=dense) options = { "trust_remote_code": True, @@ -1349,6 +1350,7 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader( assert (call["trust_remote_code"], call["attn_implementation"]) == (True, "eager") assert call["freeze"] is False assert call["hub_options"] == {"local_files_only": True} + assert call["hf_config"] is config assert distribution_plan(model) == plan assert loaders == {} assert auto_class.calls == [] From 5cbb6455a78395cb66913b5c4668dd8a632b362b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:44:10 +0000 Subject: [PATCH 11/32] Let None ask nothing of the expert-parallel load's own flags load_pretrained drops an option passed as None, from_pretrained's own default, so it asks nothing. freeze and require_lm_head default to True, though, and a None passed for either reached load_and_prepare_ep, which read it as False: require_lm_head=None built a VLM without its LM head. Drop None for them too, so their defaults hold. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 6 +++++- tests/unit/torch/utils/test_model_load_utils.py | 8 +++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index f868c7d0a61..1b07516221a 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -889,7 +889,11 @@ def handed() -> DistributionPlan | None: f"{sharded} only apply to the FSDP2 load, but {plan!r} gives every rank the whole " "model; pass a plan with dense='fsdp' to shard it." ) - # As do the options only the expert-parallel load takes. + # As do the options only the expert-parallel load takes, but for one passed as None, which + # asks nothing of either load: its default holds, rather than None read as False. + for name in _EP_ONLY_LOAD_OPTIONS: + if name in kwargs and kwargs[name] is None: + del kwargs[name] if route is None and (sharded := sorted(set(kwargs) & _EP_ONLY_LOAD_OPTIONS)): raise TypeError( f"{sharded} only apply to the expert-parallel load, but {plan!r} gives every rank " diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 9b27824ec01..b4f833b5b37 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1325,9 +1325,9 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader( run_as, serve, loaders, auto_class, a_gpu, tmp_path, dense ): """An EP plan loads through load_and_prepare_ep, with the plan and the Hub options the load was - given, and the checkpoint's config, read once for both, and through no other load. That load - replicates the dense part, so a plan that also FSDP2-shards it is refused before anything - loads, naming the ways out.""" + given, and the checkpoint's config, read once for both, and through no other load. An option + passed as None asks nothing of it. That load replicates the dense part, so a plan that also + FSDP2-shards it is refused before anything loads, naming the ways out.""" run_as(4) config = serve(_deepseek_v3(n_routed_experts=8)) plan = DistributionPlan(4, ep_size=4, dense=dense) @@ -1335,6 +1335,7 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader( "trust_remote_code": True, "attn_implementation": "eager", "freeze": False, + "require_lm_head": None, "local_files_only": True, } @@ -1349,6 +1350,7 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader( assert (call["ckpt_path"], call["plan"]) == (str(tmp_path), plan) assert (call["trust_remote_code"], call["attn_implementation"]) == (True, "eager") assert call["freeze"] is False + assert "require_lm_head" not in call # its default, not None read as False assert call["hub_options"] == {"local_files_only": True} assert call["hf_config"] is config assert distribution_plan(model) == plan From 14d36c451bd97b6778592b7a2c9ab04d3d43f7db Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:44:32 +0000 Subject: [PATCH 12/32] Return the model unfrozen from load_pretrained's expert-parallel load load_and_prepare_ep freezes the model it loads unless told otherwise, which the FSDP2 and whole-model loads do not. load_pretrained decides which load runs from the world size and the model, so the same call returned a frozen model on some runs and not on others. Pass freeze=False unless the caller asks for freeze, as the other loads return the model. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 6 ++++-- tests/unit/torch/utils/test_model_load_utils.py | 8 ++++---- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 1b07516221a..91c514e9835 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -621,6 +621,8 @@ def _sharded_load_kwargs( hf_config = kwargs.pop("hf_config", None) if route == _EXPERT_PARALLEL: _refuse_a_rank_without_its_gpu(route, decided) + # Unfrozen unless asked, as the other loads return the model: which load runs is decided. + kwargs.setdefault("freeze", False) return { "plan": plan, "trust_remote_code": trust_remote_code, @@ -723,8 +725,8 @@ def load_pretrained( * ``ep_size > 1``: :func:`load_and_prepare_ep` shards an MoE's routed experts over the process group as they load and replicates the dense part, so a plan with ``dense="fsdp"`` as well is refused. It builds and places the model itself, so it takes its own options but no - ``auto_class``, placement or ``dtype``, and refuses a quantized checkpoint. Unlike the other - loads, it freezes the model unless ``freeze=False``. + ``auto_class``, placement or ``dtype``, and refuses a quantized checkpoint. Here it leaves the + model unfrozen, as the other loads do, unless ``freeze=True``. * ``dense="fsdp"``: :func:`parallel_load_and_prepare_fsdp2` shards a causal LM, or a text-only VLM it supports, over the process group (one process needs ``torchrun --nproc_per_node=1``). It builds and places the model itself, so it takes its own options but no ``auto_class``, diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b4f833b5b37..db133a52a29 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1325,16 +1325,16 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader( run_as, serve, loaders, auto_class, a_gpu, tmp_path, dense ): """An EP plan loads through load_and_prepare_ep, with the plan and the Hub options the load was - given, and the checkpoint's config, read once for both, and through no other load. An option - passed as None asks nothing of it. That load replicates the dense part, so a plan that also - FSDP2-shards it is refused before anything loads, naming the ways out.""" + given, and the checkpoint's config, read once for both, and through no other load. It returns + the model unfrozen, as the other loads do, and an option passed as None asks nothing of it. That + load replicates the dense part, so a plan that also FSDP2-shards it is refused before anything + loads, naming the ways out.""" run_as(4) config = serve(_deepseek_v3(n_routed_experts=8)) plan = DistributionPlan(4, ep_size=4, dense=dense) options = { "trust_remote_code": True, "attn_implementation": "eager", - "freeze": False, "require_lm_head": None, "local_files_only": True, } From cce9572dcf33456317af9e2cf0746f821232102f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:44:42 +0000 Subject: [PATCH 13/32] Say a decided expert-parallel layout was decided when it refuses a placement With neither a plan nor MODELOPT_EP_DEGREE, the plan decision can shard an MoE's experts, and the expert-parallel load then refuses options that place a whole model, such as the device_map hf_ptq passes, as if the caller had asked for it. Say the layout was decided, and how to keep the load whole, as the FSDP2 load already does when rank 0's GPU memory chose to shard: the measured flag _refuse_sharded_load_options takes now also covers a layout the decision chose for the model's routed experts. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 24 +++++++++++++++---- .../unit/torch/utils/test_model_load_utils.py | 18 ++++++++++++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 91c514e9835..50027a9eb34 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -552,7 +552,10 @@ def _whole_instead(decided: bool, route: str = _FSDP2) -> str: def _refuse_sharded_load_options( route, auto_class, kwargs, decided: bool, measured: bool = False ) -> None: - """Refuse what the sharded load ``route`` cannot honor; ``measured``: GPU memory chose it.""" + """Refuse what the sharded load ``route`` cannot honor; ``measured``: the decision chose it. + + By rank 0's GPU memory for the FSDP2 load, by the model's experts for the expert-parallel one. + """ # An option passed as None, from_pretrained's own default (device_map=None, say), asks nothing. passed = {name for name, value in kwargs.items() if value is not None} unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) @@ -560,10 +563,16 @@ def _refuse_sharded_load_options( dtypes = [name for name in ("dtype", "torch_dtype") if kwargs.get(name) not in (None, "auto")] if measured and (refused := sorted({*unknown, *dtypes})): # Say so, rather than refuse the options below as if the caller had asked for this load. + chose = ( + "no plan or MODELOPT_EP_DEGREE asked for a layout, and the plan decision shards the " + "model's routed experts over the ranks" + if route == _EXPERT_PARALLEL + else "no plan or MODELOPT_FSDP_SIZE asked for a layout, and by rank 0's GPU memory the " + f"model does not fit one GPU, so the {route} load shards it" + ) raise TypeError( - f"{refused} do not apply to a sharded load, but no plan or MODELOPT_FSDP_SIZE asked for " - f"a layout, and by rank 0's GPU memory the model does not fit one GPU, so the {route} " - f"load shards it. To load it whole on every rank: {_whole_instead(decided)}." + f"{refused} do not apply to a sharded load, but {chose}. To load it whole on every " + f"rank: {_whole_instead(decided, route)}." ) if _init_quantized_weights(): raise NotImplementedError( @@ -901,7 +910,12 @@ def handed() -> DistributionPlan | None: f"{sharded} only apply to the expert-parallel load, but {plan!r} gives every rank " "the whole model; pass a plan with ep_size > 1 to shard an MoE's routed experts." ) - measured = decided and route == _FSDP2 and _dense_from_env(plan.world_size) == "auto" + # Whether the decision chose this sharded layout with nothing asking for it: by rank 0's GPU + # memory for the FSDP2 load, by the model's routed experts for the expert-parallel one. + measured = decided and ( + (route == _FSDP2 and _dense_from_env(plan.world_size) == "auto") + or (route == _EXPERT_PARALLEL and _ep_degree(None) == "auto") + ) if route is not None and config is not None: # The sharded load builds from hf_config, a copy, as building writes to the config. del kwargs["config"] diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index db133a52a29..ed498c4165b 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1763,6 +1763,24 @@ def test_a_sharded_load_is_refused_inside_init_quantized_weights( assert loaders == {} +@needs_qwen3_moe_ep_plan +def test_an_automatic_expert_parallel_layout_says_so_when_it_refuses_a_placement( + run_as, serve, loaders, process_group, monkeypatch, tmp_path +): + """With neither a plan nor MODELOPT_EP_DEGREE, the decision shards an MoE's experts, and the + expert-parallel load refuses options that place a whole model. The refusal says the layout was + decided rather than asked for, and how to keep the load whole.""" + run_as(2, rank=1) + serve(_qwen3_moe(num_experts=8)) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) + + with pytest.raises( + TypeError, match=r"\['device_map'\] do not apply.*no plan or MODELOPT_EP_DEGREE" + ): + load_pretrained(str(tmp_path), device_map="auto") + assert loaders == {} + + @pytest.mark.parametrize( ("checkpoint", "match"), [("quantized", "dequantize"), ("mtp", "MTP head")], From dee3a0c932b4d2369f5b5db50225d331c9f75ae5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:06:26 +0000 Subject: [PATCH 14/32] Take the expert-parallel load's own options on a decided layout The refusal of a whole-model placement on a decided layout read as placement every option the FSDP2 load does not take, freeze and require_lm_head among them. So with neither a plan nor MODELOPT_EP_DEGREE, load_pretrained(..., freeze=True) or require_lm_head=False, the expert- parallel load's own options, was refused as if it placed a whole model. An option either sharded load takes places nothing: _refuse_sharded_load_options, where the placement a measured layout refuses is now computed, leaves both loads' options out of it. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 5 ++++- tests/unit/torch/utils/test_model_load_utils.py | 16 +++++++++++++--- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 50027a9eb34..ba4f597d055 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -561,7 +561,10 @@ def _refuse_sharded_load_options( unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) # torch_dtype is from_pretrained's older name for dtype, and None means "auto" to it. dtypes = [name for name in ("dtype", "torch_dtype") if kwargs.get(name) not in (None, "auto")] - if measured and (refused := sorted({*unknown, *dtypes})): + # What does not apply to a sharded load: a dtype to load in, or an option neither load takes. + taken = _SHARDED_LOAD_OPTIONS[_FSDP2] | _SHARDED_LOAD_OPTIONS[_EXPERT_PARALLEL] + refused = sorted({*(passed - taken - {"dtype", "torch_dtype"}), *dtypes}) + if measured and refused: # Say so, rather than refuse the options below as if the caller had asked for this load. chose = ( "no plan or MODELOPT_EP_DEGREE asked for a layout, and the plan decision shards the " diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index ed498c4165b..d07d27c3f14 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1764,20 +1764,30 @@ def test_a_sharded_load_is_refused_inside_init_quantized_weights( @needs_qwen3_moe_ep_plan +@pytest.mark.parametrize( + "options", + [{"device_map": "auto"}, {"freeze": True, "require_lm_head": False}], + ids=["placement", "its-own-options"], +) def test_an_automatic_expert_parallel_layout_says_so_when_it_refuses_a_placement( - run_as, serve, loaders, process_group, monkeypatch, tmp_path + run_as, serve, loaders, process_group, monkeypatch, tmp_path, options ): """With neither a plan nor MODELOPT_EP_DEGREE, the decision shards an MoE's experts, and the expert-parallel load refuses options that place a whole model. The refusal says the layout was - decided rather than asked for, and how to keep the load whole.""" + decided rather than asked for, and how to keep the load whole. The load's own options place + nothing, and reach it.""" run_as(2, rank=1) serve(_qwen3_moe(num_experts=8)) monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) + if "device_map" not in options: + load_pretrained(str(tmp_path), **options) + assert {name: loaders["ep"][name] for name in options} == options + return with pytest.raises( TypeError, match=r"\['device_map'\] do not apply.*no plan or MODELOPT_EP_DEGREE" ): - load_pretrained(str(tmp_path), device_map="auto") + load_pretrained(str(tmp_path), **options) assert loaders == {} From 526083e1c9c8a8b11783bcdcb5b8cf465eccd5b6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:06:42 +0000 Subject: [PATCH 15/32] Say what the GPU refusal adds beside the expert-parallel load's check The expert-parallel load now refuses a run in which only some ranks have a GPU of their own, so no rank is left waiting in its collectives. The refusal before it still covers a run in which none has one, and names the way to a whole model, which the load's own check does not. Signed-off-by: Shengliang Xu --- 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 ba4f597d055..1fccf09e3f5 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -677,7 +677,9 @@ def _sharded_load_kwargs( def _refuse_a_rank_without_its_gpu(route, decided: bool) -> None: """Refuse a rank without a GPU of its own before the ``route`` load, with the other ranks. - Inside the load, the others would wait for this rank in its collectives. + Inside the FSDP2 load, the others would wait for this rank in its collectives; the + expert-parallel load refuses only a run in which some ranks have one, and neither names the way + to a whole model. """ if not _has_own_gpu(): raise RuntimeError( From 9bd0edd5c39655970d2f5ace644618e723fbf515 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:39:04 +0000 Subject: [PATCH 16/32] Refuse the other sharded load's options as that load's own A sharded load refused an option only the other sharded load takes as one that places a whole model, and named the way to a whole model on every rank, where the whole-model load refuses the option again. It now says which load takes it. Which options a sharded load counts as passed now follows the one rule the whole-model load drops them by, _ASKS_NOTHING's: an FSDP2-only option asks nothing at the FSDP2 load's own default, any other at None. So broadcast_chunk_size=None, which asks the FSDP2 load to broadcast all of a source's layers at once and was dropped without a word on the expert-parallel load, is refused as the FSDP2 load's option there, and cpu_offload=False, what hf_ptq passes by default, asks nothing of it. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 20 +++++++++++++++---- .../unit/torch/utils/test_model_load_utils.py | 16 +++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 1fccf09e3f5..aa0afcbc965 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -469,8 +469,9 @@ def record_unplaced_source_keys( ), _FSDP2: _FSDP2_ONLY | {"trust_remote_code", "attn_implementation", "hf_config"}, } -# The FSDP2-only values that ask nothing of a whole load: the FSDP2 load's defaults, and None where -# that leaves the option unset (a None broadcast_chunk_size asks for all of a source's layers). +# The FSDP2-only values that ask nothing of any load but the FSDP2 one: its defaults, and None where +# that leaves the option unset (a None broadcast_chunk_size asks for all of a source's layers). Any +# other option asks nothing at None, from_pretrained's own default (device_map=None, say). _ASKS_NOTHING = {"cpu_offload": (False, None), "mp_policy": (None,), "broadcast_chunk_size": (8,)} @@ -556,8 +557,12 @@ def _refuse_sharded_load_options( By rank 0's GPU memory for the FSDP2 load, by the model's experts for the expert-parallel one. """ - # An option passed as None, from_pretrained's own default (device_map=None, say), asks nothing. - passed = {name for name, value in kwargs.items() if value is not None} + # The options that ask anything (_ASKS_NOTHING). + passed = { + name + for name, value in kwargs.items() + if (value not in _ASKS_NOTHING[name] if name in _ASKS_NOTHING else value is not None) + } unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) # torch_dtype is from_pretrained's older name for dtype, and None means "auto" to it. dtypes = [name for name in ("dtype", "torch_dtype") if kwargs.get(name) not in (None, "auto")] @@ -598,6 +603,13 @@ def _refuse_sharded_load_options( f"The {route} load reads the checkpoint's tensors as stored and keeps their dtype, so " f"{name}={kwargs[name]!r} does not apply to it{computes}." ) + other = _FSDP2 if route == _EXPERT_PARALLEL else _EXPERT_PARALLEL + if elsewhere := sorted(passed & (_SHARDED_LOAD_OPTIONS[other] - _SHARDED_LOAD_OPTIONS[route])): + # Not placement: a whole model on every rank, the way out below, would refuse them too. + raise TypeError( + f"{elsewhere} only apply to the {other} load, and the {route} load takes " + f"{sorted(_SHARDED_LOAD_OPTIONS[route])}." + ) if unknown: raise TypeError( f"{unknown} do not apply to the {route} load, which places the model itself; it " diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index d07d27c3f14..21d82bdf87d 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1652,6 +1652,20 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( TypeError, r"\['freeze'\] only apply to the expert-parallel load", ), + ( + DistributionPlan(2, dense="fsdp"), + "causal-lm", + {"require_lm_head": False}, + TypeError, + r"\['require_lm_head'\] only apply to the expert-parallel load, and the FSDP2 load", + ), + ( + DistributionPlan(2, ep_size=2), + "causal-lm", + {"broadcast_chunk_size": None}, + TypeError, + r"\['broadcast_chunk_size'\] only apply to the FSDP2 load, and the expert-parallel", + ), ], ids=[ "another-world", @@ -1662,6 +1676,8 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( "hf-config-on-a-whole-load", "fsdp2-option-on-a-whole-load", "ep-option-on-a-whole-load", + "ep-option-on-the-fsdp2-load", + "fsdp2-chunking-as-none-on-the-ep-load", ], ) def test_a_load_refuses_what_it_cannot_load( From ddb1ba870067d34cf10e4d68a414bcde77be1473 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:26:12 +0000 Subject: [PATCH 17/32] Name a way forward when a sharded load refuses the other one's options The refusal of an option only the other sharded load takes said which load takes it, but not what to do. It now says to drop the options, or to pass a plan that takes that load: one with ep_size > 1 for the expert-parallel load's, and one with ep_size=1 and dense='fsdp' for the FSDP2 load's, since an expert-parallel plan that FSDP2-shards the dense part is refused too. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 10 ++++++++-- tests/unit/torch/utils/test_model_load_utils.py | 10 ++++++++-- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index aa0afcbc965..d0194de734a 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -605,10 +605,16 @@ def _refuse_sharded_load_options( ) other = _FSDP2 if route == _EXPERT_PARALLEL else _EXPERT_PARALLEL if elsewhere := sorted(passed & (_SHARDED_LOAD_OPTIONS[other] - _SHARDED_LOAD_OPTIONS[route])): - # Not placement: a whole model on every rank, the way out below, would refuse them too. + # Not placement: a whole model on every rank, the way out below, would refuse them too. An + # expert-parallel plan that FSDP2-shards the dense part is refused as well (_route). + instead = ( + "pass a plan with ep_size > 1 to shard an MoE's routed experts" + if other == _EXPERT_PARALLEL + else "pass a plan with ep_size=1 and dense='fsdp' to shard the whole model with FSDP2" + ) raise TypeError( f"{elsewhere} only apply to the {other} load, and the {route} load takes " - f"{sorted(_SHARDED_LOAD_OPTIONS[route])}." + f"{sorted(_SHARDED_LOAD_OPTIONS[route])}; drop them, or {instead}." ) if unknown: raise TypeError( diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 21d82bdf87d..97b3f174f84 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1657,14 +1657,20 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( "causal-lm", {"require_lm_head": False}, TypeError, - r"\['require_lm_head'\] only apply to the expert-parallel load, and the FSDP2 load", + ( + r"\['require_lm_head'\] only apply to the expert-parallel load, and the FSDP2 " + r"load .*; drop them, or pass a plan with ep_size > 1" + ), ), ( DistributionPlan(2, ep_size=2), "causal-lm", {"broadcast_chunk_size": None}, TypeError, - r"\['broadcast_chunk_size'\] only apply to the FSDP2 load, and the expert-parallel", + ( + r"\['broadcast_chunk_size'\] only apply to the FSDP2 load, and the " + r"expert-parallel load .*; drop them, or pass a plan with ep_size=1 and dense='fsdp'" + ), ), ], ids=[ From 176c0a79715109c006933815657a8f65892a7b2d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:27:07 +0000 Subject: [PATCH 18/32] Load an expert-parallel plan on a launch with one GPU per process A launch that shows each process one GPU, bound to cuda:0, gives every rank a GPU of its own whatever its local rank, and the GPU refusal before the expert-parallel load takes it so through the shared predicate. A case of the expert-parallel load's routing test, on local rank 1 with one GPU visible, holds it: under the predicate's older reading, by local rank alone, that rank is refused. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 97b3f174f84..3f009cc5bb8 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1320,16 +1320,25 @@ def test_a_load_it_decides_shards_the_experts( assert distribution_plan(model) == DistributionPlan(4, ep_size=4) -@pytest.mark.parametrize("dense", ["replicate", "fsdp"]) +@pytest.mark.parametrize( + ("dense", "one_gpu_per_process"), + [("replicate", False), ("fsdp", False), ("replicate", True)], + ids=["replicate", "fsdp", "one-gpu-per-process"], +) def test_an_expert_parallel_plan_loads_through_the_ep_loader( - run_as, serve, loaders, auto_class, a_gpu, tmp_path, dense + run_as, serve, loaders, auto_class, a_gpu, monkeypatch, tmp_path, dense, one_gpu_per_process ): """An EP plan loads through load_and_prepare_ep, with the plan and the Hub options the load was given, and the checkpoint's config, read once for both, and through no other load. It returns the model unfrozen, as the other loads do, and an option passed as None asks nothing of it. That load replicates the dense part, so a plan that also FSDP2-shards it is refused before anything - loads, naming the ways out.""" - run_as(4) + loads, naming the ways out. A launch that shows each process one GPU, bound to cuda:0, gives a + rank of any local rank a GPU of its own.""" + run_as(4, rank=1 if one_gpu_per_process else 0) + if one_gpu_per_process: + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) config = serve(_deepseek_v3(n_routed_experts=8)) plan = DistributionPlan(4, ep_size=4, dense=dense) options = { From abbdb986e1e1713e0956b0aebd8bacc483a49ca2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:27:35 +0000 Subject: [PATCH 19/32] Say in the changelog that an MoE loads expert-parallel when distributed The backward-breaking entry for a distributed load_pretrained still said an MoE loads without expert parallelism until the expert-parallel load lands. This is that load: an MoE has its routed experts sharded across the ranks, and is refused where it would otherwise load whole on every rank. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1f947089d83..ba54cb16280 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -52,7 +52,7 @@ Changelog - Bump minimum required PyTorch version to 2.10. - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. -- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE loads without expert parallelism, and is refused when ``MODELOPT_EP_DEGREE`` asks for it, until the expert-parallel load lands; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. +- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE has its routed experts sharded across the ranks, and is refused where it would otherwise load whole on every rank; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. - ``load_pretrained`` now records how it distributed the model on the model it returns, even in a single process, so ``fsdp2_wrap`` refuses that model; to shard a model yourself, wrap one loaded with ``from_pretrained``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. From c6ad5950123c5cf3aab45090e9935c3e81e35f72 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:06:59 +0000 Subject: [PATCH 20/32] Shard a decided MoE whole with FSDP2 when its dense part does not fit With the EP degree left to the plan decision, an MoE whose dense part, replicated beside a rank's share of the routed experts, does not fit one GPU was planned expert-parallel with the dense part FSDP2-sharded, a layout the expert-parallel load refuses, as it replicates the dense part. Such a load, which sharded the whole model with FSDP2 before, now failed. An EP degree the decision chose now gives way to FSDP2 for the whole model there; one the caller asked for stands, and the load refuses it naming the ways out. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- modelopt/torch/utils/plugins/model_load_utils.py | 8 +++++++- tests/unit/torch/utils/test_model_load_utils.py | 8 ++++++-- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ba54cb16280..61712ae372c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -52,7 +52,7 @@ Changelog - Bump minimum required PyTorch version to 2.10. - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. -- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE has its routed experts sharded across the ranks, and is refused where it would otherwise load whole on every rank; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. +- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE has its routed experts sharded across the ranks, or the whole model FSDP2-sharded when a rank's share with the rest of it replicated would not fit one GPU, and is refused where it would otherwise load whole on every rank; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. - ``load_pretrained`` now records how it distributed the model on the model it returns, even in a single process, so ``fsdp2_wrap`` refuses that model; to shard a model yourself, wrap one loaded with ``from_pretrained``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d0194de734a..53e713edba8 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1610,6 +1610,7 @@ def _decide_distribution( f"n_routed_experts or num_routed_experts) to shard them by. Set {knob}=1 and " f"{replicate} to load it without expert parallelism, whole on every rank." ) + chosen = ep_size == "auto" # an EP degree no knob or argument set if ep_size == "auto": ep_size = _choose_ep_degree(world_size, num_experts) elif ep_size > 1 and world_size % ep_size: @@ -1632,6 +1633,10 @@ def _decide_distribution( memory = _gpu_memory() too_big = memory is not None and held > _REPLICATED_SHARE * memory dense = "fsdp" if too_big else "replicate" + if chosen and ep_size > 1 and dense == "fsdp": + # The expert-parallel load replicates the dense part, and no load shards both yet: with + # nothing asking for expert parallelism, FSDP2-shard the whole model instead. + ep_size = 1 if experts and asked_to_replicate: print_rank_0( f"{hf_config.model_type}: ep_size=1 and dense='replicate' as asked, so every rank " @@ -1697,7 +1702,8 @@ def resolve_distribution_plan( 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``. + world and the routed-expert count of a model whose config carries a ``base_model_ep_plan``, + or 1 where the dense part is FSDP2-sharded, which the expert-parallel load does not do. * ``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"`` diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 3f009cc5bb8..f95dca87e62 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -670,7 +670,9 @@ def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit( ): """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.""" + whole MoE must fit, and here it does not. The expert-parallel load replicates the dense part, + so an EP degree left to the decision gives way to FSDP2 for the whole model, where one asked + for stands, for the load to refuse.""" run_as(4) config = serve(_qwen3_moe(num_experts=8)) model = _meta(config) @@ -678,7 +680,9 @@ def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit( # 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) + decided = DistributionPlan(4, 4) if dense == "replicate" else DistributionPlan(4, dense="fsdp") + assert resolve_distribution_plan(str(tmp_path))[0] == decided + assert resolve_distribution_plan(str(tmp_path), ep_size=4)[0] == DistributionPlan(4, 4, dense) assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan( 4, dense="fsdp" ) From 6302ffcb0a8103adc4059401a403a951f4a6aa5a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:41:12 +0000 Subject: [PATCH 21/32] Plan without expert parallelism a decided MoE its load would refuse With nothing asking for expert parallelism, the plan decision committed any MoE whose config carries a base_model_ep_plan to the expert-parallel load, the ones it refuses included: a plan that shards no expert weights through grouped_gemm, or styles them by one it drops (GPT-OSS's and Llama-4's below their specs), and a checkpoint format it does not load. Its refusals named no way back, and some came only after the download. An EP degree the decision chose now gives way to none there, decided by the load's own checks, which _ep_plan_refusal now holds for both, and such an MoE is not refused for being left whole, as it cannot be expert-parallel here. An EP degree the caller asked for still meets the load's refusal. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 109 ++++++++++-------- .../unit/torch/utils/test_model_load_utils.py | 30 ++++- 2 files changed, 90 insertions(+), 49 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 53e713edba8..64c11cc8276 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1611,8 +1611,16 @@ def _decide_distribution( f"{replicate} to load it without expert parallelism, whole on every rank." ) chosen = ep_size == "auto" # an EP degree no knob or argument set + loadable = True # whether the expert-parallel load takes this model if ep_size == "auto": ep_size = _choose_ep_degree(world_size, num_experts) + if ep_size > 1 and ( + _ckpt_quant_method(hf_config) not in _EP_SOURCE_FORMATS + or _ep_plan_refusal(ep_plan) is not None + ): + # The expert-parallel load would refuse this model, for some reasons only after the + # checkpoint's download: with nothing asking for expert parallelism, plan without it. + ep_size, loadable = 1, False 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): @@ -1642,7 +1650,7 @@ def _decide_distribution( f"{hf_config.model_type}: ep_size=1 and dense='replicate' as asked, so every rank " "holds the whole MoE." ) - elif experts and ep_size == 1 and dense == "replicate": + elif experts and loadable 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. @@ -1792,6 +1800,57 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st # Not MXFP4: transformers dequantizes it into GPT-OSS's transposed expert layout, whatever the model's. _EP_SOURCE_FORMATS: frozenset[str | None] = frozenset({None}) +# The ep_plan styles the expert-parallel load swaps for its torch_a2a styles. It keeps these and +# grouped_gemm, and drops any other. +_EP_STYLE_SWAP = { + "ep_router": "torch_a2a_router", + "moe_tp_experts": "torch_a2a_experts", + "nemotron_experts_ep": "torch_a2a_experts", +} + + +def _ep_plan_refusal(ep_plan: dict[str, str]) -> str | None: + """Why the expert-parallel load refuses ``ep_plan``'s styles, or ``None``: by the plan alone.""" + if not _routed_expert_keys(ep_plan): + return ( + f"base_model_ep_plan={ep_plan} styles no routed experts ({', '.join(_EXPERTS_STYLES)} " + "or grouped_gemm), so this load would leave them whole on every rank." + ) + # Only grouped_gemm entries shard an expert's weights here: the torch_a2a styles the + # containers take for the load shard nothing themselves (any other style is refused below). + unsharded = sorted( + key + for key in _routed_expert_keys(ep_plan) + if not any( + style == "grouped_gemm" and name.startswith(key + ".") + for name, style in ep_plan.items() + ) + ) + if unsharded: + return ( + f"base_model_ep_plan={ep_plan} shards no weights of {', '.join(unsharded)} through " + "grouped_gemm, the only expert-weight style this load shards by, so every rank would " + "hold every expert." + ) + # The load drops any other style an expert weight takes, so that weight would load whole + # beside its sharded siblings, and the grouped forward would pair slices of different experts. + containers = _routed_expert_keys(ep_plan) + kept_styles = {"grouped_gemm", *_EP_STYLE_SWAP} + dropped = sorted( + name + for name, style in ep_plan.items() + if style not in kept_styles + and name not in containers + and any(name.startswith(key + ".") for key in containers) + ) + if dropped: + return ( + f"base_model_ep_plan={ep_plan} styles {', '.join(dropped)} by a style this load " + "drops, so those expert weights would load whole on every rank beside their sharded " + "siblings." + ) + return None + def _ckpt_quant_method(cfg): """Return the checkpoint's native quantization method, or ``None``. @@ -1892,14 +1951,7 @@ def load_and_prepare_ep( ) world_size, ep_degree, edp = plan.world_size, plan.ep_size, plan.edp hub_options = dict(hub_options or {}) - # The ep_plan styles this load keeps: grouped_gemm, and those it swaps for its torch_a2a styles. - # The checks and the swapped plan both read it. - swap = { - "ep_router": "torch_a2a_router", - "moe_tp_experts": "torch_a2a_experts", - "nemotron_experts_ep": "torch_a2a_experts", - } - kept_styles = {"grouped_gemm", *swap} + swap, kept_styles = _EP_STYLE_SWAP, {"grouped_gemm", *_EP_STYLE_SWAP} # Each rank checks its own transformers, the config it reads or was handed, and the plan against # it: one refusing fails every rank here, before the checkpoint's fetch and the meshes. with _refused_together("the expert-parallel load's checks", alone=False): @@ -1929,43 +1981,8 @@ def load_and_prepare_ep( ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {}) if not ep_plan: raise ValueError("expert-parallel PTQ requires the model to declare base_model_ep_plan") - if not _routed_expert_keys(ep_plan): - raise ValueError( - f"base_model_ep_plan={ep_plan} styles no routed experts ({', '.join(_EXPERTS_STYLES)} " - "or grouped_gemm), so this load would leave them whole on every rank." - ) - # Only grouped_gemm entries shard an expert's weights here: the torch_a2a styles the - # containers take for the load shard nothing themselves (any other style is refused below). - unsharded = sorted( - key - for key in _routed_expert_keys(ep_plan) - if not any( - style == "grouped_gemm" and name.startswith(key + ".") - for name, style in ep_plan.items() - ) - ) - if unsharded: - raise ValueError( - f"base_model_ep_plan={ep_plan} shards no weights of {', '.join(unsharded)} through " - "grouped_gemm, the only expert-weight style this load shards by, so every rank would " - "hold every expert." - ) - # The load drops any other style an expert weight takes, so that weight would load whole - # beside its sharded siblings, and the grouped forward would pair slices of different experts. - containers = _routed_expert_keys(ep_plan) - dropped = sorted( - name - for name, style in ep_plan.items() - if style not in kept_styles - and name not in containers - and any(name.startswith(key + ".") for key in containers) - ) - if dropped: - raise ValueError( - f"base_model_ep_plan={ep_plan} styles {', '.join(dropped)} by a style this load " - "drops, so those expert weights would load whole on every rank beside their sharded " - "siblings." - ) + if (refusal := _ep_plan_refusal(ep_plan)) is not None: + raise ValueError(refusal) # The all-to-all transport needs every rank to hold as many experts, counted as the plan's # decision counts them; a count the config does not name is checked once the experts load. num_experts = _num_experts_from_config(ep_cfg) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index f95dca87e62..6a4a626fb7b 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -600,6 +600,27 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( assert "MODELOPT_" not in str(refused.value), refused.value +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize("refused_by", ["its-plan", "its-format"]) +def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( + run_as, serve, monkeypatch, tmp_path, refused_by +): + """Left to the decision, an MoE the expert-parallel load would refuse, by a plan that shards no + expert weights through grouped_gemm or a checkpoint format it does not load, is planned without + expert parallelism rather than committed to a load that refuses it, for some reasons only after + the download. Nor is it refused for being left whole: it cannot be expert-parallel here.""" + run_as(2) + config = _qwen3_moe(num_experts=8) + if refused_by == "its-plan": + config.base_model_ep_plan = {"layers.*.mlp.experts": "moe_tp_experts"} + else: + config.quantization_config = {"quant_method": "fp8"} + serve(config) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) + + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2) + + @needs_qwen3_moe_ep_plan def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(run_as, serve, monkeypatch, tmp_path): """Asked for no expert parallelism but for dense sharding, the MoE is sharded whole by FSDP2, @@ -782,9 +803,12 @@ def test_the_expert_count_is_read_from_the_config_holding_the_plan( ): """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 - ) + plan = { + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + } + text = SimpleNamespace(base_model_ep_plan=plan, 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 From 85df01f97420693e8188f19d66fcfb7a5b6ab3b2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:41:56 +0000 Subject: [PATCH 22/32] Say in the EP+FSDP2 refusal that a decided EP degree gives way The refusal of an expert-parallel plan that FSDP2-shards the dense part said a plan left to decide shards it when a rank's share would not fit, which the decision no longer does: an EP degree it chose gives way to FSDP2 for the whole model there, and only one the caller asked for meets this refusal. Say that, keeping the 60% and the ways out. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 64c11cc8276..25e90a9eff1 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -493,9 +493,10 @@ def _route(plan: DistributionPlan) -> str | None: raise NotImplementedError( f"{plan!r} FSDP2-shards the dense part, but the expert-parallel load replicates it " "on every rank. Sharding the dense part under expert parallelism is not supported " - "yet. A plan left to decide shards it when a whole copy of it and a rank's share " - "of the routed experts would take more than " - f"{_REPLICATED_SHARE:.0%} of a GPU. Replicate it anyway (MODELOPT_FSDP_SIZE=1, or " + "yet. An EP degree left to the decision gives way to FSDP2 for the whole model when " + "a whole copy of the dense part and a rank's share of the routed experts would take " + f"more than {_REPLICATED_SHARE:.0%} of a GPU; the degree asked for stands. " + "Replicate the dense part anyway (MODELOPT_FSDP_SIZE=1, or " "a plan with dense='replicate'), or shard the whole model with FSDP2 " "(MODELOPT_EP_DEGREE=1 with MODELOPT_FSDP_SIZE=world) if the checkpoint is an " "unquantized causal LM." From fcb111d452f866f1dcb45181879cf2698d9c5017 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:44:42 +0000 Subject: [PATCH 23/32] Name the expert-parallel load in the texts #2779 wrote before it meet the expert-parallel load. A launch with no process group that sets MODELOPT_EP_DEGREE needs only the group, as one that sets MODELOPT_FSDP_SIZE does. The expert-parallel load builds the checkpoint's own class, a VLM's image-text class included, and its refusal of an auto_class names its own way to a whole model. The docstring and the decision's auto_class refusal say which class each sharded load builds. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 19 +++++++++++-------- .../unit/torch/utils/test_model_load_utils.py | 9 ++------- 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 25e90a9eff1..cb2d832cb52 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -591,10 +591,15 @@ def _refuse_sharded_load_options( ) if auto_class not in (None, AutoModelForCausalLM): name = getattr(auto_class, "__name__", repr(auto_class)) + builds = ( + "the checkpoint's own, a VLM's image-text class included" + if route == _EXPERT_PARALLEL + else "a causal LM's or a text-only VLM's image-text class" + ) raise NotImplementedError( - f"The {route} load builds its model class itself, a causal LM's or a text-only VLM's " - f"image-text class, so it takes no auto_class={name}: leave it out to have the load " - f"pick, or load the model whole with it: {_whole_instead(decided)}." + f"The {route} load builds its model class itself, {builds}, so it takes no " + f"auto_class={name}: leave it out to have the load pick, or load the model whole with " + f"it: {_whole_instead(decided, route)}." ) if dtypes: # A sharded load keeps the checkpoint's dtype. name = dtypes[0] @@ -814,16 +819,14 @@ def load_pretrained( if alone and distribution is None: # Knobs asking to shard over the launch are refused, as is a value that is no EP degree. ep, world = _ep_degree(None), int(os.environ["WORLD_SIZE"]) - experts = ep != "auto" and ep > 1 - asked = [f"MODELOPT_EP_DEGREE={env.ep_degree()}"] if experts else [] + asked = [f"MODELOPT_EP_DEGREE={env.ep_degree()}"] if ep != "auto" and ep > 1 else [] if _dense_from_env(world) == "fsdp": asked.append(f"MODELOPT_FSDP_SIZE={env.fsdp_size()}") if asked: - ep_way = ", with MODELOPT_EP_DEGREE=1: no load here shards the routed experts" raise RuntimeError( f"{' and '.join(asked)} ask to shard the model over this " f"{os.environ['WORLD_SIZE']}-rank launch, whose process group is not set up; load " - f"inside distributed_run(), which sets it up{ep_way if experts else ''}." + "inside distributed_run(), which sets it up." ) # None is these options' own default: passed explicitly, it is no option to refuse. for name in ("config", "variant", "hf_config"): @@ -1564,7 +1567,7 @@ def _decide_distribution( # Refused here, rather than build the decision's own pick of class, which may not build. name = getattr(auto_class, "__name__", repr(auto_class)) raise NotImplementedError( - f"A sharded load builds its model class itself, a causal LM's or a text-only VLM's " + f"A sharded load builds its model class itself, a causal LM's or a supported VLM's " f"image-text class, so it takes no auto_class={name} while ep_size={ep_size!r}, " f"dense={dense!r} ask to shard. To give every rank the whole model, plan ep_size=1 and " "dense='replicate' (MODELOPT_EP_DEGREE=1, MODELOPT_FSDP_SIZE=1)." diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 6a4a626fb7b..ede0527e986 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1594,7 +1594,7 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, ( {"auto_class": AutoModel}, NotImplementedError, - r"auto_class=AutoModel: leave it out.*DistributionPlan\(2, dense='replicate'\)", + r"auto_class=AutoModel: leave it out.*DistributionPlan\(2, (1, |dense=)'replicate'\)", ), ({"mtp": True}, NotImplementedError, "MTP head"), ( @@ -2015,12 +2015,7 @@ def test_a_launch_that_sets_no_group_up_loads_whole_on_each_process( RuntimeError, "MODELOPT_FSDP_SIZE=world", ), - ( - None, - {"MODELOPT_EP_DEGREE": "4"}, - RuntimeError, - r"MODELOPT_EP_DEGREE=4 .*distributed_run\(\).*MODELOPT_EP_DEGREE=1", - ), + (None, {"MODELOPT_EP_DEGREE": "4"}, RuntimeError, "MODELOPT_EP_DEGREE=4"), (None, {"MODELOPT_EP_DEGREE": "eigth"}, ValueError, "MODELOPT_EP_DEGREE='eigth'"), ], ids=["several-ranks", "sharded", "fsdp-knob", "ep-knob", "ep-knob-not-a-degree"], From af2fa662e03fefddbabf21044892bd466b9d5640 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:11:38 +0000 Subject: [PATCH 24/32] Give way for a model spec or transformers the EP load refuses, and say so The plan decision's give-way for a model the expert-parallel load would refuse read the plan's styles and the checkpoint format, but not a model spec ruling expert parallelism out (Llama-4's, GPT-OSS's, whose experts the load refuses once built) nor a transformers below the load's floor or the spec's: an EP degree it chose still committed those to a refusal, after the download. _ep_load_refusal now gathers every reason the config tells, and the decision gives way for any of them. Each give-way now says why on rank 0, as the plan's other outcomes do: the load's reason, or for an FSDP2-sharded dense part the knobs that shard the experts anyway. The changelog names only the MoEs the expert-parallel load takes as refused when left whole. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- .../torch/utils/plugins/model_load_utils.py | 41 +++++++++++++++++-- .../unit/torch/utils/test_model_load_utils.py | 20 ++++++--- 3 files changed, 52 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 61712ae372c..0b4335a203c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -52,7 +52,7 @@ Changelog - Bump minimum required PyTorch version to 2.10. - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. -- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE has its routed experts sharded across the ranks, or the whole model FSDP2-sharded when a rank's share with the rest of it replicated would not fit one GPU, and is refused where it would otherwise load whole on every rank; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. +- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE the expert-parallel load takes has its routed experts sharded across the ranks, or the whole model FSDP2-sharded when a rank's share with the rest of it replicated would not fit one GPU, and is refused where it would otherwise load whole on every rank; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. - ``load_pretrained`` now records how it distributed the model on the model it returns, even in a single process, so ``fsdp2_wrap`` refuses that model; to shard a model yourself, wrap one loaded with ``from_pretrained``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index cb2d832cb52..874c2511fff 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -33,6 +33,7 @@ import torch.nn as nn import transformers.distributed as hf_distributed from huggingface_hub import snapshot_download +from packaging.version import parse if TYPE_CHECKING: from collections.abc import Iterable @@ -53,6 +54,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.models.specs import get_spec from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import env, moe from modelopt.torch.utils.distributed import ( @@ -1618,12 +1620,13 @@ def _decide_distribution( loadable = True # whether the expert-parallel load takes this model if ep_size == "auto": ep_size = _choose_ep_degree(world_size, num_experts) - if ep_size > 1 and ( - _ckpt_quant_method(hf_config) not in _EP_SOURCE_FORMATS - or _ep_plan_refusal(ep_plan) is not None - ): + if ep_size > 1 and (refusal := _ep_load_refusal(hf_config, ep_plan)) is not None: # The expert-parallel load would refuse this model, for some reasons only after the # checkpoint's download: with nothing asking for expert parallelism, plan without it. + print_rank_0( + f"{hf_config.model_type}: planned without expert parallelism, which its load " + f"would refuse: {refusal}" + ) ep_size, loadable = 1, False elif ep_size > 1 and world_size % ep_size: raise ValueError(f"{knob}={ep_size} must divide the world size ({world_size}).") @@ -1632,6 +1635,7 @@ def _decide_distribution( f"{knob}={ep_size} must divide the routed experts {hf_config.model_type} shards " f"({num_experts})." ) + dense_chosen = dense == "auto" # a dense layout no knob or argument set if dense == "auto": # 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. @@ -1648,6 +1652,14 @@ def _decide_distribution( if chosen and ep_size > 1 and dense == "fsdp": # The expert-parallel load replicates the dense part, and no load shards both yet: with # nothing asking for expert parallelism, FSDP2-shard the whole model instead. + why = ( + "would not fit one GPU beside a rank's experts" if dense_chosen else "is to be sharded" + ) + print_rank_0( + f"{hf_config.model_type}: FSDP2-sharded whole, without expert parallelism, as the " + f"expert-parallel load replicates the dense part, which {why}. To shard the experts " + f"anyway, the dense part replicated: {knob}={ep_size} with {replicate}." + ) ep_size = 1 if experts and asked_to_replicate: print_rank_0( @@ -1856,6 +1868,27 @@ def _ep_plan_refusal(ep_plan: dict[str, str]) -> str | None: return None +def _ep_load_refusal(hf_config, ep_plan: dict[str, str]) -> str | None: + """Why the expert-parallel load refuses this model, as far as its config tells, or ``None``. + + A model spec that rules it out, the load's transformers floor or the spec's later one, a + checkpoint format the load does not take, and the plan's styles; the rest shows only as the + experts load. + """ + spec = get_spec(hf_config.model_type) + ep = spec.ep_spec if spec is not None else None + if ep is not None and ep.unsupported: + return ep.unsupported + floor = max("5.7", (ep.requires_transformers if ep is not None else None) or "5.7", key=parse) + try: + moe.require_transformers(floor, f"{hf_config.model_type}'s expert-parallel load") + except RuntimeError as e: + return str(e) + if (method := _ckpt_quant_method(hf_config)) not in _EP_SOURCE_FORMATS: + return f"The expert-parallel load cannot dequantize a {method!r} checkpoint." + return _ep_plan_refusal(ep_plan) + + def _ckpt_quant_method(cfg): """Return the checkpoint's native quantization method, or ``None``. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index ede0527e986..437c0c48236 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -72,6 +72,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist import modelopt.torch.utils.moe as moe +from modelopt.torch.models import EPSpec, ModelSpec from modelopt.torch.quantization.plugins import accelerate as accelerate_plugin from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -601,20 +602,26 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( @needs_qwen3_moe_ep_plan -@pytest.mark.parametrize("refused_by", ["its-plan", "its-format"]) +@pytest.mark.parametrize("refused_by", ["its-plan", "its-format", "its-spec", "its-transformers"]) def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( run_as, serve, monkeypatch, tmp_path, refused_by ): - """Left to the decision, an MoE the expert-parallel load would refuse, by a plan that shards no - expert weights through grouped_gemm or a checkpoint format it does not load, is planned without - expert parallelism rather than committed to a load that refuses it, for some reasons only after - the download. Nor is it refused for being left whole: it cannot be expert-parallel here.""" + """Left to the decision, an MoE the expert-parallel load would refuse -- by a plan that shards no + expert weights through grouped_gemm, a checkpoint format it does not load, a model spec ruling + it out (Llama-4's, GPT-OSS's) or a transformers below its floor -- is planned without expert + parallelism rather than committed to a load that refuses it, for some reasons only after the + download. Nor is it refused for being left whole: it cannot be expert-parallel here.""" run_as(2) config = _qwen3_moe(num_experts=8) if refused_by == "its-plan": config.base_model_ep_plan = {"layers.*.mlp.experts": "moe_tp_experts"} - else: + elif refused_by == "its-format": config.quantization_config = {"quant_method": "fp8"} + elif refused_by == "its-spec": + spec = ModelSpec(model_type="qwen3_moe", ep_spec=EPSpec(unsupported="Not yet.")) + monkeypatch.setattr(model_load_utils, "get_spec", lambda model_type: spec) + else: + monkeypatch.setattr(transformers, "__version__", "5.6.2") serve(config) monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) @@ -813,6 +820,7 @@ def test_the_expert_count_is_read_from_the_config_holding_the_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) + monkeypatch.setattr(transformers, "__version__", "5.7.0") # one the expert-parallel load takes run_as(4) assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4) From 80499d6ec41a8b2690bda6e7586d6f7159c87b9d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 01:23:08 +0000 Subject: [PATCH 25/32] Keep the expert-parallel load's transformers floor in one constant The decision's give-way and load_and_prepare_ep each spelled 5.7 out, so raising the load's floor in one place would leave the decision planning expert parallelism for a model the load then refuses. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 874c2511fff..59519c8ac4b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1812,6 +1812,10 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st ) +# The transformers release the expert-parallel load needs, for every model: 5.6 builds a model only +# with the experts implementations it ships, never torch_a2a. A model spec may need a later one. +_EP_MIN_TRANSFORMERS = "5.7" + # The checkpoint formats load_and_prepare_ep dequantizes as it loads, beside unquantized ones (None). # Not MXFP4: transformers dequantizes it into GPT-OSS's transposed expert layout, whatever the model's. _EP_SOURCE_FORMATS: frozenset[str | None] = frozenset({None}) @@ -1879,7 +1883,8 @@ def _ep_load_refusal(hf_config, ep_plan: dict[str, str]) -> str | None: ep = spec.ep_spec if spec is not None else None if ep is not None and ep.unsupported: return ep.unsupported - floor = max("5.7", (ep.requires_transformers if ep is not None else None) or "5.7", key=parse) + stated = ep.requires_transformers if ep is not None else None + floor = max(_EP_MIN_TRANSFORMERS, stated or _EP_MIN_TRANSFORMERS, key=parse) try: moe.require_transformers(floor, f"{hf_config.model_type}'s expert-parallel load") except RuntimeError as e: @@ -1994,7 +1999,8 @@ def load_and_prepare_ep( with _refused_together("the expert-parallel load's checks", alone=False): # 5.7 is the first release that builds a model with the torch_a2a experts. moe.require_transformers( - "5.7", "ModelOpt DP=EP expert-parallel PTQ (transformers MoE EP framework)" + _EP_MIN_TRANSFORMERS, + "ModelOpt DP=EP expert-parallel PTQ (transformers MoE EP framework)", ) if hf_config is None: config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code, **hub_options} From 5651400e186bc644ce40928731b138b792720dea Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 01:23:23 +0000 Subject: [PATCH 26/32] Resolve whether the EP load refuses an MoE however its EP degree was set Only an EP degree the decision picked consulted the load's refusal, so an EP degree of 1 the caller set left an MoE the load refuses (GPT-OSS by its spec, a quantized checkpoint, a plan sharding no expert weights) to the left-whole refusal, which offered expert parallelism as the way out. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 10 ++++++---- tests/unit/torch/utils/test_model_load_utils.py | 8 ++++++-- 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 59519c8ac4b..ccd27838e05 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1617,17 +1617,19 @@ def _decide_distribution( f"{replicate} to load it without expert parallelism, whole on every rank." ) chosen = ep_size == "auto" # an EP degree no knob or argument set - loadable = True # whether the expert-parallel load takes this model + # Why the expert-parallel load would refuse this MoE, or None, whatever EP degree is set: an MoE + # it refuses cannot be expert-parallel here, so is not refused below for being left whole. + refusal = _ep_load_refusal(hf_config, ep_plan) if experts else None if ep_size == "auto": ep_size = _choose_ep_degree(world_size, num_experts) - if ep_size > 1 and (refusal := _ep_load_refusal(hf_config, ep_plan)) is not None: + if ep_size > 1 and refusal is not None: # The expert-parallel load would refuse this model, for some reasons only after the # checkpoint's download: with nothing asking for expert parallelism, plan without it. print_rank_0( f"{hf_config.model_type}: planned without expert parallelism, which its load " f"would refuse: {refusal}" ) - ep_size, loadable = 1, False + ep_size = 1 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): @@ -1666,7 +1668,7 @@ def _decide_distribution( f"{hf_config.model_type}: ep_size=1 and dense='replicate' as asked, so every rank " "holds the whole MoE." ) - elif experts and loadable and ep_size == 1 and dense == "replicate": + elif experts and refusal is None 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. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 437c0c48236..50f2c16c37d 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -602,15 +602,19 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( @needs_qwen3_moe_ep_plan +@pytest.mark.parametrize("ep_degree", [None, "1"], ids=["auto", "set-to-1"]) @pytest.mark.parametrize("refused_by", ["its-plan", "its-format", "its-spec", "its-transformers"]) def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( - run_as, serve, monkeypatch, tmp_path, refused_by + run_as, serve, monkeypatch, tmp_path, refused_by, ep_degree ): """Left to the decision, an MoE the expert-parallel load would refuse -- by a plan that shards no expert weights through grouped_gemm, a checkpoint format it does not load, a model spec ruling it out (Llama-4's, GPT-OSS's) or a transformers below its floor -- is planned without expert parallelism rather than committed to a load that refuses it, for some reasons only after the - download. Nor is it refused for being left whole: it cannot be expert-parallel here.""" + download. Nor is it refused for being left whole, whether the decision or the caller made its + EP degree 1: it cannot be expert-parallel here.""" + if ep_degree is not None: + monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) run_as(2) config = _qwen3_moe(num_experts=8) if refused_by == "its-plan": From f190bfda4995114e1eedf7cf5de8d838847eee8a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 03:54:50 +0000 Subject: [PATCH 27/32] Refuse in the expert-parallel load what the plan decision gives way for The load checked a plan's styles alone, so an EP degree asked for, or a plan handed in, loaded a model whose spec rules expert parallelism out (GPT-OSS's, Llama-4's) or states a later transformers, which the decision plans without expert parallelism. The load now asks the same _ep_load_refusal, after its own floor and format checks, which keep their messages. The load tests' stand-in configs carry a model type for the spec lookup, one no spec names. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 5 +- .../unit/torch/utils/test_model_load_utils.py | 47 ++++++++++++++----- 2 files changed, 38 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ccd27838e05..b61012f3120 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2026,7 +2026,10 @@ def load_and_prepare_ep( ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {}) if not ep_plan: raise ValueError("expert-parallel PTQ requires the model to declare base_model_ep_plan") - if (refusal := _ep_plan_refusal(ep_plan)) is not None: + # The refusals the plan decision gives way for, which an EP degree asked for, or a plan + # handed in, meets here: a model spec ruling expert parallelism out or stating a later + # transformers, and the plan's styles. + if (refusal := _ep_load_refusal(hf_config, ep_plan)) is not None: raise ValueError(refusal) # The all-to-all transport needs every rank to hold as many experts, counted as the plan's # decision counts them; a count the config does not name is checked once the experts load. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 50f2c16c37d..a6cb710dff5 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2067,6 +2067,9 @@ def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( # ------------------------------------------------------------------------------- the plan +# The model type of the configs that stand in for a checkpoint's here: no model spec names it. +_NO_SPEC = "toy_moe" + _EP_PLAN = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", @@ -2184,7 +2187,7 @@ def resolve(path, **hub_options): load_and_prepare_ep( "org/model", DistributionPlan(2, ep_size=2), - hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), + hf_config=SimpleNamespace(model_type=_NO_SPEC, base_model_ep_plan=_EP_PLAN), ) assert resolved == ["org/model"] @@ -2211,7 +2214,7 @@ def resolve(path, **hub_options): load_and_prepare_ep( "org/model", DistributionPlan(4, ep_size=4), - hf_config=SimpleNamespace(text_config=text, num_experts=3), + hf_config=SimpleNamespace(model_type=_NO_SPEC, text_config=text, num_experts=3), ) @@ -2228,14 +2231,17 @@ def resolve(path, **hub_options): ), (SimpleNamespace(), ValueError, "declare base_model_ep_plan"), ( - SimpleNamespace(base_model_ep_plan={"layers.*.self_attn.q_proj": "colwise"}), + SimpleNamespace( + model_type=_NO_SPEC, base_model_ep_plan={"layers.*.self_attn.q_proj": "colwise"} + ), ValueError, "styles no routed experts", ), *( ( SimpleNamespace( - base_model_ep_plan={"layers.*.mlp.experts": "moe_tp_experts", **params} + model_type=_NO_SPEC, + base_model_ep_plan={"layers.*.mlp.experts": "moe_tp_experts", **params}, ), ValueError, r"shards no weights of layers\.\*\.mlp\.experts through grouped_gemm", @@ -2244,11 +2250,17 @@ def resolve(path, **hub_options): ), ( SimpleNamespace( - base_model_ep_plan={**_EP_PLAN, "layers.*.mlp.experts.down_proj": "packed_rowwise"} + model_type=_NO_SPEC, + base_model_ep_plan={**_EP_PLAN, "layers.*.mlp.experts.down_proj": "packed_rowwise"}, ), ValueError, r"styles layers\.\*\.mlp\.experts\.down_proj by a style this load drops", ), + ( + SimpleNamespace(model_type="gpt_oss", base_model_ep_plan=_EP_PLAN), + ValueError, + "the expert-parallel export cannot write that layout", + ), ], ids=[ "gptq", @@ -2260,6 +2272,7 @@ def resolve(path, **hub_options): "experts-container-only", "experts-sharded-otherwise", "experts-sharded-partly-otherwise", + "spec-refuses-ep", ], ) def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( @@ -2269,8 +2282,10 @@ def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( or, MXFP4, dequantize into GPT-OSS's transposed expert layout whatever the model's, or a plan that leaves the experts whole on every rank: one naming no routed experts, or none of whose entries shards their weights by grouped_gemm, the only expert-weight style the load - shards by, or one that styles some of their weights otherwise. The load refuses each before it - reads anything.""" + shards by, or one that styles some of their weights otherwise. So is a model whose spec rules + expert parallelism out (GPT-OSS's), which the plan decision plans without it: an EP degree asked + for, or a plan handed in, meets that refusal here. The load refuses each before it reads + anything.""" with pytest.raises(error, match=match): load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) assert ep_world_of_two == [] @@ -2282,7 +2297,7 @@ def test_an_ep_size_that_does_not_divide_the_experts_is_refused_before_anything_ """The all-to-all transport finds an expert's rank by dividing its id by this rank's share of the experts: with 3 experts over 2 ranks the ranks would disagree on who holds which, and send tokens to the wrong experts.""" - config = SimpleNamespace(base_model_ep_plan=_EP_PLAN, num_experts=3) + config = SimpleNamespace(model_type=_NO_SPEC, base_model_ep_plan=_EP_PLAN, num_experts=3) with pytest.raises(ValueError, match="ep_size=2 must divide the 3 routed experts"): load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) @@ -2305,7 +2320,8 @@ def from_pretrained(path, **kwargs): monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) config = SimpleNamespace( - base_model_ep_plan={"layers.*.mlp.experts.*.gate_up_proj": "grouped_gemm"} + model_type=_NO_SPEC, + base_model_ep_plan={"layers.*.mlp.experts.*.gate_up_proj": "grouped_gemm"}, ) with pytest.raises(ValueError, match="that no module of"): @@ -2342,6 +2358,7 @@ def next_check(*args, **kwargs): monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) class _Config(SimpleNamespace): + model_type = _NO_SPEC base_model_ep_plan = _EP_PLAN # 8 experts over 2 ranks: the stand-in model holds this rank's 4. @@ -2442,7 +2459,7 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", auto_class) - config = SimpleNamespace(dtype=torch.float16, base_model_ep_plan=_EP_PLAN) + config = SimpleNamespace(model_type=_NO_SPEC, dtype=torch.float16, base_model_ep_plan=_EP_PLAN) reads, downloads = [], [] monkeypatch.setattr( "transformers.AutoConfig.from_pretrained", lambda *a, **k: reads.append(k) or config @@ -2628,7 +2645,7 @@ def from_pretrained(path, **kwargs): load_and_prepare_ep( "org/model", DistributionPlan(2, ep_size=2), - hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), + hf_config=SimpleNamespace(model_type=_NO_SPEC, base_model_ep_plan=_EP_PLAN), ) given = {k: v for k, v in vars(loaded["distributed_config"]).items() if v is not None} @@ -2652,7 +2669,9 @@ def _meshes_in_step(*args, **kwargs): def _one_rank_fails_job(rank, size, ckpt, failure, failing_rank, match): path, hf_config, patches = ckpt, None, [] if failure != "unreadable-config": - hf_config = SimpleNamespace(base_model_ep_plan=dict(_EP_PLAN), num_experts=4) + hf_config = SimpleNamespace( + model_type=_NO_SPEC, base_model_ep_plan=dict(_EP_PLAN), num_experts=4 + ) if failure != "unseen-checkpoint": local = mock.patch.object( model_load_utils, "ensure_local_checkpoint", side_effect=lambda p, **k: (None, p) @@ -2735,7 +2754,9 @@ def test_a_load_failing_on_one_rank_fails_every_rank( def _gpu_layout_job(rank, size, visible, bound, shared, refused): - hf_config = SimpleNamespace(base_model_ep_plan=dict(_EP_PLAN), num_experts=4) + hf_config = SimpleNamespace( + model_type=_NO_SPEC, base_model_ep_plan=dict(_EP_PLAN), num_experts=4 + ) reached = None with bounded_collectives(), ExitStack() as stack: # This rank is on cuda:0, bound to it or not yet; the local ranks run 0 and 1. From eda4c883244c2444d1dc87b4531c29248c3ad42f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:20:01 +0000 Subject: [PATCH 28/32] Raise each expert-parallel load refusal as the error it is _ep_load_refusal now returns the error rather than its message: a model spec ruling expert parallelism out and a checkpoint format the load does not take are NotImplementedError, a transformers below the floor a spec states is the RuntimeError require_transformers raised, and a plan's styles stay a ValueError. The load raises it as it is, so a floor refusal has one type whichever floor it is below; the plan decision only prints it. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 22 +++++++++++-------- .../unit/torch/utils/test_model_load_utils.py | 2 +- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index b61012f3120..e3864945383 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1874,26 +1874,30 @@ def _ep_plan_refusal(ep_plan: dict[str, str]) -> str | None: return None -def _ep_load_refusal(hf_config, ep_plan: dict[str, str]) -> str | None: - """Why the expert-parallel load refuses this model, as far as its config tells, or ``None``. +def _ep_load_refusal(hf_config, ep_plan: dict[str, str]) -> Exception | None: + """The error the expert-parallel load refuses this model with, by its config, or ``None``. - A model spec that rules it out, the load's transformers floor or the spec's later one, a - checkpoint format the load does not take, and the plan's styles; the rest shows only as the + A model spec that rules it out (``NotImplementedError``), the load's transformers floor or the + spec's later one (``RuntimeError``), a checkpoint format the load does not take + (``NotImplementedError``), and the plan's styles (``ValueError``); the rest shows only as the experts load. """ spec = get_spec(hf_config.model_type) ep = spec.ep_spec if spec is not None else None if ep is not None and ep.unsupported: - return ep.unsupported + return NotImplementedError(ep.unsupported) stated = ep.requires_transformers if ep is not None else None floor = max(_EP_MIN_TRANSFORMERS, stated or _EP_MIN_TRANSFORMERS, key=parse) try: moe.require_transformers(floor, f"{hf_config.model_type}'s expert-parallel load") except RuntimeError as e: - return str(e) + return e if (method := _ckpt_quant_method(hf_config)) not in _EP_SOURCE_FORMATS: - return f"The expert-parallel load cannot dequantize a {method!r} checkpoint." - return _ep_plan_refusal(ep_plan) + return NotImplementedError( + f"The expert-parallel load cannot dequantize a {method!r} checkpoint." + ) + refusal = _ep_plan_refusal(ep_plan) + return ValueError(refusal) if refusal is not None else None def _ckpt_quant_method(cfg): @@ -2030,7 +2034,7 @@ def load_and_prepare_ep( # handed in, meets here: a model spec ruling expert parallelism out or stating a later # transformers, and the plan's styles. if (refusal := _ep_load_refusal(hf_config, ep_plan)) is not None: - raise ValueError(refusal) + raise refusal # The all-to-all transport needs every rank to hold as many experts, counted as the plan's # decision counts them; a count the config does not name is checked once the experts load. num_experts = _num_experts_from_config(ep_cfg) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index a6cb710dff5..65605923e27 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2258,7 +2258,7 @@ def resolve(path, **hub_options): ), ( SimpleNamespace(model_type="gpt_oss", base_model_ep_plan=_EP_PLAN), - ValueError, + NotImplementedError, "the expert-parallel export cannot write that layout", ), ], From 608b0b8bba2c3eb4f3804ef121d0b838f90d689c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:20:41 +0000 Subject: [PATCH 29/32] Check the transformers floor a model spec states above the load's own The give-way test tripped only the load's own floor. A case with a spec stating 5.13 on transformers 5.10 decides it by the spec's floor, which a string comparison would put below 5.7. Signed-off-by: Shengliang Xu --- .../unit/torch/utils/test_model_load_utils.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 65605923e27..b8056f93074 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -603,16 +603,20 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( @needs_qwen3_moe_ep_plan @pytest.mark.parametrize("ep_degree", [None, "1"], ids=["auto", "set-to-1"]) -@pytest.mark.parametrize("refused_by", ["its-plan", "its-format", "its-spec", "its-transformers"]) +@pytest.mark.parametrize( + "refused_by", + ["its-plan", "its-format", "its-spec", "its-transformers", "its-specs-transformers"], +) def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( run_as, serve, monkeypatch, tmp_path, refused_by, ep_degree ): - """Left to the decision, an MoE the expert-parallel load would refuse -- by a plan that shards no - expert weights through grouped_gemm, a checkpoint format it does not load, a model spec ruling - it out (Llama-4's, GPT-OSS's) or a transformers below its floor -- is planned without expert - parallelism rather than committed to a load that refuses it, for some reasons only after the - download. Nor is it refused for being left whole, whether the decision or the caller made its - EP degree 1: it cannot be expert-parallel here.""" + """Left to the decision, an MoE the expert-parallel load would refuse -- by a plan that + shards no expert weights through grouped_gemm, a checkpoint format it does not load, a model + spec ruling it out (Llama-4's, GPT-OSS's), or a transformers below its floor or below the + later one the spec states (MiniMax-M3-VL's) -- is planned without expert parallelism rather + than committed to a load that refuses it, for some reasons only after the download. Nor is it + refused for being left whole, whether the decision or the caller made its EP degree 1: it + cannot be expert-parallel here.""" if ep_degree is not None: monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) run_as(2) @@ -624,6 +628,11 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( elif refused_by == "its-spec": spec = ModelSpec(model_type="qwen3_moe", ep_spec=EPSpec(unsupported="Not yet.")) monkeypatch.setattr(model_load_utils, "get_spec", lambda model_type: spec) + elif refused_by == "its-specs-transformers": + # Above the load's own 5.7 by version, as 5.10 is not, though not as strings compare. + spec = ModelSpec(model_type="qwen3_moe", ep_spec=EPSpec(requires_transformers="5.13")) + monkeypatch.setattr(model_load_utils, "get_spec", lambda model_type: spec) + monkeypatch.setattr(transformers, "__version__", "5.10.0") else: monkeypatch.setattr(transformers, "__version__", "5.6.2") serve(config) From 5455a0c0355d981edc9abb271ec9cfac686616b7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:44:17 +0000 Subject: [PATCH 30/32] Describe the expert-parallel load in the distributed-loading guide #2779's guide, written while no load sharded the routed experts, said a plan with ep_size > 1 is refused. It now lists the expert-parallel route, how the decision picks an EP degree and when it gives way, the way to a whole model on that route, and what the expert-parallel load builds, takes and refuses beside the FSDP2 load. Signed-off-by: Shengliang Xu --- docs/source/guides/_distributed_loading.rst | 68 ++++++++++++++------- 1 file changed, 45 insertions(+), 23 deletions(-) diff --git a/docs/source/guides/_distributed_loading.rst b/docs/source/guides/_distributed_loading.rst index 230b0249629..30d938102e2 100644 --- a/docs/source/guides/_distributed_loading.rst +++ b/docs/source/guides/_distributed_loading.rst @@ -19,23 +19,36 @@ recorded on the model, where calibration and export read it. with distributed_run(): model = load_pretrained("meta-llama/Llama-3.1-8B") # a whole copy on every rank's GPU # With MODELOPT_FSDP_SIZE=world: FSDP2-sharded over the 8 ranks instead. + moe = load_pretrained("Qwen/Qwen3-30B-A3B") # its routed experts sharded over the 8 ranks Routes ====== +* A plan with ``ep_size > 1`` takes the expert-parallel load (``load_and_prepare_ep``): each rank + holds ``num_experts / ep_size`` of an MoE's routed experts and a replica of the dense part, and + tokens reach their experts' ranks by all-to-all. It replicates the dense part, so a plan that + FSDP2-shards it as well is refused. * A plan with ``dense="fsdp"`` takes the FSDP2 load (``parallel_load_and_prepare_fsdp2``): a causal LM, or a text-only VLM it supports (Qwen3.5), is FSDP2-sharded over the process group, each rank reading its own layers. * Any other plan gives every rank the whole model, through ``auto_class.from_pretrained`` (``auto_class`` defaults to ``AutoModelForCausalLM``). -No load shards the routed experts yet. A plan the load decides keeps the expert-parallel degree at -1, so an MoE loads without expert parallelism; an EP degree ``MODELOPT_EP_DEGREE`` asks for, or a -plan with ``ep_size > 1``, is refused. +Left to decide, the EP degree of an MoE is the largest that divides both its routed experts and the +world size. The decision plans without expert parallelism, and says why, an MoE the expert-parallel +load would refuse: one its model spec rules out (GPT-OSS's, Llama-4's), one below the transformers +release it needs, a quantized checkpoint, or a plan that shards no expert weights. It FSDP2-shards +the whole model instead when a rank's share of the experts with a whole copy of the dense part would +not fit one GPU. An EP degree a knob or argument asks for stands, and the load refuses what it +cannot take. A plan that leaves an expert-parallel MoE whole on every rank, which no knob asked +for, is refused, naming the ways out; ``MODELOPT_EP_DEGREE=1`` with ``MODELOPT_FSDP_SIZE=1`` asks +for it. A refusal on any rank raises on every rank, before the load, and says how to load the model whole: ``MODELOPT_FSDP_SIZE=1`` for a plan the load decided, or -``distribution=DistributionPlan(, dense="replicate")``. +``distribution=DistributionPlan(, dense="replicate")``; for an expert-parallel plan, +``MODELOPT_EP_DEGREE=1`` with ``MODELOPT_FSDP_SIZE=1``, or +``distribution=DistributionPlan(, 1, "replicate")``. The plan is recorded once, so ``fsdp2_wrap`` refuses a model ``load_pretrained`` returned: to shard a model yourself, load it with ``from_pretrained``. @@ -56,15 +69,23 @@ Whole-model loads Sharded loads ============= -* The FSDP2 load builds the model class and places the model itself: a text-only VLM through its - image-text class, anything else through ``AutoModelForCausalLM``. Leave ``auto_class`` out. -* Besides its own options (``cpu_offload``, ``mp_policy``, ``broadcast_chunk_size``) it takes - ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other keyword - rather than ignore it. An option passed as ``None``, ``from_pretrained``'s own default, asks - nothing and is not refused. -* It reads the checkpoint's tensors as stored and keeps their dtype, so it refuses a ``dtype`` and - a quantized checkpoint. It cannot run the model type's loading hook, so it refuses a checkpoint - the hook would build modules for, such as a Nemotron-H MTP head. +* A sharded load builds the model class and places the model itself, so leave ``auto_class`` out. + The FSDP2 load builds a text-only VLM through its image-text class and anything else through + ``AutoModelForCausalLM``; the expert-parallel load builds the checkpoint's own class, a VLM's + image-text class included. +* Besides its own options (the FSDP2 load's ``cpu_offload``, ``mp_policy`` and + ``broadcast_chunk_size``; the expert-parallel load's ``freeze`` and ``require_lm_head``) each + takes ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other + keyword rather than ignore it, the other sharded load's options as that load's own. An option + passed as ``None``, ``from_pretrained``'s own default, asks nothing and is not refused, nor is an + FSDP2 option at that load's own default. The expert-parallel load leaves the model unfrozen, as + the other loads do, unless ``freeze=True``. +* Each keeps the checkpoint's dtype, so it refuses a ``dtype``. The FSDP2 load reads the tensors as + stored, so it refuses a quantized checkpoint; the expert-parallel load loads unquantized ones + only. Neither can run the model type's loading hook, so each refuses a checkpoint the hook would + build modules for, such as a Nemotron-H MTP head. +* Each shards onto every rank's own GPU, and a rank without one is refused with the others. The + expert-parallel load runs with one GPU per process, each bound to ``cuda:0``, too. * A single process, or a group of one, reads no knob: it takes the FSDP2 load only from a plan passed as ``distribution``, and sets a group of one up for it, which needs ``torchrun --nproc_per_node=1``. @@ -76,11 +97,12 @@ With neither a plan nor ``MODELOPT_FSDP_SIZE``, rank 0's GPU memory decides whet is sharded. The model is weighed in the checkpoint's own dtype (its config's, else bfloat16), not one the load asks for: hand a load that widens it, ``dtype=torch.float32`` say, its plan. -A call that places the model (``device_map``, ``quantization_config``, a ``dtype``) loads only -where the model fits one GPU; where it does not, the refusal says the layout was decided rather -than asked for. Pass ``distribution=`` or set ``MODELOPT_FSDP_SIZE=1`` to keep the load whole. An -option at ``from_pretrained``'s own default, ``device_map=None`` included, places nothing and does -not keep the load whole. +A call that places the model (``device_map``, ``quantization_config``, a ``dtype``) loads only where +the model fits one GPU; where it does not, the refusal says the layout was decided rather than asked +for. So does the expert-parallel load's, for an MoE whose experts the decision shards with neither a +plan nor ``MODELOPT_EP_DEGREE``. Pass ``distribution=`` or set ``MODELOPT_FSDP_SIZE=1`` to keep the +load whole. An option at ``from_pretrained``'s own default, ``device_map=None`` included, places +nothing and does not keep the load whole. Configs ======= @@ -95,10 +117,10 @@ Hub options A Hub model ID is made local first, so the export can read the source checkpoint from disk; the ``revision``, ``token``, ``cache_dir`` and ``local_files_only`` options go to that fetch. The config -read and a whole-model load also 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 the model -from the config with ``from_config``, which takes none of them, so such code comes with the -environment's Hub settings. +read, a whole-model load and the expert-parallel load also 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 the model from the config with ``from_config``, which takes none of them, so +such code comes with the environment's Hub settings. Launches without a process group ================================ @@ -113,7 +135,7 @@ Inside ``init_quantized_weights`` ================================= ``init_quantized_weights`` (see :doc:`_compress_quantized_models`) quantizes and compresses the -model in the ``from_pretrained`` it patches, which the FSDP2 load never calls. Inside it, a plan +model in the ``from_pretrained`` it patches, which the sharded loads never call. Inside it, a plan the load decides gives every rank the whole model, with nothing to measure, and the load leaves the device map to ``init_quantized_weights``, which computes it and refuses one it is passed. A knob or a plan that asks to shard is refused there. From c2d18e6c4732b8e49834958dcf7978a015ba5567 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:45:55 +0000 Subject: [PATCH 31/32] Check the expert-parallel load is refused inside init_quantized_weights #2779 refuses a sharded load inside init_quantized_weights, whose from_pretrained the sharded loads never call. The expert-parallel load meets that refusal too, naming its own way to a whole model: an EP degree of 1 as well, for a plan decided from MODELOPT_EP_DEGREE or handed in. Signed-off-by: Shengliang Xu --- .../unit/torch/utils/test_model_load_utils.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b8056f93074..0baf38edc3b 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1843,6 +1843,30 @@ def test_a_sharded_load_is_refused_inside_init_quantized_weights( assert loaders == {} +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize("handed", [False, True], ids=["knob", "handed"]) +def test_an_expert_parallel_load_is_refused_inside_init_quantized_weights( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, handed +): + """The expert-parallel load never calls from_pretrained either, and its way to a whole model on + every rank sets the EP degree to 1 as well, as the decision refuses to replicate an MoE unasked.""" + run_as(2, rank=1) + serve(_qwen3_moe(num_experts=8)) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) + monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") + plan = {"distribution": DistributionPlan(2, 2)} if handed else {} + + with ( + init_quantized_weights(mtq.INT8_DEFAULT_CFG), + pytest.raises( + NotImplementedError, + match=r"expert-parallel load .*init_quantized_weights.*DistributionPlan\(2, 1, 'replicate'\)", + ), + ): + load_pretrained(str(tmp_path), **plan) + assert loaders == {} + + @needs_qwen3_moe_ep_plan @pytest.mark.parametrize( "options", From f7b11feae8e15a71b84f8687d51c76115ca41f73 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 05:20:30 +0000 Subject: [PATCH 32/32] Check an EP degree left to decide gives way to the dense sharding asked for The decided EP degree gives way to FSDP2 for the whole model either when rank 0's GPU memory chose to shard the dense part or when the caller asked for it. The test pinned the EP degree to 1 itself, so only the first was checked; it now also leaves the degree to the decision. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 0baf38edc3b..355be7c3f72 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -642,10 +642,15 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( @needs_qwen3_moe_ep_plan -def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(run_as, serve, monkeypatch, tmp_path): - """Asked for no expert parallelism but for dense sharding, the MoE is sharded whole by FSDP2, - which is not the waste of a replica on every rank.""" - monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") +@pytest.mark.parametrize("ep_degree", [None, "1"], ids=["auto", "set-to-1"]) +def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead( + run_as, serve, monkeypatch, tmp_path, ep_degree +): + """Asked for dense sharding, the MoE is sharded whole by FSDP2, which is not the waste of a + replica on every rank: an EP degree left to the decision gives way, as the expert-parallel + load replicates the dense part.""" + if ep_degree is not None: + monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") run_as(2) serve(_qwen3_moe(num_experts=8))