diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0b48641edec..0b4335a203c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +43,7 @@ Changelog *Misc* - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs. -- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now loads a model distributed under ``torchrun``, FSDP2-sharded or whole on every rank, by the plan ``resolve_distribution_plan`` decides or one passed as its new ``distribution`` argument. +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now loads a model distributed under ``torchrun``, with an MoE's routed experts sharded across the ranks, FSDP2-sharded or whole on every rank, by the plan ``resolve_distribution_plan`` decides or one passed as its new ``distribution`` argument. **Backward Breaking Changes** @@ -52,7 +52,7 @@ Changelog - Bump minimum required PyTorch version to 2.10. - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. -- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE loads without expert parallelism, and is refused when ``MODELOPT_EP_DEGREE`` asks for it, until the expert-parallel load lands; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. +- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE the expert-parallel load takes has its routed experts sharded across the ranks, or the whole model FSDP2-sharded when a rank's share with the rest of it replicated would not fit one GPU, and is refused where it would otherwise load whole on every rank; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. - ``load_pretrained`` now records how it distributed the model on the model it returns, even in a single process, so ``fsdp2_wrap`` refuses that model; to shard a model yourself, wrap one loaded with ``from_pretrained``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. diff --git a/docs/source/guides/_distributed_loading.rst b/docs/source/guides/_distributed_loading.rst index 230b0249629..30d938102e2 100644 --- a/docs/source/guides/_distributed_loading.rst +++ b/docs/source/guides/_distributed_loading.rst @@ -19,23 +19,36 @@ recorded on the model, where calibration and export read it. with distributed_run(): model = load_pretrained("meta-llama/Llama-3.1-8B") # a whole copy on every rank's GPU # With MODELOPT_FSDP_SIZE=world: FSDP2-sharded over the 8 ranks instead. + moe = load_pretrained("Qwen/Qwen3-30B-A3B") # its routed experts sharded over the 8 ranks Routes ====== +* A plan with ``ep_size > 1`` takes the expert-parallel load (``load_and_prepare_ep``): each rank + holds ``num_experts / ep_size`` of an MoE's routed experts and a replica of the dense part, and + tokens reach their experts' ranks by all-to-all. It replicates the dense part, so a plan that + FSDP2-shards it as well is refused. * A plan with ``dense="fsdp"`` takes the FSDP2 load (``parallel_load_and_prepare_fsdp2``): a causal LM, or a text-only VLM it supports (Qwen3.5), is FSDP2-sharded over the process group, each rank reading its own layers. * Any other plan gives every rank the whole model, through ``auto_class.from_pretrained`` (``auto_class`` defaults to ``AutoModelForCausalLM``). -No load shards the routed experts yet. A plan the load decides keeps the expert-parallel degree at -1, so an MoE loads without expert parallelism; an EP degree ``MODELOPT_EP_DEGREE`` asks for, or a -plan with ``ep_size > 1``, is refused. +Left to decide, the EP degree of an MoE is the largest that divides both its routed experts and the +world size. The decision plans without expert parallelism, and says why, an MoE the expert-parallel +load would refuse: one its model spec rules out (GPT-OSS's, Llama-4's), one below the transformers +release it needs, a quantized checkpoint, or a plan that shards no expert weights. It FSDP2-shards +the whole model instead when a rank's share of the experts with a whole copy of the dense part would +not fit one GPU. An EP degree a knob or argument asks for stands, and the load refuses what it +cannot take. A plan that leaves an expert-parallel MoE whole on every rank, which no knob asked +for, is refused, naming the ways out; ``MODELOPT_EP_DEGREE=1`` with ``MODELOPT_FSDP_SIZE=1`` asks +for it. A refusal on any rank raises on every rank, before the load, and says how to load the model whole: ``MODELOPT_FSDP_SIZE=1`` for a plan the load decided, or -``distribution=DistributionPlan(, dense="replicate")``. +``distribution=DistributionPlan(, dense="replicate")``; for an expert-parallel plan, +``MODELOPT_EP_DEGREE=1`` with ``MODELOPT_FSDP_SIZE=1``, or +``distribution=DistributionPlan(, 1, "replicate")``. The plan is recorded once, so ``fsdp2_wrap`` refuses a model ``load_pretrained`` returned: to shard a model yourself, load it with ``from_pretrained``. @@ -56,15 +69,23 @@ Whole-model loads Sharded loads ============= -* The FSDP2 load builds the model class and places the model itself: a text-only VLM through its - image-text class, anything else through ``AutoModelForCausalLM``. Leave ``auto_class`` out. -* Besides its own options (``cpu_offload``, ``mp_policy``, ``broadcast_chunk_size``) it takes - ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other keyword - rather than ignore it. An option passed as ``None``, ``from_pretrained``'s own default, asks - nothing and is not refused. -* It reads the checkpoint's tensors as stored and keeps their dtype, so it refuses a ``dtype`` and - a quantized checkpoint. It cannot run the model type's loading hook, so it refuses a checkpoint - the hook would build modules for, such as a Nemotron-H MTP head. +* A sharded load builds the model class and places the model itself, so leave ``auto_class`` out. + The FSDP2 load builds a text-only VLM through its image-text class and anything else through + ``AutoModelForCausalLM``; the expert-parallel load builds the checkpoint's own class, a VLM's + image-text class included. +* Besides its own options (the FSDP2 load's ``cpu_offload``, ``mp_policy`` and + ``broadcast_chunk_size``; the expert-parallel load's ``freeze`` and ``require_lm_head``) each + takes ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other + keyword rather than ignore it, the other sharded load's options as that load's own. An option + passed as ``None``, ``from_pretrained``'s own default, asks nothing and is not refused, nor is an + FSDP2 option at that load's own default. The expert-parallel load leaves the model unfrozen, as + the other loads do, unless ``freeze=True``. +* Each keeps the checkpoint's dtype, so it refuses a ``dtype``. The FSDP2 load reads the tensors as + stored, so it refuses a quantized checkpoint; the expert-parallel load loads unquantized ones + only. Neither can run the model type's loading hook, so each refuses a checkpoint the hook would + build modules for, such as a Nemotron-H MTP head. +* Each shards onto every rank's own GPU, and a rank without one is refused with the others. The + expert-parallel load runs with one GPU per process, each bound to ``cuda:0``, too. * A single process, or a group of one, reads no knob: it takes the FSDP2 load only from a plan passed as ``distribution``, and sets a group of one up for it, which needs ``torchrun --nproc_per_node=1``. @@ -76,11 +97,12 @@ With neither a plan nor ``MODELOPT_FSDP_SIZE``, rank 0's GPU memory decides whet is sharded. The model is weighed in the checkpoint's own dtype (its config's, else bfloat16), not one the load asks for: hand a load that widens it, ``dtype=torch.float32`` say, its plan. -A call that places the model (``device_map``, ``quantization_config``, a ``dtype``) loads only -where the model fits one GPU; where it does not, the refusal says the layout was decided rather -than asked for. Pass ``distribution=`` or set ``MODELOPT_FSDP_SIZE=1`` to keep the load whole. An -option at ``from_pretrained``'s own default, ``device_map=None`` included, places nothing and does -not keep the load whole. +A call that places the model (``device_map``, ``quantization_config``, a ``dtype``) loads only where +the model fits one GPU; where it does not, the refusal says the layout was decided rather than asked +for. So does the expert-parallel load's, for an MoE whose experts the decision shards with neither a +plan nor ``MODELOPT_EP_DEGREE``. Pass ``distribution=`` or set ``MODELOPT_FSDP_SIZE=1`` to keep the +load whole. An option at ``from_pretrained``'s own default, ``device_map=None`` included, places +nothing and does not keep the load whole. Configs ======= @@ -95,10 +117,10 @@ Hub options A Hub model ID is made local first, so the export can read the source checkpoint from disk; the ``revision``, ``token``, ``cache_dir`` and ``local_files_only`` options go to that fetch. The config -read and a whole-model load also take all but ``revision``: an ``auto_map`` entry can name custom -code in another Hub repo, which Transformers fetches with them. The FSDP2 load builds the model -from the config with ``from_config``, which takes none of them, so such code comes with the -environment's Hub settings. +read, a whole-model load and the expert-parallel load also take all but ``revision``: an +``auto_map`` entry can name custom code in another Hub repo, which Transformers fetches with them. +The FSDP2 load builds the model from the config with ``from_config``, which takes none of them, so +such code comes with the environment's Hub settings. Launches without a process group ================================ @@ -113,7 +135,7 @@ Inside ``init_quantized_weights`` ================================= ``init_quantized_weights`` (see :doc:`_compress_quantized_models`) quantizes and compresses the -model in the ``from_pretrained`` it patches, which the FSDP2 load never calls. Inside it, a plan +model in the ``from_pretrained`` it patches, which the sharded loads never call. Inside it, a plan the load decides gives every rank the whole model, with nothing to measure, and the load leaves the device map to ``init_quantized_weights``, which computes it and refuses one it is passed. A knob or a plan that asks to shard is refused there. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 705e21d5a93..e3864945383 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -19,6 +19,7 @@ import copy import dataclasses import logging +import math import os import re import warnings @@ -32,6 +33,7 @@ import torch.nn as nn import transformers.distributed as hf_distributed from huggingface_hub import snapshot_download +from packaging.version import parse if TYPE_CHECKING: from collections.abc import Iterable @@ -52,6 +54,7 @@ from transformers.conversion_mapping import get_model_conversion_mapping from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key +from modelopt.torch.models.specs import get_spec from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import env, moe from modelopt.torch.utils.distributed import ( @@ -451,18 +454,34 @@ def record_unplaced_source_keys( return keys +_EXPERT_PARALLEL = "expert-parallel" _FSDP2 = "FSDP2" # What each sharded load takes besides the checkpoint. These loads build the model class and place # it themselves, so from_pretrained's own arguments (device_map, max_memory, ...) do not reach them. _FSDP2_ONLY = frozenset({"cpu_offload", "mp_policy", "broadcast_chunk_size"}) _SHARDED_LOAD_OPTIONS = { - _FSDP2: _FSDP2_ONLY | {"trust_remote_code", "attn_implementation", "hf_config"} + _EXPERT_PARALLEL: frozenset( + { + "trust_remote_code", + "attn_implementation", + "hf_config", + "freeze", + "require_lm_head", + } + ), + _FSDP2: _FSDP2_ONLY | {"trust_remote_code", "attn_implementation", "hf_config"}, } -# The FSDP2-only values that ask nothing of a whole load: the FSDP2 load's defaults, and None where -# that leaves the option unset (a None broadcast_chunk_size asks for all of a source's layers). +# The FSDP2-only values that ask nothing of any load but the FSDP2 one: its defaults, and None where +# that leaves the option unset (a None broadcast_chunk_size asks for all of a source's layers). Any +# other option asks nothing at None, from_pretrained's own default (device_map=None, say). _ASKS_NOTHING = {"cpu_offload": (False, None), "mp_policy": (None,), "broadcast_chunk_size": (8,)} +# The options only the expert-parallel load takes, which a whole-model load would hand on to +# from_pretrained. +_EP_ONLY_LOAD_OPTIONS = _SHARDED_LOAD_OPTIONS[_EXPERT_PARALLEL] - _SHARDED_LOAD_OPTIONS[_FSDP2] + + def _route(plan: DistributionPlan) -> str | None: """The sharded load ``plan`` takes, or ``None`` when every rank loads the whole model.""" if plan.world_size != dist.size(): @@ -470,11 +489,21 @@ def _route(plan: DistributionPlan) -> str | None: f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}." ) if plan.is_expert_parallel: - raise NotImplementedError( - f"{plan!r} shards the routed experts, which load_pretrained has no load for. Plan " - "ep_size=1 (MODELOPT_EP_DEGREE=1) to give every rank the whole model, or shard it " - "with FSDP2." - ) + if plan.is_fsdp_dense: + # Refused before anything loads, where the remedies can be named: the expert-parallel + # load replicates the dense part, so it cannot honor this plan. + raise NotImplementedError( + f"{plan!r} FSDP2-shards the dense part, but the expert-parallel load replicates it " + "on every rank. Sharding the dense part under expert parallelism is not supported " + "yet. An EP degree left to the decision gives way to FSDP2 for the whole model when " + "a whole copy of the dense part and a rank's share of the routed experts would take " + f"more than {_REPLICATED_SHARE:.0%} of a GPU; the degree asked for stands. " + "Replicate the dense part anyway (MODELOPT_FSDP_SIZE=1, or " + "a plan with dense='replicate'), or shard the whole model with FSDP2 " + "(MODELOPT_EP_DEGREE=1 with MODELOPT_FSDP_SIZE=world) if the checkpoint is an " + "unquantized causal LM." + ) + return _EXPERT_PARALLEL return _FSDP2 if plan.is_fsdp_dense else None @@ -509,58 +538,110 @@ def _init_quantized_weights() -> bool: return getattr(AutoModelForCausalLM.from_pretrained, "_modelopt_init_quantized_weights", False) -def _whole_instead(decided: bool) -> str: - """How to have every rank load the model whole: the knob counts only for a decided plan.""" - plan = f"distribution=DistributionPlan({dist.size()}, dense='replicate')" - return f"MODELOPT_FSDP_SIZE=1 or {plan}" if decided else plan +def _whole_instead(decided: bool, route: str = _FSDP2) -> str: + """How to have every rank load the model whole: the knob counts only for a decided plan. + + Instead of the expert-parallel load it takes both knobs, as the plan decision refuses to + replicate an expert-parallel MoE unasked. + """ + if route == _EXPERT_PARALLEL: + plan = f"distribution=DistributionPlan({dist.size()}, 1, 'replicate')" + knobs = "MODELOPT_EP_DEGREE=1 with MODELOPT_FSDP_SIZE=1" + else: + plan = f"distribution=DistributionPlan({dist.size()}, dense='replicate')" + knobs = "MODELOPT_FSDP_SIZE=1" + return f"{knobs} or {plan}" if decided else plan def _refuse_sharded_load_options( route, auto_class, kwargs, decided: bool, measured: bool = False ) -> None: - """Refuse what the sharded load ``route`` cannot honor; ``measured``: GPU memory chose it.""" - # An option passed as None, from_pretrained's own default (device_map=None, say), asks nothing. - passed = {name for name, value in kwargs.items() if value is not None} + """Refuse what the sharded load ``route`` cannot honor; ``measured``: the decision chose it. + + By rank 0's GPU memory for the FSDP2 load, by the model's experts for the expert-parallel one. + """ + # The options that ask anything (_ASKS_NOTHING). + passed = { + name + for name, value in kwargs.items() + if (value not in _ASKS_NOTHING[name] if name in _ASKS_NOTHING else value is not None) + } unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) # torch_dtype is from_pretrained's older name for dtype, and None means "auto" to it. dtypes = [name for name in ("dtype", "torch_dtype") if kwargs.get(name) not in (None, "auto")] - if measured and (refused := sorted({*unknown, *dtypes})): + # What does not apply to a sharded load: a dtype to load in, or an option neither load takes. + taken = _SHARDED_LOAD_OPTIONS[_FSDP2] | _SHARDED_LOAD_OPTIONS[_EXPERT_PARALLEL] + refused = sorted({*(passed - taken - {"dtype", "torch_dtype"}), *dtypes}) + if measured and refused: # Say so, rather than refuse the options below as if the caller had asked for this load. + chose = ( + "no plan or MODELOPT_EP_DEGREE asked for a layout, and the plan decision shards the " + "model's routed experts over the ranks" + if route == _EXPERT_PARALLEL + else "no plan or MODELOPT_FSDP_SIZE asked for a layout, and by rank 0's GPU memory the " + f"model does not fit one GPU, so the {route} load shards it" + ) raise TypeError( - f"{refused} do not apply to a sharded load, but no plan or MODELOPT_FSDP_SIZE asked for " - f"a layout, and by rank 0's GPU memory the model does not fit one GPU, so the {route} " - f"load shards it. To load it whole on every rank: {_whole_instead(decided)}." + f"{refused} do not apply to a sharded load, but {chose}. To load it whole on every " + f"rank: {_whole_instead(decided, route)}." ) if _init_quantized_weights(): raise NotImplementedError( f"The {route} load builds the model without from_pretrained, so it would skip " "init_quantized_weights' quantization; load the model whole on every rank: " - f"{_whole_instead(decided)}." + f"{_whole_instead(decided, route)}." ) if auto_class not in (None, AutoModelForCausalLM): name = getattr(auto_class, "__name__", repr(auto_class)) + builds = ( + "the checkpoint's own, a VLM's image-text class included" + if route == _EXPERT_PARALLEL + else "a causal LM's or a text-only VLM's image-text class" + ) raise NotImplementedError( - f"The {route} load builds its model class itself, a causal LM's or a text-only VLM's " - f"image-text class, so it takes no auto_class={name}: leave it out to have the load " - f"pick, or load the model whole with it: {_whole_instead(decided)}." + f"The {route} load builds its model class itself, {builds}, so it takes no " + f"auto_class={name}: leave it out to have the load pick, or load the model whole with " + f"it: {_whole_instead(decided, route)}." ) if dtypes: # A sharded load keeps the checkpoint's dtype. name = dtypes[0] + # Only the FSDP2 load computes in another dtype, through its mp_policy. + computes = "; FSDP2's mp_policy sets the dtype it computes in" if route == _FSDP2 else "" raise TypeError( f"The {route} load reads the checkpoint's tensors as stored and keeps their dtype, so " - f"{name}={kwargs[name]!r} does not apply to it; FSDP2's mp_policy sets the dtype it " - "computes in." + f"{name}={kwargs[name]!r} does not apply to it{computes}." + ) + other = _FSDP2 if route == _EXPERT_PARALLEL else _EXPERT_PARALLEL + if elsewhere := sorted(passed & (_SHARDED_LOAD_OPTIONS[other] - _SHARDED_LOAD_OPTIONS[route])): + # Not placement: a whole model on every rank, the way out below, would refuse them too. An + # expert-parallel plan that FSDP2-shards the dense part is refused as well (_route). + instead = ( + "pass a plan with ep_size > 1 to shard an MoE's routed experts" + if other == _EXPERT_PARALLEL + else "pass a plan with ep_size=1 and dense='fsdp' to shard the whole model with FSDP2" + ) + raise TypeError( + f"{elsewhere} only apply to the {other} load, and the {route} load takes " + f"{sorted(_SHARDED_LOAD_OPTIONS[route])}; drop them, or {instead}." ) if unknown: raise TypeError( f"{unknown} do not apply to the {route} load, which places the model itself; it " f"takes {sorted(_SHARDED_LOAD_OPTIONS[route])}. To place the model yourself, have " - f"every rank load it whole: {_whole_instead(decided)}." + f"every rank load it whole: {_whole_instead(decided, route)}." ) def _sharded_load_kwargs( - route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided: bool, measured: bool + route, + plan, + ckpt_path, + auto_class, + model_type, + kwargs, + hub_options, + decided: bool, + measured: bool, ) -> dict: """What the sharded load ``route`` names takes, refusing what it cannot honor.""" from modelopt.torch.models.hf import checkpoint_has_mtp @@ -572,10 +653,21 @@ def _sharded_load_kwargs( # The head is built on a whole model only; carried over, it would export unquantized. raise NotImplementedError( f"The {route} load cannot build the MTP head this {model_type} checkpoint carries. To " - f"load it whole on every rank: {_whole_instead(decided)}." + f"load it whole on every rank: {_whole_instead(decided, route)}." ) trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) + if route == _EXPERT_PARALLEL: + _refuse_a_rank_without_its_gpu(route, decided) + # Unfrozen unless asked, as the other loads return the model: which load runs is decided. + kwargs.setdefault("freeze", False) + return { + "plan": plan, + "trust_remote_code": trust_remote_code, + "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 @@ -604,14 +696,23 @@ def _sharded_load_kwargs( f"The {route} load shards through a process group, even on one rank, which a process " "started without a launcher cannot set up; launch it with `torchrun --nproc_per_node=1`." ) + _refuse_a_rank_without_its_gpu(route, decided) + return {"trust_remote_code": trust_remote_code, "hf_config": hf_config, **kwargs} + + +def _refuse_a_rank_without_its_gpu(route, decided: bool) -> None: + """Refuse a rank without a GPU of its own before the ``route`` load, with the other ranks. + + Inside the FSDP2 load, the others would wait for this rank in its collectives; the + expert-parallel load refuses only a run in which some ranks have one, and neither names the way + to a whole model. + """ if not _has_own_gpu(): - # Refused with the other ranks, not inside the load, where they would wait for this one. raise RuntimeError( f"The {route} load shards onto each rank's own GPU, and this rank has none: it sees " f"{torch.cuda.device_count()} GPU(s) as local rank {dist.local_rank()}. To load it " - f"whole on every rank: {_whole_instead(decided)}." + f"whole on every rank: {_whole_instead(decided, route)}." ) - return {"trust_remote_code": trust_remote_code, "hf_config": hf_config, **kwargs} def _own_gpu_index() -> int | None: @@ -632,8 +733,10 @@ def _has_own_gpu() -> bool: return _own_gpu_index() is not None -def _load_sharded(ckpt_path, load_kwargs) -> nn.Module: - """The FSDP2 load, by the arguments :func:`_sharded_load_kwargs` checked.""" +def _load_sharded(route, ckpt_path, load_kwargs) -> nn.Module: + """The sharded load ``route`` names, by the arguments :func:`_sharded_load_kwargs` checked.""" + if route == _EXPERT_PARALLEL: + return load_and_prepare_ep(ckpt_path, **load_kwargs) if not dist.is_initialized(): # A single process's group of one, which the sharded model needs while it lives. dist.setup() @@ -659,6 +762,11 @@ def load_pretrained( ``distribution`` is the :class:`~modelopt.torch.utils.distributed.DistributionPlan` to load by: every rank takes rank 0's or, when rank 0 passes none, :func:`resolve_distribution_plan`'s. + * ``ep_size > 1``: :func:`load_and_prepare_ep` shards an MoE's routed experts over the process + group as they load and replicates the dense part, so a plan with ``dense="fsdp"`` as well is + refused. It builds and places the model itself, so it takes its own options but no + ``auto_class``, placement or ``dtype``, and refuses a quantized checkpoint. Here it leaves the + model unfrozen, as the other loads do, unless ``freeze=True``. * ``dense="fsdp"``: :func:`parallel_load_and_prepare_fsdp2` shards a causal LM, or a text-only VLM it supports, over the process group (one process needs ``torchrun --nproc_per_node=1``). It builds and places the model itself, so it takes its own options but no ``auto_class``, @@ -667,8 +775,7 @@ def load_pretrained( ``AutoModelForCausalLM``); on a distributed run every rank loads the whole model, onto its own GPU unless ``device_map`` is passed or ``init_quantized_weights`` places it. - No load here shards the routed experts: a plan with ``ep_size > 1`` is refused. A refusal on any - rank raises on every rank, before the load. The plan is recorded on the model + A refusal on any rank raises on every rank, before the load. The plan is recorded on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) for calibration and export, so :func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap` refuses it: to shard a model yourself, load it with ``from_pretrained``. :doc:`/guides/_distributed_loading` has the routing in full. @@ -677,10 +784,11 @@ def load_pretrained( * A Hub model ID is made local first (:func:`ensure_local_checkpoint`, which takes the ``revision``, ``token``, ``cache_dir`` and ``local_files_only`` ``kwargs``): the export reads the source checkpoint from disk. Under ``torch.distributed`` every rank must call this. The - config read and a whole-model load still take all but ``revision``: an ``auto_map`` entry can - name custom code in another Hub repo, which Transformers fetches with them. A sharded load - builds the model with ``from_config``, which takes none. A ``subfolder`` is part of where the - checkpoint is: the config, the loading hook and the record below read it, as the load does. + config read, a whole-model load and the expert-parallel load still take all but + ``revision``: an ``auto_map`` entry can name custom code in another Hub repo, which + Transformers fetches with them. The FSDP2 load builds the model with ``from_config``, which + takes none. A ``subfolder`` is part of where the checkpoint is: the config, the loading hook + and the record below read it, as the load does. * The model type's loading hook runs (:func:`~modelopt.torch.models.hf.prepare_model_for_loading`), e.g. to keep the MTP head a Nemotron-H checkpoint carries. ``model_type`` selects it; when omitted, the type of the @@ -713,16 +821,14 @@ def load_pretrained( if alone and distribution is None: # Knobs asking to shard over the launch are refused, as is a value that is no EP degree. ep, world = _ep_degree(None), int(os.environ["WORLD_SIZE"]) - experts = ep != "auto" and ep > 1 - asked = [f"MODELOPT_EP_DEGREE={env.ep_degree()}"] if experts else [] + asked = [f"MODELOPT_EP_DEGREE={env.ep_degree()}"] if ep != "auto" and ep > 1 else [] if _dense_from_env(world) == "fsdp": asked.append(f"MODELOPT_FSDP_SIZE={env.fsdp_size()}") if asked: - ep_way = ", with MODELOPT_EP_DEGREE=1: no load here shards the routed experts" raise RuntimeError( f"{' and '.join(asked)} ask to shard the model over this " f"{os.environ['WORLD_SIZE']}-rank launch, whose process group is not set up; load " - f"inside distributed_run(), which sets it up{ep_way if experts else ''}." + "inside distributed_run(), which sets it up." ) # None is these options' own default: passed explicitly, it is no option to refuse. for name in ("config", "variant", "hf_config"): @@ -749,9 +855,9 @@ def handed() -> DistributionPlan | None: decided = plan is None and dist.size() > 1 # Each rank checks its own arguments; a refusal on one fails every rank before the fetch. with _refused_together("load arguments", alone): - # No load here shards the routed experts yet, so an EP degree left to the decision is 1. - # Parsed in this check, a knob one rank cannot parse fails every rank. - ep_size = 1 if decided and _ep_degree(None) == "auto" else None + if decided: + # Parsed in this check, a knob one rank cannot parse fails every rank. + _ep_degree(None) if kwargs.pop("output_loading_info", False): raise ValueError( "load_pretrained returns the model alone and does not take " @@ -803,7 +909,6 @@ def handed() -> DistributionPlan | None: config=config, auto_class=auto_class, trust_remote_code=trust_remote_code, - ep_size=ep_size, **hub_options, ) if _route(plan) is not None: @@ -824,7 +929,22 @@ def handed() -> DistributionPlan | None: f"{sharded} only apply to the FSDP2 load, but {plan!r} gives every rank the whole " "model; pass a plan with dense='fsdp' to shard it." ) - measured = decided and route == _FSDP2 and _dense_from_env(plan.world_size) == "auto" + # As do the options only the expert-parallel load takes, but for one passed as None, which + # asks nothing of either load: its default holds, rather than None read as False. + for name in _EP_ONLY_LOAD_OPTIONS: + if name in kwargs and kwargs[name] is None: + del kwargs[name] + if route is None and (sharded := sorted(set(kwargs) & _EP_ONLY_LOAD_OPTIONS)): + raise TypeError( + f"{sharded} only apply to the expert-parallel load, but {plan!r} gives every rank " + "the whole model; pass a plan with ep_size > 1 to shard an MoE's routed experts." + ) + # Whether the decision chose this sharded layout with nothing asking for it: by rank 0's GPU + # memory for the FSDP2 load, by the model's routed experts for the expert-parallel one. + measured = decided and ( + (route == _FSDP2 and _dense_from_env(plan.world_size) == "auto") + or (route == _EXPERT_PARALLEL and _ep_degree(None) == "auto") + ) if route is not None and config is not None: # The sharded load builds from hf_config, a copy, as building writes to the config. del kwargs["config"] @@ -832,8 +952,8 @@ def handed() -> DistributionPlan | None: kwargs["hf_config"] = copy.deepcopy(config) if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. - if route == _FSDP2 and hf_config is None and config is None: - # The FSDP2 load builds from the checkpoint's config too: read it once, for both. + if route is not None and hf_config is None and config is None: + # A sharded load builds from the checkpoint's config too: read it once, for both. hf_config = AutoConfig.from_pretrained( ckpt_path, trust_remote_code=trust_remote_code, **hub_options ) @@ -846,11 +966,19 @@ def handed() -> DistributionPlan | None: {} if route is None else _sharded_load_kwargs( - route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided, measured + route, + plan, + ckpt_path, + auto_class, + model_type, + kwargs, + hub_options, + decided, + measured, ) ) if route is not None: - model = _load_sharded(ckpt_path, load_kwargs) + model = _load_sharded(route, ckpt_path, load_kwargs) # A sharded load finds the unplaced keys while it places the weights. keys = model._modelopt_unplaced_source_keys else: @@ -1412,8 +1540,10 @@ def _decide_distribution( ckpt_path: str, hf_config, world_size: int, ep_size, dense, trust_remote_code: bool, auto_class ) -> DistributionPlan: """The plan :func:`resolve_distribution_plan` returns, decided by this rank alone.""" - # What the caller set the EP degree with, for a refusal to name it. + # What the caller set the EP degree and the dense layout with, for a refusal to name them. knob = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size" + replicate = "MODELOPT_FSDP_SIZE=1" if dense is None else "dense='replicate'" + fsdp = "MODELOPT_FSDP_SIZE=world" if dense is None else "dense='fsdp'" if world_size == 1: # Nothing to shard over, so nothing about the model matters, nor do the knobs: a shell still # exporting a multi-GPU run's would otherwise reshape a single process's load. The caller's @@ -1425,6 +1555,9 @@ def _decide_distribution( ) ep_size = _ep_degree(ep_size) dense = _dense_from_env(world_size) if dense is None else dense + # A replication both knobs ask for is the caller's choice, and loads; one reached from "auto" + # would leave an expert-parallel MoE whole on every rank unasked, and is refused below. + asked_to_replicate = ep_size == 1 and dense == "replicate" if ( (auto_class not in (None, AutoModelForCausalLM) or _init_quantized_weights()) and ep_size in ("auto", 1) @@ -1436,7 +1569,7 @@ def _decide_distribution( # Refused here, rather than build the decision's own pick of class, which may not build. name = getattr(auto_class, "__name__", repr(auto_class)) raise NotImplementedError( - f"A sharded load builds its model class itself, a causal LM's or a text-only VLM's " + f"A sharded load builds its model class itself, a causal LM's or a supported VLM's " f"image-text class, so it takes no auto_class={name} while ep_size={ep_size!r}, " f"dense={dense!r} ask to shard. To give every rank the whole model, plan ep_size=1 and " "dense='replicate' (MODELOPT_EP_DEGREE=1, MODELOPT_FSDP_SIZE=1)." @@ -1471,8 +1604,8 @@ def _decide_distribution( raise ValueError( f"{hf_config.model_type} has routed experts and a base_model_ep_plan, but none of the " f"plan's expert containers ({containers}) names a module of the built model, so they " - f"can be neither counted nor sharded. Set {knob}=1 to load it without expert " - "parallelism." + f"can be neither counted nor sharded. Set {knob}=1 and {replicate} to load it without " + "expert parallelism, whole on every rank." ) if experts and not num_experts and ep_size != 1: # The mirror case: the containers are there, but no count to choose a degree by, which @@ -1480,11 +1613,23 @@ def _decide_distribution( raise ValueError( f"{hf_config.model_type}'s base_model_ep_plan names routed experts the built model has, " "but its config gives no expert count (num_experts, num_local_experts, " - f"n_routed_experts or num_routed_experts) to shard them by. Set {knob}=1 to load it " - "without expert parallelism." + f"n_routed_experts or num_routed_experts) to shard them by. Set {knob}=1 and " + f"{replicate} to load it without expert parallelism, whole on every rank." ) + chosen = ep_size == "auto" # an EP degree no knob or argument set + # Why the expert-parallel load would refuse this MoE, or None, whatever EP degree is set: an MoE + # it refuses cannot be expert-parallel here, so is not refused below for being left whole. + refusal = _ep_load_refusal(hf_config, ep_plan) if experts else None if ep_size == "auto": ep_size = _choose_ep_degree(world_size, num_experts) + if ep_size > 1 and refusal is not None: + # The expert-parallel load would refuse this model, for some reasons only after the + # checkpoint's download: with nothing asking for expert parallelism, plan without it. + print_rank_0( + f"{hf_config.model_type}: planned without expert parallelism, which its load " + f"would refuse: {refusal}" + ) + ep_size = 1 elif ep_size > 1 and world_size % ep_size: raise ValueError(f"{knob}={ep_size} must divide the world size ({world_size}).") elif ep_size > 1 and (num_experts == 0 or num_experts % ep_size): @@ -1492,6 +1637,7 @@ def _decide_distribution( f"{knob}={ep_size} must divide the routed experts {hf_config.model_type} shards " f"({num_experts})." ) + dense_chosen = dense == "auto" # a dense layout no knob or argument set if dense == "auto": # A rank holds a whole copy of what expert parallelism does not shard, and its 1/ep_size # share of the routed experts it does. @@ -1505,6 +1651,50 @@ def _decide_distribution( memory = _gpu_memory() too_big = memory is not None and held > _REPLICATED_SHARE * memory dense = "fsdp" if too_big else "replicate" + if chosen and ep_size > 1 and dense == "fsdp": + # The expert-parallel load replicates the dense part, and no load shards both yet: with + # nothing asking for expert parallelism, FSDP2-shard the whole model instead. + why = ( + "would not fit one GPU beside a rank's experts" if dense_chosen else "is to be sharded" + ) + print_rank_0( + f"{hf_config.model_type}: FSDP2-sharded whole, without expert parallelism, as the " + f"expert-parallel load replicates the dense part, which {why}. To shard the experts " + f"anyway, the dense part replicated: {knob}={ep_size} with {replicate}." + ) + ep_size = 1 + if experts and asked_to_replicate: + print_rank_0( + 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": + # 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 ({knob} above 1, dividing both its {num_experts} experts and " + "the world size), shard" + ) + # 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'))." + ) return DistributionPlan(world_size, ep_size, dense) @@ -1538,7 +1728,8 @@ def resolve_distribution_plan( hang. Collective: every rank must call it. * ``ep_size``: the expert-parallel degree, or ``"auto"``, the largest that divides both the - world and the routed-expert count of a model whose config carries a ``base_model_ep_plan``. + world and the routed-expert count of a model whose config carries a ``base_model_ep_plan``, + or 1 where the dense part is FSDP2-sharded, which the expert-parallel load does not do. * ``dense``: ``"replicate"`` or ``"fsdp"`` for the rest of the model, or ``"auto"``, which FSDP2-shards it when a whole copy of it and the rank's share of the routed experts would take more than 60% of one GPU's memory. That GPU is rank 0's, by its total memory, so ``"auto"`` @@ -1561,6 +1752,10 @@ def resolve_distribution_plan( is the class the load builds with: with one the sharded loads refuse, or inside ``init_quantized_weights``, and nothing asking to shard, the plan loads the model whole, decided without a build. + + A plan that leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for + exactly that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and + ``MODELOPT_FSDP_SIZE=1``) replicates it, noting that every rank holds the whole MoE. """ # Each rank reads the checkpoint for itself. One that cannot must fail every rank here, rather # than leave the others waiting in a broadcast below for a rank that already raised. @@ -1619,10 +1814,91 @@ 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}) +# The ep_plan styles the expert-parallel load swaps for its torch_a2a styles. It keeps these and +# grouped_gemm, and drops any other. +_EP_STYLE_SWAP = { + "ep_router": "torch_a2a_router", + "moe_tp_experts": "torch_a2a_experts", + "nemotron_experts_ep": "torch_a2a_experts", +} + + +def _ep_plan_refusal(ep_plan: dict[str, str]) -> str | None: + """Why the expert-parallel load refuses ``ep_plan``'s styles, or ``None``: by the plan alone.""" + if not _routed_expert_keys(ep_plan): + return ( + f"base_model_ep_plan={ep_plan} styles no routed experts ({', '.join(_EXPERTS_STYLES)} " + "or grouped_gemm), so this load would leave them whole on every rank." + ) + # Only grouped_gemm entries shard an expert's weights here: the torch_a2a styles the + # containers take for the load shard nothing themselves (any other style is refused below). + unsharded = sorted( + key + for key in _routed_expert_keys(ep_plan) + if not any( + style == "grouped_gemm" and name.startswith(key + ".") + for name, style in ep_plan.items() + ) + ) + if unsharded: + return ( + f"base_model_ep_plan={ep_plan} shards no weights of {', '.join(unsharded)} through " + "grouped_gemm, the only expert-weight style this load shards by, so every rank would " + "hold every expert." + ) + # The load drops any other style an expert weight takes, so that weight would load whole + # beside its sharded siblings, and the grouped forward would pair slices of different experts. + containers = _routed_expert_keys(ep_plan) + kept_styles = {"grouped_gemm", *_EP_STYLE_SWAP} + dropped = sorted( + name + for name, style in ep_plan.items() + if style not in kept_styles + and name not in containers + and any(name.startswith(key + ".") for key in containers) + ) + if dropped: + return ( + f"base_model_ep_plan={ep_plan} styles {', '.join(dropped)} by a style this load " + "drops, so those expert weights would load whole on every rank beside their sharded " + "siblings." + ) + return None + + +def _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. + """ + 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." + ) + refusal = _ep_plan_refusal(ep_plan) + return ValueError(refusal) if refusal is not None else None + def _ckpt_quant_method(cfg): """Return the checkpoint's native quantization method, or ``None``. @@ -1683,6 +1959,7 @@ def load_and_prepare_ep( freeze: bool = True, hf_config=None, require_lm_head: bool = True, + hub_options: dict | None = None, ) -> nn.Module: """Load a MoE causal LM expert-parallel for PTQ: each rank holds its share of the routed experts. @@ -1700,7 +1977,9 @@ def load_and_prepare_ep( The model builds with its LM head where an auto class builds one; ``require_lm_head=False`` builds a VLM without it. Without ``hf_config`` every rank reads the checkpoint's config itself; - on many ranks, pass it. The caller records the plan on the model + on many ranks, pass it. ``hub_options`` (``token``, ``cache_dir``, ``local_files_only``, as + :func:`load_pretrained` takes them) go to the checkpoint's download, its config read and + ``from_pretrained``. The caller records the plan on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), as :func:`load_pretrained` does; :func:`~modelopt.torch.utils.moe.expert_parallel_mesh` then finds its mesh. """ @@ -1719,23 +1998,18 @@ def load_and_prepare_ep( f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}." ) world_size, ep_degree, edp = plan.world_size, plan.ep_size, plan.edp - # The ep_plan styles this load keeps: grouped_gemm, and those it swaps for its torch_a2a styles. - # The checks and the swapped plan both read it. - swap = { - "ep_router": "torch_a2a_router", - "moe_tp_experts": "torch_a2a_experts", - "nemotron_experts_ep": "torch_a2a_experts", - } - kept_styles = {"grouped_gemm", *swap} + hub_options = dict(hub_options or {}) + swap, kept_styles = _EP_STYLE_SWAP, {"grouped_gemm", *_EP_STYLE_SWAP} # Each rank checks its own transformers, the config it reads or was handed, and the plan against # it: one refusing fails every rank here, before the checkpoint's fetch and the meshes. with _refused_together("the expert-parallel load's checks", alone=False): # 5.7 is the first release that builds a model with the torch_a2a experts. moe.require_transformers( - "5.7", "ModelOpt DP=EP expert-parallel PTQ (transformers MoE EP framework)" + _EP_MIN_TRANSFORMERS, + "ModelOpt DP=EP expert-parallel PTQ (transformers MoE EP framework)", ) if hf_config is None: - config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code} + 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) @@ -1756,43 +2030,11 @@ def load_and_prepare_ep( ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {}) if not ep_plan: raise ValueError("expert-parallel PTQ requires the model to declare base_model_ep_plan") - if not _routed_expert_keys(ep_plan): - raise ValueError( - f"base_model_ep_plan={ep_plan} styles no routed experts ({', '.join(_EXPERTS_STYLES)} " - "or grouped_gemm), so this load would leave them whole on every rank." - ) - # Only grouped_gemm entries shard an expert's weights here: the torch_a2a styles the - # containers take for the load shard nothing themselves (any other style is refused below). - unsharded = sorted( - key - for key in _routed_expert_keys(ep_plan) - if not any( - style == "grouped_gemm" and name.startswith(key + ".") - for name, style in ep_plan.items() - ) - ) - if unsharded: - raise ValueError( - f"base_model_ep_plan={ep_plan} shards no weights of {', '.join(unsharded)} through " - "grouped_gemm, the only expert-weight style this load shards by, so every rank would " - "hold every expert." - ) - # The load drops any other style an expert weight takes, so that weight would load whole - # beside its sharded siblings, and the grouped forward would pair slices of different experts. - containers = _routed_expert_keys(ep_plan) - dropped = sorted( - name - for name, style in ep_plan.items() - if style not in kept_styles - and name not in containers - and any(name.startswith(key + ".") for key in containers) - ) - if dropped: - raise ValueError( - f"base_model_ep_plan={ep_plan} styles {', '.join(dropped)} by a style this load " - "drops, so those expert weights would load whole on every rank beside their sharded " - "siblings." - ) + # The refusals the plan decision gives way for, which an EP degree asked for, or a plan + # handed in, meets here: a model spec ruling expert parallelism out or stating a later + # transformers, and the plan's styles. + if (refusal := _ep_load_refusal(hf_config, ep_plan)) is not None: + raise refusal # The all-to-all transport needs every rank to hold as many experts, counted as the plan's # decision counts them; a count the config does not name is checked once the experts load. num_experts = _num_experts_from_config(ep_cfg) @@ -1844,7 +2086,7 @@ def gpu() -> tuple[bool, bool, str | None]: # The load and the export both read the local snapshot, made only once the checks pass, so a # refused load downloads nothing. with _refused_together("the local checkpoint", alone=False): - _, ckpt_path = ensure_local_checkpoint(ckpt_path) + _, ckpt_path = ensure_local_checkpoint(ckpt_path, **hub_options) # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over # the whole world (``dp``). @@ -1871,6 +2113,7 @@ def gpu() -> tuple[bool, bool, str | None]: "device_mesh": ep_submesh, "distributed_config": distributed_config, "trust_remote_code": trust_remote_code, + **hub_options, } if not _dc_has_tp_plan: from_pretrained_kwargs["tp_plan"] = "auto" diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py index 54b96461485..eb2685c68b2 100644 --- a/tests/gpu/torch/utils/test_model_load_utils.py +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -22,10 +22,12 @@ import pytest import torch +import transformers from _test_utils.torch.transformers_models import create_tiny_llama_dir +from packaging.version import Version from safetensors.torch import load_file from torch.distributed.tensor import DTensor -from transformers import AutoModelForCausalLM, LlamaConfig +from transformers import AutoModelForCausalLM, LlamaConfig, Qwen3MoeConfig from modelopt.torch.export.unified_export_hf import export_hf_checkpoint from modelopt.torch.utils.distributed import ( @@ -187,6 +189,7 @@ def test_carry_over_unplaced_weights(dist_workers, tmp_path): # ------------------------------------------------------------------ load_pretrained, left to decide +NUM_EXPERTS = 8 _TINY = { "hidden_size": 64, "intermediate_size": 128, @@ -210,9 +213,10 @@ def _full(tensor): def _load_and_check(rank, size, ckpt_dir, knobs, expected): """Load on every rank of the run with the plan left to decide, and check what each rank holds. - Left to decide, ``load_pretrained`` FSDP2-shards a dense model when asked to, and otherwise gives - every rank the whole model on its own GPU. Whichever load runs, every rank ends up with the - checkpoint's weights and the load's record of what it could not place. + Left to decide, ``load_pretrained`` shards an expert-parallel MoE's experts, FSDP2-shards a dense + model when asked to, and otherwise gives every rank the whole model on its own GPU. Whichever + load runs, every rank ends up with the checkpoint's weights and the load's record of what it + could not place. """ device = torch.device(f"cuda:{rank}") os.environ.update(knobs) @@ -227,14 +231,31 @@ def _load_and_check(rank, size, ckpt_dir, knobs, expected): if expected == "replicated": assert not sharded assert all(param.device == device for param in params.values()) - else: + elif expected == "fsdp2": assert {name for name in params if ".layers." in name} <= sharded + else: + # Each rank holds its contiguous share of the routed experts, and the rest whole. + experts = {name for name in params if ".mlp.experts." in name} + assert experts and not sharded - experts + per_rank = NUM_EXPERTS // size + for name in experts: + local = params.pop(name) + local = local.to_local() if isinstance(local, DTensor) else local + assert local.shape[0] == per_rank + owned = whole_params[name][rank * per_rank : (rank + 1) * per_rank] + torch.testing.assert_close(local.cpu(), owned) for name, param in params.items(): # the same order on every rank: full_tensor gathers torch.testing.assert_close(_full(param).cpu(), whole_params[name]) assert model._modelopt_unplaced_source_keys == [] assert model._modelopt_source_checkpoint == ckpt_dir - dense = "replicate" if expected == "replicated" else "fsdp" - assert distribution_plan(model) == DistributionPlan(size, dense=dense) + assert ( + distribution_plan(model) + == { + "replicated": DistributionPlan(size), + "fsdp2": DistributionPlan(size, dense="fsdp"), + "expert-parallel": DistributionPlan(size, ep_size=size), + }[expected] + ) finally: for knob in knobs: del os.environ[knob] @@ -248,3 +269,16 @@ def test_a_dense_model_loads_whole_on_every_rank(need_2_gpus, dist_workers, tmp_ def test_a_dense_model_asked_to_shard_loads_fsdp2_sharded(need_2_gpus, dist_workers, tmp_path): _save(LlamaConfig(**_TINY), tmp_path) dist_workers.run(_load_and_check, str(tmp_path), {"MODELOPT_FSDP_SIZE": "world"}, "fsdp2") + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("5.7"), + reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7", +) +def test_an_expert_parallel_moe_loads_with_its_experts_sharded(need_2_gpus, dist_workers, tmp_path): + if NUM_EXPERTS % dist_workers.world_size: + pytest.skip(f"{NUM_EXPERTS} experts do not shard over {dist_workers.world_size} GPUs") + config = Qwen3MoeConfig(moe_intermediate_size=32, num_experts=NUM_EXPERTS, **_TINY) + config.num_experts_per_tok = 2 + _save(config, tmp_path) + dist_workers.run(_load_and_check, str(tmp_path), {}, "expert-parallel") diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 5f07787ddee..355be7c3f72 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -54,6 +54,7 @@ AutoModelForImageTextToText, AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq, + DeepseekV3Config, FineGrainedFP8Config, LlamaConfig, LlamaForCausalLM, @@ -71,6 +72,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist import modelopt.torch.utils.moe as moe +from modelopt.torch.models import EPSpec, ModelSpec from modelopt.torch.quantization.plugins import accelerate as accelerate_plugin from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -233,6 +235,24 @@ def _qwen3_moe(num_experts=8): return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY) +def _deepseek_v3(n_routed_experts=8): + """DeepSeek-V3 without a plan of its own, so only its model spec can make it EP-capable. + + Its one layer is an MoE layer: DeepSeek-V3's first layers are dense by default. + """ + config = DeepseekV3Config( + n_routed_experts=n_routed_experts, + first_k_dense_replace=0, + moe_intermediate_size=8, + n_group=1, + topk_group=1, + num_experts_per_tok=1, + **_TINY, + ) + config.base_model_ep_plan = None # holds even if a later transformers ships a plan + return config + + def _bytes(module): return sum(p.numel() * p.element_size() for p in module.parameters()) @@ -440,14 +460,6 @@ def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, c assert ("Distribution:" in capsys.readouterr().out) == (rank == 0) -@needs_qwen3_moe_ep_plan -def test_no_ep_degree_dividing_the_experts_leaves_them_whole(run_as, serve, tmp_path): - run_as(2) - serve(_qwen3_moe(num_experts=3)) - - assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2) - - def test_an_moe_without_an_ep_plan_is_not_expert_parallel(run_as, serve, tmp_path): """The config counts 8 experts, but with no base_model_ep_plan nothing could shard them.""" run_as(4) @@ -464,17 +476,22 @@ def test_an_ep_plan_naming_no_module_of_the_model_is_refused( ): """A plan whose expert containers match no module of the built model (a renamed submodule, a hand-written plan) leaves no expert to count or shard: the automatic degree would replicate - every expert on every rank without a word. A degree set to 1 asks for exactly that.""" + every expert on every rank without a word. A degree of 1 with the dense part replicated asks + for exactly that.""" run_as(4) config = _qwen3_moe(num_experts=8) config.base_model_ep_plan = {"layers.*.mlp.renamed_experts": "moe_tp_experts"} serve(config) if degree is None: - with pytest.raises(ValueError, match=r"renamed_experts\) names .* MODELOPT_EP_DEGREE=1"): + with pytest.raises( + ValueError, + match=r"renamed_experts\) names .* MODELOPT_EP_DEGREE=1 and MODELOPT_FSDP_SIZE=1", + ): resolve_distribution_plan(str(tmp_path)) else: monkeypatch.setenv("MODELOPT_EP_DEGREE", degree) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "1") assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) @@ -484,7 +501,7 @@ def test_routed_experts_the_config_does_not_count_are_refused( ): """The mirror case: the plan's containers are in the built model, but the config counts its experts under no field the decision reads, which would leave every expert replicated on every - rank without a word. A degree set to 1 asks for exactly that.""" + rank without a word. A degree of 1 with the dense part replicated asks for exactly that.""" plan = {"layers.*.mlp.experts": "moe_tp_experts"} serve(SimpleNamespace(model_type="moe", architectures=[], base_model_ep_plan=plan)) built = model_with_decoder_at("model.") @@ -493,10 +510,13 @@ def test_routed_experts_the_config_does_not_count_are_refused( run_as(4) if degree is None: - with pytest.raises(ValueError, match=r"no expert count .* MODELOPT_EP_DEGREE=1"): + with pytest.raises( + ValueError, match=r"no expert count .* MODELOPT_EP_DEGREE=1 and MODELOPT_FSDP_SIZE=1" + ): resolve_distribution_plan(str(tmp_path)) else: monkeypatch.setenv("MODELOPT_EP_DEGREE", degree) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "1") assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) @@ -551,6 +571,107 @@ 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"), + pytest.param(_qwen3_moe(num_experts=8), 1, id="explicit_degree_one_as_an_argument"), + ], +) +def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( + run_as, serve, monkeypatch, tmp_path, config, ep_degree +): + """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none. With + nothing sharded either, every rank would hold the whole MoE: that must fail, not run. The ways + out are named as the caller set the plan, so arguments, which override the knobs, are named.""" + if isinstance(ep_degree, str): + monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) + arguments = {"ep_size": ep_degree, "dense": "auto"} if isinstance(ep_degree, int) else {} + run_as(2) + serve(config) + + with pytest.raises( + ValueError, match="can be expert-parallel, but the plan resolved" + ) as refused: + resolve_distribution_plan(str(tmp_path), **arguments) + if arguments: # a knob the arguments override would lead straight back here + assert "ep_size=1 with dense='replicate'" in str(refused.value), refused.value + assert "MODELOPT_" not in str(refused.value), refused.value + + +@needs_qwen3_moe_ep_plan +@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"], +) +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.""" + if ep_degree is not None: + monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) + run_as(2) + config = _qwen3_moe(num_experts=8) + if refused_by == "its-plan": + 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) + + 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 +): + """Asked for dense sharding, the MoE is sharded whole by FSDP2, which is not the waste of a + replica on every rank: an EP degree left to the decision gives way, as the expert-parallel + load replicates the dense part.""" + if ep_degree is not None: + monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + run_as(2) + serve(_qwen3_moe(num_experts=8)) + + 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", @@ -595,7 +716,9 @@ def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit( ): """At EP 4 a rank holds the rest of the model whole and a quarter of the routed experts, so the dense part is sharded when that does not fit one GPU, however small the rest; with no EP the - whole MoE must fit, and here it does not.""" + whole MoE must fit, and here it does not. The expert-parallel load replicates the dense part, + so an EP degree left to the decision gives way to FSDP2 for the whole model, where one asked + for stands, for the load to refuse.""" run_as(4) config = serve(_qwen3_moe(num_experts=8)) model = _meta(config) @@ -603,7 +726,9 @@ def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit( # A GPU 60% of which holds the rest of the model and that share of the experts. gpu_memory(int((_bytes(model) - experts + experts_fitting * experts) / 0.6)) - assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4, dense) + decided = DistributionPlan(4, 4) if dense == "replicate" else DistributionPlan(4, dense="fsdp") + assert resolve_distribution_plan(str(tmp_path))[0] == decided + assert resolve_distribution_plan(str(tmp_path), ep_size=4)[0] == DistributionPlan(4, 4, dense) assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan( 4, dense="fsdp" ) @@ -703,13 +828,17 @@ def test_the_expert_count_is_read_from_the_config_holding_the_plan( ): """A VLM whose nested config holds the MoE and its plan counts the experts there, not by an expert-ish field on its root config.""" - text = SimpleNamespace( - base_model_ep_plan={"layers.*.mlp.experts": "moe_tp_experts"}, num_experts=8 - ) + plan = { + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + } + text = SimpleNamespace(base_model_ep_plan=plan, num_experts=8) serve(SimpleNamespace(model_type="vlm", architectures=[], text_config=text, num_experts=3)) built = model_with_decoder_at("model.") built.tie_weights = lambda: None # a PreTrainedModel's; nothing here is tied monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: built) + monkeypatch.setattr(transformers, "__version__", "5.7.0") # one the expert-parallel load takes run_as(4) assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4) @@ -903,6 +1032,11 @@ def stub(*args, **kwargs): return stub + monkeypatch.setattr( + model_load_utils, + "load_and_prepare_ep", + recording("ep", model_load_utils.load_and_prepare_ep), + ) monkeypatch.setattr( model_load_utils, "parallel_load_and_prepare_fsdp2", @@ -1224,36 +1358,65 @@ def test_a_model_class_given_by_name_is_refused_by_name( @needs_qwen3_moe_ep_plan @pytest.mark.parametrize("knob", [None, "4"], ids=["left-to-decide", "asked-for"]) -def test_a_load_it_decides_keeps_the_experts_whole( - run_as, serve, loaders, monkeypatch, tmp_path, knob +def test_a_load_it_decides_shards_the_experts( + run_as, serve, loaders, a_gpu, monkeypatch, tmp_path, knob ): - """No load here shards the routed experts yet, so an MoE whose EP degree is left to the decision - is decided without expert parallelism and loads, where an automatic degree would be one the - load refuses; a degree MODELOPT_EP_DEGREE asks for is still refused, not ignored.""" + """An MoE whose EP degree is left to the decision is decided expert-parallel and loads through + the expert-parallel load, as one MODELOPT_EP_DEGREE asks for does.""" run_as(4) serve(_qwen3_moe()) if knob is not None: monkeypatch.setenv("MODELOPT_EP_DEGREE", knob) - monkeypatch.setattr( - transformers.AutoModelForCausalLM, - "from_pretrained", - classmethod(lambda cls, *a, **k: torch.nn.Linear(1, 1)), - ) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) - if knob is not None: - with pytest.raises(NotImplementedError, match=r"ep_size=4.*MODELOPT_EP_DEGREE=1"): - load_pretrained(str(tmp_path)) - return - assert distribution_plan(load_pretrained(str(tmp_path))) == DistributionPlan(4) + model = load_pretrained(str(tmp_path)) + assert loaders["ep"]["plan"] == DistributionPlan(4, ep_size=4) + assert distribution_plan(model) == DistributionPlan(4, ep_size=4) -def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, tmp_path): - """No load here shards the routed experts: an EP plan fails before anything loads.""" - run_as(2) - serve(Qwen3MoeConfig(num_experts=8, moe_intermediate_size=8, **_TINY)) +@pytest.mark.parametrize( + ("dense", "one_gpu_per_process"), + [("replicate", False), ("fsdp", False), ("replicate", True)], + ids=["replicate", "fsdp", "one-gpu-per-process"], +) +def test_an_expert_parallel_plan_loads_through_the_ep_loader( + run_as, serve, loaders, auto_class, a_gpu, 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.""" + run_as(4, rank=1 if one_gpu_per_process else 0) + if one_gpu_per_process: + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + config = serve(_deepseek_v3(n_routed_experts=8)) + plan = DistributionPlan(4, ep_size=4, dense=dense) + options = { + "trust_remote_code": True, + "attn_implementation": "eager", + "require_lm_head": None, + "local_files_only": True, + } - with pytest.raises(NotImplementedError, match=r"ep_size=2.*MODELOPT_EP_DEGREE=1"): - load_pretrained(str(tmp_path), auto_class, distribution=DistributionPlan(2, ep_size=2)) + if dense == "fsdp": + with pytest.raises( + NotImplementedError, match=r"load replicates it on every rank\..*\(MODELOPT_FSDP_SIZE=1" + ): + load_pretrained(str(tmp_path), distribution=plan, **options) + else: + model = load_pretrained(str(tmp_path), distribution=plan, **options) + call = loaders.pop("ep") + assert (call["ckpt_path"], call["plan"]) == (str(tmp_path), plan) + assert (call["trust_remote_code"], call["attn_implementation"]) == (True, "eager") + assert call["freeze"] is False + assert "require_lm_head" not in call # its default, not None read as False + assert call["hub_options"] == {"local_files_only": True} + assert call["hf_config"] is config + assert distribution_plan(model) == plan assert loaders == {} assert auto_class.calls == [] @@ -1443,6 +1606,7 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, cpu_dist_workers.run(_one_rank_fails_to_load_job, ckpt=str(tmp_path)) +@pytest.mark.parametrize("sharding", ["expert-parallel", "fsdp2"]) @pytest.mark.parametrize( ("change", "error", "match"), [ @@ -1456,7 +1620,7 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, ( {"auto_class": AutoModel}, NotImplementedError, - r"auto_class=AutoModel: leave it out.*DistributionPlan\(2, dense='replicate'\)", + r"auto_class=AutoModel: leave it out.*DistributionPlan\(2, (1, |dense=)'replicate'\)", ), ({"mtp": True}, NotImplementedError, "MTP head"), ( @@ -1464,27 +1628,37 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, TypeError, "config and hf_config", ), + ({"gpus": 0}, RuntimeError, "this rank has none"), ], - ids=["placement", "dtype", "torch_dtype", "model_class", "mtp_head", "two_configs"], + ids=["placement", "dtype", "torch_dtype", "model_class", "mtp_head", "two_configs", "no_gpu"], ) def test_a_sharded_load_refuses_what_it_cannot_honor( - run_as, serve, loaders, process_group, monkeypatch, tmp_path, 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.""" + """Placement, a model class, a head only a whole model builds, two configs to load with, a rank + without a GPU of its own: none reaches a sharded load, which fails before it starts rather than + load something other than what was asked.""" run_as(2) - 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) + if "gpus" in change: # a host that shows this rank no GPU + monkeypatch.setattr(torch.cuda, "device_count", lambda: change["gpus"]) + plan = DistributionPlan(2, ep_size=2) if expert_parallel else DistributionPlan(2, dense="fsdp") - with pytest.raises(error, match=match): + with pytest.raises(error, match=match) as refused: 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", {}) ) + if "device_map" in change.get("kwargs", {}) or "gpus" in change: + # The way to a whole model it names, for the plan handed in: an expert-parallel MoE's is a + # plan with an EP degree of 1. + whole_moe = "DistributionPlan(2, 1, 'replicate')" + assert (whole_moe in str(refused.value)) == expert_parallel, refused.value + if {"dtype", "torch_dtype"} & set(change.get("kwargs", {})): + # Only the FSDP2 load takes an mp_policy to compute in another dtype. + assert ("mp_policy" in str(refused.value)) != expert_parallel, refused.value assert loaders == {} @@ -1534,6 +1708,33 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( TypeError, r"\['cpu_offload'\] only apply", ), + ( + DistributionPlan(2), + "causal-lm", + {"freeze": False}, + TypeError, + r"\['freeze'\] only apply to the expert-parallel load", + ), + ( + DistributionPlan(2, dense="fsdp"), + "causal-lm", + {"require_lm_head": False}, + TypeError, + ( + r"\['require_lm_head'\] only apply to the expert-parallel load, and the FSDP2 " + r"load .*; drop them, or pass a plan with ep_size > 1" + ), + ), + ( + DistributionPlan(2, ep_size=2), + "causal-lm", + {"broadcast_chunk_size": None}, + TypeError, + ( + r"\['broadcast_chunk_size'\] only apply to the FSDP2 load, and the " + r"expert-parallel load .*; drop them, or pass a plan with ep_size=1 and dense='fsdp'" + ), + ), ], ids=[ "another-world", @@ -1543,6 +1744,9 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( "not-a-plan", "hf-config-on-a-whole-load", "fsdp2-option-on-a-whole-load", + "ep-option-on-a-whole-load", + "ep-option-on-the-fsdp2-load", + "fsdp2-chunking-as-none-on-the-ep-load", ], ) def test_a_load_refuses_what_it_cannot_load( @@ -1562,8 +1766,8 @@ def test_a_load_refuses_what_it_cannot_load( read as stored, even when only a VLM's language model is quantized and its remote code builds it as a causal LM; a VLM but a text-only one it supports, which FSDP2 would build through AutoModelForCausalLM, refused or as its language model alone; anything but a plan; and the - sharded loads' config, or an option only the FSDP2 load takes, which from_pretrained would take - for a model argument and fail far from here, or drop.""" + sharded loads' config, or an option only one sharded load takes, which from_pretrained would + take for a model argument and fail far from here, or drop.""" run_as(2) config = serve(transformers.LlavaConfig() if checkpoint == "vlm" else LlamaConfig(**_TINY)) if checkpoint == "quantized": @@ -1644,6 +1848,58 @@ def test_a_sharded_load_is_refused_inside_init_quantized_weights( assert loaders == {} +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize("handed", [False, True], ids=["knob", "handed"]) +def test_an_expert_parallel_load_is_refused_inside_init_quantized_weights( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, handed +): + """The expert-parallel load never calls from_pretrained either, and its way to a whole model on + every rank sets the EP degree to 1 as well, as the decision refuses to replicate an MoE unasked.""" + run_as(2, rank=1) + serve(_qwen3_moe(num_experts=8)) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) + monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") + plan = {"distribution": DistributionPlan(2, 2)} if handed else {} + + with ( + init_quantized_weights(mtq.INT8_DEFAULT_CFG), + pytest.raises( + NotImplementedError, + match=r"expert-parallel load .*init_quantized_weights.*DistributionPlan\(2, 1, 'replicate'\)", + ), + ): + load_pretrained(str(tmp_path), **plan) + assert loaders == {} + + +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize( + "options", + [{"device_map": "auto"}, {"freeze": True, "require_lm_head": False}], + ids=["placement", "its-own-options"], +) +def test_an_automatic_expert_parallel_layout_says_so_when_it_refuses_a_placement( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, options +): + """With neither a plan nor MODELOPT_EP_DEGREE, the decision shards an MoE's experts, and the + expert-parallel load refuses options that place a whole model. The refusal says the layout was + decided rather than asked for, and how to keep the load whole. The load's own options place + nothing, and reach it.""" + run_as(2, rank=1) + serve(_qwen3_moe(num_experts=8)) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) + + if "device_map" not in options: + load_pretrained(str(tmp_path), **options) + assert {name: loaders["ep"][name] for name in options} == options + return + with pytest.raises( + TypeError, match=r"\['device_map'\] do not apply.*no plan or MODELOPT_EP_DEGREE" + ): + load_pretrained(str(tmp_path), **options) + assert loaders == {} + + @pytest.mark.parametrize( ("checkpoint", "match"), [("quantized", "dequantize"), ("mtp", "MTP head")], @@ -1809,12 +2065,7 @@ def test_a_launch_that_sets_no_group_up_loads_whole_on_each_process( RuntimeError, "MODELOPT_FSDP_SIZE=world", ), - ( - None, - {"MODELOPT_EP_DEGREE": "4"}, - RuntimeError, - r"MODELOPT_EP_DEGREE=4 .*distributed_run\(\).*MODELOPT_EP_DEGREE=1", - ), + (None, {"MODELOPT_EP_DEGREE": "4"}, RuntimeError, "MODELOPT_EP_DEGREE=4"), (None, {"MODELOPT_EP_DEGREE": "eigth"}, ValueError, "MODELOPT_EP_DEGREE='eigth'"), ], ids=["several-ranks", "sharded", "fsdp-knob", "ep-knob", "ep-knob-not-a-degree"], @@ -1854,6 +2105,9 @@ def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( # ------------------------------------------------------------------------------- the plan +# The model type of the configs that stand in for a checkpoint's here: no model spec names it. +_NO_SPEC = "toy_moe" + _EP_PLAN = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", @@ -1971,7 +2225,7 @@ def resolve(path, **hub_options): load_and_prepare_ep( "org/model", DistributionPlan(2, ep_size=2), - hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), + hf_config=SimpleNamespace(model_type=_NO_SPEC, base_model_ep_plan=_EP_PLAN), ) assert resolved == ["org/model"] @@ -1998,7 +2252,7 @@ def resolve(path, **hub_options): load_and_prepare_ep( "org/model", DistributionPlan(4, ep_size=4), - hf_config=SimpleNamespace(text_config=text, num_experts=3), + hf_config=SimpleNamespace(model_type=_NO_SPEC, text_config=text, num_experts=3), ) @@ -2015,14 +2269,17 @@ def resolve(path, **hub_options): ), (SimpleNamespace(), ValueError, "declare base_model_ep_plan"), ( - SimpleNamespace(base_model_ep_plan={"layers.*.self_attn.q_proj": "colwise"}), + SimpleNamespace( + model_type=_NO_SPEC, base_model_ep_plan={"layers.*.self_attn.q_proj": "colwise"} + ), ValueError, "styles no routed experts", ), *( ( SimpleNamespace( - base_model_ep_plan={"layers.*.mlp.experts": "moe_tp_experts", **params} + model_type=_NO_SPEC, + base_model_ep_plan={"layers.*.mlp.experts": "moe_tp_experts", **params}, ), ValueError, r"shards no weights of layers\.\*\.mlp\.experts through grouped_gemm", @@ -2031,11 +2288,17 @@ def resolve(path, **hub_options): ), ( SimpleNamespace( - base_model_ep_plan={**_EP_PLAN, "layers.*.mlp.experts.down_proj": "packed_rowwise"} + model_type=_NO_SPEC, + base_model_ep_plan={**_EP_PLAN, "layers.*.mlp.experts.down_proj": "packed_rowwise"}, ), ValueError, r"styles layers\.\*\.mlp\.experts\.down_proj by a style this load drops", ), + ( + SimpleNamespace(model_type="gpt_oss", base_model_ep_plan=_EP_PLAN), + NotImplementedError, + "the expert-parallel export cannot write that layout", + ), ], ids=[ "gptq", @@ -2047,6 +2310,7 @@ def resolve(path, **hub_options): "experts-container-only", "experts-sharded-otherwise", "experts-sharded-partly-otherwise", + "spec-refuses-ep", ], ) def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( @@ -2056,8 +2320,10 @@ def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( or, MXFP4, dequantize into GPT-OSS's transposed expert layout whatever the model's, or a plan that leaves the experts whole on every rank: one naming no routed experts, or none of whose entries shards their weights by grouped_gemm, the only expert-weight style the load - shards by, or one that styles some of their weights otherwise. The load refuses each before it - reads anything.""" + shards by, or one that styles some of their weights otherwise. So is a model whose spec rules + expert parallelism out (GPT-OSS's), which the plan decision plans without it: an EP degree asked + for, or a plan handed in, meets that refusal here. The load refuses each before it reads + anything.""" with pytest.raises(error, match=match): load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) assert ep_world_of_two == [] @@ -2069,7 +2335,7 @@ def test_an_ep_size_that_does_not_divide_the_experts_is_refused_before_anything_ """The all-to-all transport finds an expert's rank by dividing its id by this rank's share of the experts: with 3 experts over 2 ranks the ranks would disagree on who holds which, and send tokens to the wrong experts.""" - config = SimpleNamespace(base_model_ep_plan=_EP_PLAN, num_experts=3) + config = SimpleNamespace(model_type=_NO_SPEC, base_model_ep_plan=_EP_PLAN, num_experts=3) with pytest.raises(ValueError, match="ep_size=2 must divide the 3 routed experts"): load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) @@ -2092,7 +2358,8 @@ def from_pretrained(path, **kwargs): monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) config = SimpleNamespace( - base_model_ep_plan={"layers.*.mlp.experts.*.gate_up_proj": "grouped_gemm"} + model_type=_NO_SPEC, + base_model_ep_plan={"layers.*.mlp.experts.*.gate_up_proj": "grouped_gemm"}, ) with pytest.raises(ValueError, match="that no module of"): @@ -2129,6 +2396,7 @@ def next_check(*args, **kwargs): monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) class _Config(SimpleNamespace): + model_type = _NO_SPEC base_model_ep_plan = _EP_PLAN # 8 experts over 2 ranks: the stand-in model holds this rank's 4. @@ -2206,7 +2474,9 @@ def next_check(*args, **kwargs): def test_the_weights_load_in_the_checkpoints_dtype_through_its_class(ep_world_of_two, monkeypatch): """The checkpoint's ``dtype`` reaches ``from_pretrained``, and the class that builds it is chosen - with the trust the caller gave its code.""" + with the trust the caller gave its code. The Hub options reach the checkpoint's download, the + config read and ``from_pretrained``, which fetch the custom code an ``auto_map`` names in another + repo.""" class _LoadedError(Exception): """Stops the load at from_pretrained.""" @@ -2227,13 +2497,28 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", auto_class) - config = SimpleNamespace(dtype=torch.float16, base_model_ep_plan=_EP_PLAN) + config = SimpleNamespace(model_type=_NO_SPEC, dtype=torch.float16, base_model_ep_plan=_EP_PLAN) + reads, downloads = [], [] + monkeypatch.setattr( + "transformers.AutoConfig.from_pretrained", lambda *a, **k: reads.append(k) or config + ) + monkeypatch.setattr( + model_load_utils, + "ensure_local_checkpoint", + lambda p, **k: downloads.append(k) or (None, p), + ) with pytest.raises(_LoadedError): load_and_prepare_ep( - "org/model", DistributionPlan(2, ep_size=2), trust_remote_code=True, hf_config=config + "org/model", + DistributionPlan(2, ep_size=2), + trust_remote_code=True, + hub_options={"local_files_only": True}, ) assert loaded["dtype"] is torch.float16 + assert [read["local_files_only"] for read in reads] == [True] + assert downloads == [{"local_files_only": True}] + assert loaded["local_files_only"] is True assert classes == [True] @@ -2398,7 +2683,7 @@ def from_pretrained(path, **kwargs): load_and_prepare_ep( "org/model", DistributionPlan(2, ep_size=2), - hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), + hf_config=SimpleNamespace(model_type=_NO_SPEC, base_model_ep_plan=_EP_PLAN), ) given = {k: v for k, v in vars(loaded["distributed_config"]).items() if v is not None} @@ -2422,7 +2707,9 @@ def _meshes_in_step(*args, **kwargs): def _one_rank_fails_job(rank, size, ckpt, failure, failing_rank, match): path, hf_config, patches = ckpt, None, [] if failure != "unreadable-config": - hf_config = SimpleNamespace(base_model_ep_plan=dict(_EP_PLAN), num_experts=4) + hf_config = SimpleNamespace( + model_type=_NO_SPEC, base_model_ep_plan=dict(_EP_PLAN), num_experts=4 + ) if failure != "unseen-checkpoint": local = mock.patch.object( model_load_utils, "ensure_local_checkpoint", side_effect=lambda p, **k: (None, p) @@ -2505,7 +2792,9 @@ def test_a_load_failing_on_one_rank_fails_every_rank( def _gpu_layout_job(rank, size, visible, bound, shared, refused): - hf_config = SimpleNamespace(base_model_ep_plan=dict(_EP_PLAN), num_experts=4) + hf_config = SimpleNamespace( + model_type=_NO_SPEC, base_model_ep_plan=dict(_EP_PLAN), num_experts=4 + ) reached = None with bounded_collectives(), ExitStack() as stack: # This rank is on cuda:0, bound to it or not yet; the local ranks run 0 and 1.