diff --git a/modelopt/torch/utils/env.py b/modelopt/torch/utils/env.py new file mode 100644 index 00000000000..7e27e7857c3 --- /dev/null +++ b/modelopt/torch/utils/env.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The ``MODELOPT_*`` environment-variable knobs of distributed HuggingFace model loading. + +Each knob is a small accessor that reads ``os.environ`` *at call time*, not at import: a launcher sets +them before it forks the distributed workers, and tests may override them, so the live value is what +matters. Callers do ``from modelopt.torch.utils import env`` then ``env.ep_degree()``. + +Distribution plan (see ``resolve_distribution_plan``): + +- ``MODELOPT_EP_DEGREE``: expert-parallel degree; ``""``/unset -> ``"auto"``. +- ``MODELOPT_FSDP_SIZE``: the dense layout. ``"1"`` replicates it; ``"world"``, or the world size as a + number, FSDP2-shards it over every rank; ``""``/unset -> ``"auto"``. +""" + +import os + +__all__ = [ + "ep_degree", + "fsdp_size", +] + + +def ep_degree() -> str: + """Expert-parallel degree knob for ``resolve_distribution_plan``; ``""``/unset -> ``"auto"``.""" + return os.environ.get("MODELOPT_EP_DEGREE") or "auto" + + +def fsdp_size() -> str: + """Dense-sharding knob for ``resolve_distribution_plan``; ``""``/unset -> ``"auto"``.""" + return os.environ.get("MODELOPT_FSDP_SIZE") or "auto" diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index fc62fe9f71c..98c4bc2453c 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -15,7 +15,8 @@ """HuggingFace-coupled MoE expert-parallel (EP) helpers. -The DP=EP all-to-all expert transport (``torch_a2a``). +Finding a model's routed experts, choosing the EP degree they are sharded over, and the DP=EP +all-to-all expert transport (``torch_a2a``). The all-to-all transport is plain ``torch.distributed.all_to_all_single`` (NCCL) wrapped around transformers' ``_grouped_linear``. ``register_torch_a2a()`` registers it into @@ -24,11 +25,103 @@ """ import warnings +from typing import TYPE_CHECKING import torch import torch.distributed as dist from torch.distributed.nn.functional import all_to_all_single as autograd_all_to_all_single +if TYPE_CHECKING: + from torch import nn + + +def _ep_config(hf_config): + """Return the config that carries the EP plan / expert count. + + For a plain MoE causal LM this is the root config. For a VLM whose nested text submodel + is the MoE (e.g. ``KimiK25.text_config`` -> deepseek_v3, or ``NemotronH_Omni.llm_config`` + -> nemotron_h), this is that nested config. + """ + if getattr(hf_config, "base_model_ep_plan", None): + return hf_config + for _attr in ("text_config", "llm_config"): + sub = getattr(hf_config, _attr, None) + if sub is not None and getattr(sub, "base_model_ep_plan", None): + return sub + # Fall through: no plan anywhere; let callers decide. + return hf_config + + +def _num_experts_from_config(hf_config) -> int: + """Return the MoE expert count from an HF config (0 if the model is not MoE). + + Recurses into ``.text_config`` / ``.llm_config`` so VLMs whose nested text submodel is the MoE + (Kimi-K2.7, Llama-4 multimodal, NemotronH-Omni, ...) report the LM's expert count. + """ + for cfg in ( + hf_config, + getattr(hf_config, "text_config", None), + getattr(hf_config, "llm_config", None), + ): + if cfg is None: + continue + for attr in ("num_experts", "num_local_experts", "n_routed_experts", "num_routed_experts"): + value = getattr(cfg, attr, None) + if value: + return int(value) + return 0 + + +def _choose_ep_degree(world_size: int, num_experts: int) -> int: + """The largest EP degree dividing both ``world_size`` and ``num_experts``; ``1`` if none does. + + EP must divide ``num_experts`` (whole experts per rank) and ``world_size`` (a valid mesh + factorization). Each expert shard is then replicated over the remaining ranks + (``DistributionPlan.edp``). The largest degree holds the fewest experts per rank, which leaves + the most memory for activations, i.e. for calibration batches; any degree above 1 pays the same + two all-to-alls per MoE layer. Heuristic -- callers may override. + """ + if num_experts <= 0: + return 1 + return next( + (ep for ep in range(world_size, 0, -1) if world_size % ep == 0 and num_experts % ep == 0), + 1, + ) + + +# The styles a base_model_ep_plan gives a layer's routed-experts container. A plan that styles the +# experts' grouped-GEMM parameters names the containers through their parents too, so a plan that +# styles some containers one way and others the other names them all. +_EXPERTS_STYLES = ("moe_tp_experts", "nemotron_experts_ep") + + +def _routed_expert_keys(ep_plan: dict[str, str]) -> set[str]: + """The keys of ``ep_plan``, a ``base_model_ep_plan``, that name a routed-experts container.""" + return {key for key, style in ep_plan.items() if style in _EXPERTS_STYLES} | { + key.rsplit(".", 1)[0] for key, style in ep_plan.items() if style == "grouped_gemm" + } + + +def _routed_expert_modules(model: "nn.Module", ep_plan: dict[str, str]) -> "dict[str, nn.Module]": + """``model``'s routed-expert containers, ``{name: module}``, as its ``base_model_ep_plan`` names them. + + A module is one when its name, with layer indices wildcarded as transformers does, is a container + key or ends with one: the keys are relative to the base model, so they omit the ``model.`` or + ``language_model.model.`` a causal LM or a VLM puts in front. Deciding a model's distribution and + loading it expert-parallel both find the experts this way, so the two cannot disagree. + """ + # Local: transformers is the optional [hf] extra, and this module must import without it. + from transformers.integrations.tensor_parallel import replace_layer_number_by_wildcard + + keys = _routed_expert_keys(ep_plan) + found = {} + for name, module in model.named_modules(): + generic = replace_layer_number_by_wildcard(name) + if any(generic == key or generic.endswith("." + key) for key in keys): + found[name] = module + return found + + # --------------------------------------------------------------------------------------------------- # DP=EP all-to-all expert transport (`torch_a2a`) # diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 45a8e3ecff5..babf7f11054 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""HuggingFace model loading: :func:`load_pretrained`, and parallel FSDP2 loading.""" +"""HuggingFace model loading: :func:`load_pretrained`, and the plan to distribute a model by.""" import logging import os @@ -33,17 +33,35 @@ from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict from torch.distributed.fsdp import FSDPModule from torch.distributed.tensor import DTensor -from transformers import AutoConfig, AutoModel, AutoModelForCausalLM +from transformers import ( + AutoConfig, + AutoModel, + AutoModelForCausalLM, + AutoModelForImageTextToText, + AutoModelForSeq2SeqLM, + AutoModelForSpeechSeq2Seq, + PretrainedConfig, +) 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.utils import distributed as dist +from modelopt.torch.utils import env from modelopt.torch.utils.distributed import ( + DistributionPlan, barrier, broadcast_state_dict, fsdp2_wrap, is_initialized, ) from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 +from modelopt.torch.utils.moe import ( + _choose_ep_degree, + _ep_config, + _num_experts_from_config, + _routed_expert_keys, + _routed_expert_modules, +) from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( checkpoint_tensor_files, ensure_local_checkpoint, @@ -194,8 +212,19 @@ def build_meta_causal_lm( trust_remote_code: bool, attn_implementation: str | None, hf_config=None, + auto_class: Any = None, ): - """Build a meta-init causal LM (no real storage allocated).""" + """Build a meta-init causal LM (no real storage allocated), or ``auto_class``'s model. + + It leaves ``auto_class`` as it found it, and the auto classes a remote model's ``__init__`` + builds its parts through (``AutoModel``, ``AutoModelForCausalLM``). Building a remote-code + model through an auto class registers its class there, as it is, and a later load through that + auto class builds the registered class, without the ``GenerationMixin`` a first load adds: a + rank that built one more meta model than its peers, as rank 0 deciding a distribution does, + would load its model without the ``generate`` theirs have. That assumes a remote model registers + its classes while it is built: one its module registered when imported is undone as well, and a + later load in the process, which finds the module imported, does not register it again. + """ if hf_config is None: config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code} if attn_implementation is not None: @@ -207,16 +236,27 @@ def build_meta_causal_lm( dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 from accelerate import init_empty_weights # only real callers of this function need it - model_factory = AutoModelForCausalLM - if is_fsdp2_text_only_vlm(hf_config): - # Only newer Transformers versions supporting these models expose this factory. - from transformers import AutoModelForImageTextToText - - model_factory = AutoModelForImageTextToText - with init_empty_weights(include_buffers=False): - model = model_factory.from_config( - hf_config, dtype=dtype, trust_remote_code=trust_remote_code - ) + model_factory = auto_class + if model_factory is None: + model_factory = AutoModelForCausalLM + if is_fsdp2_text_only_vlm(hf_config): + model_factory = AutoModelForImageTextToText + # Where transformers records the class a remote-code build registers: in the auto class built + # through, and in those a remote VLM builds its towers through. Read without a fallback: if a + # release moves it, this fails here rather than leave one rank's model without generate. + kept = [ + (cls._model_mapping._extra_content, dict(cls._model_mapping._extra_content)) + for cls in (model_factory, AutoModel, AutoModelForCausalLM) + ] + try: + with init_empty_weights(include_buffers=False): + model = model_factory.from_config( + hf_config, dtype=dtype, trust_remote_code=trust_remote_code + ) + finally: + for registered, saved in kept: + registered.clear() + registered.update(saved) model.eval() return model @@ -882,3 +922,249 @@ def parallel_load_and_prepare_fsdp2( if hasattr(model, "tie_weights"): model.tie_weights() return model + + +def _auto_class(hf_config, trust_remote_code: bool = False) -> Any: + """The auto class that builds ``hf_config``'s whole model, LM head included. + + A causal LM builds through ``AutoModelForCausalLM``. A VLM -- a nested text config, a vision + config, or a ``*ForConditionalGeneration`` architecture -- builds through the remote + ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` there, + its code is trusted and it registers no image-text class, else through + ``AutoModelForImageTextToText`` if that recognizes the config or, its code trusted, the + checkpoint registers one, else through ``AutoModelForCausalLM`` if that recognizes it (as it + does Phi-4-multimodal), else through ``AutoModel``, which leaves the head out. An + encoder-decoder model (``is_encoder_decoder``, whether or not it names its architecture, which + reads like a VLM's) builds through its seq2seq class. + """ + architectures = getattr(hf_config, "architectures", None) or [] + is_vlm = ( + getattr(hf_config, "text_config", None) is not None + or getattr(hf_config, "vision_config", None) is not None + or any("conditionalgeneration" in str(arch).lower() for arch in architectures) + ) + if not is_vlm and not getattr(hf_config, "is_encoder_decoder", False): + return AutoModelForCausalLM + auto_map = getattr(hf_config, "auto_map", None) or {} + remote_with_head = ( + trust_remote_code + and "conditionalgeneration" in str(auto_map.get("AutoModelForCausalLM", "")).lower() + ) + if remote_with_head and "AutoModelForImageTextToText" not in auto_map: + return AutoModelForCausalLM + for seq2seq in (AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq): + if type(hf_config) in seq2seq._model_mapping: + return seq2seq + if type(hf_config) in AutoModelForImageTextToText._model_mapping or ( + trust_remote_code and "AutoModelForImageTextToText" in auto_map + ): + return AutoModelForImageTextToText + if type(hf_config) in AutoModelForCausalLM._model_mapping: + return AutoModelForCausalLM + return AutoModel + + +# The share of one GPU's memory a rank may take with the dense part replicated: that whole copy and +# the rank's share of the routed experts. The rest is left to calibration's activations and +# communication buffers. +_REPLICATED_SHARE = 0.6 + + +def _gpu_memory() -> int | None: + """This rank's GPU memory in bytes, or ``None`` without a GPU.""" + if not torch.cuda.is_available(): + return None + return torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory + + +def _dense_from_env(world_size: int) -> str: + """``MODELOPT_FSDP_SIZE`` as a dense layout, or ``"auto"``.""" + value = env.fsdp_size().lower() + if value == "auto": + return "auto" + if value == "1": + return "replicate" + if value in ("world", str(world_size)): + return "fsdp" + raise ValueError( + f"MODELOPT_FSDP_SIZE={value!r}: expected auto, 1 (replicate the dense part) or world " + f"(FSDP2-shard it over all {world_size} ranks)." + ) + + +def _ep_degree(ep_size) -> int | str: + """``ep_size`` as an EP degree or ``"auto"``; ``None`` reads ``MODELOPT_EP_DEGREE``.""" + value = env.ep_degree() if ep_size is None else ep_size + if str(value).lower() in ("auto", "0"): + return "auto" + try: + return int(value) + except ValueError: + name = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size" + raise ValueError(f"{name}={value!r}: expected auto, 0 (auto) or an EP degree.") from None + + +def _decide_distribution( + ckpt_path: str, hf_config, world_size: int, ep_size, dense, trust_remote_code: bool +) -> 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. + knob = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size" + 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 + # own arguments still count, as dense="fsdp" for FSDP2 on one rank. + single = 1 if ep_size is None else _ep_degree(ep_size) + return DistributionPlan( + ep_size=single if isinstance(single, int) else 1, + dense="replicate" if dense in (None, "auto") else dense, + ) + ep_size = _ep_degree(ep_size) + dense = _dense_from_env(world_size) if dense is None else dense + try: + model = build_meta_causal_lm( + ckpt_path, + trust_remote_code, + None, + hf_config, + auto_class=_auto_class(hf_config, trust_remote_code), + ) + # init_empty_weights re-wraps every parameter it registers, which unties what from_config + # tied: re-tie, so a tied embedding and LM head count once, as the loaded model holds them. + model.tie_weights() + except Exception as e: + raise RuntimeError( + f"Deciding how to distribute {hf_config.model_type} needs its model built on the meta " + f"device, which failed: {type(e).__name__}: {e}" + ) from e + ep_plan = getattr(_ep_config(hf_config), "base_model_ep_plan", None) or {} + 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 + # without them would replicate every expert on every rank, and the expert-parallel load + # finds them the same way. + containers = ", ".join(sorted(_routed_expert_keys(ep_plan))) or "none" + 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." + ) + 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 + # would leave every expert replicated on every rank. + 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." + ) + if ep_size == "auto": + ep_size = _choose_ep_degree(world_size, num_experts) + 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): + raise ValueError( + f"{knob}={ep_size} must divide the routed experts {hf_config.model_type} shards " + f"({num_experts})." + ) + 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. + sharded = ( + {id(p) for m in experts.values() for p in m.parameters()} if ep_size > 1 else set() + ) + held = sum( + p.numel() * p.element_size() // (ep_size if id(p) in sharded else 1) + for p in model.parameters() + ) + memory = _gpu_memory() + too_big = memory is not None and held > _REPLICATED_SHARE * memory + dense = "fsdp" if too_big else "replicate" + return DistributionPlan(world_size, ep_size, dense) + + +def resolve_distribution_plan( + ckpt_path: str, + *, + trust_remote_code: bool = False, + ep_size: int | str | None = None, + dense: str | None = None, + revision: str | None = None, + token: str | bool | None = None, + cache_dir: str | os.PathLike | None = None, + local_files_only: bool = False, + subfolder: str | None = None, +) -> tuple[DistributionPlan, PretrainedConfig]: + """Decide how this run distributes ``ckpt_path``'s model over its ranks. + + The one place a run decides it: a distributed load loads the model by the plan this returns and + records it on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), and + every later step reads it back from there + (:func:`~modelopt.torch.utils.distributed.distribution_plan`). Returns ``(plan, hf_config)``, + ``hf_config`` being this rank's own copy of the config (rank 0 decides the plan from its copy), + for a sharded load to build from. + + Rank 0 decides, from the model built on the meta device: from its parameters rather than + estimates of them, and from its routed experts as an expert-parallel load finds them. Every rank + takes rank 0's plan (:func:`~modelopt.torch.utils.distributed.broadcast_from`): the inputs are + read locally, ranks deciding for themselves could disagree, and ranks loading different plans + 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``. + * ``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"`` + assumes the run's GPUs are alike; a rank 0 without a GPU replicates it. + + ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` (``1`` replicates; ``world``, + or the world size as a number, shards) from rank 0's environment, both defaulting to + ``"auto"``. A single process reads neither and holds the whole model, unless ``dense="fsdp"`` + asks for FSDP2 on one rank. + + ``revision``, ``token``, ``cache_dir``, ``local_files_only`` and ``subfolder`` locate the + checkpoint as :func:`load_pretrained` takes them, so the plan is decided for what the load reads: + all but ``subfolder`` resolve a Hub model ID (:func:`ensure_local_checkpoint`), and the config + read takes all but ``revision``, for custom code its ``auto_map`` names in another Hub repo. It + returns no path, so a caller that loads the checkpoint resolves it first, with + :func:`ensure_local_checkpoint`, and passes the local directory: the plan is then decided for + the snapshot the load reads. + """ + # 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. + read: list[tuple[str, PretrainedConfig]] = [] + + def checkpoint_read() -> bool: + _, local_path = ensure_local_checkpoint( + ckpt_path, + revision=revision, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + ) + if subfolder: + local_path = os.path.join(local_path, subfolder) + config = AutoConfig.from_pretrained( + local_path, + trust_remote_code=trust_remote_code, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + ) + read.append((local_path, config)) + return True + + dist.all_reduce("checkpoint read", checkpoint_read, all) + ckpt_path, hf_config = read[0] + world_size = dist.size() + plan = dist.broadcast_from( + "plan", + lambda: _decide_distribution( + ckpt_path, hf_config, world_size, ep_size, dense, trust_remote_code + ), + ) + if plan != DistributionPlan(): + print_rank_0(f"Distribution: {plan!r}") + return plan, hf_config diff --git a/tests/_test_utils/torch/moe_models.py b/tests/_test_utils/torch/moe_models.py new file mode 100644 index 00000000000..709944f7b34 --- /dev/null +++ b/tests/_test_utils/torch/moe_models.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Plain-module stand-ins for an MoE model: a decoder whose layers hold routed experts, under any +prefix a causal LM or a VLM puts it. +""" + +import torch +from torch import nn + + +class MoeLayer(nn.Module): + """A decoder layer: attention, and an MLP of a router, routed experts and shared experts.""" + + def __init__(self): + super().__init__() + self.self_attn = nn.Linear(2, 2) + self.mlp = nn.Module() + self.mlp.gate = nn.Linear(2, 4, bias=False) # the router + self.mlp.experts = nn.Module() + self.mlp.experts.gate_up_proj = nn.Parameter(torch.zeros(4, 4, 2)) + self.mlp.experts.down_proj = nn.Parameter(torch.zeros(4, 2, 2)) + self.mlp.shared_experts = nn.Linear(2, 4, bias=False) + + +def model_with_decoder_at(prefix): + """A two-layer MoE whose decoder (``layers.``) sits under ``prefix`` (``""``: the decoder).""" + decoder = nn.Module() + decoder.embed_tokens = nn.Embedding(8, 2) + decoder.layers = nn.ModuleList(MoeLayer() for _ in range(2)) + if not prefix: + return decoder + root = parent = nn.Module() + *path, leaf = prefix.rstrip(".").split(".") + for name in path: + parent.add_module(name, nn.Module()) + parent = getattr(parent, name) + parent.add_module(leaf, decoder) + return root diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index c2a90c33fb0..2c225dafc14 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -13,28 +13,61 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Pure-function tests for ``modelopt.torch.utils.plugins.model_load_utils``.""" +"""Tests for ``modelopt.torch.utils.plugins.model_load_utils``. + +Nothing here needs a GPU. The tests of ``resolve_distribution_plan`` answer the process group's rank +and size as a torchrun launch would, and serve tiny real transformers configs from memory, so the +decision builds their real models on the meta device; only the last ones use a process group, two +gloo ranks. +""" + +import math +import os +from contextlib import contextmanager +from types import SimpleNamespace +from unittest import mock import pytest import torch +from _test_utils.torch.distributed.utils import bounded_collectives +from _test_utils.torch.moe_models import model_with_decoder_at from safetensors.torch import save_file pytest.importorskip("accelerate") -pytest.importorskip("transformers") +transformers = pytest.importorskip("transformers") from _test_utils.torch.transformers_models import ( _TinyQwen35Tokenizer, get_tiny_qwen3_5_moe_vl, get_tiny_qwen3_5_vl, ) -from transformers import AutoConfig +from packaging.version import Version +from transformers import ( + AutoConfig, + AutoModel, + AutoModelForCausalLM, + AutoModelForImageTextToText, + AutoModelForSeq2SeqLM, + AutoModelForSpeechSeq2Seq, + LlamaConfig, + Phi4MultimodalConfig, + Qwen2_5_VLConfig, + Qwen3MoeConfig, + T5Config, + WhisperConfig, +) +import modelopt.torch.utils.distributed as dist from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector +from modelopt.torch.utils.distributed import DistributionPlan +from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import ( + _auto_class, _conversion_plan, _convert_keys, _resolve_target, build_meta_causal_lm, record_unplaced_source_keys, + resolve_distribution_plan, unplaced_source_keys, ) @@ -153,3 +186,607 @@ def test_meta_qwen_vlm_preserves_architecture_and_vision(model_type): == config.text_config.num_hidden_layers ) assert all(parameter.is_meta for parameter in model.parameters()) + + +# -------------------------------------------------------------- deciding how a model is distributed + + +_TINY = { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "vocab_size": 32, +} + +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", +) + + +def _qwen3_moe(num_experts=8): + """Qwen3-MoE, whose base_model_ep_plan ships with transformers itself.""" + return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY) + + +def _bytes(module): + return sum(p.numel() * p.element_size() for p in module.parameters()) + + +def _meta(config): + """``config``'s model as the decision builds it, in the same dtype.""" + return model_load_utils.build_meta_causal_lm(None, False, None, config) + + +@pytest.fixture(autouse=True) +def _no_distribution_knobs(monkeypatch): + """A shell that exports the knobs would otherwise steer every decision below.""" + monkeypatch.delenv("MODELOPT_EP_DEGREE", raising=False) + monkeypatch.delenv("MODELOPT_FSDP_SIZE", raising=False) + + +def _gathered_alike(gathered, mine, group=None): + """Answer a gather as if every rank handed it what this rank did.""" + gathered[:] = [mine] * len(gathered) + + +@pytest.fixture +def run_as(monkeypatch): + """Run as rank ``rank`` of a ``world_size``-rank launch, deciding for itself. + + Every rank reads the checkpoint as this one does. How the ranks agree on one plan is tested at + the end of this file. + """ + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None) + monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike) + + def enter(world_size, rank=0): + monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1) + monkeypatch.setattr(dist, "size", lambda group=None: world_size) + monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size) + + return enter + + +@pytest.fixture +def serve(monkeypatch): + """Answer every config read with ``config`` rather than a checkpoint on disk.""" + + def serve(config): + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **k: config) + return config + + return serve + + +@pytest.fixture +def gpu_memory(monkeypatch): + """Give this rank a GPU of ``size`` bytes.""" + + def give(size): + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: size) + + return give + + +# ------------------------------------------------------------------------------ model class + + +def _vlm_config(**kwargs): + return SimpleNamespace(text_config=SimpleNamespace(), architectures=[], **kwargs) + + +def _tiny_qwen2_5_vl(): + return Qwen2_5_VLConfig( + text_config={**_TINY, "max_position_embeddings": 16}, + vision_config={ + "depth": 1, + "hidden_size": 16, + "intermediate_size": 16, + "num_heads": 2, + "out_hidden_size": 16, + "patch_size": 2, + "spatial_merge_size": 1, + "temporal_patch_size": 1, + "fullatt_block_indexes": [0], + }, + ) + + +_REMOTE_VLM = {"AutoModelForCausalLM": "m.KimiK25ForConditionalGeneration"} +_REMOTE_IMAGE_TEXT_VLM = { + **_REMOTE_VLM, + "AutoModelForImageTextToText": "m.KimiK25ForConditionalGeneration", +} + + +@pytest.mark.parametrize( + ("config", "trust_remote_code", "expected"), + [ + (SimpleNamespace(architectures=["LlamaForCausalLM"]), False, AutoModelForCausalLM), + (_vlm_config(auto_map=_REMOTE_VLM), True, AutoModelForCausalLM), + (_vlm_config(auto_map=_REMOTE_VLM), False, AutoModel), + (_vlm_config(auto_map=_REMOTE_IMAGE_TEXT_VLM), True, AutoModelForImageTextToText), + (_tiny_qwen2_5_vl(), False, AutoModelForImageTextToText), + (Phi4MultimodalConfig(), False, AutoModelForCausalLM), + (_vlm_config(), False, AutoModel), + (T5Config(architectures=["T5ForConditionalGeneration"]), False, AutoModelForSeq2SeqLM), + (WhisperConfig(), False, AutoModelForSpeechSeq2Seq), + ], + ids=[ + "causal-lm", + "remote-vlm-with-head", + "remote-vlm-untrusted", + "remote-image-text-vlm", + "native-vlm-with-head", + "native-multimodal-causal-lm", + "unknown-vlm", + "encoder-decoder", + "speech-encoder-decoder", + ], +) +def test_the_class_that_builds_the_whole_model(config, trust_remote_code, expected): + """A VLM's AutoModel leaves the LM head out: it builds through the class that carries the head, + keeping the headless backbone only when no class recognizes the config. A remote class counts + only when the checkpoint's code is trusted. An encoder-decoder model's architecture reads like a + VLM's, but its seq2seq class carries its head, also for a config that names no architecture.""" + assert _auto_class(config, trust_remote_code) is expected + + +# -------------------------------------------------------------------------------- the decision + + +@pytest.mark.parametrize( + "knob", + [ + None, + ("MODELOPT_EP_DEGREE", "4"), + ("MODELOPT_FSDP_SIZE", "world"), + ("MODELOPT_FSDP_SIZE", "2"), + ], + ids=["no_knob", "ep_degree", "fsdp_world", "fsdp_part"], +) +@pytest.mark.parametrize("world_size", [None, 1], ids=["no_process_group", "one_rank"]) +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.""" + if world_size is not None: + run_as(world_size) + if knob is not None: + monkeypatch.setenv(*knob) + served = serve(_qwen3_moe()) + monkeypatch.setattr( + model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model") + ) + + assert resolve_distribution_plan(str(tmp_path)) == (DistributionPlan(), served) + assert resolve_distribution_plan(str(tmp_path), dense="fsdp")[0] == DistributionPlan( + dense="fsdp" + ) + + +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize("rank", [0, 1]) +def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, capsys, rank): + """Auto picks the largest EP degree dividing both the world and the expert count, hands back + the config it decided from, and only rank 0 reports the decision.""" + run_as(4, rank) + served = serve(_qwen3_moe(num_experts=8)) + + plan, hf_config = resolve_distribution_plan(str(tmp_path)) + + assert plan == DistributionPlan(4, ep_size=4, dense="replicate") + assert hf_config is served + 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) + config = _qwen3_moe(num_experts=8) + config.base_model_ep_plan = None + serve(config) + + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) + + +@pytest.mark.parametrize("degree", [None, "1"], ids=["auto", "set-to-1"]) +def test_an_ep_plan_naming_no_module_of_the_model_is_refused( + run_as, serve, monkeypatch, tmp_path, degree +): + """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.""" + 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"): + resolve_distribution_plan(str(tmp_path)) + else: + monkeypatch.setenv("MODELOPT_EP_DEGREE", degree) + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) + + +@pytest.mark.parametrize("degree", [None, "1"], ids=["auto", "set-to-1"]) +def test_routed_experts_the_config_does_not_count_are_refused( + run_as, serve, monkeypatch, tmp_path, degree +): + """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.""" + plan = {"layers.*.mlp.experts": "moe_tp_experts"} + serve(SimpleNamespace(model_type="moe", architectures=[], base_model_ep_plan=plan)) + built = model_with_decoder_at("model.") + built.tie_weights = lambda: None # a PreTrainedModel's; nothing here is tied + monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: built) + run_as(4) + + if degree is None: + with pytest.raises(ValueError, match=r"no expert count .* MODELOPT_EP_DEGREE=1"): + resolve_distribution_plan(str(tmp_path)) + else: + monkeypatch.setenv("MODELOPT_EP_DEGREE", degree) + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) + + +@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 +): + """Every rank belongs to one EP group, so the degree must divide the world; the refusal names + the knob or argument the degree came from, as the one for the experts does.""" + run_as(4) + serve(LlamaConfig(**_TINY)) + arguments = {} + if source == "ep_size": + arguments["ep_size"] = 3 + else: + monkeypatch.setenv(source, "3") + + with pytest.raises(ValueError, match=rf"{source}=3 must divide the world size \(4\)"): + resolve_distribution_plan(str(tmp_path), **arguments) + + +@pytest.mark.parametrize( + ("world_size", "source"), + [(4, "MODELOPT_EP_DEGREE"), (1, "ep_size")], + ids=["knob", "argument-in-a-single-process"], +) +def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from( + run_as, serve, monkeypatch, tmp_path, world_size, source +): + """A single process takes its own argument as given, but parses it as the decision does.""" + run_as(world_size) + serve(LlamaConfig(**_TINY)) + arguments = {} + if source == "ep_size": + arguments["ep_size"] = "world" + else: + monkeypatch.setenv(source, "world") + + with pytest.raises(ValueError, match=f"{source}='world'"): + resolve_distribution_plan(str(tmp_path), **arguments) + + +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize("value", ["0", "", "AUTO"], ids=["zero", "exported-empty", "capitals"]) +def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value): + """And so is one exported empty, and auto in any case: the plan decides.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", value) + run_as(4) + serve(_qwen3_moe(num_experts=8)) + + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4) + + +@pytest.mark.parametrize("tied", [False, True], ids=["untied", "tied-embeddings"]) +@pytest.mark.parametrize( + ("share", "dense"), [(0.59, "replicate"), (0.61, "fsdp")], ids=["fits", "too-big"] +) +def test_a_dense_model_is_fsdp_sharded_only_when_too_big_for_one_gpu( + run_as, serve, gpu_memory, tmp_path, share, dense, tied +): + """Taking more than 60% of a GPU, the whole model would leave too little to calibration. It is + sized as it loads, a tied embedding and LM head counting once, though its meta build unties + them.""" + run_as(4) + config = serve(LlamaConfig(**_TINY, tie_word_embeddings=tied)) + loaded = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16) + gpu_memory(math.ceil(_bytes(loaded) / share)) + + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, dense=dense) + + +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize( + ("experts_fitting", "dense"), + [(1 / 2, "replicate"), (1 / 8, "fsdp")], + ids=["fits", "the-experts-share-does-not"], +) +def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit( + run_as, serve, gpu_memory, tmp_path, experts_fitting, dense +): + """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.""" + run_as(4) + config = serve(_qwen3_moe(num_experts=8)) + model = _meta(config) + experts = sum(_bytes(layer.mlp.experts) for layer in model.model.layers) + # 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) + assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan( + 4, dense="fsdp" + ) + + +@needs_qwen3_moe_ep_plan +def test_the_env_knobs_override_the_automatic_plan(run_as, serve, monkeypatch, tmp_path): + monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + run_as(4) + serve(_qwen3_moe(num_experts=8)) + + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 2, "fsdp") + + +@needs_qwen3_moe_ep_plan +def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path): + monkeypatch.setenv("MODELOPT_EP_DEGREE", "4") + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + run_as(4) + serve(_qwen3_moe(num_experts=8)) + + plan, _ = resolve_distribution_plan(str(tmp_path), ep_size=2, dense="replicate") + + assert plan == DistributionPlan(4, 2, "replicate") + + +@pytest.mark.parametrize( + ("value", "dense"), + [ + ("1", "replicate"), + ("4", "fsdp"), + ("WORLD", "fsdp"), + ("", "fsdp"), + ("2", None), + ], + ids=["1", "4", "WORLD", "exported-empty", "part-of-the-world"], +) +def test_the_dense_knob_replicates_or_shards_over_the_whole_world( + run_as, serve, gpu_memory, monkeypatch, tmp_path, value, dense +): + """Exported empty, the knob is auto, which shards a model too big for one GPU. Sharding over + part of the world is refused: no layout does that.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", value) + run_as(4) + serve(LlamaConfig(**_TINY)) + gpu_memory(1) # too small for the whole model, so only auto tells "" from "1" + + if dense is None: + with pytest.raises(ValueError, match=r"expected auto, 1 .* or world"): + resolve_distribution_plan(str(tmp_path)) + else: + assert resolve_distribution_plan(str(tmp_path))[0].dense == dense + + +@pytest.mark.parametrize( + ("config", "ep_size", "experts"), + [(LlamaConfig(**_TINY), 2, 0), (_qwen3_moe(num_experts=6), 4, 6)], + ids=["dense-model", "experts-not-divisible"], +) +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") + run_as(4) + serve(config) + + with pytest.raises(ValueError, match=rf"ep_size={ep_size} must divide .* \({experts}\)"): + resolve_distribution_plan(str(tmp_path), ep_size=ep_size) + + +@pytest.mark.parametrize( + "config", + [ + _tiny_qwen2_5_vl(), + T5Config(d_model=16, d_ff=32, d_kv=8, num_layers=1, num_heads=2, vocab_size=32), + ], + ids=["vlm", "encoder-decoder-naming-no-architecture"], +) +def test_a_model_is_decided_from_its_whole_model(run_as, serve, gpu_memory, tmp_path, config): + """AutoModelForCausalLM builds neither every VLM nor an encoder-decoder model, whose config need + not name its architecture: the decision builds what the load builds.""" + run_as(2) + serve(config) + gpu_memory(1) + + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp") + + +def test_the_expert_count_is_read_from_the_config_holding_the_plan( + run_as, serve, monkeypatch, tmp_path +): + """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 + ) + 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) + run_as(4) + + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 4) + + +def test_a_model_that_cannot_be_built_on_meta_names_itself(run_as, serve, monkeypatch, tmp_path): + def unbuildable(*args, **kwargs): + raise NotImplementedError("custom init touches the weights") + + run_as(2) + serve(LlamaConfig(**_TINY)) + monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", unbuildable) + + with pytest.raises(RuntimeError, match=r"distribute llama .* custom init touches the weights"): + resolve_distribution_plan(str(tmp_path)) + + +def test_a_meta_build_leaves_the_auto_class_a_part_registers_in_as_it_found_it(monkeypatch): + """A remote VLM's ``__init__`` can build its text model through ``AutoModelForCausalLM`` while + the whole model builds through another auto class. That registers the text model's class in + ``AutoModelForCausalLM`` too, which would leave rank 0's later load of it without the + ``generate`` its peers' have.""" + + class ProbeConfig(transformers.PretrainedConfig): + model_type = "meta_build_part_probe" + + class ProbeModel(transformers.PreTrainedModel): + config_class = ProbeConfig + + for auto in (AutoModel, AutoModelForCausalLM): + mapping = auto._model_mapping + monkeypatch.setattr(mapping, "_extra_content", dict(mapping._extra_content)) + from_config = AutoModel.from_config + + def remote_vlm_from_config(config, **kwargs): + # What a remote VLM's __init__ does for its text model before the whole model is built. + AutoModelForCausalLM.register(ProbeConfig, ProbeModel, exist_ok=True) + return from_config(config, **kwargs) + + monkeypatch.setattr(AutoModel, "from_config", remote_vlm_from_config) + + model_load_utils.build_meta_causal_lm( + None, False, None, LlamaConfig(**_TINY), auto_class=AutoModel + ) + + assert ProbeConfig not in AutoModelForCausalLM._model_mapping + + +def test_a_meta_build_leaves_the_auto_class_as_it_found_it(monkeypatch): + """Building a remote-code model through an auto class registers its class there, as it is, and a + later load through that auto class builds the registered class, without the ``generate`` a first + load adds. Rank 0 alone builds one to decide a distribution: its load alone would lack it.""" + + class ProbeConfig(transformers.PretrainedConfig): + model_type = "meta_build_probe" + + class ProbeModel(transformers.PreTrainedModel): + config_class = ProbeConfig + + mapping = AutoModelForCausalLM._model_mapping + monkeypatch.setattr(mapping, "_extra_content", dict(mapping._extra_content)) + from_config = AutoModelForCausalLM.from_config + + def remote_from_config(config, **kwargs): + # What transformers' from_config does for a remote-code config before it builds. + AutoModelForCausalLM.register(ProbeConfig, ProbeModel, exist_ok=True) + return from_config(config, **kwargs) + + monkeypatch.setattr(AutoModelForCausalLM, "from_config", remote_from_config) + + model_load_utils.build_meta_causal_lm(None, False, None, LlamaConfig(**_TINY)) + + assert ProbeConfig not in mapping + + +def test_the_hub_options_locate_the_checkpoint_the_load_reads(tmp_path): + """``revision``, ``cache_dir``, ``local_files_only`` and ``subfolder`` find the checkpoint as + ``load_pretrained`` takes them, so the plan and its config are those of what the load reads.""" + repo = tmp_path / "models--org--model" + (repo / "refs").mkdir(parents=True) + for revision, commit, hidden_size in (("main", "a" * 40, 16), ("v2", "b" * 40, 32)): + (repo / "refs" / revision).write_text(commit) + config = LlamaConfig(**{**_TINY, "hidden_size": hidden_size}) + config.save_pretrained(repo / "snapshots" / commit / "nested") + + _, config = resolve_distribution_plan( + "org/model", + revision="v2", + cache_dir=str(tmp_path), + local_files_only=True, + subfolder="nested", + ) + + assert config.hidden_size == 32 + + +# ------------------------------------------------------------------------------ rank 0 decides + + +@contextmanager +def _knobs(**values): + """This rank's distribution knobs, exactly ``values``, within the block. + + A pooled worker keeps the environment it was spawned with, so the knobs are cleared first. The + environment is restored on the way out. + """ + with mock.patch.dict(os.environ): + for name in ("MODELOPT_EP_DEGREE", "MODELOPT_FSDP_SIZE"): + os.environ.pop(name, None) + os.environ.update(values) + yield + + +def _every_rank_takes_rank0s_plan_job(rank, size, ckpt, knobs, expected): + with bounded_collectives(), _knobs(**knobs[rank]): + plan, _ = resolve_distribution_plan(ckpt) + assert plan == expected, (rank, plan) + + +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) + cpu_dist_workers.run( + _every_rank_takes_rank0s_plan_job, + ckpt=str(tmp_path), + knobs=[{}, {"MODELOPT_FSDP_SIZE": "world"}], + expected=DistributionPlan(2), + ) + + +def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt): + isdir = os.path.isdir + with ( + bounded_collectives(), + mock.patch("os.path.isdir", lambda path: (rank == 0 or path != ckpt) and isdir(path)), + ): + try: + resolve_distribution_plan(ckpt) + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + # Rank 1's own error, carried to rank 0, rather than rank 0 timing out in the plan's broadcast. + assert "which this rank cannot see" in raised, raised + + +def test_a_rank_that_cannot_see_the_checkpoint_fails_every_rank(cpu_dist_workers, tmp_path): + """Before the plan's broadcast, which rank 0 would otherwise wait in for a rank that raised.""" + LlamaConfig(**_TINY).save_pretrained(tmp_path) + cpu_dist_workers.run(_one_rank_cannot_see_the_checkpoint_job, ckpt=str(tmp_path)) diff --git a/tests/unit/torch/utils/test_moe_a2a.py b/tests/unit/torch/utils/test_moe_a2a.py index a53c334ec51..7a872183040 100644 --- a/tests/unit/torch/utils/test_moe_a2a.py +++ b/tests/unit/torch/utils/test_moe_a2a.py @@ -23,19 +23,30 @@ The routing-weight gradient check guards that the routing weights keep their gradient: they never travel, and weight their tokens' expert results on the rank that holds them. + +The tests at the end cover the module's helpers that count a model's routed experts, choose the EP +degree they are sharded over, and find them by the model's ``base_model_ep_plan``. """ import copy import warnings from datetime import timedelta +from types import SimpleNamespace import pytest import torch import torch.distributed as dist +from _test_utils.torch.moe_models import MoeLayer, model_with_decoder_at from packaging.version import Version +from torch import nn import modelopt.torch.utils.moe as moe -from modelopt.torch.utils.moe import a2a_moe +from modelopt.torch.utils.moe import ( + _choose_ep_degree, + _num_experts_from_config, + _routed_expert_modules, + a2a_moe, +) # Small, fully-deterministic problem. E must be divisible by every world size under test. E, TOP_K, H, T = 4, 2, 8, 6 @@ -401,3 +412,94 @@ def test_fp32_routing_weights_give_grouped_mms_result_at_ep1(cpu_dist_workers): if Version(transformers.__version__) < Version("5.6"): pytest.skip("expert parallelism needs transformers >= 5.6") cpu_dist_workers.run(_fp32_routing_weights_job) + + +# ------------------------------------------------------------------------------- expert counting + + +@pytest.mark.parametrize( + ("config", "experts"), + [ + (SimpleNamespace(hidden_size=128), 0), + (SimpleNamespace(num_experts=16), 16), + (SimpleNamespace(num_local_experts=16), 16), + (SimpleNamespace(n_routed_experts=16), 16), + (SimpleNamespace(num_routed_experts=16), 16), + (SimpleNamespace(text_config=SimpleNamespace(n_routed_experts=32)), 32), + (SimpleNamespace(llm_config=SimpleNamespace(n_routed_experts=32)), 32), + ], + ids=[ + "dense", + "num_experts", + "num_local_experts", + "n_routed_experts", + "num_routed_experts", + "vlm-text_config", + "vlm-llm_config", + ], +) +def test_the_expert_count_is_read_under_every_spelling(config, experts): + """Architectures disagree on the field name, and a VLM's MoE is its nested LM, whose root + config has no expert count of its own: missing either would silently disable EP.""" + assert _num_experts_from_config(config) == experts + + +# ------------------------------------------------------------------------------- degree selection + + +@pytest.mark.parametrize( + ("world", "experts", "degree"), + [(8, 64, 8), (8, 6, 2), (8, 5, 1), (8, 0, 1), (1, 64, 1)], + ids=["largest-valid", "divides-the-experts", "nothing-divides", "dense", "single-rank"], +) +def test_the_largest_ep_degree_dividing_the_world_and_the_experts_is_chosen(world, experts, degree): + """EP must divide the expert count, so every rank holds whole experts; with nothing dividing, a + dense model, or a single rank, there is no expert parallelism.""" + assert _choose_ep_degree(world, experts) == degree + + +# ---------------------------------------------------------------------------- routed experts + + +_STOCK_PLAN = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_tp_experts", +} + + +@pytest.mark.parametrize( + "plan", + [ + _STOCK_PLAN, + {"layers.*.mlp.experts": "moe_tp_experts"}, + {"layers.*.mlp.experts": "nemotron_experts_ep"}, + {k: v for k, v in _STOCK_PLAN.items() if v == "grouped_gemm"}, + { + "layers.*.block_sparse_moe.experts": "moe_tp_experts", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + }, + ], + ids=["stock", "experts-style-only", "nemotron-style-only", "grouped-gemm-only", "hybrid"], +) +@pytest.mark.parametrize( + "prefix", + ["", "model.", "model.language_model.", "language_model.model."], + ids=["base-model", "causal-lm", "vlm", "legacy-vlm"], +) +def test_routed_experts_are_the_plans_containers_under_any_prefix(plan, prefix): + """The base-relative plan keys match the decoder wherever the model nests it, at every layer + index, the container named by its style or as its grouped-GEMM parameters' parent, also in a + plan that names some containers one way and some the other. The router, the shared experts, and + a module whose name merely ends in a key's text are not routed experts.""" + pytest.importorskip("transformers.integrations.tensor_parallel") + model = model_with_decoder_at(prefix) + model.add_module("sublayers", nn.ModuleList([MoeLayer()])) # "sublayers.0.mlp.experts" + + found = _routed_expert_modules(model, plan) + + modules = dict(model.named_modules()) + assert found == { + name: modules[name] for name in [f"{prefix}layers.{i}.mlp.experts" for i in (0, 1)] + }