From 38c9b594ef006336addc6bc2e55e768cff8837f5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 01:45:33 +0000 Subject: [PATCH 01/11] Load an expert-parallel plan through load_pretrained A DistributionPlan that shards the routed experts now loads through load_and_prepare_ep, the dense part replicated or FSDP2-sharded as fsdp_size says; its freeze, require_lm_head and dense_strategy options pass through. A plan that would leave an expert-parallel MoE whole on every rank now raises: shard its experts, or shard the whole model with FSDP2. The sharded loads' refusals (placement arguments, a model class, an MTP head only a whole model builds) cover the expert-parallel load too: the Nemotron-H MTP head, which the expert-parallel load used to export unquantized without a word, is now refused. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- .../torch/utils/plugins/model_load_utils.py | 25 ++++++++++++- .../test_load_pretrained_distribution.py | 36 ++++++++++++++++--- .../test_load_pretrained_distribution.py | 16 +++++++-- 4 files changed, 70 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7be5cc2db0c..fde979af3e1 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -37,7 +37,7 @@ Changelog *Misc* - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs. -- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it FSDP2-shards the model or gives every rank a whole copy, by the plan the new ``resolve_distribution_plan`` decides from the model and the ``MODELOPT_FSDP_SIZE`` variable, or by a plan passed as its new ``distribution`` argument. The loaded model records the plan (``modelopt.torch.utils.distributed.distribution_plan``). Until the expert-parallel load lands, an MoE planned expert-parallel, as the default ``MODELOPT_EP_DEGREE=auto`` plans one whose config has a ``base_model_ep_plan``, is refused before it loads: set ``MODELOPT_EP_DEGREE=1`` (with ``MODELOPT_FSDP_SIZE=world`` to FSDP2-shard it) or pass a plan with ``ep_size=1``. +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it shards an MoE's routed experts across the ranks, FSDP2-shards the model, or gives every rank a whole copy, by the plan the new ``resolve_distribution_plan`` decides from the model and the ``MODELOPT_EP_DEGREE`` / ``MODELOPT_FSDP_SIZE`` variables, or by a plan passed as its new ``distribution`` argument. The loaded model records the plan (``modelopt.torch.utils.distributed.distribution_plan``). **Backward Breaking Changes** diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 61c438485b0..1ff80a3a177 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -419,10 +419,21 @@ def record_unplaced_source_keys( return keys +_EXPERT_PARALLEL = "expert-parallel" _FSDP2 = "FSDP2" # What each sharded load takes besides the checkpoint. These loads build the model class and place # it themselves, so from_pretrained's own arguments (device_map, max_memory, ...) do not reach them. _SHARDED_LOAD_OPTIONS = { + _EXPERT_PARALLEL: frozenset( + { + "trust_remote_code", + "attn_implementation", + "hf_config", + "freeze", + "require_lm_head", + "dense_strategy", + } + ), _FSDP2: frozenset( { "trust_remote_code", @@ -460,7 +471,7 @@ def _route(plan: DistributionPlan) -> str | None: return _FSDP2 if plan.is_fsdp_dense else None -def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) -> nn.Module: +def _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_options) -> nn.Module: """The sharded load ``route`` names, refusing what it cannot honor before it starts.""" from modelopt.torch.models.hf import checkpoint_has_mtp from modelopt.torch.utils import distributed as dist @@ -501,6 +512,18 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) ) trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) + if route == _EXPERT_PARALLEL: + from modelopt.torch.utils.moe import load_and_prepare_ep + + return load_and_prepare_ep( + ckpt_path, + dist.get_active_mo_session(), + plan.ep_size, + trust_remote_code=trust_remote_code, + hf_config=hf_config, + fsdp_dense=plan.is_fsdp_dense, + **kwargs, + ) if hf_config is None: hf_config = AutoConfig.from_pretrained( ckpt_path, trust_remote_code=trust_remote_code, **hub_options diff --git a/tests/gpu/torch/utils/test_load_pretrained_distribution.py b/tests/gpu/torch/utils/test_load_pretrained_distribution.py index 7f9eb7a24f3..d8b98211175 100644 --- a/tests/gpu/torch/utils/test_load_pretrained_distribution.py +++ b/tests/gpu/torch/utils/test_load_pretrained_distribution.py @@ -15,24 +15,26 @@ """``load_pretrained`` in a distributed run, on real GPUs. -Left to decide, it FSDP2-shards a dense model when asked to, and otherwise gives every rank the whole -model on its own GPU. Whichever load runs, every rank ends up with the checkpoint's weights and the -load's record of what it could not place. +Left to decide, it shards an expert-parallel MoE's experts, FSDP2-shards a dense model when asked to, +and otherwise gives every rank the whole model on its own GPU. Whichever load runs, every rank ends +up with the checkpoint's weights and the load's record of what it could not place. """ import os import pytest import torch +from packaging.version import Version transformers = pytest.importorskip("transformers") from torch.distributed.tensor import DTensor -from transformers import AutoModelForCausalLM, LlamaConfig +from transformers import AutoModelForCausalLM, LlamaConfig, Qwen3MoeConfig from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins.model_load_utils import load_pretrained +NUM_EXPERTS = 8 _TINY = { "hidden_size": 64, "intermediate_size": 128, @@ -68,8 +70,19 @@ 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 == [] @@ -89,3 +102,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_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index cd0bd663959..2bafeae5b9a 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -113,6 +113,7 @@ def stub(*args, **kwargs): return stub + monkeypatch.setattr(moe, "load_and_prepare_ep", recording("ep", moe.load_and_prepare_ep)) monkeypatch.setattr( model_load_utils, "parallel_load_and_prepare_fsdp2", @@ -323,6 +324,7 @@ def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path, cpu_dist_workers.run(_every_rank_loads_rank0s_plan_job, ckpt=str(tmp_path), handed=handed) +@pytest.mark.parametrize("distribution", ["ep", "fsdp2"]) @pytest.mark.parametrize( ("change", "error", "match"), [ @@ -344,15 +346,25 @@ def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path, 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, + distribution, + 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)) + serve(_deepseek_v3(n_routed_experts=8) if distribution == "ep" else LlamaConfig(**_TINY)) if change.get("mtp"): monkeypatch.setattr(models_hf, "checkpoint_has_mtp", lambda *a: True) + plan = DistributionPlan(2, ep_size=2) if distribution == "ep" else "fsdp2" with pytest.raises(error, match=match): load_pretrained( From 688734a749b792aa778f84463a0209d1c38b8358 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 17:34:53 +0000 Subject: [PATCH 02/11] Route an expert-parallel plan without the session: the EP loader reads the live process group Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 1ff80a3a177..48376798192 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -517,7 +517,6 @@ def _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_op return load_and_prepare_ep( ckpt_path, - dist.get_active_mo_session(), plan.ep_size, trust_remote_code=trust_remote_code, hf_config=hf_config, From d7012a39e65036583083e2ffa1a6c87c5a1ecb27 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 01:46:30 +0000 Subject: [PATCH 03/11] Route an expert-parallel plan through load_pretrained An expert-parallel DistributionPlan loads through load_and_prepare_ep with the plan itself, so the loader places the dense part as the plan says; dense_strategy no longer passes through. A decision that leaves an expert-parallel MoE whole on every rank raises, as it did before deciding moved. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 45 ++++++------ .../test_load_pretrained_distribution.py | 10 ++- .../test_load_pretrained_distribution.py | 70 ++++++++++++------- .../utils/test_resolve_distribution_plan.py | 42 ++++++++--- 4 files changed, 111 insertions(+), 56 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 48376798192..4a4cf2e1910 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -431,7 +431,6 @@ def record_unplaced_source_keys( "hf_config", "freeze", "require_lm_head", - "dense_strategy", } ), _FSDP2: frozenset( @@ -463,11 +462,7 @@ def _route(plan: DistributionPlan) -> str | None: f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}." ) if plan.is_expert_parallel: - raise NotImplementedError( - f"{plan!r} shards the routed experts, which load_pretrained has no load for. Plan " - "ep_size=1 (MODELOPT_EP_DEGREE=1) to give every rank the whole model, or shard it " - "with FSDP2." - ) + return _EXPERT_PARALLEL return _FSDP2 if plan.is_fsdp_dense else None @@ -516,12 +511,7 @@ def _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_op from modelopt.torch.utils.moe import load_and_prepare_ep return load_and_prepare_ep( - ckpt_path, - plan.ep_size, - trust_remote_code=trust_remote_code, - hf_config=hf_config, - fsdp_dense=plan.is_fsdp_dense, - **kwargs, + ckpt_path, plan, trust_remote_code=trust_remote_code, hf_config=hf_config, **kwargs ) if hf_config is None: hf_config = AutoConfig.from_pretrained( @@ -575,22 +565,24 @@ def load_pretrained( the model over the run's ranks by. Every rank takes rank 0's, or, when rank 0 passes none, :func:`resolve_distribution_plan`'s. By the plan: - * ``dense="fsdp"``: a causal LM FSDP2-sharded over the whole process group, each rank reading its - own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``, ``mp_policy`` and - ``broadcast_chunk_size`` ``kwargs`` may set). A single process launched with + * ``ep_size > 1``: the routed experts sharded as they load, the dense part replicated or + FSDP2-sharded as the plan says (:func:`~modelopt.torch.utils.moe.load_and_prepare_ep`, whose + ``freeze`` and ``require_lm_head`` ``kwargs`` may set); + * otherwise ``dense="fsdp"``: a causal LM FSDP2-sharded over the whole process group, each rank + reading its own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``, + ``mp_policy`` and ``broadcast_chunk_size`` ``kwargs`` may set). A single process launched with ``torchrun --nproc_per_node=1`` has no group to shard through, so the load sets up one of one rank, which :func:`~modelopt.torch.utils.distributed.distributed_run` destroys with the run; a process started without a launcher cannot set one up and is refused. It reads the - checkpoint's tensors as stored, so it refuses a quantized checkpoint. + checkpoint's tensors as stored, so it refuses a quantized checkpoint; * otherwise ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to ``AutoModelForCausalLM``), placed as ``kwargs`` say. Every rank of a distributed run gets the whole model, ``device_map`` defaulting to the rank's GPU. - A plan with ``ep_size > 1`` raises: no load here shards the routed experts. On a distributed run - every rank loads the plan rank 0 was handed. A launch whose processes set no process group up (a - launcher that only exports ``WORLD_SIZE``) has each process load as a single process does, the - whole model with no plan to agree on; a plan for several ranks and a sharded load need the group, - so there they are refused: set it up first, with + On a distributed run every rank loads the plan rank 0 was handed. A launch whose processes set + no process group up (a launcher that only exports ``WORLD_SIZE``) has each process load as a + single process does, the whole model with no plan to agree on; a plan for several ranks and a + sharded load need the group, so there they are refused: set it up first, with :func:`~modelopt.torch.utils.distributed.distributed_run`. The plan is recorded on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), where calibration and export read it back. @@ -757,7 +749,7 @@ def handed() -> DistributionPlan | None: ).model_type ) if route is not None: - model = _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) + model = _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_options) # A sharded load finds the unplaced keys while it places the weights. keys = model._modelopt_unplaced_source_keys else: @@ -1524,6 +1516,13 @@ def _decide_distribution( raise NotImplementedError( f"The dense part of {hf_config.model_type} cannot be sharded {unshardable}.{remedy}" ) + if experts and ep_size == 1 and dense == "replicate": + 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 its experts " + "(MODELOPT_EP_DEGREE above 1, dividing both the expert count and the world size), or " + "shard the whole model with FSDP2." + ) return DistributionPlan(world_size, ep_size, dense) @@ -1584,6 +1583,8 @@ def resolve_distribution_plan( ``config``, the config the load will use when the caller has one, is decided from instead of the checkpoint's, which is then not read: each rank takes a copy, as building writes to it. + + A plan that leaves an expert-parallel MoE whole on every rank raises. """ from modelopt.torch.utils import distributed as dist diff --git a/tests/gpu/torch/utils/test_load_pretrained_distribution.py b/tests/gpu/torch/utils/test_load_pretrained_distribution.py index d8b98211175..69726595ac1 100644 --- a/tests/gpu/torch/utils/test_load_pretrained_distribution.py +++ b/tests/gpu/torch/utils/test_load_pretrained_distribution.py @@ -87,8 +87,14 @@ def _load_and_check(rank, size, ckpt_dir, knobs, expected): 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] diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 2bafeae5b9a..d4b93329560 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -29,10 +29,11 @@ transformers = pytest.importorskip("transformers") -from transformers import AutoModel, LlamaConfig, LlamaForCausalLM, Qwen3MoeConfig +from transformers import AutoModel, DeepseekV3Config, LlamaConfig, LlamaForCausalLM import modelopt.torch.models.hf as models_hf import modelopt.torch.utils.distributed as dist +import modelopt.torch.utils.moe as moe from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins import model_load_utils @@ -48,6 +49,24 @@ } +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 + + @pytest.fixture(autouse=True) def _no_distribution_knobs(monkeypatch): """A shell that exports the knobs would otherwise steer every decision below.""" @@ -282,13 +301,27 @@ class may not build (an encoder-decoder loaded through AutoModel).""" assert loaders == {} -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)) +def test_an_expert_parallel_plan_loads_through_the_ep_loader( + run_as, serve, loaders, auto_class, tmp_path +): + """An EP plan loads through load_and_prepare_ep, with the plan, and through no other load.""" + run_as(4) + serve(_deepseek_v3(n_routed_experts=8)) + plan = DistributionPlan(4, ep_size=4, dense="fsdp") + + model = load_pretrained( + str(tmp_path), + distribution=plan, + trust_remote_code=True, + attn_implementation="eager", + freeze=False, + ) - with pytest.raises(NotImplementedError, match="shards the routed experts"): - load_pretrained(str(tmp_path), auto_class, distribution=DistributionPlan(2, ep_size=2)) + 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 distribution_plan(model) == plan assert loaders == {} assert auto_class.calls == [] @@ -324,7 +357,7 @@ def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path, cpu_dist_workers.run(_every_rank_loads_rank0s_plan_job, ckpt=str(tmp_path), handed=handed) -@pytest.mark.parametrize("distribution", ["ep", "fsdp2"]) +@pytest.mark.parametrize("sharding", ["expert-parallel", "fsdp2"]) @pytest.mark.parametrize( ("change", "error", "match"), [ @@ -346,32 +379,21 @@ def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path, 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, - distribution, - 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(_deepseek_v3(n_routed_experts=8) if distribution == "ep" else 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 distribution == "ep" else "fsdp2" + 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 == {} diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index beab6ae696b..3a0ca6e8ba5 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -388,14 +388,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) @@ -465,6 +457,40 @@ 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 @pytest.mark.parametrize("value", ["0", "", "AUTO"], ids=["zero", "exported-empty", "capitals"]) def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value): From 39e34776aad1594df91593cfa7982e8e497555ea Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 23:43:35 +0000 Subject: [PATCH 04/11] Route an expert-parallel plan to the loader in the same module load_and_prepare_ep lives in model_load_utils with load_pretrained, so the EP route calls it there, and the routing test records the call there. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 4 +--- .../unit/torch/utils/test_load_pretrained_distribution.py | 7 +++++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 4a4cf2e1910..a43a2e8bf97 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -508,8 +508,6 @@ def _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_op trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) if route == _EXPERT_PARALLEL: - from modelopt.torch.utils.moe import load_and_prepare_ep - return load_and_prepare_ep( ckpt_path, plan, trust_remote_code=trust_remote_code, hf_config=hf_config, **kwargs ) @@ -566,7 +564,7 @@ def load_pretrained( :func:`resolve_distribution_plan`'s. By the plan: * ``ep_size > 1``: the routed experts sharded as they load, the dense part replicated or - FSDP2-sharded as the plan says (:func:`~modelopt.torch.utils.moe.load_and_prepare_ep`, whose + FSDP2-sharded as the plan says (:func:`load_and_prepare_ep`, whose ``freeze`` and ``require_lm_head`` ``kwargs`` may set); * otherwise ``dense="fsdp"``: a causal LM FSDP2-sharded over the whole process group, each rank reading its own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``, diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index d4b93329560..bef24728f3c 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -33,7 +33,6 @@ import modelopt.torch.models.hf as models_hf import modelopt.torch.utils.distributed as dist -import modelopt.torch.utils.moe as moe from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins import model_load_utils @@ -132,7 +131,11 @@ def stub(*args, **kwargs): return stub - monkeypatch.setattr(moe, "load_and_prepare_ep", recording("ep", moe.load_and_prepare_ep)) + 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", From 437e6199d4afac332444bccb22103eef1855edea Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 20:33:57 +0000 Subject: [PATCH 05/11] Name the expert count and every way out when a plan would replicate a whole MoE The refusal now names the experts the EP degree must divide, so a world no degree fits shows at once, and every route that still loads the model: FSDP2 for an unquantized checkpoint, a single process, or a replicated plan handed to load_pretrained, which skips the decision. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index a43a2e8bf97..3aa80b45b93 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1518,8 +1518,10 @@ 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 its experts " - "(MODELOPT_EP_DEGREE above 1, dividing both the expert count and the world size), or " - "shard the whole model with FSDP2." + f"(MODELOPT_EP_DEGREE above 1, dividing both its {num_experts} experts and the world " + "size), 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 with " + f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))." ) return DistributionPlan(world_size, ep_size, dense) From a6622bc0b847af01d921bdbb32067825e2940c55 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:28:59 +0000 Subject: [PATCH 06/11] Say why no EP degree can shard experts whose count the config hides With the routed-expert count spelled none of the ways ModelOpt reads it, the refusal of a whole-on-every-rank MoE named its 0 experts and asked for an EP degree dividing them, which none does. It now says the count is unreadable and names the escapes that remain. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 3aa80b45b93..af9c83afc21 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1515,12 +1515,19 @@ def _decide_distribution( f"The dense part of {hf_config.model_type} cannot be sharded {unshardable}.{remedy}" ) if experts and ep_size == 1 and dense == "replicate": + # A count the config spells none of the ways ModelOpt reads leaves no EP degree to name. + shard = ( + f"Shard its experts (MODELOPT_EP_DEGREE above 1, dividing both its {num_experts} " + "experts and the world size), shard" + if num_experts + else "Its config names its routed-expert count none of the ways ModelOpt reads, so no " + "EP degree can shard its experts. 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 its experts " - f"(MODELOPT_EP_DEGREE above 1, dividing both its {num_experts} experts and the world " - "size), 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 with " + 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 with " f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))." ) return DistributionPlan(world_size, ep_size, dense) From 79fa1a8325ed4e6aad43771d7a7c70428ec00997 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 06:51:26 +0000 Subject: [PATCH 07/11] Load an MoE both distribution knobs ask to replicate MODELOPT_EP_DEGREE=1 with MODELOPT_FSDP_SIZE=1, or ep_size=1 with dense="replicate", is a replication the caller chose: the decision now returns it for an expert-parallel MoE, noting that every rank holds the whole MoE, rather than refusing it. A plan that reaches the replication from auto is still refused, and the refusal now names that knob pair as a way out. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 18 +++++++++++++++--- .../utils/test_resolve_distribution_plan.py | 14 ++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index af9c83afc21..0b6c8836e75 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1402,6 +1402,9 @@ def _decide_distribution( ) ep_size = _ep_degree(ep_size) dense = _dense_from_env(world_size) if dense is None else dense + # A replication both knobs ask for is the caller's choice, and loads; one reached from "auto" + # would leave an expert-parallel MoE whole on every rank unasked, and is refused below. + asked_to_replicate = ep_size == 1 and dense == "replicate" if ( auto_class not in (None, AutoModelForCausalLM) and ep_size in ("auto", 1) @@ -1514,7 +1517,13 @@ def _decide_distribution( raise NotImplementedError( f"The dense part of {hf_config.model_type} cannot be sharded {unshardable}.{remedy}" ) - if experts and ep_size == 1 and dense == "replicate": + if experts and asked_to_replicate: + print( + f"{hf_config.model_type}: ep_size=1 and dense='replicate' as asked, so every rank " + "holds the whole MoE.", + flush=True, + ) + elif experts and ep_size == 1 and dense == "replicate": # A count the config spells none of the ways ModelOpt reads leaves no EP degree to name. shard = ( f"Shard its experts (MODELOPT_EP_DEGREE above 1, dividing both its {num_experts} " @@ -1527,7 +1536,8 @@ def _decide_distribution( 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 with " + "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) @@ -1591,7 +1601,9 @@ def resolve_distribution_plan( ``config``, the config the load will use when the caller has one, is decided from instead of the checkpoint's, which is then not read: each rank takes a copy, as building writes to it. - A plan that leaves an expert-parallel MoE whole on every rank raises. + A plan that leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for + exactly that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and + ``MODELOPT_FSDP_SIZE=1``) replicates it, noting that every rank holds the whole MoE. """ from modelopt.torch.utils import distributed as dist diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 3a0ca6e8ba5..d42f05e893b 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -491,6 +491,20 @@ def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(run_as, serve, monkeypat 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", ["0", "", "AUTO"], ids=["zero", "exported-empty", "capitals"]) def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value): From 1120f93efb58764715d02f393dfe3713b824b48a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 06:51:46 +0000 Subject: [PATCH 08/11] Offer no EP degree when none divides the experts and the world When the expert count shares no factor above 1 with the world size, the refusal of an MoE left whole on every rank offered an EP degree above 1 dividing both, which does not exist. It now says that no EP degree can shard the experts on that many ranks, and names only the other ways out. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 27 +++++++++++++------ 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 0b6c8836e75..80099cecb44 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -20,6 +20,7 @@ import inspect import json import logging +import math import os import re from collections import defaultdict @@ -1524,14 +1525,24 @@ def _decide_distribution( flush=True, ) elif experts and ep_size == 1 and dense == "replicate": - # A count the config spells none of the ways ModelOpt reads leaves no EP degree to name. - shard = ( - f"Shard its experts (MODELOPT_EP_DEGREE above 1, dividing both its {num_experts} " - "experts and the world size), shard" - if num_experts - else "Its config names its routed-expert count none of the ways ModelOpt reads, so no " - "EP degree can shard its experts. Shard" - ) + # 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 " From 7861d6374db077fda3d99d620db96488249f7e1a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 06:52:03 +0000 Subject: [PATCH 09/11] Pass the load's Hub options to the expert-parallel load load_pretrained hands token, cache_dir and local_files_only to the FSDP2 load's config read, but not to the expert-parallel load, which reads the config itself when handed a plan and builds through from_pretrained: custom code an auto_map names in another Hub repo was fetched without them. load_and_prepare_ep now takes them as hub_options, for both its config read and from_pretrained. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 27 +++++++++++++------ tests/unit/torch/utils/test_ep_loader.py | 14 ++++++++-- .../test_load_pretrained_distribution.py | 5 +++- 3 files changed, 35 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 80099cecb44..34c92108c36 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -510,7 +510,12 @@ def _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_op hf_config = kwargs.pop("hf_config", None) if route == _EXPERT_PARALLEL: return load_and_prepare_ep( - ckpt_path, plan, trust_remote_code=trust_remote_code, hf_config=hf_config, **kwargs + ckpt_path, + plan, + trust_remote_code=trust_remote_code, + hf_config=hf_config, + hub_options=hub_options, + **kwargs, ) if hf_config is None: hf_config = AutoConfig.from_pretrained( @@ -603,11 +608,12 @@ def load_pretrained( * A Hub model ID is made local first (:func:`ensure_local_checkpoint`, which takes the ``revision``, ``token``, ``cache_dir`` and ``local_files_only`` ``kwargs``): the export reads the source checkpoint from disk. Under ``torch.distributed`` every rank must call this. The - config read and the whole-model load still take all but ``revision``: an ``auto_map`` entry - can name custom code in another Hub repo, which Transformers fetches with them. A sharded - load builds its model from the config with ``from_config``, which takes none, so such code - comes with the environment's Hub settings. A ``subfolder`` is part of where the checkpoint - is: the config, the loading hook and the record below read it, as the load does. + config read, the whole-model load and the expert-parallel load still take all but + ``revision``: an ``auto_map`` entry can name custom code in another Hub repo, which + Transformers fetches with them. The FSDP2 load builds its model from the config with + ``from_config``, which takes none, so such code comes with the environment's Hub settings. + A ``subfolder`` is part of where the checkpoint is: the config, the loading hook and the + record below read it, as the load does. * The model type's loading hook runs (:func:`~modelopt.torch.models.hf.prepare_model_for_loading`), e.g. to keep the MTP head a Nemotron-H checkpoint carries. ``model_type`` selects it; when omitted, the type of the @@ -2381,6 +2387,7 @@ def load_and_prepare_ep( freeze: bool = True, hf_config=None, require_lm_head: bool = True, + hub_options: dict | None = None, ) -> nn.Module: """Load + shard a MoE causal LM for expert-parallel PTQ via the DP=EP all-to-all transport. @@ -2389,7 +2396,9 @@ def load_and_prepare_ep( does (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), as :func:`load_pretrained` does for every load. The model builds with its LM head where an auto class builds one; ``require_lm_head=False`` builds a VLM without it, through ``AutoModel`` (a - causal LM always keeps its head). + causal LM always keeps its head). ``hub_options`` (``token``, ``cache_dir``, + ``local_files_only``, as :func:`load_pretrained` takes them) go to its config read and + ``from_pretrained``, which fetch the custom code an ``auto_map`` names in another Hub repo. Builds an ``(edp, ep)`` device mesh and loads experts EP-sharded over the ``ep`` submesh (``enable_expert_parallel=True`` -- each rank owns ``num_experts / ep`` experts as grouped-GEMM @@ -2474,8 +2483,9 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A from transformers.distributed import DistributedConfig world_size, ep_degree, edp = plan.world_size, plan.ep_size, plan.edp + hub_options = dict(hub_options or {}) if hf_config is None: - config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code} + config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code, **hub_options} if attn_implementation is not None: config_kwargs["attn_implementation"] = attn_implementation hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) @@ -2577,6 +2587,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A "device_mesh": ep_submesh, "distributed_config": distributed_config, "trust_remote_code": trust_remote_code, + **hub_options, } if not _dc_has_tp_plan: from_pretrained_kwargs["tp_plan"] = "auto" diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 7971645d6d2..1cb7a0dff2f 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -518,7 +518,8 @@ def next_check(experts): def test_the_weights_load_in_the_checkpoints_dtype_through_its_class(ep_world_of_two, monkeypatch): """The checkpoint's ``dtype`` reaches ``from_pretrained``, and the class that builds it is chosen - with the trust the caller gave its code.""" + with the trust the caller gave its code. The Hub options reach the config read and + ``from_pretrained``, which fetch the custom code an ``auto_map`` names in another repo.""" class _LoadedError(Exception): """Stops the load at from_pretrained.""" @@ -540,12 +541,21 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): ) monkeypatch.setattr(model_load_utils, "_auto_class", auto_class) config = SimpleNamespace(dtype=torch.float16, base_model_ep_plan=_EP_PLAN) + reads = [] + monkeypatch.setattr( + "transformers.AutoConfig.from_pretrained", lambda *a, **k: reads.append(k) or config + ) with pytest.raises(_LoadedError): load_and_prepare_ep( - "org/model", DistributionPlan(2, ep_size=2), trust_remote_code=True, hf_config=config + "org/model", + DistributionPlan(2, ep_size=2), + trust_remote_code=True, + hub_options={"local_files_only": True}, ) assert loaded["dtype"] is torch.float16 + assert [read["local_files_only"] for read in reads] == [True] + assert loaded["local_files_only"] is True assert classes == [True] diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index bef24728f3c..4cd7672b7be 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -307,7 +307,8 @@ class may not build (an encoder-decoder loaded through AutoModel).""" def test_an_expert_parallel_plan_loads_through_the_ep_loader( run_as, serve, loaders, auto_class, tmp_path ): - """An EP plan loads through load_and_prepare_ep, with the plan, and through no other load.""" + """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.""" run_as(4) serve(_deepseek_v3(n_routed_experts=8)) plan = DistributionPlan(4, ep_size=4, dense="fsdp") @@ -318,12 +319,14 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader( trust_remote_code=True, attn_implementation="eager", freeze=False, + local_files_only=True, ) 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 == [] From 12075cf32f24245e627ef9348584cb93834b53aa Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 18:54:39 +0000 Subject: [PATCH 10/11] Pass the expert-parallel load's Hub options to its download too load_and_prepare_ep took the Hub options for its config read and from_pretrained, but resolved a Hub ID through snapshot_download with none: a token for a gated repo, a cache_dir or local_files_only were dropped by the one call that fetches the weights. load_pretrained hands it a local snapshot, so only a direct caller met it. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 13 +++++++------ tests/unit/torch/utils/test_ep_loader.py | 11 ++++++++--- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 34c92108c36..c8c1720c375 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -74,15 +74,15 @@ logger = logging.getLogger(__name__) -def _resolve_checkpoint_dir(ckpt_path: str, rank: int) -> str: +def _resolve_checkpoint_dir(ckpt_path: str, rank: int, **hub_options) -> str: """Local dir for ``ckpt_path``; resolves an HF Hub ID (rank 0 downloads, others wait).""" if os.path.isdir(ckpt_path): return ckpt_path if rank == 0: - snapshot_download(ckpt_path) + snapshot_download(ckpt_path, **hub_options) if is_initialized(): barrier() - return snapshot_download(ckpt_path) + return snapshot_download(ckpt_path, **hub_options) def _materialize_meta_model(model: nn.Module, device: torch.device) -> None: @@ -2397,8 +2397,9 @@ def load_and_prepare_ep( :func:`load_pretrained` does for every load. The model builds with its LM head where an auto class builds one; ``require_lm_head=False`` builds a VLM without it, through ``AutoModel`` (a causal LM always keeps its head). ``hub_options`` (``token``, ``cache_dir``, - ``local_files_only``, as :func:`load_pretrained` takes them) go to its config read and - ``from_pretrained``, which fetch the custom code an ``auto_map`` names in another Hub repo. + ``local_files_only``, as :func:`load_pretrained` takes them) go to the checkpoint's download, its + config read and ``from_pretrained``; the last two fetch the custom code an ``auto_map`` names in + another Hub repo. Builds an ``(edp, ep)`` device mesh and loads experts EP-sharded over the ``ep`` submesh (``enable_expert_parallel=True`` -- each rank owns ``num_experts / ep`` experts as grouped-GEMM @@ -2542,7 +2543,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A # The weights load, and the checkpoint recorded for the export to read the unplaced tensors back # from, take the local snapshot: a Hub ID recorded there would leave the export nothing to read. # Resolved once the plan and the config pass, so a refused load downloads nothing. - ckpt_path = _resolve_checkpoint_dir(ckpt_path, torch.distributed.get_rank()) + ckpt_path = _resolve_checkpoint_dir(ckpt_path, torch.distributed.get_rank(), **hub_options) # Two INDEPENDENT meshes (NeMo-AutoModel style), named by the model part they govern -- two # factorizations of the same world: diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 1cb7a0dff2f..93c1d969c12 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -518,8 +518,9 @@ def next_check(experts): def test_the_weights_load_in_the_checkpoints_dtype_through_its_class(ep_world_of_two, monkeypatch): """The checkpoint's ``dtype`` reaches ``from_pretrained``, and the class that builds it is chosen - with the trust the caller gave its code. The Hub options reach the config read and - ``from_pretrained``, which fetch the custom code an ``auto_map`` names in another repo.""" + 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.""" @@ -541,10 +542,13 @@ 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 = [] + reads, downloads = [], [] monkeypatch.setattr( "transformers.AutoConfig.from_pretrained", lambda *a, **k: reads.append(k) or config ) + monkeypatch.setattr( + model_load_utils, "_resolve_checkpoint_dir", lambda p, r, **k: downloads.append(k) or p + ) with pytest.raises(_LoadedError): load_and_prepare_ep( @@ -555,6 +559,7 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): ) 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 c3c1c5ab66f7104853fbee8a38ced8950ba2a0c8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 18:55:20 +0000 Subject: [PATCH 11/11] Print the replicated-MoE note through print_rank_0 The rest of the plan decision reports through the rank-aware helpers; this note printed on whichever rank ran it, rank 0 only because the decision runs inside broadcast_from's callable. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c8c1720c375..01ad8618a8d 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1525,10 +1525,9 @@ def _decide_distribution( f"The dense part of {hf_config.model_type} cannot be sharded {unshardable}.{remedy}" ) if experts and asked_to_replicate: - print( + print_rank_0( f"{hf_config.model_type}: ep_size=1 and dense='replicate' as asked, so every rank " - "holds the whole MoE.", - flush=True, + "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