diff --git a/docs/source/guides/_distributed_loading.rst b/docs/source/guides/_distributed_loading.rst index 30d938102e2..bac542f7d8f 100644 --- a/docs/source/guides/_distributed_loading.rst +++ b/docs/source/guides/_distributed_loading.rst @@ -110,7 +110,8 @@ Configs ``hf_config`` is the config ``resolve_distribution_plan`` returned with the plan handed in as ``distribution``: the load builds from it as given, so without that plan it is refused. A ``config`` passed for the load, or a path to one, is the one the plan is decided from, and a sharded -load takes it as its ``hf_config``. +load takes it as its ``hf_config``: the expert-parallel load a copy of it the model spec prepares +(``apply_distributed_support``), as deciding does. Hub options =========== diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 4599813e12a..fba161de33d 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -211,7 +211,11 @@ class EPSpec(SpecSection): config to load. Reserved for checkpoints carrying their own modeling code, where the work is patching classes or installing on-the-fly weight decompression. Runs after the declarative fields above have been applied, so a ``prepare`` that returns another config than the one it - was handed must carry over the ``base_model_ep_plan`` they attached.""" + was handed must carry over the ``base_model_ep_plan`` they attached. + + It runs before the run's plan is decided, wherever the plan could still be expert-parallel, so + a run whose EP degree comes out at 1 has run it too: it must leave a load of the whole model, + on every rank or FSDP2-sharded, as it would be without it, and may run more than once.""" prepared_ep_plan: dict[str, str] | None = None """The ``base_model_ep_plan`` ``prepare`` attaches, for a model whose plan only ``prepare`` diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 5dd6ce764ea..20d5ef7be04 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -156,6 +156,30 @@ def _ep_spec(model_type: str): return getattr(get_spec(model_type), "ep_spec", None) +def _ep_floor(model_type: str) -> str: + """The transformers release ``model_type``'s expert parallelism needs. + + The expert-parallel load's own, or its spec's when that is later. A model whose spec states no + release, or that has no spec, still needs the load's: auto must pick no EP degree below it, + rather than reach the load and fail there. + """ + ep = _ep_spec(model_type) + stated = ep.requires_transformers if ep is not None else None + return max(_EP_MIN_TRANSFORMERS, stated or _EP_MIN_TRANSFORMERS, key=parse) + + +def _fsdp2_way_out(fsdp: str = "MODELOPT_FSDP_SIZE=world") -> str: + """FSDP2 as a refusal names it, the way around expert parallelism ``fsdp`` asks for.""" + # The FSDP2 load builds a causal LM and reads its tensors as stored, so a quantized or + # multimodal checkpoint (GPT-OSS's MXFP4, Llama-4's) needs a single process, as does one with + # an MTP head the load cannot build (--use_fsdp2 refuses Nemotron-H's). + return ( + f"(--use_fsdp2 in examples/hf_ptq/hf_ptq.py, or {fsdp}) if the FSDP2 load supports the " + "checkpoint, which takes an unquantized causal LM with no MTP head it cannot build (such " + "as Nemotron-H's)" + ) + + def check_distributed_support(model_type: str) -> None: """Refuse ``model_type``'s expert parallelism if its spec does, or below its transformers floor. @@ -167,26 +191,16 @@ def check_distributed_support(model_type: str) -> None: expert-parallel plan handed in, meets these errors; both name the way around expert parallelism. """ ep = _ep_spec(model_type) - # The FSDP2 load builds a causal LM and reads its tensors as stored, so a quantized or - # multimodal checkpoint (GPT-OSS's MXFP4, Llama-4's) needs a single process, as does one with - # an MTP head the load cannot build (--use_fsdp2 refuses Nemotron-H's). without_ep = ( - "without expert parallelism: in a single process, or FSDP2-sharded (--use_fsdp2 in " - "examples/hf_ptq/hf_ptq.py, or MODELOPT_FSDP_SIZE=world) if the FSDP2 load supports the " - "checkpoint, which takes an unquantized causal LM with no MTP head it cannot build (such " - "as Nemotron-H's)" + f"without expert parallelism: in a single process, or FSDP2-sharded {_fsdp2_way_out()}" ) if ep is not None and ep.unsupported: raise NotImplementedError( f"{model_type} cannot be quantized expert-parallel yet: {ep.unsupported} Quantize it " f"{without_ep}." ) - # A model whose spec states no release, or that has no spec, still needs the load's floor: auto - # must pick no EP degree below it, rather than reach the load and fail there. - stated = ep.requires_transformers if ep is not None else None - floor = max(_EP_MIN_TRANSFORMERS, stated or _EP_MIN_TRANSFORMERS, key=parse) try: - require_transformers(floor, f"{model_type} expert-parallel PTQ") + require_transformers(_ep_floor(model_type), f"{model_type} expert-parallel PTQ") except RuntimeError as e: raise RuntimeError(f"{e} Or quantize it {without_ep}.") from None diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c9dfe4295b0..122455daaa0 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -18,6 +18,7 @@ import contextlib import copy import dataclasses +import json import logging import math import os @@ -33,7 +34,6 @@ 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 @@ -54,7 +54,6 @@ 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 ( @@ -68,12 +67,21 @@ from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 from modelopt.torch.utils.moe import ( + _EP_MIN_TRANSFORMERS, _EXPERTS_STYLES, + _SUB_CONFIG_ATTRS, _choose_ep_degree, _ep_config, + _ep_floor, + _ep_spec, + _fsdp2_way_out, _num_experts_from_config, _routed_expert_keys, _routed_expert_modules, + _spec_ep_plan, + _spec_refuses_ep, + apply_distributed_support, + check_distributed_support, register_torch_a2a, ) from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -642,8 +650,15 @@ def _sharded_load_kwargs( hub_options, decided: bool, measured: bool, + prepared: bool = False, ) -> dict: - """What the sharded load ``route`` names takes, refusing what it cannot honor.""" + """What the sharded load ``route`` names takes, refusing what it cannot honor. + + ``prepared`` is whether an ``hf_config`` in ``kwargs`` is prepared by the model's spec already: + one the caller or the plan decision handed in is; the checkpoint's own config read for the load, + or a copy of a ``config`` passed for it, is not. Left out, it is taken as not, as preparing + again is harmless (:attr:`~modelopt.torch.models.EPSpec.prepare`) and skipping it is not. + """ from modelopt.torch.models.hf import checkpoint_has_mtp _refuse_sharded_load_options(route, auto_class, kwargs, decided, measured) @@ -659,6 +674,18 @@ def _sharded_load_kwargs( hf_config = kwargs.pop("hf_config", None) if route == _EXPERT_PARALLEL: _refuse_a_rank_without_its_gpu(route, decided) + if hf_config is None or not prepared: + if hf_config is None: + hf_config = AutoConfig.from_pretrained( + ckpt_path, trust_remote_code=trust_remote_code, **hub_options + ) + hf_config = apply_distributed_support( + model_type, hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code + ) + else: + # Prepared with its plan, and preparing can patch modeling code, so it is not prepared + # again; a model its spec refuses is refused all the same. + check_distributed_support(model_type) # Unfrozen unless asked, as the other loads return the model: which load runs is decided. kwargs.setdefault("freeze", False) return { @@ -945,11 +972,15 @@ def handed() -> DistributionPlan | None: (route == _FSDP2 and _dense_from_env(plan.world_size) == "auto") or (route == _EXPERT_PARALLEL and _ep_degree(None) == "auto") ) + # Whether hf_config is prepared by the model's spec: handed in with its plan, it is; read or + # copied here, the expert-parallel load prepares it, once its own checks pass. + prepared = True 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"] if "hf_config" not in kwargs: kwargs["hf_config"] = copy.deepcopy(config) + prepared = False if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. if route is not None and hf_config is None and config is None: @@ -958,6 +989,7 @@ def handed() -> DistributionPlan | None: ckpt_path, trust_remote_code=trust_remote_code, **hub_options ) kwargs["hf_config"] = hf_config + prepared = False chosen = hf_config if hf_config is not None else config model_type = getattr(chosen, "model_type", None) or ( _read_config(ckpt_path, auto_class, trust_remote_code, **hub_options).model_type @@ -975,6 +1007,7 @@ def handed() -> DistributionPlan | None: hub_options, decided, measured, + prepared=prepared, ) ) if route is not None: @@ -1517,11 +1550,20 @@ def _dense_from_env(world_size: int) -> str: ) -def _ep_degree(ep_size) -> int | str: - """``ep_size`` as an EP degree or ``"auto"``; ``None`` reads ``MODELOPT_EP_DEGREE``.""" +# How ep_size and MODELOPT_EP_DEGREE spell "auto", for every reading of them to agree. +_EP_AUTO = ("auto", "0") + + +def _ep_degree(ep_size, model_type: str | None = None) -> int | str: + """``ep_size`` as an EP degree or ``"auto"``; ``None`` reads ``MODELOPT_EP_DEGREE``. + + Auto is 1 for a model whose spec refuses expert parallelism on the installed transformers, for + good or below the release it needs: only an EP degree asked for explicitly meets that refusal. + Without ``model_type``, where no config is read yet, auto stays auto. + """ value = env.ep_degree() if ep_size is None else ep_size - if str(value).lower() in ("auto", "0"): - return "auto" + if str(value).lower() in _EP_AUTO: + return 1 if model_type is not None and _spec_refuses_ep(model_type) else "auto" try: degree = int(value) except (ValueError, TypeError): # TypeError: an argument that is no number at all, as [2] @@ -1551,11 +1593,14 @@ def _decide_distribution( ep_size=single if isinstance(single, int) else 1, dense="replicate" if dense in (None, "auto") else dense, ) - ep_size = _ep_degree(ep_size) + # Read before _ep_degree, which takes auto as 1 for a model whose spec refuses expert + # parallelism: only an EP degree the caller set can ask for a replication. + ep_asked = str(env.ep_degree() if ep_size is None else ep_size).strip().lower() not in _EP_AUTO + ep_size = _ep_degree(ep_size, hf_config.model_type) 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" + asked_to_replicate = ep_asked and ep_size == 1 and dense == "replicate" if ( (auto_class not in (None, AutoModelForCausalLM) or _init_quantized_weights()) and ep_size in ("auto", 1) @@ -1593,6 +1638,12 @@ def _decide_distribution( ) from e ep_plan = getattr(_ep_config(hf_config), "base_model_ep_plan", None) or {} experts = _routed_expert_modules(model, ep_plan) + if not experts and ep_size == 1: + # An EP degree of 1 skips the model's spec (resolve_distribution_plan), so a model whose plan + # only its spec supplies shows no routed experts: look for them through the plan the spec + # would supply, or such an MoE would load whole on every rank where one with its own refuses. + ep_plan = _spec_ep_plan(hf_config.model_type) + experts = _routed_expert_modules(model, ep_plan) num_experts = _num_experts_from_config(_ep_config(hf_config)) if experts else 0 if ep_plan and not experts and ep_size != 1 and _num_experts_from_config(_ep_config(hf_config)): # The plan names routed experts the built model has none of under those names: deciding @@ -1666,32 +1717,72 @@ 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 refusal is None and ep_size == 1 and dense == "replicate": + elif ( + experts + and ep_size == 1 + and dense == "replicate" + # One whose spec refuses expert parallelism is refused below saying so, whatever else the + # expert-parallel load would refuse it for. + and (refusal is None or _spec_refuses_ep(hf_config.model_type)) + ): + # The ways out besides sharding the experts, which a spec refusing EP rules out. + others = ( + f"shard the whole model with FSDP2 {_fsdp2_way_out(fsdp)}, 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'))." + ) + ep = _ep_spec(hf_config.model_type) + if ep is not None and ep.unsupported: + # NotImplementedError, as check_distributed_support refuses the same spec. + raise NotImplementedError( + f"{hf_config.model_type} cannot be quantized expert-parallel yet: {ep.unsupported} " + f"The plan replicates it at world_size={world_size}, so every rank would hold the " + f"whole MoE. Instead, {others}" + ) + if ( + _spec_refuses_ep(hf_config.model_type) + and num_experts + and math.gcd(num_experts, world_size) > 1 + ): + # Below the release its expert parallelism needs: an EP degree above 1 is refused too. + # Where no EP degree could shard the experts anyway, the arms below say so instead, as + # upgrading would not change it. An EP degree the caller set to 1 stays 1 after an + # upgrade, so it has to be raised as well. + degree = ( + f" and set {knob} above 1, dividing both its {num_experts} experts and the world " + "size," + if ep_asked + else "" + ) + raise RuntimeError( + f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 " + f"at world_size={world_size}, so every rank would hold the whole MoE, and this " + "transformers cannot shard its experts. Upgrade transformers to " + f"{_ep_floor(hf_config.model_type)} or later{degree} to shard them, or {others}" + ) # 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 = ( + way_out = ( "Its config names its routed-expert count none of the ways ModelOpt reads, so " - "no EP degree can shard its experts. Shard" + f"no EP degree can shard its experts. {others[:1].upper()}{others[1:]}" ) elif math.gcd(num_experts, world_size) == 1: - shard = ( + way_out = ( 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" + f"so none can shard its experts on {world_size} ranks. " + f"{others[:1].upper()}{others[1:]}" ) else: - shard = ( + way_out = ( f"Shard its experts ({knob} above 1, dividing both its {num_experts} experts and " - "the world size), shard" + f"the world size), {others}" ) # 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 " - f"process, or replicate it on every rank outright: {knob}=1 with {replicate}, or " - f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))." + f"world_size={world_size}, so every rank would hold the whole MoE. {way_out}" ) return DistributionPlan(world_size, ep_size, dense) @@ -1727,7 +1818,8 @@ def resolve_distribution_plan( * ``ep_size``: the expert-parallel degree, or ``"auto"``, the largest that divides both the world and the routed-expert count of a model whose config carries a ``base_model_ep_plan``, - or 1 where the dense part is FSDP2-sharded, which the expert-parallel load does not do. + or 1 where the dense part is FSDP2-sharded, which the expert-parallel load does not do, and + for a model whose spec refuses expert parallelism on the installed transformers. * ``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"`` @@ -1751,8 +1843,17 @@ def resolve_distribution_plan( ``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 + Unless the call or rank 0's knobs rule expert parallelism out (a single process, an + ``auto_class`` the sharded loads refuse, or an EP degree of 1, which auto is for a model whose + spec refuses expert parallelism on the installed transformers), every rank applies the model's + spec to a copy of the config first + (:func:`~modelopt.torch.utils.moe.apply_distributed_support`): for several models it is what + supplies the expert-parallel plan, and an expert-parallel plan returns that copy. Any other + plan returns the config without it, which the export would otherwise write out; the spec's + ``prepare`` has run for it all the same, which its contract allows + (:attr:`~modelopt.torch.models.EPSpec.prepare`). 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 @@ -1786,15 +1887,43 @@ def checkpoint_read() -> bool: dist.all_reduce("checkpoint read", checkpoint_read, all) ckpt_path, hf_config = read[0] world_size = dist.size() + # Whether the run rules expert parallelism out is rank 0's call, as the plan is: a rank that + # skipped the spec while the others applied it would load without the expert-parallel plan + # they shard by, and fail once the mesh is up, leaving them waiting in a collective. A model + # class the sharded loads refuse rules it out too: the decision then plans the whole model on + # every rank, or refuses, without the spec, whose prepare could fail, or patch the process, + # for a load that never shards. + prepared = [hf_config] + if dist.broadcast_from( + "spec", + lambda: ( + world_size > 1 + and auto_class in (None, AutoModelForCausalLM) + and _ep_degree(ep_size, hf_config.model_type) != 1 + ), + ): + # Before deciding: deciding without the spec would read an expert-parallel MoE as dense. A + # rank that cannot apply it fails every rank, as a checkpoint read does above. + def spec_applied() -> bool: + prepared[0] = apply_distributed_support( + hf_config.model_type, + copy.deepcopy(hf_config), + ckpt_path=ckpt_path, + trust_remote_code=trust_remote_code, + ) + return True + + dist.all_reduce("model spec", spec_applied, all) plan = dist.broadcast_from( "plan", lambda: _decide_distribution( - ckpt_path, hf_config, world_size, ep_size, dense, trust_remote_code, auto_class + ckpt_path, prepared[0], world_size, ep_size, dense, trust_remote_code, auto_class ), ) if plan != DistributionPlan(): print_rank_0(f"Distribution: {plan!r}") - return plan, hf_config + # Only the expert-parallel load shards by the plan the spec attached. + return plan, prepared[0] if plan.is_expert_parallel else hf_config def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_state) -> None: @@ -1812,10 +1941,6 @@ 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}) @@ -1875,21 +2000,11 @@ def _ep_plan_refusal(ep_plan: dict[str, str]) -> str | 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 (``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. + A checkpoint format the load does not take (``NotImplementedError``) and the plan's styles + (``ValueError``); the rest shows only as the experts load. A model spec ruling expert + parallelism out, or the transformers release it needs, is not asked here: the plan decision and + the load ask the spec first (:func:`~modelopt.torch.utils.moe.check_distributed_support`). """ - 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 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 e if (method := _ckpt_quant_method(hf_config)) not in _EP_SOURCE_FORMATS: return NotImplementedError( f"The expert-parallel load cannot dequantize a {method!r} checkpoint." @@ -1914,6 +2029,43 @@ def _ckpt_quant_method(cfg): return None +_UNDECLARED = object() + + +def _declared_ep_plan(ckpt_dir: str, hf_config, ep_cfg): + """The plan the checkpoint's ``config.json`` declares for ``ep_cfg``, or ``_UNDECLARED``. + + ``ep_cfg`` is ``hf_config`` or one of its nested configs, and the plan is read from the same + place in the file. Read from the file, not the config: a plan the config carries besides is + ModelOpt's -- the one a model spec attaches + (:func:`~modelopt.torch.utils.moe.apply_distributed_support`) -- the config class's, or one set + on a config passed to the load, and the loaded model keeps none of them: an export writes out + what the checkpoint declares, if anything. A file this rank cannot read as a JSON object + declares none, with a warning on this rank, which may be the only one that cannot: what the + checkpoint declared is lost to its model's config then. + """ + path = os.path.join(ckpt_dir, "config.json") + keeps_none = ( + f"so rank {dist.rank()}'s model keeps no expert-parallel plan of its own, and an export " + "from it writes none" + ) + try: + with open(path, encoding="utf-8") as f: + source = json.load(f) + except (OSError, ValueError) as e: + warnings.warn(f"Cannot read {path} ({e}), {keeps_none}.") + return _UNDECLARED + if not isinstance(source, dict): + warnings.warn(f"{path} is not a JSON object, {keeps_none}.") + return _UNDECLARED + if ep_cfg is not hf_config: + attr = next((a for a in _SUB_CONFIG_ATTRS if getattr(hf_config, a, None) is ep_cfg), None) + source = source.get(attr) if attr else None + if isinstance(source, dict) and "base_model_ep_plan" in source: + return source["base_model_ep_plan"] + return _UNDECLARED + + def _transformers_mesh_builder(): """The module whose ``initialize_distributed_mesh`` ``from_pretrained`` calls, or ``None``. @@ -2011,6 +2163,10 @@ def load_and_prepare_ep( if attn_implementation is not None: config_kwargs["attn_implementation"] = attn_implementation hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) + # A model spec ruling expert parallelism out, or stating a later transformers, refuses it + # first, as load_pretrained does, for an EP degree asked for or a plan handed in: no other + # checkpoint of the model, nor a plan of its own, would load. + check_distributed_support(hf_config.model_type) quant_method = _ckpt_quant_method(hf_config) if quant_method not in _EP_SOURCE_FORMATS: # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights. @@ -2028,9 +2184,7 @@ 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") - # 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. + # The plan's styles, which the plan decision gives way for too. if (refusal := _ep_load_refusal(hf_config, ep_plan)) is not None: raise refusal # The all-to-all transport needs every rank to hold as many experts, counted as the plan's @@ -2127,14 +2281,16 @@ def gpu() -> tuple[bool, bool, str | None]: # some plans style only their parameters, grouped_gemm. swapped_plan.update(dict.fromkeys(_routed_expert_keys(ep_plan), "torch_a2a_experts")) from_pretrained_kwargs["experts_implementation"] = "torch_a2a" - # The checkpoint's own plan, for the config the model keeps (below); None: a class default. - declared_plan = vars(ep_cfg).get("base_model_ep_plan") ep_cfg.base_model_ep_plan = swapped_plan from_pretrained_kwargs["config"] = hf_config print_rank_0(f"[ep] ep_plan={swapped_plan} experts_implementation=torch_a2a") # A rank failing in the load fails every rank here, before the warmup's collectives: # from_pretrained runs no collective, so the others finish it and meet the failure in the gather. with _refused_together("the expert-parallel load", alone=False): + # The checkpoint's own plan, for the config the model keeps (below). The read raises on no + # failure of the file's, but sits in this block regardless: a rank failing there alone + # would leave the others in the load. + declared_plan = _declared_ep_plan(ckpt_path, hf_config, ep_cfg) # The class the plan's decision chose, with its LM head where an auto class builds one. loader = _auto_class(hf_config, require_lm_head, trust_remote_code=trust_remote_code) print_rank_0(f"[ep] loader={loader.__name__} (require_lm_head={require_lm_head})") @@ -2146,7 +2302,7 @@ def gpu() -> tuple[bool, bool, str | None]: kept_config = getattr(model, "config", None) if kept_config is not None: kept_ep_config = _ep_config(kept_config) - if declared_plan is None: + if declared_plan is _UNDECLARED: vars(kept_ep_config).pop("base_model_ep_plan", None) else: kept_ep_config.base_model_ep_plan = declared_plan diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py index eb2685c68b2..e5be242cdd2 100644 --- a/tests/gpu/torch/utils/test_model_load_utils.py +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -27,7 +27,7 @@ from packaging.version import Version from safetensors.torch import load_file from torch.distributed.tensor import DTensor -from transformers import AutoModelForCausalLM, LlamaConfig, Qwen3MoeConfig +from transformers import AutoModelForCausalLM, DeepseekV3Config, LlamaConfig, Qwen3MoeConfig from modelopt.torch.export.unified_export_hf import export_hf_checkpoint from modelopt.torch.utils.distributed import ( @@ -36,6 +36,7 @@ distribution_plan, set_distribution_plan, ) +from modelopt.torch.utils.moe import _EP_MIN_TRANSFORMERS from modelopt.torch.utils.plugins.model_load_utils import ( load_pretrained, parallel_load_and_prepare_fsdp2, @@ -216,7 +217,8 @@ def _load_and_check(rank, size, ckpt_dir, knobs, expected): 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. + could not place. An expert-parallel model also computes what the whole model does: each rank's + tokens reach the ranks holding their experts and come back. """ device = torch.device(f"cuda:{rank}") os.environ.update(knobs) @@ -256,6 +258,15 @@ def _load_and_check(rank, size, ckpt_dir, knobs, expected): "expert-parallel": DistributionPlan(size, ep_size=size), }[expected] ) + if expected == "expert-parallel": + # Every rank sends its own batch through the experts, as calibration does. + whole.to(device) + generator = torch.Generator().manual_seed(rank) + input_ids = torch.randint(0, _TINY["vocab_size"], (2, 16), generator=generator) + with torch.no_grad(): + sharded = model(input_ids.to(device)).logits + reference = whole(input_ids.to(device)).logits + torch.testing.assert_close(sharded.float(), reference.float(), atol=2e-2, rtol=2e-2) finally: for knob in knobs: del os.environ[knob] @@ -271,14 +282,46 @@ def test_a_dense_model_asked_to_shard_loads_fsdp2_sharded(need_2_gpus, dist_work dist_workers.run(_load_and_check, str(tmp_path), {"MODELOPT_FSDP_SIZE": "world"}, "fsdp2") +def _qwen3_moe(): + """Qwen3-MoE, whose config ships its own expert-parallel plan from transformers 5.7.""" + config = Qwen3MoeConfig(moe_intermediate_size=32, num_experts=NUM_EXPERTS, **_TINY) + config.num_experts_per_tok = 2 + return config + + +def _deepseek_v3(): + """DeepSeek-V3 with no plan of its own, as before transformers 5.13, so only its model spec + makes it expert-parallel. Every layer is an MoE layer, its attention shrunk to the tiny width.""" + config = DeepseekV3Config( + n_routed_experts=NUM_EXPERTS, + num_experts_per_tok=2, + moe_intermediate_size=32, + first_k_dense_replace=0, + n_group=1, + topk_group=1, + q_lora_rank=32, + kv_lora_rank=16, + qk_rope_head_dim=8, + qk_nope_head_dim=8, + v_head_dim=16, + **_TINY, + ) + config.base_model_ep_plan = None # saved in config.json: no plan of its own on any transformers + return config + + @pytest.mark.skipif( - Version(transformers.__version__) < Version("5.7"), - reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7", + Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS), + reason=f"expert parallelism needs transformers >= {_EP_MIN_TRANSFORMERS}", +) +@pytest.mark.parametrize( + "config", [_qwen3_moe, _deepseek_v3], ids=["its-own-plan", "its-spec-plan"] ) -def test_an_expert_parallel_moe_loads_with_its_experts_sharded(need_2_gpus, dist_workers, tmp_path): +def test_an_expert_parallel_moe_loads_with_its_experts_sharded( + need_2_gpus, dist_workers, tmp_path, config +): + """Whether the plan comes with the model's config or only from its model spec.""" 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) + _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 c090a8052f1..63894086116 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -25,10 +25,11 @@ import copy import dataclasses import inspect +import json import math import os import warnings -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack, contextmanager, nullcontext from types import SimpleNamespace from unittest import mock @@ -57,6 +58,7 @@ AutoModelForSpeechSeq2Seq, DeepseekV3Config, FineGrainedFP8Config, + GptOssConfig, LlamaConfig, LlamaForCausalLM, Phi4MultimodalConfig, @@ -73,11 +75,12 @@ 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.models import get_spec 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 from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan +from modelopt.torch.utils.moe import _EP_MIN_TRANSFORMERS from modelopt.torch.utils.plugins import hf_checkpoint_utils, model_load_utils from modelopt.torch.utils.plugins.model_load_utils import ( _auto_class, @@ -225,6 +228,12 @@ def test_meta_qwen_vlm_preserves_architecture_and_vision(model_type): "vocab_size": 32, } +# Applying a model's spec enforces the transformers floor the expert-parallel load needs. +needs_ep_transformers = pytest.mark.skipif( + Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS), + reason=f"expert parallelism needs transformers >= {_EP_MIN_TRANSFORMERS}", +) + needs_qwen3_moe_ep_plan = pytest.mark.skipif( Version(transformers.__version__) < Version("5.7"), reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7", @@ -428,9 +437,9 @@ class counts only when the checkpoint's code is trusted. An encoder-decoder mode def test_a_single_process_holds_the_whole_model( run_as, serve, monkeypatch, tmp_path, world_size, knob ): - """With no ranks to shard over, even an MoE loads whole, and no model is built to decide it. A - shell still exporting a multi-GPU run's knobs must not reshape that: only the caller's own - arguments ask for FSDP2 on one rank.""" + """With no ranks to shard over, even an MoE loads whole, and no model is built, nor its spec + consulted, to decide it. A shell still exporting a multi-GPU run's knobs must not reshape that: + only the caller's own arguments ask for FSDP2 on one rank.""" if world_size is not None: run_as(world_size) if knob is not None: @@ -439,8 +448,12 @@ def test_a_single_process_holds_the_whole_model( monkeypatch.setattr( model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model") ) + monkeypatch.setattr( + model_load_utils, "_spec_refuses_ep", lambda *a: pytest.fail("consulted the spec") + ) assert resolve_distribution_plan(str(tmp_path)) == (DistributionPlan(), served) + assert resolve_distribution_plan(str(tmp_path), ep_size="auto")[0] == DistributionPlan() assert resolve_distribution_plan(str(tmp_path), dense="fsdp")[0] == DistributionPlan( dense="fsdp" ) @@ -457,7 +470,7 @@ def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, c plan, hf_config = resolve_distribution_plan(str(tmp_path)) assert plan == DistributionPlan(4, ep_size=4, dense="replicate") - assert hf_config is served + assert hf_config.base_model_ep_plan == served.base_model_ep_plan assert ("Distribution:" in capsys.readouterr().out) == (rank == 0) @@ -471,7 +484,10 @@ def test_an_moe_without_an_ep_plan_is_not_expert_parallel(run_as, serve, tmp_pat assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) -@pytest.mark.parametrize("degree", [None, "1"], ids=["auto", "set-to-1"]) +@pytest.mark.parametrize( + "degree", + [pytest.param(None, id="auto", marks=needs_ep_transformers), pytest.param("1", id="set-to-1")], +) def test_an_ep_plan_naming_no_module_of_the_model_is_refused( run_as, serve, monkeypatch, tmp_path, degree ): @@ -496,7 +512,10 @@ def test_an_ep_plan_naming_no_module_of_the_model_is_refused( assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) -@pytest.mark.parametrize("degree", [None, "1"], ids=["auto", "set-to-1"]) +@pytest.mark.parametrize( + "degree", + [pytest.param(None, id="auto", marks=needs_ep_transformers), pytest.param("1", id="set-to-1")], +) def test_routed_experts_the_config_does_not_count_are_refused( run_as, serve, monkeypatch, tmp_path, degree ): @@ -521,6 +540,7 @@ def test_routed_experts_the_config_does_not_count_are_refused( assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) +@needs_ep_transformers @pytest.mark.parametrize("source", ["MODELOPT_EP_DEGREE", "ep_size"]) def test_an_ep_degree_that_does_not_divide_the_world_names_where_it_came_from( run_as, serve, monkeypatch, tmp_path, source @@ -572,31 +592,199 @@ 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 +@needs_ep_transformers +def test_the_spec_is_applied_before_the_decision(run_as, serve, tmp_path): + """DeepSeek-V3's plan comes only from its spec. Deciding before applying it would read the MoE + as dense; applied first, it shards, and the config handed on carries the spec's plan.""" + run_as(4) + serve(_deepseek_v3(n_routed_experts=8)) + + plan, hf_config = resolve_distribution_plan(str(tmp_path)) + + assert plan == DistributionPlan(4, 4) + assert hf_config.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + + +def test_only_an_expert_parallel_plan_carries_the_spec_plan_on(run_as, serve, tmp_path): + """No EP degree divides 3 experts over 2 ranks, so the whole model is FSDP2-sharded: the config + handed on keeps no expert-parallel plan, which the export would write into config.json.""" + run_as(2) + serve(_deepseek_v3(n_routed_experts=3)) + + plan, hf_config = resolve_distribution_plan(str(tmp_path), dense="fsdp") + + assert plan == DistributionPlan(2, dense="fsdp") + assert hf_config.base_model_ep_plan is None + + @pytest.mark.parametrize( - ("config", "ep_degree"), + ("config", "floor", "version", "error", "refusal"), [ - 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"), + pytest.param( + GptOssConfig(num_local_experts=4, **_TINY), + None, + None, + NotImplementedError, + "cannot be quantized expert-parallel yet", + id="spec-refuses-ep", + ), + pytest.param( + _qwen3_moe(num_experts=4), + "99", + None, + RuntimeError, + "requires transformers >= 99", + id="below-the-specs-floor", + ), + pytest.param( + _deepseek_v3(n_routed_experts=4), + None, + "5.6.2", + RuntimeError, + r"requires transformers >= 5\.7", + id="below-the-expert-parallel-loads-floor", + ), + ], +) +def test_a_run_without_expert_parallelism_needs_no_spec( + run_as, serve, monkeypatch, tmp_path, config, floor, version, error, refusal +): + """GPT-OSS cannot be expert-parallel, which its spec says by refusing an EP degree above 1, nor + can a model on a transformers older than its spec, or the expert-parallel load, asks for: a + plan the spec supplies does not lift the load's floor. A run that rules expert parallelism out + does not ask the spec, nor does one leaving the EP degree to auto, which picks none for + either: asking for FSDP2 alone shards it whole.""" + if floor is not None: + monkeypatch.setattr(get_spec(config.model_type).ep_spec, "requires_transformers", floor) + if version is not None: + monkeypatch.setattr(transformers, "__version__", version) + run_as(2) + serve(config) + + with pytest.raises(error, match=refusal): + resolve_distribution_plan(str(tmp_path), ep_size=2) + plan, _ = resolve_distribution_plan(str(tmp_path), ep_size=1, dense="fsdp") + assert plan == DistributionPlan(2, dense="fsdp") + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + assert resolve_distribution_plan(str(tmp_path))[0] == plan + monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") + assert resolve_distribution_plan(str(tmp_path))[0] == plan + + +# The refusals of an MoE left whole on every rank, by whether its spec allows expert parallelism. +_CAN_BE_EP = "can be expert-parallel, but the plan resolved .* divid(?:es|ing) both its {} experts" +_NO_EP_YET = ( + "cannot be quantized expert-parallel yet: .* The plan replicates .* FSDP2 .* if the " + "FSDP2 load supports the checkpoint, which takes .*, load it in a single process, or " + r"replicate .*DistributionPlan\(2, 1, 'replicate'\)" +) +_NEEDS_LATER = ( + "can be expert-parallel, but the plan resolved .* this transformers cannot shard its experts. " + r"Upgrade transformers to 5\.7 or later to shard them, or shard the whole model with FSDP2" +) +# The same, for an EP degree the caller set to 1, which an upgrade alone leaves at 1. +_NEEDS_LATER_AND_A_DEGREE = ( + "can be expert-parallel, but the plan resolved .* this transformers cannot shard its experts. " + r"Upgrade transformers to 5\.7 or later and set MODELOPT_EP_DEGREE above 1, dividing both its " + "8 experts and the world size, to shard them, or shard the whole model with FSDP2" +) + + +@pytest.mark.parametrize( + ("config", "ep_degree", "refusal", "version"), + [ + pytest.param( + _deepseek_v3(n_routed_experts=3), + None, + _CAN_BE_EP.format(3), + None, + id="auto_finds_no_degree", + ), + pytest.param( + _qwen3_moe(num_experts=8), "1", _CAN_BE_EP.format(8), None, id="explicit_degree_one" + ), + pytest.param( + _deepseek_v3(n_routed_experts=8), + "1", + _CAN_BE_EP.format(8), + None, + id="explicit_degree_one_spec_plan", + ), + pytest.param( + # MXFP4, as GPT-OSS's checkpoints are, which the expert-parallel load refuses too. + GptOssConfig( + num_local_experts=4, quantization_config={"quant_method": "mxfp4"}, **_TINY + ), + None, + _NO_EP_YET, + None, + id="spec_refuses_ep", + ), + pytest.param( + _deepseek_v3(n_routed_experts=8), None, _NEEDS_LATER, "5.6.2", id="below_the_ep_floor" + ), + pytest.param( + _qwen3_moe(num_experts=8), + 1, + _CAN_BE_EP.format(8), + None, + id="explicit_degree_one_as_an_argument", + ), + pytest.param( + _deepseek_v3(n_routed_experts=3), + None, + _CAN_BE_EP.format(3), + "5.6.2", + id="below_the_ep_floor_no_degree", + ), + pytest.param( + _deepseek_v3(n_routed_experts=8), + "1", + _NEEDS_LATER_AND_A_DEGREE, + "5.6.2", + id="below_the_ep_floor_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 + run_as, serve, monkeypatch, tmp_path, config, ep_degree, refusal, version ): - """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. 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) + """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none, + whether the plan comes with the config (Qwen3-MoE) or only from the model's spec, which an EP + degree of 1 skips (DeepSeek-V3). With nothing sharded either, every rank would hold the whole + MoE: that must fail, not run, and name the expert count an EP degree has to divide. GPT-OSS + carries a plan too, but its spec refuses expert parallelism, so auto picks no EP degree for it + and this check refuses it as an MoE that cannot be expert-parallel yet, whatever else the + expert-parallel load refuses it for, offering every other way out (FSDP2 only for an + unquantized causal LM); the spec's own refusal, a + NotImplementedError and so a RuntimeError too, words it differently. Below the transformers + release expert parallelism needs, no EP degree is offered, since none would load: upgrading + is, where an EP degree divides the experts; where none does, upgrading would not shard them. + The ways out are named as the caller set the plan, so arguments, which override the knobs, are + named.""" + if version is not None: + monkeypatch.setattr(transformers, "__version__", version) + experts = getattr(config, "n_routed_experts", None) or getattr(config, "num_experts", 0) + below = Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS) + if version is None and refusal.startswith("can be") and below and experts % 2 == 0: + # On a transformers below the floor, upgrading is the way out where an EP degree divides + # the experts, with a degree set to 1 raised as well. + refusal = _NEEDS_LATER_AND_A_DEGREE if ep_degree is not None else _NEEDS_LATER + if ep_degree is not None: + if config.model_type == "qwen3_moe" and Version(transformers.__version__) < Version("5.7"): + pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7") + 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( - ValueError, match="can be expert-parallel, but the plan resolved" - ) as refused: + with pytest.raises((RuntimeError, ValueError), match=refusal) as refused: resolve_distribution_plan(str(tmp_path), **arguments) + if refusal is _NO_EP_YET: # as the load refuses the same spec + assert refused.type is NotImplementedError + if refusal.startswith(_CAN_BE_EP[: _CAN_BE_EP.index("{")]): # what its inputs resolve to + assert refused.type is ValueError, refused.type 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 @@ -604,20 +792,18 @@ 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", "its-specs-transformers"], -) +@pytest.mark.parametrize("refused_by", ["its-plan", "its-format", "its-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 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.""" + """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, whether the decision or the + caller made its EP degree 1: it cannot be expert-parallel here. Below the floor its spec + states, though, the decision's own refusal names the upgrade, as one would shard it, and an EP + degree the caller set to 1 to raise; a spec ruling expert parallelism out has a refusal of its + own (GPT-OSS's, in test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error).""" if ep_degree is not None: monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) run_as(2) @@ -626,23 +812,19 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( config.base_model_ep_plan = {"layers.*.mlp.experts": "moe_tp_experts"} 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) - 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) monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) + if refused_by == "its-transformers": + refusal = _NEEDS_LATER if ep_degree is None else _NEEDS_LATER_AND_A_DEGREE + with pytest.raises(RuntimeError, match=refusal): + resolve_distribution_plan(str(tmp_path)) + return assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2) -@needs_qwen3_moe_ep_plan @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 @@ -654,7 +836,7 @@ def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead( monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") run_as(2) - serve(_qwen3_moe(num_experts=8)) + serve(_deepseek_v3(n_routed_experts=8)) assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp") @@ -673,6 +855,19 @@ def test_an_ep_capable_moe_is_replicated_when_both_knobs_ask_for_it( assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2) +def test_an_ep_degree_of_1_auto_picks_does_not_ask_for_the_replication( + run_as, serve, monkeypatch, tmp_path +): + """Auto is an EP degree of 1 for GPT-OSS, whose spec refuses expert parallelism; with the dense + part asked to be replicated, that is still a replication the plan reached, and is refused.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "1") + run_as(2) + serve(GptOssConfig(num_local_experts=4, **_TINY)) + + with pytest.raises(RuntimeError, match="cannot be quantized expert-parallel yet"): + resolve_distribution_plan(str(tmp_path)) + + @needs_qwen3_moe_ep_plan @pytest.mark.parametrize( "value", @@ -787,6 +982,7 @@ def test_the_dense_knob_replicates_or_shards_over_the_whole_world( assert resolve_distribution_plan(str(tmp_path))[0].dense == dense +@needs_ep_transformers @pytest.mark.parametrize( ("config", "ep_size", "experts"), [(LlamaConfig(**_TINY), 2, 0), (_qwen3_moe(num_experts=6), 4, 6)], @@ -796,9 +992,7 @@ def test_an_ep_degree_that_cannot_shard_the_experts_is_refused( run_as, serve, tmp_path, config, ep_size, experts ): """Every rank of an EP group must hold whole experts, and a model without routed experts has - none to shard.""" - if experts and Version(transformers.__version__) < Version("5.7"): - pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7") + none to shard. (Below the expert-parallel load's transformers, any EP degree is refused first.)""" run_as(4) serve(config) @@ -824,6 +1018,7 @@ def test_a_model_is_decided_from_its_whole_model(run_as, serve, gpu_memory, tmp_ assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp") +@needs_ep_transformers def test_the_expert_count_is_read_from_the_config_holding_the_plan( run_as, serve, monkeypatch, tmp_path ): @@ -964,21 +1159,57 @@ def _knobs(**values): yield -def _every_rank_takes_rank0s_plan_job(rank, size, ckpt, knobs, expected): +def _every_rank_takes_rank0s_plan_job(rank, size, ckpt, knobs, expected, spec_applied): with bounded_collectives(), _knobs(**knobs[rank]): - plan, _ = resolve_distribution_plan(ckpt) + plan, hf_config = resolve_distribution_plan(ckpt) assert plan == expected, (rank, plan) + if spec_applied is not None: + spec_plan = get_spec(hf_config.model_type).ep_spec.base_model_ep_plan + applied = getattr(hf_config, "base_model_ep_plan", None) == spec_plan + assert applied is spec_applied, (rank, applied) -def test_every_rank_takes_the_plan_rank0_decided(cpu_dist_workers, tmp_path): - """Over a real broadcast. Rank 1's own knobs, FSDP2 over the world, would decide another plan, - and ranks loading different plans would build mismatched meshes and hang.""" - LlamaConfig(**_TINY).save_pretrained(tmp_path) +@pytest.mark.parametrize( + ("config", "knobs", "expected", "spec_applied"), + [ + pytest.param( + LlamaConfig(**_TINY), + [{}, {"MODELOPT_FSDP_SIZE": "world"}], + DistributionPlan(2), + None, + id="rank1-would-shard", + ), + pytest.param( + _deepseek_v3(n_routed_experts=8), + [{}, {"MODELOPT_EP_DEGREE": "1"}], + DistributionPlan(2, 2), + True, + marks=needs_ep_transformers, + id="rank1-would-skip-the-spec", + ), + pytest.param( + _deepseek_v3(n_routed_experts=8), + [{"MODELOPT_EP_DEGREE": "1", "MODELOPT_FSDP_SIZE": "world"}, {}], + DistributionPlan(2, dense="fsdp"), + False, + marks=needs_ep_transformers, + id="rank1-would-apply-the-spec", + ), + ], +) +def test_every_rank_takes_the_plan_rank0_decided( + cpu_dist_workers, tmp_path, config, knobs, expected, spec_applied +): + """Over real broadcasts. Rank 1's own knobs would decide another plan, or apply the model's spec + where rank 0 skipped it, or skip it where rank 0 applied it: ranks loading different plans, or + one without the expert-parallel plan the others shard by, build mismatched meshes and hang.""" + config.save_pretrained(tmp_path) cpu_dist_workers.run( _every_rank_takes_rank0s_plan_job, ckpt=str(tmp_path), - knobs=[{}, {"MODELOPT_FSDP_SIZE": "world"}], - expected=DistributionPlan(2), + knobs=knobs, + expected=expected, + spec_applied=spec_applied, ) @@ -1318,9 +1549,10 @@ def test_a_model_class_of_the_callers_own_loads_whole( """The sharded loads build their own model class and refuse another, so with nothing asking to shard, a class of the caller's own loads whole on every rank, and with a knob asking to shard it is refused. Neither builds the model on the meta device, which the decision's own pick of - class may not build (an encoder-decoder loaded through AutoModel). The decision reads the - checkpoint's config as the class does, through its own config class, which AutoConfig need not - know.""" + class may not build (an encoder-decoder loaded through AutoModel), nor applies the model's spec, + whose prepare could fail every rank, or patch the process, for a load that is never + expert-parallel. The decision reads the checkpoint's config as the class does, through its own + config class, which AutoConfig need not know.""" run_as(2) if own_config: _OwnConfig(**_TINY).save_pretrained(tmp_path) @@ -1330,6 +1562,11 @@ class may not build (an encoder-decoder loaded through AutoModel). The decision monkeypatch.setattr( model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model") ) + monkeypatch.setattr( + model_load_utils, "apply_distributed_support", lambda *a, **k: pytest.fail("applied a spec") + ) + # As on a transformers new enough for expert parallelism, whose floor would rule it out here. + monkeypatch.setattr(model_load_utils, "_spec_refuses_ep", lambda model_type: False) if fsdp_size is None: model = load_pretrained(str(tmp_path), auto_class) assert distribution_plan(model) == DistributionPlan(2) @@ -1375,6 +1612,7 @@ def test_a_load_it_decides_shards_the_experts( assert distribution_plan(model) == DistributionPlan(4, ep_size=4) +@needs_ep_transformers @pytest.mark.parametrize( ("dense", "one_gpu_per_process"), [("replicate", False), ("fsdp", False), ("replicate", True)], @@ -1383,19 +1621,22 @@ def test_a_load_it_decides_shards_the_experts( def test_an_expert_parallel_plan_loads_through_the_ep_loader( 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. A launch that shows each process one GPU, bound to cuda:0, gives a - rank of any local rank a GPU of its own.""" + """Decided automatically, an EP plan loads through load_and_prepare_ep -- with the plan, the + spec-prepared config and the Hub options the load was given -- 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 asks for expert parallelism and + FSDP2-shards the dense part too is refused before anything 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.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", {"replicate": "1", "fsdp": "world"}[dense]) + if dense == "fsdp": # left to the decision, the EP degree would give way to FSDP2 + monkeypatch.setenv("MODELOPT_EP_DEGREE", "4") 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) + served = serve(_deepseek_v3(n_routed_experts=8)) options = { "trust_remote_code": True, "attn_implementation": "eager", @@ -1407,21 +1648,94 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader( with pytest.raises( NotImplementedError, match=r"load replicates it on every rank\..*\(MODELOPT_FSDP_SIZE=1" ): - load_pretrained(str(tmp_path), distribution=plan, **options) + load_pretrained(str(tmp_path), **options) else: - model = load_pretrained(str(tmp_path), distribution=plan, **options) + model = load_pretrained(str(tmp_path), **options) call = loaders.pop("ep") + plan = DistributionPlan(4, ep_size=4, dense=dense) 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 ( + call["hf_config"].base_model_ep_plan + == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + ) + assert served.base_model_ep_plan is None assert distribution_plan(model) == plan assert loaders == {} assert auto_class.calls == [] +@needs_ep_transformers +@pytest.mark.parametrize( + ("planned", "passed_as", "preparations"), + [(False, None, 1), (True, None, 1), (True, "hf_config", 0), (True, "config", 1)], + ids=[ + "decided_by_the_load", + "planned", + "planned_with_its_config", + "planned_with_a_config_for_the_load", + ], +) +def test_a_load_prepares_its_config_once( + run_as, loaders, monkeypatch, tmp_path, planned, passed_as, preparations +): + """A model's EP preparation can patch its modeling code, so a load runs it once, and not at all + when handed the config resolve_distribution_plan prepared. A config passed for the load is the + caller's own, and the one deciding reads may go on to a load that is not expert-parallel, so + either is prepared as a copy. Every config it reads takes the Hub options the load was given + (custom code an ``auto_map`` names in another repo).""" + run_as(4) + # A GPU of its own, which the expert-parallel load needs; the decision weighs no GPU memory. + monkeypatch.setattr(model_load_utils, "_has_own_gpu", lambda: True) + config, reads = _deepseek_v3(n_routed_experts=8), [] + as_passed = config.to_dict() + monkeypatch.setattr( + transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(k) or config + ) + prepared = [] + + def apply_distributed_support(model_type, hf_config, **kwargs): + prepared.append(model_type) + hf_config.base_model_ep_plan = get_spec("deepseek_v3").ep_spec.base_model_ep_plan + return hf_config + + monkeypatch.setattr(model_load_utils, "apply_distributed_support", apply_distributed_support) + kwargs = {"distribution": DistributionPlan(4, ep_size=4)} if planned else {} + if passed_as: + kwargs[passed_as] = config + + load_pretrained(str(tmp_path), local_files_only=True, **kwargs) + + assert len(prepared) == preparations + loaded = loaders["ep"]["hf_config"] + if passed_as == "config" or not planned: + assert loaded is not config and loaded.base_model_ep_plan + assert config.to_dict() == as_passed + else: + assert loaded is config + if planned and passed_as is None: # read once, for the model type and for the load + assert len(reads) == 1, reads + assert all(read.get("local_files_only") is True for read in reads) + + +def test_a_model_its_spec_refuses_is_refused_with_its_config_handed_in( + run_as, loaders, a_gpu, tmp_path +): + """Handed the config prepared with its plan, the load does not prepare it again, but a model + whose spec refuses expert parallelism is refused all the same, before anything loads.""" + run_as(4) + config = GptOssConfig(**_TINY, num_local_experts=4) + + with pytest.raises(NotImplementedError, match="gpt_oss cannot be quantized expert-parallel"): + load_pretrained( + str(tmp_path), distribution=DistributionPlan(4, ep_size=4), hf_config=config + ) + assert loaders == {} + + def _every_rank_loads_rank0s_plan_job(rank, size, ckpt, handed, placed): model, refused = None, None # Ranks without GPUs, whatever the host has: the pooled workers keep no CUDA context. @@ -1643,6 +1957,51 @@ def test_a_placement_warning_raised_on_one_rank_fails_every_rank(cpu_dist_worker cpu_dist_workers.run(_one_rank_warns_as_an_error_job, ckpt=str(tmp_path)) +def _one_rank_cannot_apply_the_spec_job(rank, size, ckpt, step): + apply = model_load_utils.apply_distributed_support + + def apply_on_rank0_only(*args, **kwargs): + if rank == 1: + raise OSError("this rank cannot import the checkpoint's code") + return apply(*args, **kwargs) + + with ExitStack() as stack: + stack.enter_context(bounded_collectives()) + stack.enter_context( + mock.patch.object(model_load_utils, "apply_distributed_support", apply_on_rank0_only) + ) + if step == "load": + # Every rank has a GPU of its own, without which the load refuses it first. + stack.enter_context(mock.patch.object(torch.cuda, "is_available", return_value=True)) + stack.enter_context(mock.patch.object(torch.cuda, "device_count", return_value=size)) + stack.enter_context( + mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step) + ) + stack.enter_context( + mock.patch.object(model_load_utils, "load_and_prepare_ep", _loaded_in_step) + ) + try: + if step == "decision": + resolve_distribution_plan(ckpt) + else: + load_pretrained(ckpt, distribution=DistributionPlan(2, ep_size=2)) + except (OSError, RuntimeError) as error: + raised = str(error) + else: + raised = "no error" + assert "cannot import the checkpoint's code" in raised, raised + + +@needs_ep_transformers +@pytest.mark.parametrize("step", ["decision", "load"]) +def test_a_rank_that_cannot_apply_the_spec_fails_every_rank(cpu_dist_workers, tmp_path, step): + """A spec's preparation reads the checkpoint and imports its code on each rank, which can fail + on one alone; the others must not wait for it, in the plan's broadcast when deciding, nor in the + load a handed expert-parallel plan starts.""" + _deepseek_v3(n_routed_experts=8).save_pretrained(tmp_path) + cpu_dist_workers.run(_one_rank_cannot_apply_the_spec_job, ckpt=str(tmp_path), step=step) + + @pytest.mark.parametrize("sharding", ["expert-parallel", "fsdp2"]) @pytest.mark.parametrize( ("change", "error", "match"), @@ -1666,28 +2025,44 @@ def test_a_placement_warning_raised_on_one_rank_fails_every_rank(cpu_dist_worker "config and hf_config", ), ({"gpus": 0}, RuntimeError, "this rank has none"), + ({"gpus": 0, "config": True}, RuntimeError, "this rank has none"), + ], + ids=[ + "placement", + "dtype", + "torch_dtype", + "model_class", + "mtp_head", + "two_configs", + "no_gpu", + "no_gpu_config_passed", ], - 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, 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.""" + load something other than what was asked. A rank without a GPU is refused before the model's + spec prepares its config, which can import the checkpoint's code, whichever config it loads + with.""" run_as(2) expert_parallel = sharding == "expert-parallel" - serve(_deepseek_v3() if expert_parallel else LlamaConfig(**_TINY)) + config = 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"]) + monkeypatch.setattr( + model_load_utils, + "apply_distributed_support", + lambda *a, **k: pytest.fail("prepared a rank that cannot load"), + ) plan = DistributionPlan(2, ep_size=2) if expert_parallel else DistributionPlan(2, dense="fsdp") + kwargs = {**change.get("kwargs", {}), **({"config": config} if change.get("config") else {})} with pytest.raises(error, match=match) as refused: - load_pretrained( - str(tmp_path), change.get("auto_class"), distribution=plan, **change.get("kwargs", {}) - ) + load_pretrained(str(tmp_path), change.get("auto_class"), distribution=plan, **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. @@ -2298,13 +2673,13 @@ def resolve(path, **hub_options): [ *( ( - SimpleNamespace(quantization_config={"quant_method": method}), + SimpleNamespace(model_type=_NO_SPEC, quantization_config={"quant_method": method}), NotImplementedError, f"cannot dequantize a '{method}'", ) for method in ("gptq", "modelopt", "compressed-tensors", "mxfp4") ), - (SimpleNamespace(), ValueError, "declare base_model_ep_plan"), + (SimpleNamespace(model_type=_NO_SPEC), ValueError, "declare base_model_ep_plan"), ( SimpleNamespace( model_type=_NO_SPEC, base_model_ep_plan={"layers.*.self_attn.q_proj": "colwise"} @@ -2332,9 +2707,10 @@ def resolve(path, **hub_options): r"styles layers\.\*\.mlp\.experts\.down_proj by a style this load drops", ), ( - SimpleNamespace(model_type="gpt_oss", base_model_ep_plan=_EP_PLAN), + # An MXFP4 checkpoint with no plan, refused by its spec before either. + SimpleNamespace(model_type="gpt_oss", quantization_config={"quant_method": "mxfp4"}), NotImplementedError, - "the expert-parallel export cannot write that layout", + "gpt_oss cannot be quantized expert-parallel yet: .* cannot write that layout", ), ], ids=[ @@ -2358,9 +2734,9 @@ def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( 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. 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.""" + expert parallelism out (GPT-OSS's), as load_pretrained refuses it, before its checkpoint's + format or plan is looked at: a plan handed to this load 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 == [] @@ -2403,13 +2779,26 @@ def from_pretrained(path, **kwargs): load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) -@pytest.mark.parametrize("declared", [True, False], ids=["on-the-config", "a-class-default"]) +@pytest.mark.parametrize( + ("source", "on_the_config"), + [ + ({"base_model_ep_plan": _EP_PLAN}, True), + ({}, False), + ({}, True), + ('{"base_model_ep_plan": ', True), + ], + ids=["by-its-checkpoint", "a-class-default", "attached-by-its-spec", "in-an-unreadable-file"], +) def test_the_model_keeps_the_plan_its_checkpoint_declares( - ep_world_of_two, gloo_world_of_one, monkeypatch, declared + ep_world_of_two, gloo_world_of_one, monkeypatch, tmp_path, source, on_the_config ): """The load shards by the torch_a2a styles, but the config the model keeps goes back to the - checkpoint's plan: an export writes it out, and nothing without ModelOpt knows those styles. A - plan the config class ships (transformers' own, e.g. Qwen3-MoE's) stays off the instance.""" + plan its checkpoint's config.json declares: an export writes it out, and nothing without + ModelOpt knows those styles. A plan the config class ships (transformers' own, e.g. + Qwen3-MoE's) stays off the instance, and so does one the model's spec attached to the config, + which names styles only ModelOpt registers (Nemotron-H's nemotron_experts_ep). A config.json + the load cannot read declares none, with a warning, since an export then loses the plan.""" + declared = isinstance(source, dict) and "base_model_ep_plan" in source class _PastError(Exception): """Raised by the load's next check.""" @@ -2432,18 +2821,25 @@ def next_check(*args, **kwargs): monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) - class _Config(SimpleNamespace): - model_type = _NO_SPEC + class _Config(transformers.PretrainedConfig): + model_type = "_probe" base_model_ep_plan = _EP_PLAN + unreadable = isinstance(source, str) + (tmp_path / "config.json").write_text(source if unreadable else json.dumps(source)) # 8 experts over 2 ranks: the stand-in model holds this rank's 4. - config = _Config(num_experts=8, **({"base_model_ep_plan": dict(_EP_PLAN)} if declared else {})) + plan = {"base_model_ep_plan": dict(_EP_PLAN)} if on_the_config else {} + config = _Config(num_experts=8, **plan) - with pytest.raises(_PastError): - load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) + warns = pytest.warns(UserWarning, match="config.json") if unreadable else nullcontext() + with warns, pytest.raises(_PastError): + load_and_prepare_ep(str(tmp_path), DistributionPlan(2, ep_size=2), hf_config=config) assert model.config.base_model_ep_plan == _EP_PLAN assert ("base_model_ep_plan" in vars(model.config)) is declared + model.config.save_pretrained(tmp_path / "saved") + written = json.loads((tmp_path / "saved" / "config.json").read_text()) + assert ("base_model_ep_plan" in written) is declared def test_a_loaded_model_keeps_its_sharding_styles_out_of_its_saved_config( @@ -2741,9 +3137,23 @@ def _meshes_in_step(*args, **kwargs): raise _PastTheMeshesError +class _InTheLoadError(Exception): + """Raised by the weights' load stand-in: this rank got into the load.""" + + +def _load_in_step(*args, **kwargs): + """The weights' load, as the collective it is: a rank in it waits for every peer.""" + torch.distributed.barrier() + raise _InTheLoadError + + def _one_rank_fails_job(rank, size, ckpt, failure, failing_rank, match): path, hf_config, patches = ckpt, None, [] - if failure != "unreadable-config": + if failure == "unreadable-plan-provenance": + # A VLM's config, handed in: the plan to keep is read from where it sits in config.json. + plan = {"base_model_ep_plan": dict(_EP_PLAN), "num_experts": 4} + hf_config = SimpleNamespace(model_type=_NO_SPEC, text_config=SimpleNamespace(**plan)) + elif failure != "unreadable-config": hf_config = SimpleNamespace( model_type=_NO_SPEC, base_model_ep_plan=dict(_EP_PLAN), num_experts=4 ) @@ -2768,9 +3178,15 @@ def fail(*args, **kwargs): elif failure == "unseen-checkpoint" and fails_here: # Rank 0 makes the checkpoint local; a rank on a node without that storage cannot see it. patches.append(mock.patch.object(hf_checkpoint_utils.os.path, "isdir", return_value=False)) - if failure in ("load", "record"): + elif failure == "unreadable-plan-provenance" and fails_here: + # This rank reads a config.json that is no JSON object: it declares no plan, not an error. + reads = SimpleNamespace(load=lambda f: []) + patches.append(mock.patch.object(model_load_utils, "json", reads)) + if failure in ("load", "record", "unreadable-plan-provenance"): loaded = model_with_decoder_at("model.") load = fail if failure == "load" and fails_here else (lambda path, **kwargs: loaded) + if failure == "unreadable-plan-provenance": + load = _load_in_step loader = SimpleNamespace(__name__="Loader", from_pretrained=load) record = fail if failure == "record" and fails_here else None patches += [ @@ -2786,6 +3202,8 @@ def fail(*args, **kwargs): with bounded_collectives(), ExitStack() as stack: for patch in patches: stack.enter_context(patch) + caught = stack.enter_context(warnings.catch_warnings(record=True)) + warnings.simplefilter("always") try: # not pytest.raises: its failure would end the pooled worker instead of the job load_and_prepare_ep(path, DistributionPlan(2, ep_size=2), hf_config=hf_config) except Exception as error: @@ -2793,7 +3211,16 @@ def fail(*args, **kwargs): # The one rank's failure on every rank, before any rank went on alone into a collective. The GPU # check decides on every rank at once, from one gather, so no rank relays another's failure. assert refused is not None and match in refused, (rank, refused) - assert fails_here or failure == "no-gpu" or f"rank {failing_rank} failed" in refused, refused + if match != _InTheLoadError.__name__: # else no rank failed, and every rank got into the load + assert fails_here or failure == "no-gpu" or f"rank {failing_rank} failed" in refused, ( + refused + ) + else: # the plan the one rank could not read is lost to its model alone, which says so + warned = [ + str(w.message) for w in caught if "keeps no expert-parallel plan" in str(w.message) + ] + assert bool(warned) == fails_here, (rank, warned) + assert not warned or f"rank {rank}'s model" in warned[0], warned @pytest.mark.skipif( @@ -2809,6 +3236,7 @@ def fail(*args, **kwargs): ("no-gpu", 1, "rank(s) [1] have none while the others have theirs"), ("load", 1, "the GPU ran out of memory"), ("record", 1, "a weight file this rank cannot read"), + ("unreadable-plan-provenance", 1, _InTheLoadError.__name__), ], ) def test_a_load_failing_on_one_rank_fails_every_rank( @@ -2817,7 +3245,9 @@ def test_a_load_failing_on_one_rank_fails_every_rank( """One rank fails alone: reading the config, on its transformers, on its arguments, seeing the checkpoint rank 0 made local, without a GPU of its own where the others have theirs, or in the load. Every rank raises that failure, rather than leave the others waiting for it in the meshes - (stood in by a bounded barrier) or the warmup's collectives.""" + (stood in by a bounded barrier) or the warmup's collectives. A config.json one rank cannot read + as a JSON object declares no plan to keep there, so every rank goes on into the load (stood in + by a bounded barrier too).""" Qwen3MoeConfig(num_experts=4, num_hidden_layers=1, hidden_size=8).save_pretrained(tmp_path) cpu_dist_workers.run( _one_rank_fails_job,