diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d480a908903..db0d62604bf 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 dcc30d949ce..22926e63450 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 @@ -89,15 +90,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: @@ -435,10 +436,20 @@ def record_unplaced_source_keys( return keys +_EXPERT_PARALLEL = "expert-parallel" _FSDP2 = "FSDP2" # What each sharded load takes besides the checkpoint. These loads build the model class and place # it themselves, so from_pretrained's own arguments (device_map, max_memory, ...) do not reach them. _SHARDED_LOAD_OPTIONS = { + _EXPERT_PARALLEL: frozenset( + { + "trust_remote_code", + "attn_implementation", + "hf_config", + "freeze", + "require_lm_head", + } + ), _FSDP2: frozenset( { "trust_remote_code", @@ -466,11 +477,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 @@ -504,7 +511,7 @@ def _refuse_sharded_load_options(route, auto_class, kwargs) -> 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 @@ -519,6 +526,15 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) ) trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) + if route == _EXPERT_PARALLEL: + return load_and_prepare_ep( + ckpt_path, + plan, + trust_remote_code=trust_remote_code, + hf_config=hf_config, + hub_options=hub_options, + **kwargs, + ) if hf_config is None: hf_config = AutoConfig.from_pretrained( ckpt_path, trust_remote_code=trust_remote_code, **hub_options @@ -571,22 +587,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:`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. @@ -608,11 +626,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 @@ -766,7 +785,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: @@ -1406,6 +1425,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) @@ -1516,6 +1538,38 @@ def _decide_distribution( raise NotImplementedError( f"The dense part of {hf_config.model_type} cannot be sharded {unshardable}.{remedy}" ) + if experts and asked_to_replicate: + print_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) @@ -1577,6 +1631,10 @@ def resolve_distribution_plan( ``config``, the config the load will use when the caller has one, is decided from instead of the checkpoint's, which is then not read: each rank takes a copy, as building writes to it. + + A plan that leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for + exactly that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and + ``MODELOPT_FSDP_SIZE=1``) replicates it, noting that every rank holds the whole MoE. """ # 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. @@ -2335,6 +2393,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. @@ -2343,7 +2402,10 @@ 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 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 @@ -2430,8 +2492,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) @@ -2488,7 +2551,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: @@ -2533,6 +2596,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/gpu/torch/utils/test_load_pretrained_distribution.py b/tests/gpu/torch/utils/test_load_pretrained_distribution.py index 7f9eb7a24f3..69726595ac1 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,14 +70,31 @@ def _load_and_check(rank, size, ckpt_dir, knobs, expected): if expected == "replicated": assert not sharded assert all(param.device == device for param in params.values()) - else: + elif expected == "fsdp2": assert {name for name in params if ".layers." in name} <= sharded + else: + # Each rank holds its contiguous share of the routed experts, and the rest whole. + experts = {name for name in params if ".mlp.experts." in name} + assert experts and not sharded - experts + per_rank = NUM_EXPERTS // size + for name in experts: + local = params.pop(name) + local = local.to_local() if isinstance(local, DTensor) else local + assert local.shape[0] == per_rank + owned = whole_params[name][rank * per_rank : (rank + 1) * per_rank] + torch.testing.assert_close(local.cpu(), owned) for name, param in params.items(): # the same order on every rank: full_tensor gathers torch.testing.assert_close(_full(param).cpu(), whole_params[name]) assert model._modelopt_unplaced_source_keys == [] assert model._modelopt_source_checkpoint == ckpt_dir - dense = "replicate" if expected == "replicated" else "fsdp" - assert distribution_plan(model) == DistributionPlan(size, dense=dense) + assert ( + distribution_plan(model) + == { + "replicated": DistributionPlan(size), + "fsdp2": DistributionPlan(size, dense="fsdp"), + "expert-parallel": DistributionPlan(size, ep_size=size), + }[expected] + ) finally: for knob in knobs: del os.environ[knob] @@ -89,3 +108,16 @@ def test_a_dense_model_loads_whole_on_every_rank(need_2_gpus, dist_workers, tmp_ def test_a_dense_model_asked_to_shard_loads_fsdp2_sharded(need_2_gpus, dist_workers, tmp_path): _save(LlamaConfig(**_TINY), tmp_path) dist_workers.run(_load_and_check, str(tmp_path), {"MODELOPT_FSDP_SIZE": "world"}, "fsdp2") + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("5.7"), + reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7", +) +def test_an_expert_parallel_moe_loads_with_its_experts_sharded(need_2_gpus, dist_workers, tmp_path): + if NUM_EXPERTS % dist_workers.world_size: + pytest.skip(f"{NUM_EXPERTS} experts do not shard over {dist_workers.world_size} GPUs") + config = Qwen3MoeConfig(moe_intermediate_size=32, num_experts=NUM_EXPERTS, **_TINY) + config.num_experts_per_tok = 2 + _save(config, tmp_path) + dist_workers.run(_load_and_check, str(tmp_path), {}, "expert-parallel") diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 7971645d6d2..93c1d969c12 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -518,7 +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.""" + 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.""" @@ -540,12 +542,25 @@ 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, "_resolve_checkpoint_dir", lambda p, r, **k: downloads.append(k) or 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] diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 3c8ed95b137..f5f2cd10b2c 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -29,7 +29,7 @@ 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 @@ -48,6 +48,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.""" @@ -113,6 +131,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", @@ -281,13 +304,30 @@ 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 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") - with pytest.raises(NotImplementedError, match="shards the routed experts"): - load_pretrained(str(tmp_path), auto_class, distribution=DistributionPlan(2, ep_size=2)) + model = load_pretrained( + str(tmp_path), + distribution=plan, + 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 == [] @@ -323,6 +363,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("sharding", ["expert-parallel", "fsdp2"]) @pytest.mark.parametrize( ("change", "error", "match"), [ @@ -344,22 +385,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, 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 == {} diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index beab6ae696b..d42f05e893b 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,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", ["0", "", "AUTO"], ids=["zero", "exported-empty", "capitals"]) def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value):