From f2b01f7357819772f4a1dfbc937abdeee95d7e21 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:15:00 +0000 Subject: [PATCH 01/23] Load an MoE expert-parallel: shard the experts at load time load_and_prepare_ep(ckpt_path, plan) loads a model by an expert-parallel DistributionPlan, with every rank both data-parallel and expert-parallel: transformers' own expert-parallel load builds only each rank's experts, over the (edp, ep) mesh's ep submesh, with the base_model_ep_plan swapped to the torch_a2a styles for the load and put back on the config it keeps. The dense part is replicated. Every module gets the ParallelState its amax syncs over, and the meshes' communicators are created in step. The plan, the process group, the transformers floor (5.7), the checkpoint's format and the plan's routed experts are checked before anything loads; a native MXFP4 checkpoint is dequantized as it loads; on transformers 5.18+, whose from_pretrained builds a world-sized mesh of its own, the experts are still sharded over the EP group. The model keeps the mesh for calibration (moe.expert_parallel_mesh), and the checkpoint tensors it has no parameter for are recorded for the export. This is the loader #2741 introduced, without the post-load refusals, which stay in #2741 on top of it. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 40 +- .../torch/utils/plugins/model_load_utils.py | 383 ++++++++++++ .../torch/utils/test_expert_parallel_load.py | 113 ++++ tests/unit/torch/utils/test_ep_loader.py | 582 ++++++++++++++++++ 4 files changed, 1109 insertions(+), 9 deletions(-) create mode 100644 tests/gpu/torch/utils/test_expert_parallel_load.py create mode 100644 tests/unit/torch/utils/test_ep_loader.py diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 2aff08e6891..29282c5f4cf 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -15,25 +15,25 @@ """HuggingFace-coupled MoE expert-parallel (EP) helpers. -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``). +Finding a model's routed experts, choosing the EP degree they are sharded over, the DP=EP +all-to-all expert transport (``torch_a2a``). The expert-parallel model loader is +:func:`~modelopt.torch.utils.plugins.model_load_utils.load_and_prepare_ep`. The all-to-all transport is plain ``torch.distributed.all_to_all_single`` (NCCL) wrapped around -transformers' quantizable ``_grouped_linear``. ``register_torch_a2a()`` registers it into -transformers' ``ExpertsInterface`` lazily -- never at import time, so importing this module has no -transformers-``ExpertsInterface`` side effect. +transformers' quantizable ``_grouped_linear``. ``register_torch_a2a()`` registers +it into transformers' ``ExpertsInterface`` (lazily, from ``load_and_prepare_ep``'s a2a path -- never at +import time, so importing this module has no transformers-``ExpertsInterface`` side effect). """ import warnings -from typing import TYPE_CHECKING import torch import torch.distributed as dist +import torch.nn as nn from packaging.version import parse from torch.distributed.nn.functional import all_to_all_single as autograd_all_to_all_single -if TYPE_CHECKING: - from torch import nn +from .distributed import distribution_plan def require_transformers(min_version: str, feature: str) -> None: @@ -406,7 +406,9 @@ def register_torch_a2a(): These two styles are functionally transformers' ``megamoe_router`` / ``megamoe_experts`` (added in 5.10.4 for DeepGEMM Mega-MoE), reimplemented here so the transport works below 5.10.4 too. The base classes (``RouterParallel``, ``TensorParallelLayer``) are stock wherever the EP infra - exists (transformers >= 5.6); ``grouped_gemm`` (expert weight sharding) is stock too. + exists (transformers >= 5.6); ``grouped_gemm`` (expert weight sharding) is stock too. A model + builds with ``torch_a2a`` only from 5.7, the floor ``load_and_prepare_ep`` enforces: 5.6 accepts + only the experts implementations it ships. """ global _TORCH_A2A_REGISTERED if _TORCH_A2A_REGISTERED: @@ -473,3 +475,23 @@ def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None): ALL_PARALLEL_STYLES.register("torch_a2a_router", TorchA2ARouter()) ALL_PARALLEL_STYLES.register("torch_a2a_experts", TorchA2AExperts()) _TORCH_A2A_REGISTERED = True + + +def expert_parallel_mesh(model: nn.Module): + """The mesh ``model``'s routed experts are sharded over, or ``None`` if its plan shards none. + + :func:`~modelopt.torch.utils.plugins.model_load_utils.load_and_prepare_ep` loads the model over + its expert-parallel submesh, and keeps that submesh on the model as ``_device_mesh``. + Calibration takes the mesh to re-install the all-to-all forwards + (:func:`~modelopt.torch.utils.distributed.distributed_forward_loop`). + """ + if not distribution_plan(model).is_expert_parallel: + return None + mesh = getattr(model, "_device_mesh", None) + if mesh is None: + raise RuntimeError( + f"{type(model).__name__} is planned expert-parallel but carries no device mesh. Pass " + "the model load_and_prepare_ep returned, not a submodule of it: the loader keeps the " + "mesh it shards the experts over on that model alone." + ) + return mesh diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 094a2d79504..3581fd284d9 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -15,7 +15,9 @@ """HuggingFace model loading: :func:`load_pretrained`, by a plan to distribute the model.""" +import contextlib import copy +import dataclasses import logging import os import re @@ -73,6 +75,17 @@ read_safetensors_subset, ) +__all__ = [ + "build_meta_causal_lm", + "fsdp2_wrap", + "load_and_prepare_ep", + "load_pretrained", + "parallel_load_and_prepare_fsdp2", + "record_unplaced_source_keys", + "resolve_distribution_plan", + "unplaced_source_keys", +] + logger = logging.getLogger(__name__) @@ -1583,3 +1596,373 @@ def checkpoint_read() -> bool: if plan != DistributionPlan(): print_rank_0(f"Distribution: {plan!r}") return plan, hf_config + + +def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_state) -> None: + """Seed every module's ``_parallel_state`` for calibration's amax sync. + + The routed-expert containers named in ``expert_names`` and everything under them get + ``expert_state``; every other module gets ``dense_state``. Seeding the plain modules works because + quantization keeps a ``_parallel_state`` that is already present when it converts a module -- the same + way Megatron's expert MLPs seed theirs. + """ + prefixes = tuple(name + "." for name in expert_names) + for name, module in model.named_modules(): + module._parallel_state = ( + expert_state if prefixes and (name + ".").startswith(prefixes) else dense_state + ) + + +# The checkpoint formats load_and_prepare_ep dequantizes as it loads, beside unquantized ones (None). +# mxfp4 is for experts stored in the standard layout: GPT-OSS's are transposed, which the load refuses. +_EP_SOURCE_FORMATS = frozenset({None, "mxfp4"}) + + +def _ckpt_quant_method(cfg): + """Return the checkpoint's native quantization method, or ``None``. + + Return the checkpoint's native quantization method (e.g. ``"fp8"`` for block-fp8 DeepSeek-V3), + or ``None`` if the checkpoint isn't natively quantized. Checks the top config, then a VLM's + ``text_config`` or ``llm_config``. + """ + for c in (cfg, getattr(cfg, "text_config", None), getattr(cfg, "llm_config", None)): + qc = getattr(c, "quantization_config", None) + m = qc.get("quant_method") if isinstance(qc, dict) else getattr(qc, "quant_method", None) + if m: + # transformers' config objects carry a QuantizationMethod enum, whose str() is not its value. + return str(getattr(m, "value", m)) + return None + + +def _transformers_mesh_builder(): + """The module whose ``initialize_distributed_mesh`` ``from_pretrained`` calls, or ``None``. + + Newer transformers drops ``from_pretrained``'s ``device_mesh`` once ``distributed_config`` is + given: it builds its own world-sized ``tp`` mesh from the config's sizes. That shards the experts + over every rank instead of this rank's EP group whenever the EP degree is below the world size. + Older releases honor ``device_mesh`` and have no such builder. + """ + try: + import transformers.distributed.mixin as mixin + except ImportError: + return None + return mixin if hasattr(mixin, "initialize_distributed_mesh") else None + + +@contextlib.contextmanager +def _experts_sharded_over(ep_mesh): + """Make ``from_pretrained`` shard the experts over ``ep_mesh`` on releases that build their own mesh.""" + mixin = _transformers_mesh_builder() + if mixin is None: + yield + return + + def _ep_mesh(*args, **kwargs): + index = torch.cuda.current_device() if ep_mesh.device_type == "cuda" else None + return torch.device(ep_mesh.device_type, index), ep_mesh + + original = mixin.initialize_distributed_mesh + mixin.initialize_distributed_mesh = _ep_mesh + try: + yield + finally: + mixin.initialize_distributed_mesh = original + + +def load_and_prepare_ep( + ckpt_path: str, + plan: DistributionPlan, + trust_remote_code: bool = False, + attn_implementation: str | None = None, + freeze: bool = True, + hf_config=None, + require_lm_head: bool = True, +) -> nn.Module: + """Load + shard a MoE causal LM for expert-parallel PTQ via the DP=EP all-to-all transport. + + ``plan`` is the expert-parallel :class:`~modelopt.torch.utils.distributed.DistributionPlan` to load + by (``ep_size > 1``), for the whole process group. It records no plan on the model: the caller + does (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), as + :func:`load_pretrained` does for every load. The model builds with its LM head where an auto + class builds one; ``require_lm_head=False`` builds a VLM without it, through ``AutoModel`` (a + causal LM always keeps its head). + + Builds an ``(edp, ep)`` device mesh and loads experts EP-sharded over the ``ep`` submesh + (``enable_expert_parallel=True`` -- each rank owns ``num_experts / ep`` experts as grouped-GEMM + shards), while ATTENTION and dense params stay REPLICATED (so EP has no KV-head-divisibility cap). + The model's ``base_model_ep_plan`` is swapped to the ``torch_a2a`` styles: the router emits GLOBAL + unmasked expert ids and the experts forward routes each token to its expert-owning rank by a NCCL + all-to-all (dispatch) -> quantizable ``_grouped_linear`` -> all-to-all (combine), rather than every + rank computing all tokens followed by an EP all-reduce. Every rank is data-parallel (sees DIFFERENT + tokens), so calibration MUST shard the data per-rank. + + The ``edp`` dimension (present when ``plan.ep_size < world_size``) is additional data parallelism + across EP groups: non-expert params stay REPLICATED per rank and each EP group calibrates on + different data; experts stay EP-sharded (resident). No FSDP and no concurrent ``edp``/``ep`` + collectives -> no deadlock by construction. Costs more per-rank memory (non-experts + experts + replicated across edp) but is the reliable path at scale. + + EP CANNOT be applied to a meta-init model (transformers' ``GroupedGemmParallel.shard_tensor`` + slices materialized tensors), so this loads real weights via ``from_pretrained`` over the ``ep`` + submesh. Requires transformers >= 5.7: the MoE EP framework (ExpertsInterface / grouped-GEMM EP + styles / enable_expert_parallel) is in 5.6, but 5.6 builds a model only with the experts + implementations it ships, never ``torch_a2a``. (A model whose expert-parallel support landed + later states the release it needs in its spec.) + + Assumes the caller has already initialized (and device-bound) the process group -- distributed + bring-up is the launcher's job, not this loader's (hf_ptq.main() does it via ``distributed_run``). + The expert-parallel mesh stays recoverable from the model + (:func:`~modelopt.torch.utils.moe.expert_parallel_mesh`), once the caller has recorded the plan. + """ + # Imported here, so looked up when the load runs: its tests stand in a process group and a + # transformers release by patching these two modules. + from modelopt.torch.utils.distributed import ParallelState, is_initialized, size + from modelopt.torch.utils.moe import ( + _EXPERTS_STYLES, + _ep_config, + _num_experts_from_config, + _routed_expert_keys, + _routed_expert_modules, + register_torch_a2a, + require_transformers, + ) + + if not plan.is_expert_parallel: + raise ValueError(f"{plan!r} shards no routed experts, which is what this load is for.") + if plan.is_fsdp_dense: + raise NotImplementedError(f"{plan!r}: this load replicates the dense part on every rank.") + # This loader assumes the caller has ALREADY initialized (and device-bound) the process group -- + # distributed bring-up belongs to the launcher, not this library function. The example launcher does + # it in hf_ptq.main() via distributed_run() -> modelopt.torch.utils.distributed.setup() before + # the load. A live, device-bound PG is required here so the mesh subgroups built below create their + # NCCL comms eagerly (the synchronized warmup all-reduce per subgroup further down) instead of lazily + # on first collective (which deadlocks at DP>1 because ranks reach the lazy init in different phases). + if not is_initialized(): + raise RuntimeError( + "load_and_prepare_ep requires a live, device-bound process group; set it up first, e.g. " + "with modelopt.torch.utils.distributed.distributed_run() under torchrun." + ) + if plan.world_size != size(): + raise ValueError(f"{plan!r} is for {plan.world_size} ranks, but this run has {size()}.") + world_size, ep_degree, edp = plan.world_size, plan.ep_size, plan.edp + # Hard floor: below 5.6 there is no ExpertsInterface / grouped-GEMM EP styles, and + # enable_expert_parallel silently shards nothing (no error); 5.6 checks a model's experts + # implementation against the ones it ships, so from_pretrained would refuse torch_a2a. + require_transformers( + "5.7", "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} + if attn_implementation is not None: + config_kwargs["attn_implementation"] = attn_implementation + hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) + quant_method = _ckpt_quant_method(hf_config) + if quant_method not in _EP_SOURCE_FORMATS: + # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights. + dequantized = ", ".join(sorted(f for f in _EP_SOURCE_FORMATS if f)) + raise NotImplementedError( + f"The expert-parallel load cannot dequantize a {quant_method!r} checkpoint: it loads " + f"unquantized checkpoints and dequantizes {dequantized}. Start from an unquantized copy " + "of the model." + ) + # The swapped plan is this load's own: the caller's config keeps the plan its checkpoint declares. + hf_config = copy.deepcopy(hf_config) + ep_cfg = _ep_config(hf_config) # for VLMs the plan lives on .text_config + 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." + ) + # The all-to-all transport finds an expert's rank by dividing its id by this rank's share of + # the experts, which holds only when every rank holds as many. A count the config names none of + # the ways this reads is checked once the experts load, from their weights. + num_experts = _num_experts_from_config(hf_config) + if num_experts % ep_degree: + raise ValueError( + f"{plan!r}: ep_size={ep_degree} must divide the {num_experts} routed experts, which the " + "all-to-all transport shares out equally over the expert-parallel ranks." + ) + # transformers 5.6 on, which the floor above makes sure of. + from torch.distributed.device_mesh import init_device_mesh + from transformers.distributed import DistributedConfig + + dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 + # The weights load, and the checkpoint recorded for the export to read the unplaced tensors back + # from, take the local snapshot: a Hub ID recorded there would leave the export nothing to read. + # Resolved once the plan and the config pass, so a refused load downloads nothing. + ckpt_path = _resolve_checkpoint_dir(ckpt_path, torch.distributed.get_rank()) + + # Two INDEPENDENT meshes (NeMo-AutoModel style), named by the model part they govern -- two + # factorizations of the same world: + # sparse_moe_mesh = (edp, ep): experts are SHARDED over the ``ep`` axis and REPLICATED over ``edp``; + # the ``ep`` submesh also drives from_pretrained's expert sharding + the a2a transport. + # dense_mesh = (dp,): dense/non-expert params are data-parallel-REPLICATED over the WHOLE world. + # This subsumes the old a2a-vs-partial split: a2a is simply edp == 1, and dense amax + # always reduces over the whole world (fixing the DP>1 case, where the old dp submesh missed the ep dim). + sparse_moe_mesh = init_device_mesh("cuda", (edp, ep_degree), mesh_dim_names=("edp", "ep")) + dense_mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) + ep_submesh = sparse_moe_mesh["ep"] + + # NO TENSOR PARALLELISM IS REQUESTED HERE, despite the name. transformers implements expert + # parallelism THROUGH the TP application path: with enable_expert_parallel=True the model's + # ``tp_plan`` PROPERTY returns ``_ep_plan`` (distributed/mixin.py), so what actually gets applied + # is the expert plan. ``tp_plan="auto"`` is only the SWITCH that makes ``tp_size > 1`` in + # DistributedConfig.__post_init__ and therefore reaches ``apply_tensor_parallelism`` at all -- + # drop it and __post_init__ early-returns with tp_size=None, which then TypeErrors on + # ``tp_size > 1`` in prepare_distribute_model. Dense params stay REPLICATED (= our DP) because + # ``swapped_plan`` below keeps only expert entries; attention/embeds match no plan key. + # + # WHERE it goes is version-dependent: from transformers 5.17 from_pretrained refuses tp_plan + # beside distributed_config ("Set tensor-parallel options on DistributedConfig when using it"), + # while DistributedConfig has no tp_plan field before 5.10, nor in 5.11 and 5.12. Passing + # tp_plan ALONE is not an option either -- the 5.17+ back-compat path rebuilds the config as + # DistributedConfig(tp_plan=..., tp_size=...) and DROPS enable_expert_parallel, silently losing + # EP. So each argument goes on DistributedConfig where it has a field for it, and tp_plan + # otherwise to from_pretrained: the dataclass fields tell them apart on any release. + _dc_fields = {f.name for f in dataclasses.fields(DistributedConfig)} + _dc_has_tp_plan = "tp_plan" in _dc_fields + dc_kwargs: dict[str, Any] = {"enable_expert_parallel": True} + if _dc_has_tp_plan: + dc_kwargs["tp_plan"] = "auto" + if _transformers_mesh_builder() is not None and "tp_size" in _dc_fields: + # This release sizes its own mesh from tp_size; _experts_sharded_over swaps in the EP mesh. + dc_kwargs["tp_size"] = ep_degree + distributed_config = DistributedConfig(**dc_kwargs) + from_pretrained_kwargs: dict[str, Any] = { + "dtype": dtype, + "device_mesh": ep_submesh, + "distributed_config": distributed_config, + "trust_remote_code": trust_remote_code, + } + if not _dc_has_tp_plan: + from_pretrained_kwargs["tp_plan"] = "auto" + if attn_implementation is not None: + from_pretrained_kwargs["attn_implementation"] = attn_implementation + if quant_method == "mxfp4": + # Native mxfp4 checkpoints (e.g. gpt-oss: fused `gate_up_proj_blocks`/`_scales` + biases) must be + # DEQUANTIZED to bf16 on load so ModelOpt can insert nvfp4/fp8 quantizers on the fused bf16 experts + # (the a2a grouped_gemm expects bf16). transformers dequantizes blocks->bf16 during from_pretrained. + # NB: the single-process loader needed device_map="sequential" to dodge a CUDA illegal-access in the + # mxfp4 dequant under multi-GPU device_map="auto"; the EP path uses device_mesh (no device_map), so + # watch for an EP-shard-awareness issue in the dequant convert-op. + from transformers import Mxfp4Config + + from_pretrained_kwargs["quantization_config"] = Mxfp4Config(dequantize=True) + print_rank_0( + "[ep] native mxfp4 checkpoint (e.g. gpt-oss) -> Mxfp4Config(dequantize=True) -> bf16" + ) + # DP=EP all-to-all: swap the ep_plan to our torch_a2a styles (torch_a2a_router emits GLOBAL ids + # unmasked; torch_a2a_experts injects the EP process_group + skips the all-reduce) and select the + # `torch_a2a` experts forward (NCCL all-to-all dispatch/combine around the quantizable + # _grouped_linear). Tokens are routed to the expert-owning rank by all-to-all instead of every rank + # computing all tokens + all-reduce. Requires per-rank (data-parallel) calib data. nemotron_h's + # logits-only gate has no ep_router entry, so its nemotron_experts_ep also maps to torch_a2a_experts. + register_torch_a2a() # idempotent; register the experts forward + the two styles, right before use + swap = { + "ep_router": "torch_a2a_router", + "moe_tp_experts": "torch_a2a_experts", + "nemotron_experts_ep": "torch_a2a_experts", + } + # Keep expert-related entries (grouped_gemm expert params + the router/experts styles); DROP the TP + # styles (a2a builds no TP group for them -> would shard attention across the wrong mesh). + swapped_plan = { + k: swap.get(v, v) for k, v in ep_plan.items() if v in swap or v == "grouped_gemm" + } + # Tag every routed-experts CONTAINER torch_a2a_experts (injects the EP process_group + marks the + # routed expert), found as the distribution plan's decision finds them: some plans tag only the + # expert PARAMS grouped_gemm, and name the container through their parents. + swapped_plan.update(dict.fromkeys(_routed_expert_keys(ep_plan), "torch_a2a_experts")) + from_pretrained_kwargs["experts_implementation"] = "torch_a2a" + # The checkpoint's own plan, for the config the model keeps (below); None: a class default. + declared_plan = vars(ep_cfg).get("base_model_ep_plan") + ep_cfg.base_model_ep_plan = swapped_plan + from_pretrained_kwargs["config"] = hf_config + print_rank_0(f"[ep] ep_plan={swapped_plan} experts_implementation=torch_a2a") + # The class that builds the whole checkpoint -- with its lm_head where an auto class builds one, + # so the head is calibrated and exported and the generation preview works -- chosen as the + # distribution plan's decision chose it, unless the caller asked for a VLM without its head + # (require_lm_head=False). The per-model plugin (models/) has already installed any on-the-fly + # weight decompression as a transformers monkeypatch. + loader = _auto_class(hf_config, require_lm_head, trust_remote_code=trust_remote_code) + print_rank_0(f"[ep] loader={loader.__name__} (require_lm_head={require_lm_head})") + with _experts_sharded_over(ep_submesh): + model = loader.from_pretrained(ckpt_path, **from_pretrained_kwargs) + model.eval() + # The swapped styles drove the sharding, and the model keeps the plan it was sharded by + # (``model.tp_plan``, copied when it was built). The config it keeps goes back to the + # checkpoint's plan: an export writes the config out, and nothing without ModelOpt knows the + # torch_a2a styles. + kept_config = getattr(model, "config", None) + if kept_config is not None: + kept_ep_config = _ep_config(kept_config) + if declared_plan is None: + vars(kept_ep_config).pop("base_model_ep_plan", None) + else: + kept_ep_config.base_model_ep_plan = declared_plan + # Force the experts forward to dispatch to `torch_a2a`: from_pretrained's `experts_implementation` + # kwarg resolves non-fp8 impls back to "grouped_mm", but @use_experts_implementation reads + # `config._experts_implementation` at FORWARD time, so set it on every (sub)config now. (Without + # this the torch_a2a_experts style passes a 5th process_group arg to grouped_mm_experts_forward -> + # TypeError.) Unlike the plan it stays set, and an export still writes no torch_a2a: transformers + # keeps it as ``_experts_implementation_internal``, which it leaves out of a serialized config. + for _m in model.modules(): + _cfg = getattr(_m, "config", None) + if _cfg is not None and hasattr(_cfg, "_experts_implementation"): + _cfg._experts_implementation = "torch_a2a" + + # The checkpoint tensors the model has no parameter for (an MTP head), worked out once, at load, + # for the export to carry. + record_unplaced_source_keys(model, ckpt_path) + + # DP=EP all-to-all: no FSDP. Non-experts stay replicated per rank (from_pretrained already placed + # them on this rank's GPU); experts stay EP-sharded on the ep submesh, which the model carries as + # ``_device_mesh`` for calibration to re-install the all-to-all forwards over (expert_parallel_mesh). + # transformers sets it from the mesh this load hands it, through internals that differ by release; + # set here, the model carries it on every release. + model._device_mesh = ep_submesh + if hasattr(model, "tie_weights"): + model.tie_weights() + if freeze: + model.requires_grad_(False) + + # The routed-expert containers, from the SAME base_model_ep_plan that drove the sharding. A MoE model + # loaded WITHOUT EP (dense DP) never reaches this loader, so its experts stay dense. + expert_modules = _routed_expert_modules(model, ep_plan) + + # Each module's ParallelState for calibration's amax sync: dense modules reduce over the whole world; + # routed experts reduce over their expert-data-parallel replicas, and share one input amax across the + # layer's experts over the expert-parallel group. A size-1 axis is left unused (-1): no collective. + def _axis(mesh): + return mesh.get_group() if mesh.size() > 1 else -1 + + _seed_parallel_states( + model, + expert_modules, + dense_state=ParallelState(data_parallel_group=_axis(dense_mesh["dp"])), + expert_state=ParallelState( + data_parallel_group=_axis(sparse_moe_mesh["edp"]), + expert_model_parallel_group=_axis(sparse_moe_mesh["ep"]), + ), + ) + print_rank_0( + f"[ep] seeded the expert-parallel state of {len(expert_modules)} routed-expert " + "module(s) from base_model_ep_plan" + ) + # Eagerly create each amax collective's NCCL comm in a globally-consistent order (barrier around + # each), so the sub-communicators are built in lockstep rather than lazily mid-calibration -- which + # deadlocks when submesh groups are phase-skewed at DP>1. NO forward warmup (variable per-group time + # would desync ranks before comm creation). Skip size-1 axes (no collective). The EP group is already + # live from from_pretrained's expert sharding; re-warming it is a harmless redundant all-reduce. Same + # order on every rank -> lockstep. + warm_device = next(model.parameters()).device + for _sub in (dense_mesh["dp"], sparse_moe_mesh["ep"], sparse_moe_mesh["edp"]): + if _sub.size() <= 1: + continue + torch.distributed.barrier() + torch.distributed.all_reduce(torch.ones(1, device=warm_device), group=_sub.get_group()) + torch.cuda.synchronize() + torch.distributed.barrier() + return model diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py new file mode 100644 index 00000000000..849305f4f6a --- /dev/null +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -0,0 +1,113 @@ +# 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. + +"""Expert-parallel loading and the all-to-all expert transport on real GPUs. + +``load_and_prepare_ep`` shards an MoE's routed experts over the ranks, and the ``torch_a2a`` experts +forward sends each token to the rank that holds its expert. Every rank runs its own batch through the +sharded model, and the result must match the whole model run on that batch alone. +""" + +import pytest +import torch +from packaging.version import Version + +transformers = pytest.importorskip("transformers") +if Version(transformers.__version__) < Version("5.7"): + pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7", allow_module_level=True) + +from torch.distributed.tensor import DTensor +from transformers import AutoConfig, AutoModelForCausalLM, Qwen3MoeConfig + +from modelopt.torch.utils.distributed import ( + DistributionPlan, + distribution_plan, + set_distribution_plan, +) +from modelopt.torch.utils.moe import expert_parallel_mesh +from modelopt.torch.utils.plugins.model_load_utils import load_and_prepare_ep + +NUM_EXPERTS = 8 + + +def _save_tiny_qwen3_moe(ckpt_dir): + config = Qwen3MoeConfig( + hidden_size=64, + intermediate_size=128, + moe_intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=NUM_EXPERTS, + num_experts_per_tok=2, + vocab_size=128, + max_position_embeddings=64, + ) + torch.manual_seed(0) + AutoModelForCausalLM.from_config(config).to(torch.bfloat16).save_pretrained(ckpt_dir) + + +def _local(tensor): + return tensor.to_local() if isinstance(tensor, DTensor) else tensor + + +def _sharded_model_matches_the_whole_one(rank, size, ckpt_dir, ep_degree): + device = torch.device(f"cuda:{rank}") + config = AutoConfig.from_pretrained(ckpt_dir) # carries qwen3_moe's own base_model_ep_plan + plan = DistributionPlan(size, ep_size=ep_degree) + model = load_and_prepare_ep(ckpt_dir, plan, hf_config=config).eval() + # The loader swaps the plan on its own copy: the caller's config keeps the checkpoint's. + assert config.base_model_ep_plan == AutoConfig.from_pretrained(ckpt_dir).base_model_ep_plan + # The loader records no plan; loading through it directly, the caller records one. + with pytest.raises(RuntimeError, match="no distribution plan recorded"): + distribution_plan(model) + set_distribution_plan(model, plan) + # Calibration re-installs the all-to-all over the mesh the loader keeps on the model. + assert expert_parallel_mesh(model).size() == ep_degree + whole = AutoModelForCausalLM.from_pretrained(ckpt_dir, dtype=torch.bfloat16).to(device).eval() + + # Each rank holds its contiguous share of every layer's experts within its EP group, whose + # ranks are consecutive. + per_rank = NUM_EXPERTS // ep_degree + ep_rank = rank % ep_degree + owned = slice(ep_rank * per_rank, (ep_rank + 1) * per_rank) + for layer, whole_layer in zip(model.model.layers, whole.model.layers): + for name in ("gate_up_proj", "down_proj"): + local = _local(getattr(layer.mlp.experts, name)) + assert local.shape[0] == per_rank + torch.testing.assert_close(local, getattr(whole_layer.mlp.experts, name)[owned]) + # The load recorded what it could not place, for the export to carry: here, nothing. + assert model._modelopt_unplaced_source_keys == [] + assert model._modelopt_source_checkpoint == ckpt_dir + + # Every rank calibrates on different tokens, so each sends its own batch through the experts. + generator = torch.Generator().manual_seed(rank) + input_ids = torch.randint(0, 128, (2, 16), generator=generator).to(device) + with torch.no_grad(): + sharded = model(input_ids).logits + reference = whole(input_ids).logits + torch.testing.assert_close(sharded.float(), reference.float(), atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("ep_share", [1, 2], ids=["ep-world", "ep-half"]) +def test_expert_parallel_model_matches_the_whole_model( + need_2_gpus, dist_workers, tmp_path, ep_share +): + """With the experts sharded over every rank, or over half of them and replicated across the halves.""" + ep_degree = dist_workers.world_size // ep_share + if ep_degree < 2 or dist_workers.world_size % ep_share or NUM_EXPERTS % ep_degree: + pytest.skip(f"{NUM_EXPERTS} experts do not shard over EP {ep_degree} on these GPUs") + _save_tiny_qwen3_moe(tmp_path) + dist_workers.run(_sharded_model_matches_the_whole_one, str(tmp_path), ep_degree) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py new file mode 100644 index 00000000000..6f27ea1e17c --- /dev/null +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -0,0 +1,582 @@ +# 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. + +"""What the expert-parallel loader reads off the loaded model and its config. + +The loader takes an expert-parallel plan, ``_seed_parallel_states`` gives every module the +parallel state its amax syncs over, +``expert_parallel_mesh`` hands calibration the mesh an expert-parallel plan's experts are sharded +over, and ``_ckpt_quant_method`` reads the checkpoint's native quantization. Everything runs in one +CPU process. +""" + +import copy +import dataclasses +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +pytest.importorskip("transformers") + +from transformers import FineGrainedFP8Config + +import modelopt.torch.utils.distributed as dist +from modelopt.torch.utils import moe +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan +from modelopt.torch.utils.moe import expert_parallel_mesh +from modelopt.torch.utils.plugins import model_load_utils +from modelopt.torch.utils.plugins.model_load_utils import ( + _ckpt_quant_method, + _experts_sharded_over, + _seed_parallel_states, + _transformers_mesh_builder, + load_and_prepare_ep, +) + +# ----------------------------------------------------------------------------------- tiny MoE + + +class _Experts(nn.Module): + """A layer's routed experts as one grouped-GEMM container (every expert in fused tensors).""" + + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(4, 4, 2)) + self.down_proj = nn.Parameter(torch.zeros(4, 2, 2)) + self.weight_quantizer = nn.Identity() # quantizer insertion adds children like this + + def forward(self, hidden_states, top_k_index, top_k_weights, process_group=None): + return hidden_states + + +class _SharedExperts(nn.Module): + """The always-on shared expert MLP: dense, replicated on every rank.""" + + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Linear(2, 4, bias=False) + + def forward(self, hidden_states): + return hidden_states + + +class _MoeLayer(nn.Module): + 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 = _Experts() + self.mlp.shared_experts = _SharedExperts() + + +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 + + +# ------------------------------------------------------------------------------- the plan + +_EP_PLAN = {"layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts": "moe_tp_experts"} + + +class _Mesh: + """A device mesh, as far as the loader reads one; its groups run no collective (``-1``).""" + + device_type = "cpu" + + def __getitem__(self, name): + return self + + def size(self): + return 2 + + def get_group(self): + return -1 + + +@pytest.fixture +def ep_world_of_two(monkeypatch): + """A two-rank expert-parallel world as far as the loader's checks see, and the checkpoint paths + it resolves; nothing else is set up.""" + resolved = [] + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + monkeypatch.setattr( + model_load_utils, "_resolve_checkpoint_dir", lambda p, r: resolved.append(p) or p + ) + return resolved + + +@pytest.mark.parametrize( + ("plan", "error", "match"), + [ + (DistributionPlan(2), ValueError, "shards no routed experts"), + (DistributionPlan(2, ep_size=2, dense="fsdp"), NotImplementedError, "replicates the dense"), + ], + ids=["no-expert-parallelism", "sharded-dense"], +) +def test_the_loader_refuses_a_plan_it_cannot_load_by(plan, error, match): + """Before anything else: no process group, transformers release or checkpoint is looked at.""" + with pytest.raises(error, match=match): + load_and_prepare_ep("unused", plan) + + +def test_the_loader_refuses_a_plan_for_another_world(gloo_world_of_one): + with pytest.raises(ValueError, match="is for 2 ranks, but this run has 1"): + load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2)) + + +def test_the_loader_refuses_a_transformers_that_builds_no_model_with_torch_a2a(monkeypatch): + """transformers 5.6 checks a model's experts implementation against the ones it ships, so its + from_pretrained would refuse torch_a2a once the load had set everything up: refused first.""" + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr("transformers.__version__", "5.6.2") + + with pytest.raises(RuntimeError, match=r"requires transformers >= 5\.7"): + load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2), hf_config=SimpleNamespace()) + + +def test_a_hub_id_is_made_local_before_the_weights_load(monkeypatch): + """The export reads the tensors the load could not place back from the checkpoint it records, + which must be a local directory: a Hub ID there leaves the export nothing to read. The load + resolves it once the plan and the config pass, before anything reads a weight.""" + + class _ResolvedError(Exception): + """Stops the load where it resolves the checkpoint.""" + + resolved = [] + + def resolve(path, rank): + resolved.append(path) + raise _ResolvedError + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + monkeypatch.setattr(model_load_utils, "_resolve_checkpoint_dir", resolve) + + with pytest.raises(_ResolvedError): + load_and_prepare_ep( + "org/model", + DistributionPlan(2, ep_size=2), + hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), + ) + + assert resolved == ["org/model"] + + +@pytest.mark.parametrize( + ("config", "error", "match"), + [ + *( + ( + SimpleNamespace(quantization_config={"quant_method": method}), + NotImplementedError, + f"cannot dequantize a '{method}'", + ) + for method in ("gptq", "modelopt", "compressed-tensors") + ), + (SimpleNamespace(), ValueError, "declare base_model_ep_plan"), + ( + SimpleNamespace(base_model_ep_plan={"layers.*.self_attn.q_proj": "colwise"}), + ValueError, + "styles no routed experts", + ), + ], + ids=["gptq", "modelopt", "compressed-tensors", "no-ep-plan", "no-routed-experts"], +) +def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( + ep_world_of_two, config, error, match +): + """A format transformers' own quantizer would load as stored, for PTQ to treat as plain weights, + or a plan that leaves the experts whole on every rank; the load refuses either 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 == [] + + +def test_an_ep_size_that_does_not_divide_the_experts_is_refused_before_anything_loads( + ep_world_of_two, +): + """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) + + 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) + assert ep_world_of_two == [] + + +@pytest.mark.parametrize("declared", [True, False], ids=["on-the-config", "a-class-default"]) +def test_the_model_keeps_the_plan_its_checkpoint_declares( + ep_world_of_two, gloo_world_of_one, monkeypatch, declared +): + """The load shards by the torch_a2a styles, but the config the model keeps goes back to the + checkpoint's plan: an export writes it out, and nothing without ModelOpt knows those styles. A + plan the config class ships (transformers' own, e.g. Qwen3-MoE's) stays off the instance.""" + + class _PastError(Exception): + """Raised by the load's next check.""" + + model = _model_with_decoder_at("model.") + + class _Loader: + @staticmethod + def from_pretrained(path, config, **kwargs): + model.config = copy.deepcopy(config) # transformers keeps a copy of the config + return model + + def next_check(*args, **kwargs): + raise _PastError + + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) + + class _Config(SimpleNamespace): + base_model_ep_plan = _EP_PLAN + + # 8 experts over 2 ranks: the stand-in model holds this rank's 4. + config = _Config(num_experts=8, **({"base_model_ep_plan": dict(_EP_PLAN)} if declared else {})) + + with pytest.raises(_PastError): + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) + + assert model.config.base_model_ep_plan == _EP_PLAN + assert ("base_model_ep_plan" in vars(model.config)) is declared + + +def test_a_loaded_model_keeps_its_sharding_styles_out_of_its_saved_config( + ep_world_of_two, gloo_world_of_one, monkeypatch +): + """What the loader leaves on a transformers model, as transformers treats it. Built with the + torch_a2a styles, the model keeps them as ``model.tp_plan``, which calibration re-installs the + all-to-all forwards by, and the expert-parallel mesh it re-installs them over, which this stand-in + from_pretrained does not set; while its config goes back to the checkpoint's plan; and the experts + implementation its forwards dispatch on, torch_a2a, stays out of the config it serializes, + which an export writes.""" + from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM + from transformers.distributed import DistributedConfig + + if not Qwen3MoeConfig.base_model_ep_plan: + pytest.skip("this transformers ships no expert-parallel plan for Qwen3-MoE") + + class _PastError(Exception): + """Raised by the load's next check.""" + + built = [] + + class _Loader: + @staticmethod + def from_pretrained(path, config, **kwargs): + # As transformers' expert-parallel load builds it: this rank's 2 of the 4 experts. + config.distributed_config = DistributedConfig(enable_expert_parallel=True) + share = copy.deepcopy(config) + share.num_experts //= 2 + with torch.device("meta"): + model = Qwen3MoeForCausalLM(share) + model.config = config + built.append(model) + return model + + def next_check(*args, **kwargs): + raise _PastError + + mesh = _Mesh() + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: mesh + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) + config = Qwen3MoeConfig( + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + vocab_size=32, + num_experts=4, + num_experts_per_tok=1, + ) + + with pytest.raises(_PastError): + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) + + (model,) = built + assert model.config.base_model_ep_plan == Qwen3MoeConfig.base_model_ep_plan + assert "torch_a2a_experts" in set(model.tp_plan.values()) + assert model._device_mesh is mesh + assert model.config._experts_implementation == "torch_a2a" + assert not any("experts_implementation" in key for key in model.config.to_dict()) + + +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.""" + + class _LoadedError(Exception): + """Stops the load at from_pretrained.""" + + loaded, classes = {}, [] + + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + loaded.update(kwargs) + raise _LoadedError + + def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): + classes.append(trust_remote_code) + return _Loader + + monkeypatch.setattr( + "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) + + with pytest.raises(_LoadedError): + load_and_prepare_ep( + "org/model", DistributionPlan(2, ep_size=2), trust_remote_code=True, hf_config=config + ) + assert loaded["dtype"] is torch.float16 + assert classes == [True] + + +# ----------------------------------------------------------------------- _seed_parallel_states + + +def test_expert_containers_and_their_children_get_the_expert_state(): + """Everything else, the root and the containers' own parents included, gets the dense state; so + does a sibling whose name merely starts with a container's (Qwen2-MoE keeps + ``shared_expert_gate`` beside ``shared_expert``).""" + model = _model_with_decoder_at("model.") + model.model.layers[0].mlp.experts_gate = nn.Linear(2, 1) + experts = {f"model.layers.{i}.mlp.experts": model.model.layers[i].mlp.experts for i in (0, 1)} + dense, expert = object(), object() + + _seed_parallel_states(model, experts, dense_state=dense, expert_state=expert) + + states = {name: module._parallel_state for name, module in model.named_modules()} + assert {name for name, state in states.items() if state is expert} == { + "model.layers.0.mlp.experts", + "model.layers.0.mlp.experts.weight_quantizer", + "model.layers.1.mlp.experts", + "model.layers.1.mlp.experts.weight_quantizer", + } + assert all(state is dense for state in states.values() if state is not expert) + + +# ------------------------------------------------------------------------ expert_parallel_mesh + + +@pytest.mark.parametrize( + ("plan", "is_expert_parallel"), + [(DistributionPlan(2, ep_size=2), True), (None, False), (DistributionPlan(2), False)], + ids=["expert-parallel", "no-plan", "replicated"], +) +def test_the_expert_parallel_mesh_comes_with_an_expert_parallel_plan_only(plan, is_expert_parallel): + """An expert-parallel plan: the mesh transformers kept on the model. Any other model has none, + even with a device mesh on it, as a load through from_pretrained's own TP can leave.""" + model = _model_with_decoder_at("model.") + if plan is not None: + set_distribution_plan(model, plan) + model._device_mesh = mesh = object() + + assert expert_parallel_mesh(model) is (mesh if is_expert_parallel else None) + + +def test_an_expert_parallel_model_without_its_mesh_is_refused(): + """A model recorded expert-parallel that load_and_prepare_ep did not load has no mesh for + calibration to re-install the all-to-all forwards over; it is refused by name.""" + model = _model_with_decoder_at("model.") + set_distribution_plan(model, DistributionPlan(2, ep_size=2)) + + with pytest.raises(RuntimeError, match="carries no device mesh"): + expert_parallel_mesh(model) + + +# --------------------------------------------------------------------------- _ckpt_quant_method + + +@pytest.mark.parametrize( + ("config", "expected"), + [ + (SimpleNamespace(quantization_config={"quant_method": "fp8"}), "fp8"), + (SimpleNamespace(quantization_config=SimpleNamespace(quant_method="mxfp4")), "mxfp4"), + ( + SimpleNamespace( + text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}) + ), + "fp8", + ), + ( + SimpleNamespace( + llm_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}) + ), + "fp8", + ), + ( + SimpleNamespace( + quantization_config={"quant_method": "mxfp4"}, + text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}), + ), + "mxfp4", + ), + (SimpleNamespace(quantization_config=FineGrainedFP8Config()), "fp8"), + (SimpleNamespace(quantization_config={"weight_block_size": [128, 128]}), None), + (SimpleNamespace(text_config=None), None), + ], + ids=[ + "dict", + "object", + "vlm-text-config", + "vlm-llm-config", + "top-level-wins", + "transformers-config-object", + "no-method", + "unquantized", + ], +) +def test_quant_method_comes_from_the_config_or_a_vlms_text_config(config, expected): + """The top-level config first, else a VLM's text or LLM config. A loaded model's config holds + transformers' config object, not the JSON dict, and callers compare against ``"fp8"`` / + ``"mxfp4"``: it reports the plain method name, not its enum.""" + assert _ckpt_quant_method(config) == expected + + +def test_from_pretrained_builds_its_mesh_from_the_ep_mesh_it_is_given(): + """Newer transformers ignores ``from_pretrained``'s ``device_mesh`` and builds a world-sized one, + which shards the experts over every rank when EP is below the world size. During the load it gets + the EP mesh instead, and its own builder is back afterwards, even when the load raises.""" + pytest.importorskip("transformers") + mixin = _transformers_mesh_builder() + if mixin is None: + pytest.skip("this transformers honors from_pretrained's device_mesh") + from transformers import PreTrainedModel + from transformers.distributed import DistributedConfig + + builder = mixin.initialize_distributed_mesh + ep_mesh = SimpleNamespace(device_type="cpu") + config = DistributedConfig(enable_expert_parallel=True, tp_plan="auto", tp_size=2) + with _experts_sharded_over(ep_mesh): + _, device_map, device_mesh = PreTrainedModel.prepare_distribute_model(config) + + assert device_mesh is ep_mesh + assert device_map == torch.device("cpu") + assert mixin.initialize_distributed_mesh is builder + with pytest.raises(RuntimeError), _experts_sharded_over(ep_mesh): + raise RuntimeError("the load failed") + assert mixin.initialize_distributed_mesh is builder + + +@pytest.mark.parametrize( + ("fields", "builds_its_mesh", "config", "tp_plan"), + [ + # transformers 5.6-5.9, 5.11, 5.12 + (["enable_expert_parallel"], False, {"enable_expert_parallel": True}, "auto"), + # 5.10, 5.13-5.17 + ( + ["tp_size", "tp_plan", "enable_expert_parallel"], + False, + {"enable_expert_parallel": True, "tp_plan": "auto"}, + None, + ), + # 5.18 on + ( + ["tp_size", "tp_plan", "enable_expert_parallel"], + True, + {"enable_expert_parallel": True, "tp_plan": "auto", "tp_size": 2}, + None, + ), + # no release so far + ( + ["tp_size", "enable_expert_parallel"], + True, + {"enable_expert_parallel": True, "tp_size": 2}, + "auto", + ), + ], + ids=["no-tp-plan-field", "tp-plan-field", "own-mesh-builder", "own-mesh-builder-no-tp-plan"], +) +def test_each_load_argument_goes_where_this_transformers_takes_it( + ep_world_of_two, monkeypatch, fields, builds_its_mesh, config, tp_plan +): + """``tp_plan="auto"`` goes on ``DistributedConfig`` where it has the field, else to + ``from_pretrained``, never both; a release that builds its own mesh gets it sized to the EP + degree. Each is decided by its own field, so a release with one and not the other loads too.""" + + class _LoadedError(Exception): + """Stops the load at from_pretrained.""" + + loaded = {} + + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + loaded.update(kwargs) + raise _LoadedError + + import transformers.distributed as hf_distributed + + distributed_config = dataclasses.make_dataclass( + "DistributedConfig", [(name, object, dataclasses.field(default=None)) for name in fields] + ) + monkeypatch.setattr(hf_distributed, "DistributedConfig", distributed_config) + monkeypatch.setattr( + model_load_utils, + "_transformers_mesh_builder", + lambda: SimpleNamespace(initialize_distributed_mesh=None) if builds_its_mesh else None, + ) + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + + with pytest.raises(_LoadedError): + load_and_prepare_ep( + "org/model", + DistributionPlan(2, ep_size=2), + hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), + ) + + given = {k: v for k, v in vars(loaded["distributed_config"]).items() if v is not None} + assert given == config + assert loaded.get("tp_plan") == tp_plan From 184675cbb1c94c215058961b7e67b047154d87c2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:15:23 +0000 Subject: [PATCH 02/23] Fail every rank when one rank's expert-parallel load fails Each rank checks its own transformers, reads or is handed its own config, sees the checkpoint where rank 0 made it local, and loads its own share, so a rank could fail alone and leave the others waiting for it in the meshes, the load's collectives or the warmup's: a node with an older transformers or without the checkpoint's storage, one rank handed another config, a rank out of memory in the load or unable to read a weight file. Those steps now run inside _refused_together blocks, whose failure-carrying all_reduce raises a failure on every rank before the next collective: the checks before anything loads, the checkpoint, now made local through ensure_local_checkpoint (rank 0 resolves it for every rank, on a group whose timeout outlasts a download), and the load through the routed experts' lookup. A test runs the load on two real gloo ranks, one failing alone at each of those points, with the meshes' collective stood in by a bounded barrier: every rank raises that rank's failure. Without the blocks every case fails, the other rank waiting at the barrier or the warmup. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 192 +++++++++--------- tests/unit/torch/utils/test_ep_loader.py | 116 ++++++++++- 2 files changed, 209 insertions(+), 99 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 3581fd284d9..8da1487ea7e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1713,6 +1713,8 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A bring-up is the launcher's job, not this loader's (hf_ptq.main() does it via ``distributed_run``). The expert-parallel mesh stays recoverable from the model (:func:`~modelopt.torch.utils.moe.expert_parallel_mesh`), once the caller has recorded the plan. + A refusal on any rank -- of its transformers, the config, the checkpoint, or in the load -- + raises on every rank, so no rank is left waiting in a collective for one that already failed. """ # Imported here, so looked up when the load runs: its tests stand in a process group and a # transformers release by patching these two modules. @@ -1745,46 +1747,49 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A if plan.world_size != size(): raise ValueError(f"{plan!r} is for {plan.world_size} ranks, but this run has {size()}.") world_size, ep_degree, edp = plan.world_size, plan.ep_size, plan.edp - # Hard floor: below 5.6 there is no ExpertsInterface / grouped-GEMM EP styles, and - # enable_expert_parallel silently shards nothing (no error); 5.6 checks a model's experts - # implementation against the ones it ships, so from_pretrained would refuse torch_a2a. - require_transformers( - "5.7", "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} - if attn_implementation is not None: - config_kwargs["attn_implementation"] = attn_implementation - hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) - quant_method = _ckpt_quant_method(hf_config) - if quant_method not in _EP_SOURCE_FORMATS: - # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights. - dequantized = ", ".join(sorted(f for f in _EP_SOURCE_FORMATS if f)) - raise NotImplementedError( - f"The expert-parallel load cannot dequantize a {quant_method!r} checkpoint: it loads " - f"unquantized checkpoints and dequantizes {dequantized}. Start from an unquantized copy " - "of the model." - ) - # The swapped plan is this load's own: the caller's config keeps the plan its checkpoint declares. - hf_config = copy.deepcopy(hf_config) - ep_cfg = _ep_config(hf_config) # for VLMs the plan lives on .text_config - 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." - ) - # The all-to-all transport finds an expert's rank by dividing its id by this rank's share of - # the experts, which holds only when every rank holds as many. A count the config names none of - # the ways this reads is checked once the experts load, from their weights. - num_experts = _num_experts_from_config(hf_config) - if num_experts % ep_degree: - raise ValueError( - f"{plan!r}: ep_size={ep_degree} must divide the {num_experts} routed experts, which the " - "all-to-all transport shares out equally over the expert-parallel ranks." + # 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): + # Hard floor: below 5.6 there is no ExpertsInterface / grouped-GEMM EP styles, and + # enable_expert_parallel silently shards nothing (no error); 5.6 checks a model's experts + # implementation against the ones it ships, so from_pretrained would refuse torch_a2a. + require_transformers( + "5.7", "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} + if attn_implementation is not None: + config_kwargs["attn_implementation"] = attn_implementation + hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) + quant_method = _ckpt_quant_method(hf_config) + if quant_method not in _EP_SOURCE_FORMATS: + # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights. + dequantized = ", ".join(sorted(f for f in _EP_SOURCE_FORMATS if f)) + raise NotImplementedError( + f"The expert-parallel load cannot dequantize a {quant_method!r} checkpoint: it loads " + f"unquantized checkpoints and dequantizes {dequantized}. Start from an unquantized copy " + "of the model." + ) + # The swapped plan is this load's own: the caller's config keeps the plan its checkpoint declares. + hf_config = copy.deepcopy(hf_config) + ep_cfg = _ep_config(hf_config) # for VLMs the plan lives on .text_config + 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." + ) + # The all-to-all transport finds an expert's rank by dividing its id by this rank's share of + # the experts, which holds only when every rank holds as many. A count the config names none of + # the ways this reads is checked once the experts load, from their weights. + num_experts = _num_experts_from_config(hf_config) + if num_experts % ep_degree: + raise ValueError( + f"{plan!r}: ep_size={ep_degree} must divide the {num_experts} routed experts, which the " + "all-to-all transport shares out equally over the expert-parallel ranks." + ) # transformers 5.6 on, which the floor above makes sure of. from torch.distributed.device_mesh import init_device_mesh from transformers.distributed import DistributedConfig @@ -1792,8 +1797,10 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 # The weights load, and the checkpoint recorded for the export to read the unplaced tensors back # from, take the local snapshot: a Hub ID recorded there would leave the export nothing to read. - # Resolved once the plan and the config pass, so a refused load downloads nothing. - ckpt_path = _resolve_checkpoint_dir(ckpt_path, torch.distributed.get_rank()) + # Resolved once the plan and the config pass, so a refused load downloads nothing. Rank 0 makes + # it local; a rank that cannot see where fails every rank. + with _refused_together("the local checkpoint", alone=False): + _, ckpt_path = ensure_local_checkpoint(ckpt_path) # Two INDEPENDENT meshes (NeMo-AutoModel style), named by the model part they govern -- two # factorizations of the same world: @@ -1881,56 +1888,59 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A ep_cfg.base_model_ep_plan = swapped_plan from_pretrained_kwargs["config"] = hf_config print_rank_0(f"[ep] ep_plan={swapped_plan} experts_implementation=torch_a2a") - # The class that builds the whole checkpoint -- with its lm_head where an auto class builds one, - # so the head is calibrated and exported and the generation preview works -- chosen as the - # distribution plan's decision chose it, unless the caller asked for a VLM without its head - # (require_lm_head=False). The per-model plugin (models/) has already installed any on-the-fly - # weight decompression as a transformers monkeypatch. - loader = _auto_class(hf_config, require_lm_head, trust_remote_code=trust_remote_code) - print_rank_0(f"[ep] loader={loader.__name__} (require_lm_head={require_lm_head})") - with _experts_sharded_over(ep_submesh): - model = loader.from_pretrained(ckpt_path, **from_pretrained_kwargs) - model.eval() - # The swapped styles drove the sharding, and the model keeps the plan it was sharded by - # (``model.tp_plan``, copied when it was built). The config it keeps goes back to the - # checkpoint's plan: an export writes the config out, and nothing without ModelOpt knows the - # torch_a2a styles. - kept_config = getattr(model, "config", None) - if kept_config is not None: - kept_ep_config = _ep_config(kept_config) - if declared_plan is None: - vars(kept_ep_config).pop("base_model_ep_plan", None) - else: - kept_ep_config.base_model_ep_plan = declared_plan - # Force the experts forward to dispatch to `torch_a2a`: from_pretrained's `experts_implementation` - # kwarg resolves non-fp8 impls back to "grouped_mm", but @use_experts_implementation reads - # `config._experts_implementation` at FORWARD time, so set it on every (sub)config now. (Without - # this the torch_a2a_experts style passes a 5th process_group arg to grouped_mm_experts_forward -> - # TypeError.) Unlike the plan it stays set, and an export still writes no torch_a2a: transformers - # keeps it as ``_experts_implementation_internal``, which it leaves out of a serialized config. - for _m in model.modules(): - _cfg = getattr(_m, "config", None) - if _cfg is not None and hasattr(_cfg, "_experts_implementation"): - _cfg._experts_implementation = "torch_a2a" - - # The checkpoint tensors the model has no parameter for (an MTP head), worked out once, at load, - # for the export to carry. - record_unplaced_source_keys(model, ckpt_path) - - # DP=EP all-to-all: no FSDP. Non-experts stay replicated per rank (from_pretrained already placed - # them on this rank's GPU); experts stay EP-sharded on the ep submesh, which the model carries as - # ``_device_mesh`` for calibration to re-install the all-to-all forwards over (expert_parallel_mesh). - # transformers sets it from the mesh this load hands it, through internals that differ by release; - # set here, the model carries it on every release. - model._device_mesh = ep_submesh - if hasattr(model, "tie_weights"): - model.tie_weights() - if freeze: - model.requires_grad_(False) - - # The routed-expert containers, from the SAME base_model_ep_plan that drove the sharding. A MoE model - # loaded WITHOUT EP (dense DP) never reaches this loader, so its experts stay dense. - expert_modules = _routed_expert_modules(model, ep_plan) + # A rank failing in the load (out of memory, a file it cannot read) fails every rank here, + # before the warmup's collectives. + with _refused_together("the expert-parallel load", alone=False): + # The class that builds the whole checkpoint -- with its lm_head where an auto class builds one, + # so the head is calibrated and exported and the generation preview works -- chosen as the + # distribution plan's decision chose it, unless the caller asked for a VLM without its head + # (require_lm_head=False). The per-model plugin (models/) has already installed any on-the-fly + # weight decompression as a transformers monkeypatch. + loader = _auto_class(hf_config, require_lm_head, trust_remote_code=trust_remote_code) + print_rank_0(f"[ep] loader={loader.__name__} (require_lm_head={require_lm_head})") + with _experts_sharded_over(ep_submesh): + model = loader.from_pretrained(ckpt_path, **from_pretrained_kwargs) + model.eval() + # The swapped styles drove the sharding, and the model keeps the plan it was sharded by + # (``model.tp_plan``, copied when it was built). The config it keeps goes back to the + # checkpoint's plan: an export writes the config out, and nothing without ModelOpt knows the + # torch_a2a styles. + kept_config = getattr(model, "config", None) + if kept_config is not None: + kept_ep_config = _ep_config(kept_config) + if declared_plan is None: + vars(kept_ep_config).pop("base_model_ep_plan", None) + else: + kept_ep_config.base_model_ep_plan = declared_plan + # Force the experts forward to dispatch to `torch_a2a`: from_pretrained's `experts_implementation` + # kwarg resolves non-fp8 impls back to "grouped_mm", but @use_experts_implementation reads + # `config._experts_implementation` at FORWARD time, so set it on every (sub)config now. (Without + # this the torch_a2a_experts style passes a 5th process_group arg to grouped_mm_experts_forward -> + # TypeError.) Unlike the plan it stays set, and an export still writes no torch_a2a: transformers + # keeps it as ``_experts_implementation_internal``, which it leaves out of a serialized config. + for _m in model.modules(): + _cfg = getattr(_m, "config", None) + if _cfg is not None and hasattr(_cfg, "_experts_implementation"): + _cfg._experts_implementation = "torch_a2a" + + # The checkpoint tensors the model has no parameter for (an MTP head), worked out once, at load, + # for the export to carry. + record_unplaced_source_keys(model, ckpt_path) + + # DP=EP all-to-all: no FSDP. Non-experts stay replicated per rank (from_pretrained already placed + # them on this rank's GPU); experts stay EP-sharded on the ep submesh, which the model carries as + # ``_device_mesh`` for calibration to re-install the all-to-all forwards over (expert_parallel_mesh). + # transformers sets it from the mesh this load hands it, through internals that differ by release; + # set here, the model carries it on every release. + model._device_mesh = ep_submesh + if hasattr(model, "tie_weights"): + model.tie_weights() + if freeze: + model.requires_grad_(False) + + # The routed-expert containers, from the SAME base_model_ep_plan that drove the sharding. A MoE model + # loaded WITHOUT EP (dense DP) never reaches this loader, so its experts stay dense. + expert_modules = _routed_expert_modules(model, ep_plan) # Each module's ParallelState for calibration's amax sync: dense modules reduce over the whole world; # routed experts reduce over their expert-data-parallel replicas, and share one input amax across the diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 6f27ea1e17c..ff1efba7d30 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -24,21 +24,26 @@ import copy import dataclasses +from contextlib import ExitStack from types import SimpleNamespace +from unittest import mock import pytest import torch +from _test_utils.torch.distributed.utils import bounded_collectives +from packaging.version import Version from torch import nn pytest.importorskip("transformers") -from transformers import FineGrainedFP8Config +import transformers +from transformers import FineGrainedFP8Config, Qwen3MoeConfig import modelopt.torch.utils.distributed as dist from modelopt.torch.utils import moe from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan from modelopt.torch.utils.moe import expert_parallel_mesh -from modelopt.torch.utils.plugins import model_load_utils +from modelopt.torch.utils.plugins import hf_checkpoint_utils, model_load_utils from modelopt.torch.utils.plugins.model_load_utils import ( _ckpt_quant_method, _experts_sharded_over, @@ -121,7 +126,7 @@ def get_group(self): @pytest.fixture -def ep_world_of_two(monkeypatch): +def ep_world_of_two(monkeypatch, gloo_world_of_one): """A two-rank expert-parallel world as far as the loader's checks see, and the checkpoint paths it resolves; nothing else is set up.""" resolved = [] @@ -130,7 +135,7 @@ def ep_world_of_two(monkeypatch): monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) monkeypatch.setattr(moe, "require_transformers", lambda *args: None) monkeypatch.setattr( - model_load_utils, "_resolve_checkpoint_dir", lambda p, r: resolved.append(p) or p + model_load_utils, "ensure_local_checkpoint", lambda p, **k: resolved.append(p) or (None, p) ) return resolved @@ -154,7 +159,9 @@ def test_the_loader_refuses_a_plan_for_another_world(gloo_world_of_one): load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2)) -def test_the_loader_refuses_a_transformers_that_builds_no_model_with_torch_a2a(monkeypatch): +def test_the_loader_refuses_a_transformers_that_builds_no_model_with_torch_a2a( + monkeypatch, gloo_world_of_one +): """transformers 5.6 checks a model's experts implementation against the ones it ships, so its from_pretrained would refuse torch_a2a once the load had set everything up: refused first.""" monkeypatch.setattr(dist, "is_initialized", lambda: True) @@ -165,7 +172,7 @@ def test_the_loader_refuses_a_transformers_that_builds_no_model_with_torch_a2a(m load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2), hf_config=SimpleNamespace()) -def test_a_hub_id_is_made_local_before_the_weights_load(monkeypatch): +def test_a_hub_id_is_made_local_before_the_weights_load(monkeypatch, gloo_world_of_one): """The export reads the tensors the load could not place back from the checkpoint it records, which must be a local directory: a Hub ID there leaves the export nothing to read. The load resolves it once the plan and the config pass, before anything reads a weight.""" @@ -175,7 +182,7 @@ class _ResolvedError(Exception): resolved = [] - def resolve(path, rank): + def resolve(path, **hub_options): resolved.append(path) raise _ResolvedError @@ -183,7 +190,7 @@ def resolve(path, rank): monkeypatch.setattr(dist, "size", lambda group=None: 2) monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) monkeypatch.setattr(moe, "require_transformers", lambda *args: None) - monkeypatch.setattr(model_load_utils, "_resolve_checkpoint_dir", resolve) + monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", resolve) with pytest.raises(_ResolvedError): load_and_prepare_ep( @@ -580,3 +587,96 @@ def from_pretrained(path, **kwargs): given = {k: v for k, v in vars(loaded["distributed_config"]).items() if v is not None} assert given == config assert loaded.get("tp_plan") == tp_plan + + +# ----------------------------------------------------------------- one rank failing alone + + +class _PastTheMeshesError(Exception): + """Raised by the meshes' stand-in: this rank got past the load's checks.""" + + +def _meshes_in_step(*args, **kwargs): + """The meshes' subgroups, as the collective they are: a rank in it waits for every peer.""" + torch.distributed.barrier() + raise _PastTheMeshesError + + +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) + if failure != "unseen-checkpoint": + local = mock.patch.object( + model_load_utils, "ensure_local_checkpoint", side_effect=lambda p, **k: (None, p) + ) + patches.append(local) + fails_here = rank == failing_rank + + def fail(*args, **kwargs): + raise (MemoryError if failure == "load" else OSError)(match) + + if failure == "unreadable-config" and fails_here: + path = ckpt + "-not-on-this-node" + elif failure == "transformers" and fails_here: + patches.append(mock.patch("transformers.__version__", "5.6.2")) + elif failure == "quantized-config" and fails_here: + hf_config.quantization_config = {"quant_method": "gptq"} + elif failure == "unseen-checkpoint" and fails_here: + # Rank 0 makes the checkpoint local; a rank on a node without that storage cannot see it. + patches.append(mock.patch.object(hf_checkpoint_utils.os.path, "isdir", return_value=False)) + if failure in ("load", "record"): + loaded = _model_with_decoder_at("model.") + load = fail if failure == "load" and fails_here else (lambda path, **kwargs: loaded) + loader = SimpleNamespace(__name__="Loader", from_pretrained=load) + record = fail if failure == "record" and fails_here else None + patches += [ + mock.patch("torch.distributed.device_mesh.init_device_mesh", return_value=_Mesh()), + mock.patch.object(model_load_utils, "_auto_class", return_value=loader), + mock.patch.object(model_load_utils, "record_unplaced_source_keys", side_effect=record), + ] + else: + patches.append( + mock.patch("torch.distributed.device_mesh.init_device_mesh", _meshes_in_step) + ) + refused = None + with bounded_collectives(), ExitStack() as stack: + for patch in patches: + stack.enter_context(patch) + try: # not pytest.raises: its failure would end the pooled worker instead of the job + load_and_prepare_ep(path, DistributionPlan(2, ep_size=2), hf_config=hf_config) + except Exception as error: + refused = f"{type(error).__name__}: {error}" + # The one rank's failure on every rank, before any rank went on alone into a collective. + assert refused is not None and match in refused, (rank, refused) + assert fails_here or f"rank {failing_rank} failed" in refused, refused + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("5.7"), reason="the EP load needs transformers 5.7" +) +@pytest.mark.parametrize( + ("failure", "failing_rank", "match"), + [ + ("unreadable-config", 1, "not-on-this-node"), + ("transformers", 1, "requires transformers >= 5.7, but 5.6.2"), + ("quantized-config", 1, "cannot dequantize a 'gptq'"), + ("unseen-checkpoint", 1, "which this rank cannot see"), + ("load", 1, "the GPU ran out of memory"), + ("record", 1, "a weight file this rank cannot read"), + ], +) +def test_a_load_failing_on_one_rank_fails_every_rank( + cpu_dist_workers, tmp_path, failure, failing_rank, match +): + """One rank fails alone: reading the config, on its transformers, on its arguments, seeing the + checkpoint rank 0 made local, or in the load. Every rank raises that failure, rather than leave the others + waiting for it in the meshes (stood in by a bounded barrier) or the warmup's collectives.""" + Qwen3MoeConfig(num_experts=4, num_hidden_layers=1, hidden_size=8).save_pretrained(tmp_path) + cpu_dist_workers.run( + _one_rank_fails_job, + ckpt=str(tmp_path), + failure=failure, + failing_rank=failing_rank, + match=match, + ) From 40ab9af79b07c61d291824491360edddb45e8a24 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:36:08 +0000 Subject: [PATCH 03/23] Refuse an EP plan whose experts no grouped_gemm entry shards The load shards an expert's weights only through the plan's grouped_gemm entries under its container: the torch_a2a styles the containers take for the load shard nothing themselves, and packed_* entries are dropped. A plan that styles only the container, or shards the expert weights another way, passed every check, and every rank loaded every expert while the all-to-all transport routed tokens as if each held its share. It is now refused with the other plan checks, before anything loads. Every plan transformers 5.18 ships, and every one ModelOpt's model specs supply, shards its experts through grouped_gemm. The tests' stand-in plan gains the grouped_gemm entries the real plans have, and two rows pin the refusal: a container-only plan and one sharding the experts by packed_colwise. Both fail without the check. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 27 +++++++++++++--- tests/unit/torch/utils/test_ep_loader.py | 32 ++++++++++++++++--- 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 8da1487ea7e..e9a91481db0 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -78,6 +78,7 @@ __all__ = [ "build_meta_causal_lm", "fsdp2_wrap", + "is_fsdp2_text_only_vlm", "load_and_prepare_ep", "load_pretrained", "parallel_load_and_prepare_fsdp2", @@ -1614,7 +1615,7 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st # The checkpoint formats load_and_prepare_ep dequantizes as it loads, beside unquantized ones (None). -# mxfp4 is for experts stored in the standard layout: GPT-OSS's are transposed, which the load refuses. +# mxfp4 is for experts stored in the standard layout. _EP_SOURCE_FORMATS = frozenset({None, "mxfp4"}) @@ -1715,6 +1716,8 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A (:func:`~modelopt.torch.utils.moe.expert_parallel_mesh`), once the caller has recorded the plan. A refusal on any rank -- of its transformers, the config, the checkpoint, or in the load -- raises on every rank, so no rank is left waiting in a collective for one that already failed. + Without ``hf_config`` every rank reads the checkpoint's config itself, from the Hub for a Hub ID, + before rank 0 makes the checkpoint local: on many ranks, pass it (as ``hf_ptq`` does). """ # Imported here, so looked up when the load runs: its tests stand in a process group and a # transformers release by patching these two modules. @@ -1781,6 +1784,22 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A 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 (packed_* entries are dropped). + 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 all-to-all transport finds an expert's rank by dividing its id by this rank's share of # the experts, which holds only when every rank holds as many. A count the config names none of # the ways this reads is checked once the experts load, from their weights. @@ -1849,7 +1868,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A if attn_implementation is not None: from_pretrained_kwargs["attn_implementation"] = attn_implementation if quant_method == "mxfp4": - # Native mxfp4 checkpoints (e.g. gpt-oss: fused `gate_up_proj_blocks`/`_scales` + biases) must be + # Native mxfp4 checkpoints (fused `gate_up_proj_blocks`/`_scales` + biases) must be # DEQUANTIZED to bf16 on load so ModelOpt can insert nvfp4/fp8 quantizers on the fused bf16 experts # (the a2a grouped_gemm expects bf16). transformers dequantizes blocks->bf16 during from_pretrained. # NB: the single-process loader needed device_map="sequential" to dodge a CUDA illegal-access in the @@ -1858,9 +1877,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A from transformers import Mxfp4Config from_pretrained_kwargs["quantization_config"] = Mxfp4Config(dequantize=True) - print_rank_0( - "[ep] native mxfp4 checkpoint (e.g. gpt-oss) -> Mxfp4Config(dequantize=True) -> bf16" - ) + print_rank_0("[ep] native mxfp4 checkpoint -> Mxfp4Config(dequantize=True) -> bf16") # DP=EP all-to-all: swap the ep_plan to our torch_a2a styles (torch_a2a_router emits GLOBAL ids # unmasked; torch_a2a_experts injects the EP process_group + skips the all-reduce) and select the # `torch_a2a` experts forward (NCCL all-to-all dispatch/combine around the quantizable diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index ff1efba7d30..fbec66a5256 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -107,7 +107,12 @@ def _model_with_decoder_at(prefix): # ------------------------------------------------------------------------------- the plan -_EP_PLAN = {"layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts": "moe_tp_experts"} +_EP_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", +} class _Mesh: @@ -219,15 +224,34 @@ def resolve(path, **hub_options): ValueError, "styles no routed experts", ), + *( + ( + SimpleNamespace( + base_model_ep_plan={"layers.*.mlp.experts": "moe_tp_experts", **params} + ), + ValueError, + r"shards no weights of layers\.\*\.mlp\.experts through grouped_gemm", + ) + for params in ({}, {"layers.*.mlp.experts.gate_up_proj": "packed_colwise"}) + ), + ], + ids=[ + "gptq", + "modelopt", + "compressed-tensors", + "no-ep-plan", + "no-routed-experts", + "experts-container-only", + "experts-sharded-otherwise", ], - ids=["gptq", "modelopt", "compressed-tensors", "no-ep-plan", "no-routed-experts"], ) def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( ep_world_of_two, config, error, match ): """A format transformers' own quantizer would load as stored, for PTQ to treat as plain weights, - or a plan that leaves the experts whole on every rank; the load refuses either before it reads - anything.""" + 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. 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 == [] From b67fa11d78a94f1beecc6741840ded4cce7212df Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:28:39 +0000 Subject: [PATCH 04/23] Refuse an EP plan that matches no experts or styles some otherwise A plan whose routed-expert containers no module of the loaded model matches left every module the dense state, so calibration reduced each expert's weight amax over ranks holding different experts: the load now refuses it, inside its guard. The refusal moves here from the loader's [3/3], word for word, with its test. A plan that shards one expert weight by grouped_gemm and styles another by a style this load drops passed the grouped_gemm check: the other weight loaded whole beside its sharded sibling, and the grouped forward paired slices of different experts. Such a plan is refused before anything loads. Every base_model_ep_plan transformers 5.18 ships and every ModelOpt spec plan passes the check. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 22 ++++++++++++ tests/unit/torch/utils/test_ep_loader.py | 34 ++++++++++++++++++- 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index e9a91481db0..ca08a78bdc4 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1800,6 +1800,22 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A "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 ("grouped_gemm", "ep_router", *_EXPERTS_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 all-to-all transport finds an expert's rank by dividing its id by this rank's share of # the experts, which holds only when every rank holds as many. A count the config names none of # the ways this reads is checked once the experts load, from their weights. @@ -1958,6 +1974,12 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A # The routed-expert containers, from the SAME base_model_ep_plan that drove the sharding. A MoE model # loaded WITHOUT EP (dense DP) never reaches this loader, so its experts stay dense. expert_modules = _routed_expert_modules(model, ep_plan) + if not expert_modules: + raise ValueError( + f"base_model_ep_plan={ep_plan} names routed-expert containers that no module of " + f"{type(model).__name__} matches, so calibration would reduce the experts' amax over " + "ranks holding different experts. Quantize the model without expert parallelism." + ) # Each module's ParallelState for calibration's amax sync: dense modules reduce over the whole world; # routed experts reduce over their expert-data-parallel replicas, and share one input amax across the diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index fbec66a5256..fb4e221fdb6 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -234,6 +234,13 @@ def resolve(path, **hub_options): ) for params in ({}, {"layers.*.mlp.experts.gate_up_proj": "packed_colwise"}) ), + ( + SimpleNamespace( + 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", + ), ], ids=[ "gptq", @@ -243,6 +250,7 @@ def resolve(path, **hub_options): "no-routed-experts", "experts-container-only", "experts-sharded-otherwise", + "experts-sharded-partly-otherwise", ], ) def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( @@ -251,7 +259,8 @@ def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( """A format transformers' own quantizer would load as stored, for PTQ to treat as plain weights, 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. The load refuses each before it reads anything.""" + shards by, or one that styles some of their weights otherwise. 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 == [] @@ -270,6 +279,29 @@ def test_an_ep_size_that_does_not_divide_the_experts_is_refused_before_anything_ assert ep_world_of_two == [] +def test_a_plan_whose_expert_containers_match_no_module_is_refused(ep_world_of_two, monkeypatch): + """A plan that styles the grouped-GEMM parameters a level below the experts names a container no + module is. Loaded anyway, every module would get the dense state, and the experts' amax would + reduce over ranks holding different experts.""" + + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + return _model_with_decoder_at("model.") + + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + config = SimpleNamespace( + base_model_ep_plan={"layers.*.mlp.experts.*.gate_up_proj": "grouped_gemm"} + ) + + with pytest.raises(ValueError, match="that no module of"): + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) + + @pytest.mark.parametrize("declared", [True, False], ids=["on-the-config", "a-class-default"]) def test_the_model_keeps_the_plan_its_checkpoint_declares( ep_world_of_two, gloo_world_of_one, monkeypatch, declared From 5546c2a79499b6cbf361dbf24c89aa39a1eabb61 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:34:22 +0000 Subject: [PATCH 05/23] Stand in a list for the unplaced keys in the no-module plan test record_unplaced_source_keys returns a sorted list, which a later check iterates; the new test's stand-in returned None. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_ep_loader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index fb4e221fdb6..8d93b9d2bdb 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -293,7 +293,7 @@ def from_pretrained(path, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + 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"} ) From 22cf2bedf8f889ab8711b37b41c09562ea541c92 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:02:05 +0000 Subject: [PATCH 06/23] Count the routed experts in the config that holds the expert plan The load read the expert count from the root config first, while the distribution plan's decision reads the config holding the plan: a VLM whose root config carries an expert-ish field of its own beside a nested MoE was planned by one count and refused by the other. The load now counts them where the decision does. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 5 ++-- tests/unit/torch/utils/test_ep_loader.py | 25 +++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ca08a78bdc4..4a2b5da80d5 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1817,9 +1817,10 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A "siblings." ) # The all-to-all transport finds an expert's rank by dividing its id by this rank's share of - # the experts, which holds only when every rank holds as many. A count the config names none of + # the experts, which holds only when every rank holds as many. Counted in the config that holds + # the plan, as the distribution plan's decision counts them. A count the config names none of # the ways this reads is checked once the experts load, from their weights. - num_experts = _num_experts_from_config(hf_config) + num_experts = _num_experts_from_config(ep_cfg) if num_experts % ep_degree: raise ValueError( f"{plan!r}: ep_size={ep_degree} must divide the {num_experts} routed experts, which the " diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 8d93b9d2bdb..ebf58eb1977 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -207,6 +207,31 @@ def resolve(path, **hub_options): assert resolved == ["org/model"] +def test_the_experts_are_counted_in_the_config_holding_the_plan(monkeypatch, gloo_world_of_one): + """As the distribution plan's decision counts them: a VLM whose nested config holds the MoE and + its plan has its experts counted there, not by an expert-ish field on its root config.""" + + class _CheckedError(Exception): + """Stops the load once its checks pass.""" + + def resolve(path, **hub_options): + raise _CheckedError + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 4) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", resolve) + text = SimpleNamespace(base_model_ep_plan=_EP_PLAN, num_experts=8) + + with pytest.raises(_CheckedError): + load_and_prepare_ep( + "org/model", + DistributionPlan(4, ep_size=4), + hf_config=SimpleNamespace(text_config=text, num_experts=3), + ) + + @pytest.mark.parametrize( ("config", "error", "match"), [ From 07105a21ac7bbf30e416e0e6f351f70e0d7946aa Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:04:18 +0000 Subject: [PATCH 07/23] Refuse an expert-parallel load in which some ranks have no GPU of their own Each rank builds the load's device meshes on its own GPU, and a live process group (gloo, say) does not give every rank one: a rank without it failed building them while the others waited in their subgroups' creation. Every rank now learns whether any rank has its GPU first, and a rank without one, where some other has its own, refuses the load on every rank, naming its device count and local rank. A run in which none has one still fails building the meshes on every rank alike. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 13 +++++++++++++ tests/unit/torch/utils/test_ep_loader.py | 9 +++++++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 4a2b5da80d5..b22daf4dc3f 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1826,6 +1826,19 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A f"{plan!r}: ep_size={ep_degree} must divide the {num_experts} routed experts, which the " "all-to-all transport shares out equally over the expert-parallel ranks." ) + # Each rank builds the meshes below on its own GPU. A live process group (gloo, say) does not + # give every rank one, and a rank without it would fail building them while the others wait in + # their subgroups' creation: refused on every rank first. A run in which no rank has one fails + # building them on every rank alike. + some_rank_has_one = dist.all_reduce("each rank's GPU", _has_own_gpu, combine=any) + with _refused_together("this rank's GPU", alone=False): + if some_rank_has_one and not _has_own_gpu(): + raise RuntimeError( + "The expert-parallel load builds its meshes on each rank's own GPU, and this rank " + f"has none, though other ranks have theirs: it sees {torch.cuda.device_count()} " + f"GPU(s) as local rank {dist.local_rank()}. Run one process per GPU, as " + "torchrun --nproc_per_node= does." + ) # transformers 5.6 on, which the floor above makes sure of. from torch.distributed.device_mesh import init_device_mesh from transformers.distributed import DistributedConfig diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index ebf58eb1977..4faafd895ef 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -706,6 +706,9 @@ def fail(*args, **kwargs): elif failure == "unseen-checkpoint" and fails_here: # Rank 0 makes the checkpoint local; a rank on a node without that storage cannot see it. patches.append(mock.patch.object(hf_checkpoint_utils.os.path, "isdir", return_value=False)) + elif failure == "no-gpu" and not fails_here: + # No rank has a GPU here: the others are given theirs. + patches.append(mock.patch.object(model_load_utils, "_has_own_gpu", return_value=True)) if failure in ("load", "record"): loaded = _model_with_decoder_at("model.") load = fail if failure == "load" and fails_here else (lambda path, **kwargs: loaded) @@ -743,6 +746,7 @@ def fail(*args, **kwargs): ("transformers", 1, "requires transformers >= 5.7, but 5.6.2"), ("quantized-config", 1, "cannot dequantize a 'gptq'"), ("unseen-checkpoint", 1, "which this rank cannot see"), + ("no-gpu", 1, "this rank has none, though other ranks have theirs"), ("load", 1, "the GPU ran out of memory"), ("record", 1, "a weight file this rank cannot read"), ], @@ -751,8 +755,9 @@ def test_a_load_failing_on_one_rank_fails_every_rank( cpu_dist_workers, tmp_path, failure, failing_rank, match ): """One rank fails alone: reading the config, on its transformers, on its arguments, seeing the - checkpoint rank 0 made local, or in the load. Every rank raises that failure, rather than leave the others - waiting for it in the meshes (stood in by a bounded barrier) or the warmup's collectives.""" + checkpoint rank 0 made local, without a GPU of its own where the others have theirs, or in the + load. Every rank raises that failure, rather than leave the others waiting for it in the meshes + (stood in by a bounded barrier) or the warmup's collectives.""" Qwen3MoeConfig(num_experts=4, num_hidden_layers=1, hidden_size=8).save_pretrained(tmp_path) cpu_dist_workers.run( _one_rank_fails_job, From 9a50ecc340c3ce0f3f2e5de77b2a2b59c6ac5495 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:06:56 +0000 Subject: [PATCH 08/23] Import the expert-parallel load's helpers at module scope The load imported its distributed and MoE helpers, the device-mesh builder and transformers' DistributedConfig inside the function, so its tests' stand-ins took effect. They are module-scope imports now, and the load reaches what the tests stand in for through its module (dist.is_initialized, dist.size, moe.require_transformers, device_mesh.init_device_mesh, hf_distributed.DistributedConfig). The loader tests' local transformers imports move to the top of the file, as every supported transformers has them. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 44 ++++++++----------- tests/unit/torch/utils/test_ep_loader.py | 13 ++---- 2 files changed, 21 insertions(+), 36 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index b22daf4dc3f..d5ba62bc8cc 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -30,10 +30,12 @@ import torch import torch.distributed.checkpoint as dcp import torch.nn as nn +import transformers.distributed as hf_distributed from huggingface_hub import snapshot_download if TYPE_CHECKING: from collections.abc import Iterable +from torch.distributed import device_mesh from torch.distributed.checkpoint import HuggingFaceStorageReader from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict from torch.distributed.fsdp import FSDPModule @@ -51,9 +53,10 @@ 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 import env, moe from modelopt.torch.utils.distributed import ( DistributionPlan, + ParallelState, barrier, broadcast_state_dict, is_initialized, @@ -62,11 +65,13 @@ from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 from modelopt.torch.utils.moe import ( + _EXPERTS_STYLES, _choose_ep_degree, _ep_config, _num_experts_from_config, _routed_expert_keys, _routed_expert_modules, + register_torch_a2a, ) from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( checkpoint_tensor_files, @@ -1719,19 +1724,6 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A Without ``hf_config`` every rank reads the checkpoint's config itself, from the Hub for a Hub ID, before rank 0 makes the checkpoint local: on many ranks, pass it (as ``hf_ptq`` does). """ - # Imported here, so looked up when the load runs: its tests stand in a process group and a - # transformers release by patching these two modules. - from modelopt.torch.utils.distributed import ParallelState, is_initialized, size - from modelopt.torch.utils.moe import ( - _EXPERTS_STYLES, - _ep_config, - _num_experts_from_config, - _routed_expert_keys, - _routed_expert_modules, - register_torch_a2a, - require_transformers, - ) - if not plan.is_expert_parallel: raise ValueError(f"{plan!r} shards no routed experts, which is what this load is for.") if plan.is_fsdp_dense: @@ -1742,13 +1734,15 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A # the load. A live, device-bound PG is required here so the mesh subgroups built below create their # NCCL comms eagerly (the synchronized warmup all-reduce per subgroup further down) instead of lazily # on first collective (which deadlocks at DP>1 because ranks reach the lazy init in different phases). - if not is_initialized(): + if not dist.is_initialized(): raise RuntimeError( "load_and_prepare_ep requires a live, device-bound process group; set it up first, e.g. " "with modelopt.torch.utils.distributed.distributed_run() under torchrun." ) - if plan.world_size != size(): - raise ValueError(f"{plan!r} is for {plan.world_size} ranks, but this run has {size()}.") + if plan.world_size != dist.size(): + raise ValueError( + 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 # 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. @@ -1756,7 +1750,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A # Hard floor: below 5.6 there is no ExpertsInterface / grouped-GEMM EP styles, and # enable_expert_parallel silently shards nothing (no error); 5.6 checks a model's experts # implementation against the ones it ships, so from_pretrained would refuse torch_a2a. - require_transformers( + moe.require_transformers( "5.7", "ModelOpt DP=EP expert-parallel PTQ (transformers MoE EP framework)" ) if hf_config is None: @@ -1839,10 +1833,6 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A f"GPU(s) as local rank {dist.local_rank()}. Run one process per GPU, as " "torchrun --nproc_per_node= does." ) - # transformers 5.6 on, which the floor above makes sure of. - from torch.distributed.device_mesh import init_device_mesh - from transformers.distributed import DistributedConfig - dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 # The weights load, and the checkpoint recorded for the export to read the unplaced tensors back # from, take the local snapshot: a Hub ID recorded there would leave the export nothing to read. @@ -1858,8 +1848,10 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A # dense_mesh = (dp,): dense/non-expert params are data-parallel-REPLICATED over the WHOLE world. # This subsumes the old a2a-vs-partial split: a2a is simply edp == 1, and dense amax # always reduces over the whole world (fixing the DP>1 case, where the old dp submesh missed the ep dim). - sparse_moe_mesh = init_device_mesh("cuda", (edp, ep_degree), mesh_dim_names=("edp", "ep")) - dense_mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) + sparse_moe_mesh = device_mesh.init_device_mesh( + "cuda", (edp, ep_degree), mesh_dim_names=("edp", "ep") + ) + dense_mesh = device_mesh.init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) ep_submesh = sparse_moe_mesh["ep"] # NO TENSOR PARALLELISM IS REQUESTED HERE, despite the name. transformers implements expert @@ -1878,7 +1870,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A # DistributedConfig(tp_plan=..., tp_size=...) and DROPS enable_expert_parallel, silently losing # EP. So each argument goes on DistributedConfig where it has a field for it, and tp_plan # otherwise to from_pretrained: the dataclass fields tell them apart on any release. - _dc_fields = {f.name for f in dataclasses.fields(DistributedConfig)} + _dc_fields = {f.name for f in dataclasses.fields(hf_distributed.DistributedConfig)} _dc_has_tp_plan = "tp_plan" in _dc_fields dc_kwargs: dict[str, Any] = {"enable_expert_parallel": True} if _dc_has_tp_plan: @@ -1886,7 +1878,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A if _transformers_mesh_builder() is not None and "tp_size" in _dc_fields: # This release sizes its own mesh from tp_size; _experts_sharded_over swaps in the EP mesh. dc_kwargs["tp_size"] = ep_degree - distributed_config = DistributedConfig(**dc_kwargs) + distributed_config = hf_distributed.DistributedConfig(**dc_kwargs) from_pretrained_kwargs: dict[str, Any] = { "dtype": dtype, "device_mesh": ep_submesh, diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 4faafd895ef..dad6123d0e7 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -37,7 +37,9 @@ pytest.importorskip("transformers") import transformers -from transformers import FineGrainedFP8Config, Qwen3MoeConfig +import transformers.distributed as hf_distributed +from transformers import FineGrainedFP8Config, PreTrainedModel, Qwen3MoeConfig, Qwen3MoeForCausalLM +from transformers.distributed import DistributedConfig import modelopt.torch.utils.distributed as dist from modelopt.torch.utils import moe @@ -378,9 +380,6 @@ def test_a_loaded_model_keeps_its_sharding_styles_out_of_its_saved_config( from_pretrained does not set; while its config goes back to the checkpoint's plan; and the experts implementation its forwards dispatch on, torch_a2a, stays out of the config it serializes, which an export writes.""" - from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM - from transformers.distributed import DistributedConfig - if not Qwen3MoeConfig.base_model_ep_plan: pytest.skip("this transformers ships no expert-parallel plan for Qwen3-MoE") @@ -574,13 +573,9 @@ def test_from_pretrained_builds_its_mesh_from_the_ep_mesh_it_is_given(): """Newer transformers ignores ``from_pretrained``'s ``device_mesh`` and builds a world-sized one, which shards the experts over every rank when EP is below the world size. During the load it gets the EP mesh instead, and its own builder is back afterwards, even when the load raises.""" - pytest.importorskip("transformers") mixin = _transformers_mesh_builder() if mixin is None: pytest.skip("this transformers honors from_pretrained's device_mesh") - from transformers import PreTrainedModel - from transformers.distributed import DistributedConfig - builder = mixin.initialize_distributed_mesh ep_mesh = SimpleNamespace(device_type="cpu") config = DistributedConfig(enable_expert_parallel=True, tp_plan="auto", tp_size=2) @@ -642,8 +637,6 @@ def from_pretrained(path, **kwargs): loaded.update(kwargs) raise _LoadedError - import transformers.distributed as hf_distributed - distributed_config = dataclasses.make_dataclass( "DistributedConfig", [(name, object, dataclasses.field(default=None)) for name in fields] ) From baf6deacc2fa65e4c8cd6adccb893deee8370b02 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:45 +0000 Subject: [PATCH 09/23] Say what the load's failure guarantee rests on A rank failing inside from_pretrained fails every rank only because the other ranks finish the load and meet it in the refusal's gather, which holds because transformers' expert-parallel load runs no collective: before 5.16 it slices each rank's experts into plain tensors, from 5.16 it builds DTensors locally, and the one mesh it would build (5.18) is replaced by the load's own. The docstring and the guard's comment now say so. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d5ba62bc8cc..06097342086 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1721,6 +1721,9 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A (:func:`~modelopt.torch.utils.moe.expert_parallel_mesh`), once the caller has recorded the plan. A refusal on any rank -- of its transformers, the config, the checkpoint, or in the load -- raises on every rank, so no rank is left waiting in a collective for one that already failed. + In the load that rests on ``from_pretrained`` running no collective, as transformers' + expert-parallel load does not: each rank reads and slices its own experts, so the others finish + it and meet a rank that failed in it in the refusal's gather. Without ``hf_config`` every rank reads the checkpoint's config itself, from the Hub for a Hub ID, before rank 0 makes the checkpoint local: on many ranks, pass it (as ``hf_ptq`` does). """ @@ -1928,7 +1931,9 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A from_pretrained_kwargs["config"] = hf_config print_rank_0(f"[ep] ep_plan={swapped_plan} experts_implementation=torch_a2a") # A rank failing in the load (out of memory, a file it cannot read) fails every rank here, - # before the warmup's collectives. + # before the warmup's collectives. from_pretrained runs no collective (transformers shards each + # rank's experts locally, and builds no mesh but the one handed to it), so the other ranks finish + # it and meet the failed one in this guard's gather. with _refused_together("the expert-parallel load", alone=False): # The class that builds the whole checkpoint -- with its lm_head where an auto class builds one, # so the head is calibrated and exported and the generation preview works -- chosen as the From b795ccc7e6a3b59010ca77cabc0fb67fd05b515f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:45:47 +0000 Subject: [PATCH 10/23] Refuse an MXFP4 checkpoint in the expert-parallel load transformers dequantizes MXFP4 experts into GPT-OSS's transposed (num_experts, in, out) layout whatever the model's own: a standard-layout MoE loaded its experts transposed, without an error, and failed on its first forward (or, with square shapes, computed garbage). GPT-OSS, the model whose layout it is, the load does not quantize either. So the load now refuses an MXFP4 checkpoint as it refuses other quantized formats, pointing to an unquantized copy of the model. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 21 ++++++------------- tests/unit/torch/utils/test_ep_loader.py | 6 ++++-- 2 files changed, 10 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 06097342086..0b871b606e9 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1620,8 +1620,8 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st # The checkpoint formats load_and_prepare_ep dequantizes as it loads, beside unquantized ones (None). -# mxfp4 is for experts stored in the standard layout. -_EP_SOURCE_FORMATS = frozenset({None, "mxfp4"}) +# Not MXFP4: transformers dequantizes it into GPT-OSS's transposed expert layout, whatever the model's. +_EP_SOURCE_FORMATS: frozenset[str | None] = frozenset({None}) def _ckpt_quant_method(cfg): @@ -1765,10 +1765,12 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A if quant_method not in _EP_SOURCE_FORMATS: # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights. dequantized = ", ".join(sorted(f for f in _EP_SOURCE_FORMATS if f)) + loads = "unquantized checkpoints" + ( + f" and dequantizes {dequantized}" if dequantized else "" + ) raise NotImplementedError( f"The expert-parallel load cannot dequantize a {quant_method!r} checkpoint: it loads " - f"unquantized checkpoints and dequantizes {dequantized}. Start from an unquantized copy " - "of the model." + f"{loads}. Start from an unquantized copy of the model." ) # The swapped plan is this load's own: the caller's config keeps the plan its checkpoint declares. hf_config = copy.deepcopy(hf_config) @@ -1892,17 +1894,6 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A from_pretrained_kwargs["tp_plan"] = "auto" if attn_implementation is not None: from_pretrained_kwargs["attn_implementation"] = attn_implementation - if quant_method == "mxfp4": - # Native mxfp4 checkpoints (fused `gate_up_proj_blocks`/`_scales` + biases) must be - # DEQUANTIZED to bf16 on load so ModelOpt can insert nvfp4/fp8 quantizers on the fused bf16 experts - # (the a2a grouped_gemm expects bf16). transformers dequantizes blocks->bf16 during from_pretrained. - # NB: the single-process loader needed device_map="sequential" to dodge a CUDA illegal-access in the - # mxfp4 dequant under multi-GPU device_map="auto"; the EP path uses device_mesh (no device_map), so - # watch for an EP-shard-awareness issue in the dequant convert-op. - from transformers import Mxfp4Config - - from_pretrained_kwargs["quantization_config"] = Mxfp4Config(dequantize=True) - print_rank_0("[ep] native mxfp4 checkpoint -> Mxfp4Config(dequantize=True) -> bf16") # DP=EP all-to-all: swap the ep_plan to our torch_a2a styles (torch_a2a_router emits GLOBAL ids # unmasked; torch_a2a_experts injects the EP process_group + skips the all-reduce) and select the # `torch_a2a` experts forward (NCCL all-to-all dispatch/combine around the quantizable diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index dad6123d0e7..a07ceea5353 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -243,7 +243,7 @@ def resolve(path, **hub_options): NotImplementedError, f"cannot dequantize a '{method}'", ) - for method in ("gptq", "modelopt", "compressed-tensors") + for method in ("gptq", "modelopt", "compressed-tensors", "mxfp4") ), (SimpleNamespace(), ValueError, "declare base_model_ep_plan"), ( @@ -273,6 +273,7 @@ def resolve(path, **hub_options): "gptq", "modelopt", "compressed-tensors", + "mxfp4", "no-ep-plan", "no-routed-experts", "experts-container-only", @@ -284,7 +285,8 @@ def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( ep_world_of_two, config, error, match ): """A format transformers' own quantizer would load as stored, for PTQ to treat as plain weights, - or a plan that leaves the experts whole on every rank: one naming no routed experts, or none of + 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.""" From 1f30540b39b551ba0e508814d3b9f94159332746 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:49:04 +0000 Subject: [PATCH 11/23] Name the transformers releases the load's failure guarantee was checked on The docstring stated its premise, that from_pretrained runs no collective, for transformers in general; it was checked on 5.7 through 5.18, and the load accepts later releases too. It now says so, and how each release range shards the experts without one. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 0b871b606e9..0deab9096b7 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1722,8 +1722,10 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A A refusal on any rank -- of its transformers, the config, the checkpoint, or in the load -- raises on every rank, so no rank is left waiting in a collective for one that already failed. In the load that rests on ``from_pretrained`` running no collective, as transformers' - expert-parallel load does not: each rank reads and slices its own experts, so the others finish - it and meet a rank that failed in it in the refusal's gather. + expert-parallel load does not on 5.7 through 5.18: before 5.16 it slices each rank's experts + into plain tensors, and from 5.16 it builds DTensors locally (``run_check=False``, + ``src_data_rank=None``). So the others finish it and meet a rank that failed in it in the + refusal's gather. Without ``hf_config`` every rank reads the checkpoint's config itself, from the Hub for a Hub ID, before rank 0 makes the checkpoint local: on many ranks, pass it (as ``hf_ptq`` does). """ From cb5495d788744ddc21cf74d9feef778a47744777 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:50:14 +0000 Subject: [PATCH 12/23] Decide the expert-parallel load's GPU check from one gather Every rank already receives every rank's answer from the gather, so the mixed case is decided on every rank alike without a second, guarded round, and the refusal names the ranks without a GPU of their own. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 26 +++++++++++-------- tests/unit/torch/utils/test_ep_loader.py | 7 ++--- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 0deab9096b7..f0d598bb578 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1829,17 +1829,21 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A ) # Each rank builds the meshes below on its own GPU. A live process group (gloo, say) does not # give every rank one, and a rank without it would fail building them while the others wait in - # their subgroups' creation: refused on every rank first. A run in which no rank has one fails - # building them on every rank alike. - some_rank_has_one = dist.all_reduce("each rank's GPU", _has_own_gpu, combine=any) - with _refused_together("this rank's GPU", alone=False): - if some_rank_has_one and not _has_own_gpu(): - raise RuntimeError( - "The expert-parallel load builds its meshes on each rank's own GPU, and this rank " - f"has none, though other ranks have theirs: it sees {torch.cuda.device_count()} " - f"GPU(s) as local rank {dist.local_rank()}. Run one process per GPU, as " - "torchrun --nproc_per_node= does." - ) + # their subgroups' creation: every rank learns which ranks have theirs, from one gather, and a + # run in which some do and some do not is refused on every rank alike. A run in which none does + # fails building the meshes on every rank alike. + own_gpu = dist.all_reduce( + "each rank's GPU", + lambda: [_has_own_gpu()], + combine=lambda ranks: [has for flags in ranks for has in flags], + ) + if any(own_gpu) and not all(own_gpu): + without = [rank for rank, has in enumerate(own_gpu) if not has] + raise RuntimeError( + "The expert-parallel load builds its meshes on each rank's own GPU, and rank(s) " + f"{without} have none while the others have theirs. Run one process per GPU, as " + "torchrun --nproc_per_node= does." + ) dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 # The weights load, and the checkpoint recorded for the export to read the unplaced tensors back # from, take the local snapshot: a Hub ID recorded there would leave the export nothing to read. diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index a07ceea5353..e2422dbd788 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -726,9 +726,10 @@ def fail(*args, **kwargs): load_and_prepare_ep(path, DistributionPlan(2, ep_size=2), hf_config=hf_config) except Exception as error: refused = f"{type(error).__name__}: {error}" - # The one rank's failure on every rank, before any rank went on alone into a collective. + # The one rank's failure on every rank, before any rank went on alone into a collective. The GPU + # check decides on every rank at once, from one gather, so no rank relays another's failure. assert refused is not None and match in refused, (rank, refused) - assert fails_here or f"rank {failing_rank} failed" in refused, refused + assert fails_here or failure == "no-gpu" or f"rank {failing_rank} failed" in refused, refused @pytest.mark.skipif( @@ -741,7 +742,7 @@ def fail(*args, **kwargs): ("transformers", 1, "requires transformers >= 5.7, but 5.6.2"), ("quantized-config", 1, "cannot dequantize a 'gptq'"), ("unseen-checkpoint", 1, "which this rank cannot see"), - ("no-gpu", 1, "this rank has none, though other ranks have theirs"), + ("no-gpu", 1, "rank(s) [1] have none while the others have theirs"), ("load", 1, "the GPU ran out of memory"), ("record", 1, "a weight file this rank cannot read"), ], From 02552e6ca221ae85d1987910e390f196f5129995 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:51:22 +0000 Subject: [PATCH 13/23] Keep the no-TP-linear warning off a model loaded expert-parallel The TP-linear converter warns when it finds a transformers mesh on the model but no TP-sharded Linear under it. An expert-parallel load keeps its experts' mesh there (calibration reads it back), and shards no Linear by design, so every expert-parallel calibration warned. The warning now skips a model whose recorded distribution plan is expert- parallel. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 6 ++++-- .../quantization/plugins/test_huggingface.py | 21 +++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 2223608c283..83cd2ba117c 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -41,7 +41,7 @@ from modelopt.torch.kernels.quantization.gemm import IS_AVAILABLE as IS_TRITON_AVAILABLE from modelopt.torch.opt.dynamic import DynamicModule from modelopt.torch.utils import warn_rank_0 -from modelopt.torch.utils.distributed import ParallelState +from modelopt.torch.utils.distributed import ParallelState, recorded_distribution_plan from ..algorithms import AutoQuantizeGradientSearcher from ..conversion import register @@ -596,7 +596,9 @@ def convert_hf_parallel_linears_on_the_fly(model): elif HFRowParallelLinear.is_compatible(module, tp_mesh): HFRowParallelLinear.convert(module) converted = True - if tp_mesh is not None and not converted: + # An expert-parallel load keeps its experts' mesh there too, and shards no Linear by design. + plan = recorded_distribution_plan(model) + if tp_mesh is not None and not converted and not (plan is not None and plan.is_expert_parallel): warnings.warn( "Found a transformers tensor-parallel mesh but no TP-sharded linear layer; they will be" " quantized as unsharded layers." diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index 078f5dacd8d..136de332d7b 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/tests/unit/torch/quantization/plugins/test_huggingface.py @@ -36,10 +36,12 @@ from modelopt.torch.quantization.plugins.huggingface import ( _QuantHFParallelLinear, _TransposedExpertsCalibMixin, + convert_hf_parallel_linears_on_the_fly, get_homogeneous_hf_decoder_layers, is_homogeneous_hf_model, ) from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan pytest.importorskip("transformers") @@ -101,6 +103,25 @@ class ParallelLinear: assert linear.weight is original_weight +@pytest.mark.parametrize("expert_parallel", [False, True], ids=["no-plan", "expert-parallel"]) +def test_a_mesh_without_tp_sharded_linears_warns_unless_the_model_is_expert_parallel( + expert_parallel, +): + """A transformers mesh with no TP-sharded Linear under it is worth a warning, except on a model + loaded expert-parallel: its mesh is the experts', and no Linear is meant to be sharded.""" + model = nn.Sequential(nn.Linear(4, 4)) + model._device_mesh = object() + if expert_parallel: + set_distribution_plan(model, DistributionPlan(2, ep_size=2)) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + convert_hf_parallel_linears_on_the_fly(model) + + warned = any("no TP-sharded linear layer" in str(w.message) for w in caught) + assert warned is not expert_parallel + + def test_convert_conv1d(): set_seed() assert transformers.pytorch_utils.Conv1D in QuantModuleRegistry From 45eaaf6ab5a1fdd44785c9d511bf0ba8017e4de1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:56:05 +0000 Subject: [PATCH 14/23] Derive the load's kept plan styles from its one style swap The plan checks refused an expert weight's style unless it was grouped_gemm, ep_router or an experts style, and the swapped plan kept an entry if its style was a swap key or grouped_gemm: two lists of the same styles. The swap is now defined once, before the checks, and both read the styles it keeps from it, so a style added to or dropped from the swap changes what the checks accept and what the load keeps alike. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 27 ++++++++++--------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index f0d598bb578..b0d137996d2 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1749,6 +1749,16 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A 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 + # (torch_a2a_router emits GLOBAL ids unmasked; torch_a2a_experts injects the EP process_group + + # skips the all-reduce). nemotron_h's logits-only gate has no ep_router entry, so its + # nemotron_experts_ep also maps to torch_a2a_experts. 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} # 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): @@ -1807,7 +1817,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A dropped = sorted( name for name, style in ep_plan.items() - if style not in ("grouped_gemm", "ep_router", *_EXPERTS_STYLES) + if style not in kept_styles and name not in containers and any(name.startswith(key + ".") for key in containers) ) @@ -1900,23 +1910,14 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A from_pretrained_kwargs["tp_plan"] = "auto" if attn_implementation is not None: from_pretrained_kwargs["attn_implementation"] = attn_implementation - # DP=EP all-to-all: swap the ep_plan to our torch_a2a styles (torch_a2a_router emits GLOBAL ids - # unmasked; torch_a2a_experts injects the EP process_group + skips the all-reduce) and select the + # DP=EP all-to-all: swap the ep_plan to our torch_a2a styles (``swap`` above) and select the # `torch_a2a` experts forward (NCCL all-to-all dispatch/combine around the quantizable # _grouped_linear). Tokens are routed to the expert-owning rank by all-to-all instead of every rank - # computing all tokens + all-reduce. Requires per-rank (data-parallel) calib data. nemotron_h's - # logits-only gate has no ep_router entry, so its nemotron_experts_ep also maps to torch_a2a_experts. + # computing all tokens + all-reduce. Requires per-rank (data-parallel) calib data. register_torch_a2a() # idempotent; register the experts forward + the two styles, right before use - swap = { - "ep_router": "torch_a2a_router", - "moe_tp_experts": "torch_a2a_experts", - "nemotron_experts_ep": "torch_a2a_experts", - } # Keep expert-related entries (grouped_gemm expert params + the router/experts styles); DROP the TP # styles (a2a builds no TP group for them -> would shard attention across the wrong mesh). - swapped_plan = { - k: swap.get(v, v) for k, v in ep_plan.items() if v in swap or v == "grouped_gemm" - } + swapped_plan = {k: swap.get(v, v) for k, v in ep_plan.items() if v in kept_styles} # Tag every routed-experts CONTAINER torch_a2a_experts (injects the EP process_group + marks the # routed expert), found as the distribution plan's decision finds them: some plans tag only the # expert PARAMS grouped_gemm, and name the container through their parents. From 0d9bed6c5313f2225ce5dd73967666ae091ed88c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:56:55 +0000 Subject: [PATCH 15/23] Build the EP loader tests' stand-in MoE from the shared test model The tests kept their own copy of tests/_test_utils/torch/moe_models.py's model_with_decoder_at, differing only in the experts' quantizer child. They now use the shared model, and the parallel-state test, the one that needs a child under the experts, adds it. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_ep_loader.py | 68 ++++-------------------- 1 file changed, 9 insertions(+), 59 deletions(-) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index e2422dbd788..3b9dc851c3c 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -31,6 +31,7 @@ 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 packaging.version import Version from torch import nn @@ -54,59 +55,6 @@ load_and_prepare_ep, ) -# ----------------------------------------------------------------------------------- tiny MoE - - -class _Experts(nn.Module): - """A layer's routed experts as one grouped-GEMM container (every expert in fused tensors).""" - - def __init__(self): - super().__init__() - self.gate_up_proj = nn.Parameter(torch.zeros(4, 4, 2)) - self.down_proj = nn.Parameter(torch.zeros(4, 2, 2)) - self.weight_quantizer = nn.Identity() # quantizer insertion adds children like this - - def forward(self, hidden_states, top_k_index, top_k_weights, process_group=None): - return hidden_states - - -class _SharedExperts(nn.Module): - """The always-on shared expert MLP: dense, replicated on every rank.""" - - def __init__(self): - super().__init__() - self.gate_up_proj = nn.Linear(2, 4, bias=False) - - def forward(self, hidden_states): - return hidden_states - - -class _MoeLayer(nn.Module): - 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 = _Experts() - self.mlp.shared_experts = _SharedExperts() - - -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 - - # ------------------------------------------------------------------------------- the plan _EP_PLAN = { @@ -316,7 +264,7 @@ def test_a_plan_whose_expert_containers_match_no_module_is_refused(ep_world_of_t class _Loader: @staticmethod def from_pretrained(path, **kwargs): - return _model_with_decoder_at("model.") + return model_with_decoder_at("model.") monkeypatch.setattr( "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() @@ -342,7 +290,7 @@ def test_the_model_keeps_the_plan_its_checkpoint_declares( class _PastError(Exception): """Raised by the load's next check.""" - model = _model_with_decoder_at("model.") + model = model_with_decoder_at("model.") class _Loader: @staticmethod @@ -476,7 +424,9 @@ def test_expert_containers_and_their_children_get_the_expert_state(): """Everything else, the root and the containers' own parents included, gets the dense state; so does a sibling whose name merely starts with a container's (Qwen2-MoE keeps ``shared_expert_gate`` beside ``shared_expert``).""" - model = _model_with_decoder_at("model.") + model = model_with_decoder_at("model.") + for layer in model.model.layers: + layer.mlp.experts.weight_quantizer = nn.Identity() # a child, as quantizer insertion adds model.model.layers[0].mlp.experts_gate = nn.Linear(2, 1) experts = {f"model.layers.{i}.mlp.experts": model.model.layers[i].mlp.experts for i in (0, 1)} dense, expert = object(), object() @@ -504,7 +454,7 @@ def test_expert_containers_and_their_children_get_the_expert_state(): def test_the_expert_parallel_mesh_comes_with_an_expert_parallel_plan_only(plan, is_expert_parallel): """An expert-parallel plan: the mesh transformers kept on the model. Any other model has none, even with a device mesh on it, as a load through from_pretrained's own TP can leave.""" - model = _model_with_decoder_at("model.") + model = model_with_decoder_at("model.") if plan is not None: set_distribution_plan(model, plan) model._device_mesh = mesh = object() @@ -515,7 +465,7 @@ def test_the_expert_parallel_mesh_comes_with_an_expert_parallel_plan_only(plan, def test_an_expert_parallel_model_without_its_mesh_is_refused(): """A model recorded expert-parallel that load_and_prepare_ep did not load has no mesh for calibration to re-install the all-to-all forwards over; it is refused by name.""" - model = _model_with_decoder_at("model.") + model = model_with_decoder_at("model.") set_distribution_plan(model, DistributionPlan(2, ep_size=2)) with pytest.raises(RuntimeError, match="carries no device mesh"): @@ -705,7 +655,7 @@ def fail(*args, **kwargs): # No rank has a GPU here: the others are given theirs. patches.append(mock.patch.object(model_load_utils, "_has_own_gpu", return_value=True)) if failure in ("load", "record"): - loaded = _model_with_decoder_at("model.") + loaded = model_with_decoder_at("model.") load = fail if failure == "load" and fails_here else (lambda path, **kwargs: loaded) loader = SimpleNamespace(__name__="Loader", from_pretrained=load) record = fail if failure == "record" and fails_here else None From 636d08f694d938c0e5724be74ec1c28c4c0e3af4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:58:48 +0000 Subject: [PATCH 16/23] Check that one visible GPU per process passes the EP load's GPU check The load's GPU check read a rank's GPU as the one its local rank numbers, so a launch that shows each process one GPU, each bound to cuda:0, refused every rank but the first. The predicate it shares, fixed in the load routing below, now takes the device the rank is bound to. Two ranks, each bound to its one visible GPU under local ranks 0 and 1, now both pass the check and go on to build the meshes; the test fails with the old predicate. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_ep_loader.py | 36 ++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 3b9dc851c3c..b537d133034 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -712,3 +712,39 @@ def test_a_load_failing_on_one_rank_fails_every_rank( failing_rank=failing_rank, match=match, ) + + +def _one_gpu_per_process_job(rank, size): + hf_config = SimpleNamespace(base_model_ep_plan=dict(_EP_PLAN), num_experts=4) + reached = None + with bounded_collectives(), ExitStack() as stack: + # Each process sees one GPU and is bound to it, as cuda:0; the local ranks run 0 and 1. + for name, value in [ + ("is_available", True), + ("is_initialized", True), + ("current_device", 0), + ("device_count", 1), + ]: + stack.enter_context(mock.patch(f"torch.cuda.{name}", return_value=value)) + stack.enter_context( + mock.patch.object( + model_load_utils, "ensure_local_checkpoint", side_effect=lambda p, **k: (None, p) + ) + ) + stack.enter_context( + mock.patch("torch.distributed.device_mesh.init_device_mesh", _meshes_in_step) + ) + try: # not pytest.raises: its failure would end the pooled worker instead of the job + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=hf_config) + except Exception as error: + reached = error + assert isinstance(reached, _PastTheMeshesError), (rank, reached) + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("5.7"), reason="the EP load needs transformers 5.7" +) +def test_one_visible_gpu_per_process_gives_every_rank_its_gpu(cpu_dist_workers): + """A launch that shows each process one GPU binds every rank to cuda:0, whatever its local rank: + every rank has its GPU, and every rank goes on to build the meshes.""" + cpu_dist_workers.run(_one_gpu_per_process_job) From 2c45e20dd40722929b3563847ecfd6cc36dcbdc5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:26:04 +0000 Subject: [PATCH 17/23] Refuse an expert-parallel load on a rank not yet bound to its GPU The GPU check reads a rank's GPU as the device it is bound to once CUDA is initialized, and before that as the one its local rank numbers. A launch that shows each process one GPU but has not bound them yet (a gloo group, say) so read every rank but the first as having none, and was refused with "Run one process per GPU", which it does. Building the meshes would have failed there anyway: DeviceMesh binds such a rank by its local rank, a GPU rank 1 does not see. The check now also gathers which ranks see CUDA but are bound to no GPU yet, and refuses such a run on every rank, telling the caller to bind each rank to its GPU first, as distributed_run() does. The one-GPU-per-process test covers that launch bound and not yet bound; the not-yet-bound case fails without the check. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 28 ++++++++++++------- tests/unit/torch/utils/test_ep_loader.py | 20 ++++++++----- 2 files changed, 31 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index b0d137996d2..f53f4695597 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1837,18 +1837,26 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A f"{plan!r}: ep_size={ep_degree} must divide the {num_experts} routed experts, which the " "all-to-all transport shares out equally over the expert-parallel ranks." ) - # Each rank builds the meshes below on its own GPU. A live process group (gloo, say) does not - # give every rank one, and a rank without it would fail building them while the others wait in - # their subgroups' creation: every rank learns which ranks have theirs, from one gather, and a - # run in which some do and some do not is refused on every rank alike. A run in which none does - # fails building the meshes on every rank alike. - own_gpu = dist.all_reduce( + # Each rank builds the meshes below on the GPU it is bound to. DeviceMesh binds a rank that sees + # CUDA but is bound to none yet (a gloo group leaves it so) by its local rank, a GPU that a launch + # showing each process one does not have; and a rank without a GPU fails building the meshes. + # Either fails alone while the others wait in their subgroups' creation, so every rank learns + # from one gather which ranks are bound and which have a GPU, and all refuse alike. A run in + # which none has a GPU fails building the meshes on every rank alike. + gpus = dist.all_reduce( "each rank's GPU", - lambda: [_has_own_gpu()], - combine=lambda ranks: [has for flags in ranks for has in flags], + lambda: [(_has_own_gpu(), torch.cuda.is_available() and not torch.cuda.is_initialized())], + combine=lambda ranks: [gpu for flags in ranks for gpu in flags], ) - if any(own_gpu) and not all(own_gpu): - without = [rank for rank, has in enumerate(own_gpu) if not has] + unbound = [rank for rank, (_, waits) in enumerate(gpus) if waits] + if unbound: + raise RuntimeError( + "The expert-parallel load builds its meshes on the GPU each rank is bound to, and rank(s) " + f"{unbound} are bound to none yet. Bind each rank to its GPU before the load " + "(torch.cuda.set_device), as distributed_run() does." + ) + without = [rank for rank, (has, _) in enumerate(gpus) if not has] + if without and len(without) < len(gpus): raise RuntimeError( "The expert-parallel load builds its meshes on each rank's own GPU, and rank(s) " f"{without} have none while the others have theirs. Run one process per GPU, as " diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index b537d133034..d9efd616afc 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -714,14 +714,14 @@ def test_a_load_failing_on_one_rank_fails_every_rank( ) -def _one_gpu_per_process_job(rank, size): +def _one_gpu_per_process_job(rank, size, bound): hf_config = SimpleNamespace(base_model_ep_plan=dict(_EP_PLAN), num_experts=4) reached = None with bounded_collectives(), ExitStack() as stack: - # Each process sees one GPU and is bound to it, as cuda:0; the local ranks run 0 and 1. + # Each process sees one GPU, cuda:0, bound to it or not yet; the local ranks run 0 and 1. for name, value in [ ("is_available", True), - ("is_initialized", True), + ("is_initialized", bound), ("current_device", 0), ("device_count", 1), ]: @@ -738,13 +738,19 @@ def _one_gpu_per_process_job(rank, size): load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=hf_config) except Exception as error: reached = error - assert isinstance(reached, _PastTheMeshesError), (rank, reached) + if bound: + assert isinstance(reached, _PastTheMeshesError), (rank, reached) + else: + assert "rank(s) [0, 1] are bound to none yet" in str(reached), (rank, reached) @pytest.mark.skipif( Version(transformers.__version__) < Version("5.7"), reason="the EP load needs transformers 5.7" ) -def test_one_visible_gpu_per_process_gives_every_rank_its_gpu(cpu_dist_workers): +@pytest.mark.parametrize("bound", [True, False], ids=["bound", "not-yet-bound"]) +def test_one_visible_gpu_per_process_needs_each_rank_bound_to_it(cpu_dist_workers, bound): """A launch that shows each process one GPU binds every rank to cuda:0, whatever its local rank: - every rank has its GPU, and every rank goes on to build the meshes.""" - cpu_dist_workers.run(_one_gpu_per_process_job) + every rank has its GPU, and every rank goes on to build the meshes. Not yet bound, DeviceMesh + would bind each by its local rank, a GPU rank 1 does not see: every rank is refused, told to + bind first.""" + cpu_dist_workers.run(_one_gpu_per_process_job, bound=bound) From 94750534528634211d460a78ed2a3f998c3f4398 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:40:51 +0000 Subject: [PATCH 18/23] Pin the CUDA state the EP load's GPU check reads in its two-rank tests The one-rank-failure rows left the GPU check to read the host's own CUDA, and the gloo workers bind no GPU. On a host with a visible GPU every rank so read as not yet bound, and the load refused before the step four rows exercise (no-gpu, unseen-checkpoint, load, record); CPU-only CI hid it. A helper now pins what the check reads of torch.cuda (availability, initialization, current device, device count): no GPU for the rows the GPU plays no part in, and in the no-gpu row a bound GPU on every rank but the failing one, which the real predicate then reads instead of a patched _has_own_gpu. The one-GPU-per-process test uses the same helper. On a host faked to show GPUs that no rank has bound, the four rows failed before and pass now. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_ep_loader.py | 27 +++++++++++++++--------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index d9efd616afc..969c97f6e83 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -628,6 +628,19 @@ def _meshes_in_step(*args, **kwargs): raise _PastTheMeshesError +def _cuda(gpu, bound=True): + """Patches pinning the CUDA state the GPU check reads: one visible GPU, cuda:0, bound to it or + not yet; or, without ``gpu``, none. The check then reads the case, not the host the test runs + on (the gloo workers bind no GPU, so a host's visible one would read as unbound).""" + state = { + "is_available": gpu, + "is_initialized": gpu and bound, + "current_device": 0, + "device_count": int(gpu), + } + return [mock.patch(f"torch.cuda.{name}", return_value=value) for name, value in state.items()] + + def _one_rank_fails_job(rank, size, ckpt, failure, failing_rank, match): path, hf_config, patches = ckpt, None, [] if failure != "unreadable-config": @@ -638,6 +651,8 @@ def _one_rank_fails_job(rank, size, ckpt, failure, failing_rank, match): ) patches.append(local) fails_here = rank == failing_rank + # No rank has a GPU, but in the no-gpu row every rank other than the failing one has its own. + patches += _cuda(failure == "no-gpu" and not fails_here) def fail(*args, **kwargs): raise (MemoryError if failure == "load" else OSError)(match) @@ -651,9 +666,6 @@ def fail(*args, **kwargs): elif failure == "unseen-checkpoint" and fails_here: # Rank 0 makes the checkpoint local; a rank on a node without that storage cannot see it. patches.append(mock.patch.object(hf_checkpoint_utils.os.path, "isdir", return_value=False)) - elif failure == "no-gpu" and not fails_here: - # No rank has a GPU here: the others are given theirs. - patches.append(mock.patch.object(model_load_utils, "_has_own_gpu", return_value=True)) if failure in ("load", "record"): loaded = model_with_decoder_at("model.") load = fail if failure == "load" and fails_here else (lambda path, **kwargs: loaded) @@ -719,13 +731,8 @@ def _one_gpu_per_process_job(rank, size, bound): reached = None with bounded_collectives(), ExitStack() as stack: # Each process sees one GPU, cuda:0, bound to it or not yet; the local ranks run 0 and 1. - for name, value in [ - ("is_available", True), - ("is_initialized", bound), - ("current_device", 0), - ("device_count", 1), - ]: - stack.enter_context(mock.patch(f"torch.cuda.{name}", return_value=value)) + for patch in _cuda(True, bound): + stack.enter_context(patch) stack.enter_context( mock.patch.object( model_load_utils, "ensure_local_checkpoint", side_effect=lambda p, **k: (None, p) From 974c72cf3744dc9a2dd2568869b7398152563c98 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:40:54 +0000 Subject: [PATCH 19/23] Say the plan checks refuse any other expert-weight style, not drop it The comment over the grouped_gemm check said packed_* entries are dropped, but the check below it refuses a plan that styles a weight under a routed-expert container by any style other than grouped_gemm, so such an entry never reaches the swapped plan. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index f53f4695597..bb7dc620337 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1796,7 +1796,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A "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 (packed_* entries are dropped). + # 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) From 82c2c918250226555fe10da929f83cf45d0a2b4e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:07:05 +0000 Subject: [PATCH 20/23] Refuse ranks the meshes would put on no GPU or on one GPU together Importing modelopt initializes CUDA wherever torchao and triton are installed (torchao's mx_formats asks has_triton(), which reads the current device), so in a launch that binds no rank every rank is on cuda:0 by the time the load runs. DeviceMesh binds a rank only while CUDA is uninitialized, so it leaves them there: on a multi-GPU node they share GPU 0, and NCCL fails at the first collective on the ranks' shared device. The GPU check now gathers the UUID of the GPU each initialized rank is on and refuses, on every rank, a run in which two ranks are on the same one, telling the caller to bind each rank to its own. Claude's review: the refusal of a rank not yet bound was broader than the launch it is for. Where CUDA is still uninitialized, DeviceMesh binds such a rank by its local rank, which is right when that names a GPU the rank sees. The refusal is now kept for the rank whose local rank names none, as a launch showing each process one GPU gives every rank but the first. The GPU-layout test covers four launches: one GPU each, bound or not yet; every GPU, not yet bound; and every rank on cuda:0. The narrowing fails the not-yet-bound one-GPU and every-GPU cases without it, and the shared-GPU case fails without the new refusal. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 38 ++++++---- tests/unit/torch/utils/test_ep_loader.py | 72 ++++++++++++------- 2 files changed, 72 insertions(+), 38 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index bb7dc620337..f7e17872917 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1837,31 +1837,45 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A f"{plan!r}: ep_size={ep_degree} must divide the {num_experts} routed experts, which the " "all-to-all transport shares out equally over the expert-parallel ranks." ) - # Each rank builds the meshes below on the GPU it is bound to. DeviceMesh binds a rank that sees - # CUDA but is bound to none yet (a gloo group leaves it so) by its local rank, a GPU that a launch - # showing each process one does not have; and a rank without a GPU fails building the meshes. - # Either fails alone while the others wait in their subgroups' creation, so every rank learns - # from one gather which ranks are bound and which have a GPU, and all refuse alike. A run in - # which none has a GPU fails building the meshes on every rank alike. + + # Each rank builds the meshes on the GPU it is bound to, and DeviceMesh binds one not bound yet + # by its local rank. A rank that would so end up on no GPU, or on one another rank is on (as an + # import that initializes CUDA leaves every rank of a launch that binds none), fails alone in the + # meshes while the others wait: every rank decides from one gather, and all refuse alike. + def gpu() -> tuple[bool, bool, str | None]: + """Whether this rank has a GPU, whether it waits to be bound to one, and the bound one's UUID.""" + if not torch.cuda.is_initialized(): + return _has_own_gpu(), torch.cuda.is_available(), None + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + return _has_own_gpu(), False, str(props.uuid) + gpus = dist.all_reduce( "each rank's GPU", - lambda: [(_has_own_gpu(), torch.cuda.is_available() and not torch.cuda.is_initialized())], - combine=lambda ranks: [gpu for flags in ranks for gpu in flags], + lambda: [gpu()], + combine=lambda ranks: [g for facts in ranks for g in facts], ) - unbound = [rank for rank, (_, waits) in enumerate(gpus) if waits] + unbound = [rank for rank, (has, waits, _) in enumerate(gpus) if waits and not has] if unbound: raise RuntimeError( "The expert-parallel load builds its meshes on the GPU each rank is bound to, and rank(s) " - f"{unbound} are bound to none yet. Bind each rank to its GPU before the load " - "(torch.cuda.set_device), as distributed_run() does." + f"{unbound} are bound to none yet, nor does their local rank name one they see. Bind each " + "rank to its GPU before the load (torch.cuda.set_device), as distributed_run() does." ) - without = [rank for rank, (has, _) in enumerate(gpus) if not has] + without = [rank for rank, (has, _, _) in enumerate(gpus) if not has] if without and len(without) < len(gpus): raise RuntimeError( "The expert-parallel load builds its meshes on each rank's own GPU, and rank(s) " f"{without} have none while the others have theirs. Run one process per GPU, as " "torchrun --nproc_per_node= does." ) + uuids = [uuid for _, _, uuid in gpus] + shared = [rank for rank, uuid in enumerate(uuids) if uuid is not None and uuids.count(uuid) > 1] + if shared: + raise RuntimeError( + "The expert-parallel load builds its meshes on the GPU each rank is bound to, and rank(s) " + f"{shared} are bound to the same one. Bind each rank to its own GPU before the load " + "(torch.cuda.set_device), as distributed_run() does." + ) dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 # The weights load, and the checkpoint recorded for the export to read the unplaced tensors back # from, take the local snapshot: a Hub ID recorded there would leave the export nothing to read. diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 969c97f6e83..dea268b7191 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -80,6 +80,20 @@ def get_group(self): return -1 +def _cuda(gpu, bound=True, count=1, uuid="GPU-0"): + """Patches pinning the CUDA state the GPU check reads: ``count`` visible GPUs and this rank on + cuda:0, the one named ``uuid``, bound to it or not yet; or, without ``gpu``, none. The check then + reads the case, not the host the test runs on, whose visible GPU no test binds.""" + state = { + "is_available": gpu, + "is_initialized": gpu and bound, + "current_device": 0, + "device_count": count if gpu else 0, + "get_device_properties": SimpleNamespace(uuid=uuid), + } + return [mock.patch(f"torch.cuda.{name}", return_value=value) for name, value in state.items()] + + @pytest.fixture def ep_world_of_two(monkeypatch, gloo_world_of_one): """A two-rank expert-parallel world as far as the loader's checks see, and the checkpoint paths @@ -628,19 +642,6 @@ def _meshes_in_step(*args, **kwargs): raise _PastTheMeshesError -def _cuda(gpu, bound=True): - """Patches pinning the CUDA state the GPU check reads: one visible GPU, cuda:0, bound to it or - not yet; or, without ``gpu``, none. The check then reads the case, not the host the test runs - on (the gloo workers bind no GPU, so a host's visible one would read as unbound).""" - state = { - "is_available": gpu, - "is_initialized": gpu and bound, - "current_device": 0, - "device_count": int(gpu), - } - return [mock.patch(f"torch.cuda.{name}", return_value=value) for name, value in state.items()] - - def _one_rank_fails_job(rank, size, ckpt, failure, failing_rank, match): path, hf_config, patches = ckpt, None, [] if failure != "unreadable-config": @@ -652,7 +653,7 @@ def _one_rank_fails_job(rank, size, ckpt, failure, failing_rank, match): patches.append(local) fails_here = rank == failing_rank # No rank has a GPU, but in the no-gpu row every rank other than the failing one has its own. - patches += _cuda(failure == "no-gpu" and not fails_here) + patches += _cuda(failure == "no-gpu" and not fails_here, uuid=f"GPU-{rank}") def fail(*args, **kwargs): raise (MemoryError if failure == "load" else OSError)(match) @@ -726,12 +727,12 @@ def test_a_load_failing_on_one_rank_fails_every_rank( ) -def _one_gpu_per_process_job(rank, size, bound): +def _gpu_layout_job(rank, size, visible, bound, shared, refused): hf_config = SimpleNamespace(base_model_ep_plan=dict(_EP_PLAN), num_experts=4) reached = None with bounded_collectives(), ExitStack() as stack: - # Each process sees one GPU, cuda:0, bound to it or not yet; the local ranks run 0 and 1. - for patch in _cuda(True, bound): + # This rank is on cuda:0, bound to it or not yet; the local ranks run 0 and 1. + for patch in _cuda(True, bound, visible, "GPU-0" if shared else f"GPU-{rank}"): stack.enter_context(patch) stack.enter_context( mock.patch.object( @@ -745,19 +746,38 @@ def _one_gpu_per_process_job(rank, size, bound): load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=hf_config) except Exception as error: reached = error - if bound: + if refused is None: assert isinstance(reached, _PastTheMeshesError), (rank, reached) else: - assert "rank(s) [0, 1] are bound to none yet" in str(reached), (rank, reached) + assert isinstance(reached, RuntimeError) and refused in str(reached), (rank, reached) @pytest.mark.skipif( Version(transformers.__version__) < Version("5.7"), reason="the EP load needs transformers 5.7" ) -@pytest.mark.parametrize("bound", [True, False], ids=["bound", "not-yet-bound"]) -def test_one_visible_gpu_per_process_needs_each_rank_bound_to_it(cpu_dist_workers, bound): - """A launch that shows each process one GPU binds every rank to cuda:0, whatever its local rank: - every rank has its GPU, and every rank goes on to build the meshes. Not yet bound, DeviceMesh - would bind each by its local rank, a GPU rank 1 does not see: every rank is refused, told to - bind first.""" - cpu_dist_workers.run(_one_gpu_per_process_job, bound=bound) +@pytest.mark.parametrize( + ("visible", "bound", "shared", "refused"), + [ + (1, True, False, None), + (1, False, False, "rank(s) [1] are bound to none yet"), + (2, False, False, None), + (2, True, True, "rank(s) [0, 1] are bound to the same one"), + ], + ids=[ + "one-gpu-each-bound", + "one-gpu-each-not-yet", + "every-gpu-not-yet", + "every-gpu-all-on-cuda-0", + ], +) +def test_each_rank_builds_the_meshes_on_a_gpu_of_its_own( + cpu_dist_workers, visible, bound, shared, refused +): + """Each rank builds the meshes on the GPU it is bound to, and DeviceMesh binds one not bound yet + by its local rank. Shown one GPU each, every rank bound to cuda:0 goes on; not yet bound, rank 1's + local rank names no GPU it sees. Shown every GPU and not yet bound, DeviceMesh binds each to its + own. Every rank left on cuda:0, as an import that initializes CUDA leaves a launch that binds + none, shares one GPU. A refusal comes on every rank, from one gather.""" + cpu_dist_workers.run( + _gpu_layout_job, visible=visible, bound=bound, shared=shared, refused=refused + ) From 71cf53d84572beab524d8f390181c3fcd4b495e8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:07:15 +0000 Subject: [PATCH 21/23] Pin the single-process EP loader tests' CUDA to none The two-rank tests pin the CUDA state the GPU check reads, but the single-process ones, through gloo_world_of_one, still read the host's. The module now overrides that fixture to pin no GPU, so those tests, the Hub-resolution and nested-expert-count ones included, reach the outcomes they check whatever GPU the host has. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_ep_loader.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index dea268b7191..2f4e3232d97 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -94,6 +94,15 @@ def _cuda(gpu, bound=True, count=1, uuid="GPU-0"): return [mock.patch(f"torch.cuda.{name}", return_value=value) for name, value in state.items()] +@pytest.fixture +def gloo_world_of_one(gloo_world_of_one): + """The one-rank world, with no GPU whatever the host has (:func:`_cuda`).""" + with ExitStack() as stack: + for patch in _cuda(False): + stack.enter_context(patch) + yield gloo_world_of_one + + @pytest.fixture def ep_world_of_two(monkeypatch, gloo_world_of_one): """A two-rank expert-parallel world as far as the loader's checks see, and the checkpoint paths From 160f89b54ca9029702bd0c14b2b919a47d0375bf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:08:04 +0000 Subject: [PATCH 22/23] Keep load_and_prepare_ep's docstring to the caller contract meenchen's review: the docstring ran to some 50 lines, and the comments after it carried launch history and version-by-version root causes. The docstring now states what a caller must know (a group set up and bound to each rank's GPU, the same plan on every rank, data of each rank's own, the formats and transformers releases it loads, the collective-free load its failure guarantee rests on), and the comments say in a line or two why the code is as it is. The rationale they carried moves to the PR description. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 175 +++++------------- 1 file changed, 48 insertions(+), 127 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index f7e17872917..705e21d5a93 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1684,61 +1684,31 @@ def load_and_prepare_ep( hf_config=None, require_lm_head: bool = True, ) -> nn.Module: - """Load + shard a MoE causal LM for expert-parallel PTQ via the DP=EP all-to-all transport. - - ``plan`` is the expert-parallel :class:`~modelopt.torch.utils.distributed.DistributionPlan` to load - by (``ep_size > 1``), for the whole process group. It records no plan on the model: the caller - does (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), as - :func:`load_pretrained` does for every load. The model builds with its LM head where an auto - class builds one; ``require_lm_head=False`` builds a VLM without it, through ``AutoModel`` (a - causal LM always keeps its head). - - Builds an ``(edp, ep)`` device mesh and loads experts EP-sharded over the ``ep`` submesh - (``enable_expert_parallel=True`` -- each rank owns ``num_experts / ep`` experts as grouped-GEMM - shards), while ATTENTION and dense params stay REPLICATED (so EP has no KV-head-divisibility cap). - The model's ``base_model_ep_plan`` is swapped to the ``torch_a2a`` styles: the router emits GLOBAL - unmasked expert ids and the experts forward routes each token to its expert-owning rank by a NCCL - all-to-all (dispatch) -> quantizable ``_grouped_linear`` -> all-to-all (combine), rather than every - rank computing all tokens followed by an EP all-reduce. Every rank is data-parallel (sees DIFFERENT - tokens), so calibration MUST shard the data per-rank. - - The ``edp`` dimension (present when ``plan.ep_size < world_size``) is additional data parallelism - across EP groups: non-expert params stay REPLICATED per rank and each EP group calibrates on - different data; experts stay EP-sharded (resident). No FSDP and no concurrent ``edp``/``ep`` - collectives -> no deadlock by construction. Costs more per-rank memory (non-experts + experts - replicated across edp) but is the reliable path at scale. - - EP CANNOT be applied to a meta-init model (transformers' ``GroupedGemmParallel.shard_tensor`` - slices materialized tensors), so this loads real weights via ``from_pretrained`` over the ``ep`` - submesh. Requires transformers >= 5.7: the MoE EP framework (ExpertsInterface / grouped-GEMM EP - styles / enable_expert_parallel) is in 5.6, but 5.6 builds a model only with the experts - implementations it ships, never ``torch_a2a``. (A model whose expert-parallel support landed - later states the release it needs in its spec.) - - Assumes the caller has already initialized (and device-bound) the process group -- distributed - bring-up is the launcher's job, not this loader's (hf_ptq.main() does it via ``distributed_run``). - The expert-parallel mesh stays recoverable from the model - (:func:`~modelopt.torch.utils.moe.expert_parallel_mesh`), once the caller has recorded the plan. - A refusal on any rank -- of its transformers, the config, the checkpoint, or in the load -- - raises on every rank, so no rank is left waiting in a collective for one that already failed. - In the load that rests on ``from_pretrained`` running no collective, as transformers' - expert-parallel load does not on 5.7 through 5.18: before 5.16 it slices each rank's experts - into plain tensors, and from 5.16 it builds DTensors locally (``run_check=False``, - ``src_data_rank=None``). So the others finish it and meet a rank that failed in it in the - refusal's gather. - Without ``hf_config`` every rank reads the checkpoint's config itself, from the Hub for a Hub ID, - before rank 0 makes the checkpoint local: on many ranks, pass it (as ``hf_ptq`` does). + """Load a MoE causal LM expert-parallel for PTQ: each rank holds its share of the routed experts. + + Collective: every rank of the process group calls it with the same ``plan``, an expert-parallel + :class:`~modelopt.torch.utils.distributed.DistributionPlan` for the whole group. The caller sets + the group up and binds each rank to its own GPU first, as + :func:`~modelopt.torch.utils.distributed.distributed_run` does. Each rank holds + ``num_experts / ep_size`` routed experts and a replica of the dense part, and tokens reach their + experts' ranks by all-to-all, so every rank must calibrate on data of its own. + + It loads unquantized checkpoints, on transformers 5.7 or later. A refusal on any rank -- of its + transformers, the config, the checkpoint, its GPU, or in the load -- raises on every rank. In + the load that rests on ``from_pretrained`` running no collective, as transformers' + expert-parallel load does not on 5.7 through 5.18. + + 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 + (: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. """ if not plan.is_expert_parallel: raise ValueError(f"{plan!r} shards no routed experts, which is what this load is for.") if plan.is_fsdp_dense: raise NotImplementedError(f"{plan!r}: this load replicates the dense part on every rank.") - # This loader assumes the caller has ALREADY initialized (and device-bound) the process group -- - # distributed bring-up belongs to the launcher, not this library function. The example launcher does - # it in hf_ptq.main() via distributed_run() -> modelopt.torch.utils.distributed.setup() before - # the load. A live, device-bound PG is required here so the mesh subgroups built below create their - # NCCL comms eagerly (the synchronized warmup all-reduce per subgroup further down) instead of lazily - # on first collective (which deadlocks at DP>1 because ranks reach the lazy init in different phases). + # Bringing the group up is the launcher's job; this load needs it live and device-bound. if not dist.is_initialized(): raise RuntimeError( "load_and_prepare_ep requires a live, device-bound process group; set it up first, e.g. " @@ -1749,10 +1719,8 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A 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 - # (torch_a2a_router emits GLOBAL ids unmasked; torch_a2a_experts injects the EP process_group + - # skips the all-reduce). nemotron_h's logits-only gate has no ep_router entry, so its - # nemotron_experts_ep also maps to torch_a2a_experts. The checks and the swapped plan both read it. + # 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", @@ -1762,9 +1730,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A # 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): - # Hard floor: below 5.6 there is no ExpertsInterface / grouped-GEMM EP styles, and - # enable_expert_parallel silently shards nothing (no error); 5.6 checks a model's experts - # implementation against the ones it ships, so from_pretrained would refuse torch_a2a. + # 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)" ) @@ -1827,10 +1793,8 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A "drops, so those expert weights would load whole on every rank beside their sharded " "siblings." ) - # The all-to-all transport finds an expert's rank by dividing its id by this rank's share of - # the experts, which holds only when every rank holds as many. Counted in the config that holds - # the plan, as the distribution plan's decision counts them. A count the config names none of - # the ways this reads is checked once the experts load, from their weights. + # 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) if num_experts % ep_degree: raise ValueError( @@ -1877,42 +1841,22 @@ def gpu() -> tuple[bool, bool, str | None]: "(torch.cuda.set_device), as distributed_run() does." ) dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 - # The weights load, and the checkpoint recorded for the export to read the unplaced tensors back - # from, take the local snapshot: a Hub ID recorded there would leave the export nothing to read. - # Resolved once the plan and the config pass, so a refused load downloads nothing. Rank 0 makes - # it local; a rank that cannot see where fails every rank. + # 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) - # Two INDEPENDENT meshes (NeMo-AutoModel style), named by the model part they govern -- two - # factorizations of the same world: - # sparse_moe_mesh = (edp, ep): experts are SHARDED over the ``ep`` axis and REPLICATED over ``edp``; - # the ``ep`` submesh also drives from_pretrained's expert sharding + the a2a transport. - # dense_mesh = (dp,): dense/non-expert params are data-parallel-REPLICATED over the WHOLE world. - # This subsumes the old a2a-vs-partial split: a2a is simply edp == 1, and dense amax - # always reduces over the whole world (fixing the DP>1 case, where the old dp submesh missed the ep dim). + # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over + # the whole world (``dp``). sparse_moe_mesh = device_mesh.init_device_mesh( "cuda", (edp, ep_degree), mesh_dim_names=("edp", "ep") ) dense_mesh = device_mesh.init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) ep_submesh = sparse_moe_mesh["ep"] - # NO TENSOR PARALLELISM IS REQUESTED HERE, despite the name. transformers implements expert - # parallelism THROUGH the TP application path: with enable_expert_parallel=True the model's - # ``tp_plan`` PROPERTY returns ``_ep_plan`` (distributed/mixin.py), so what actually gets applied - # is the expert plan. ``tp_plan="auto"`` is only the SWITCH that makes ``tp_size > 1`` in - # DistributedConfig.__post_init__ and therefore reaches ``apply_tensor_parallelism`` at all -- - # drop it and __post_init__ early-returns with tp_size=None, which then TypeErrors on - # ``tp_size > 1`` in prepare_distribute_model. Dense params stay REPLICATED (= our DP) because - # ``swapped_plan`` below keeps only expert entries; attention/embeds match no plan key. - # - # WHERE it goes is version-dependent: from transformers 5.17 from_pretrained refuses tp_plan - # beside distributed_config ("Set tensor-parallel options on DistributedConfig when using it"), - # while DistributedConfig has no tp_plan field before 5.10, nor in 5.11 and 5.12. Passing - # tp_plan ALONE is not an option either -- the 5.17+ back-compat path rebuilds the config as - # DistributedConfig(tp_plan=..., tp_size=...) and DROPS enable_expert_parallel, silently losing - # EP. So each argument goes on DistributedConfig where it has a field for it, and tp_plan - # otherwise to from_pretrained: the dataclass fields tell them apart on any release. + # No tensor parallelism: tp_plan="auto" only switches on the path through which transformers + # applies the expert plan, and the swapped plan holds expert entries only. Each argument goes on + # DistributedConfig where this release has a field for it, else to from_pretrained. _dc_fields = {f.name for f in dataclasses.fields(hf_distributed.DistributedConfig)} _dc_has_tp_plan = "tp_plan" in _dc_fields dc_kwargs: dict[str, Any] = {"enable_expert_parallel": True} @@ -1932,17 +1876,14 @@ def gpu() -> tuple[bool, bool, str | None]: from_pretrained_kwargs["tp_plan"] = "auto" if attn_implementation is not None: from_pretrained_kwargs["attn_implementation"] = attn_implementation - # DP=EP all-to-all: swap the ep_plan to our torch_a2a styles (``swap`` above) and select the - # `torch_a2a` experts forward (NCCL all-to-all dispatch/combine around the quantizable - # _grouped_linear). Tokens are routed to the expert-owning rank by all-to-all instead of every rank - # computing all tokens + all-reduce. Requires per-rank (data-parallel) calib data. + # Swap the plan to the torch_a2a styles and select the torch_a2a experts forward: tokens reach + # their experts' ranks by all-to-all around the quantizable grouped GEMM. register_torch_a2a() # idempotent; register the experts forward + the two styles, right before use # Keep expert-related entries (grouped_gemm expert params + the router/experts styles); DROP the TP # styles (a2a builds no TP group for them -> would shard attention across the wrong mesh). swapped_plan = {k: swap.get(v, v) for k, v in ep_plan.items() if v in kept_styles} - # Tag every routed-experts CONTAINER torch_a2a_experts (injects the EP process_group + marks the - # routed expert), found as the distribution plan's decision finds them: some plans tag only the - # expert PARAMS grouped_gemm, and name the container through their parents. + # Tag every routed-experts container torch_a2a_experts, found as the plan's decision finds them: + # some plans style only their parameters, grouped_gemm. swapped_plan.update(dict.fromkeys(_routed_expert_keys(ep_plan), "torch_a2a_experts")) from_pretrained_kwargs["experts_implementation"] = "torch_a2a" # The checkpoint's own plan, for the config the model keeps (below); None: a class default. @@ -1950,25 +1891,17 @@ def gpu() -> tuple[bool, bool, str | None]: ep_cfg.base_model_ep_plan = swapped_plan from_pretrained_kwargs["config"] = hf_config print_rank_0(f"[ep] ep_plan={swapped_plan} experts_implementation=torch_a2a") - # A rank failing in the load (out of memory, a file it cannot read) fails every rank here, - # before the warmup's collectives. from_pretrained runs no collective (transformers shards each - # rank's experts locally, and builds no mesh but the one handed to it), so the other ranks finish - # it and meet the failed one in this guard's gather. + # A rank failing in the load fails every rank here, before the warmup's collectives: + # from_pretrained runs no collective, so the others finish it and meet the failure in the gather. with _refused_together("the expert-parallel load", alone=False): - # The class that builds the whole checkpoint -- with its lm_head where an auto class builds one, - # so the head is calibrated and exported and the generation preview works -- chosen as the - # distribution plan's decision chose it, unless the caller asked for a VLM without its head - # (require_lm_head=False). The per-model plugin (models/) has already installed any on-the-fly - # weight decompression as a transformers monkeypatch. + # The class the plan's decision chose, with its LM head where an auto class builds one. loader = _auto_class(hf_config, require_lm_head, trust_remote_code=trust_remote_code) print_rank_0(f"[ep] loader={loader.__name__} (require_lm_head={require_lm_head})") with _experts_sharded_over(ep_submesh): model = loader.from_pretrained(ckpt_path, **from_pretrained_kwargs) model.eval() - # The swapped styles drove the sharding, and the model keeps the plan it was sharded by - # (``model.tp_plan``, copied when it was built). The config it keeps goes back to the - # checkpoint's plan: an export writes the config out, and nothing without ModelOpt knows the - # torch_a2a styles. + # The config the model keeps goes back to the checkpoint's plan, as an export writes it out + # and nothing without ModelOpt knows the torch_a2a styles. kept_config = getattr(model, "config", None) if kept_config is not None: kept_ep_config = _ep_config(kept_config) @@ -1976,12 +1909,8 @@ def gpu() -> tuple[bool, bool, str | None]: vars(kept_ep_config).pop("base_model_ep_plan", None) else: kept_ep_config.base_model_ep_plan = declared_plan - # Force the experts forward to dispatch to `torch_a2a`: from_pretrained's `experts_implementation` - # kwarg resolves non-fp8 impls back to "grouped_mm", but @use_experts_implementation reads - # `config._experts_implementation` at FORWARD time, so set it on every (sub)config now. (Without - # this the torch_a2a_experts style passes a 5th process_group arg to grouped_mm_experts_forward -> - # TypeError.) Unlike the plan it stays set, and an export still writes no torch_a2a: transformers - # keeps it as ``_experts_implementation_internal``, which it leaves out of a serialized config. + # The experts forward reads the implementation from the config when it runs, which + # from_pretrained left on grouped_mm: torch_a2a on every (sub)config. An export writes none. for _m in model.modules(): _cfg = getattr(_m, "config", None) if _cfg is not None and hasattr(_cfg, "_experts_implementation"): @@ -1991,19 +1920,15 @@ def gpu() -> tuple[bool, bool, str | None]: # for the export to carry. record_unplaced_source_keys(model, ckpt_path) - # DP=EP all-to-all: no FSDP. Non-experts stay replicated per rank (from_pretrained already placed - # them on this rank's GPU); experts stay EP-sharded on the ep submesh, which the model carries as - # ``_device_mesh`` for calibration to re-install the all-to-all forwards over (expert_parallel_mesh). - # transformers sets it from the mesh this load hands it, through internals that differ by release; - # set here, the model carries it on every release. + # The model carries the ep submesh, for calibration to re-install the all-to-all forwards + # over (expert_parallel_mesh); transformers sets it differently by release. model._device_mesh = ep_submesh if hasattr(model, "tie_weights"): model.tie_weights() if freeze: model.requires_grad_(False) - # The routed-expert containers, from the SAME base_model_ep_plan that drove the sharding. A MoE model - # loaded WITHOUT EP (dense DP) never reaches this loader, so its experts stay dense. + # The routed-expert containers, from the plan that drove the sharding. expert_modules = _routed_expert_modules(model, ep_plan) if not expert_modules: raise ValueError( @@ -2031,12 +1956,8 @@ def _axis(mesh): f"[ep] seeded the expert-parallel state of {len(expert_modules)} routed-expert " "module(s) from base_model_ep_plan" ) - # Eagerly create each amax collective's NCCL comm in a globally-consistent order (barrier around - # each), so the sub-communicators are built in lockstep rather than lazily mid-calibration -- which - # deadlocks when submesh groups are phase-skewed at DP>1. NO forward warmup (variable per-group time - # would desync ranks before comm creation). Skip size-1 axes (no collective). The EP group is already - # live from from_pretrained's expert sharding; re-warming it is a harmless redundant all-reduce. Same - # order on every rank -> lockstep. + # Build each amax collective's NCCL comm now, in the same order on every rank, rather than lazily + # mid-calibration, where subgroups phase-skewed at DP>1 deadlock. Size-1 axes run none. warm_device = next(model.parameters()).device for _sub in (dense_mesh["dp"], sparse_moe_mesh["ep"], sparse_moe_mesh["edp"]): if _sub.size() <= 1: From 930bb0f0e1c4ab3a4ec0b5d6a065ce2cf27959ac Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:20:12 +0000 Subject: [PATCH 23/23] Move the EP loader tests to the test files of the modules they cover meenchen's review: a test file is named after the module it covers. The expert-parallel load's tests move from test_ep_loader.py into test_model_load_utils.py, as a section of their own at its end, and the two expert_parallel_mesh cases into test_moe_a2a.py, beside moe.py's other tests. The tests themselves are unchanged, and the CPU suites collect and pass the same cases as before. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_ep_loader.py | 792 ------------------ .../unit/torch/utils/test_model_load_utils.py | 732 +++++++++++++++- tests/unit/torch/utils/test_moe_a2a.py | 31 + 3 files changed, 761 insertions(+), 794 deletions(-) delete mode 100644 tests/unit/torch/utils/test_ep_loader.py diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py deleted file mode 100644 index 2f4e3232d97..00000000000 --- a/tests/unit/torch/utils/test_ep_loader.py +++ /dev/null @@ -1,792 +0,0 @@ -# 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. - -"""What the expert-parallel loader reads off the loaded model and its config. - -The loader takes an expert-parallel plan, ``_seed_parallel_states`` gives every module the -parallel state its amax syncs over, -``expert_parallel_mesh`` hands calibration the mesh an expert-parallel plan's experts are sharded -over, and ``_ckpt_quant_method`` reads the checkpoint's native quantization. Everything runs in one -CPU process. -""" - -import copy -import dataclasses -from contextlib import ExitStack -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 packaging.version import Version -from torch import nn - -pytest.importorskip("transformers") - -import transformers -import transformers.distributed as hf_distributed -from transformers import FineGrainedFP8Config, PreTrainedModel, Qwen3MoeConfig, Qwen3MoeForCausalLM -from transformers.distributed import DistributedConfig - -import modelopt.torch.utils.distributed as dist -from modelopt.torch.utils import moe -from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan -from modelopt.torch.utils.moe import expert_parallel_mesh -from modelopt.torch.utils.plugins import hf_checkpoint_utils, model_load_utils -from modelopt.torch.utils.plugins.model_load_utils import ( - _ckpt_quant_method, - _experts_sharded_over, - _seed_parallel_states, - _transformers_mesh_builder, - load_and_prepare_ep, -) - -# ------------------------------------------------------------------------------- the plan - -_EP_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", -} - - -class _Mesh: - """A device mesh, as far as the loader reads one; its groups run no collective (``-1``).""" - - device_type = "cpu" - - def __getitem__(self, name): - return self - - def size(self): - return 2 - - def get_group(self): - return -1 - - -def _cuda(gpu, bound=True, count=1, uuid="GPU-0"): - """Patches pinning the CUDA state the GPU check reads: ``count`` visible GPUs and this rank on - cuda:0, the one named ``uuid``, bound to it or not yet; or, without ``gpu``, none. The check then - reads the case, not the host the test runs on, whose visible GPU no test binds.""" - state = { - "is_available": gpu, - "is_initialized": gpu and bound, - "current_device": 0, - "device_count": count if gpu else 0, - "get_device_properties": SimpleNamespace(uuid=uuid), - } - return [mock.patch(f"torch.cuda.{name}", return_value=value) for name, value in state.items()] - - -@pytest.fixture -def gloo_world_of_one(gloo_world_of_one): - """The one-rank world, with no GPU whatever the host has (:func:`_cuda`).""" - with ExitStack() as stack: - for patch in _cuda(False): - stack.enter_context(patch) - yield gloo_world_of_one - - -@pytest.fixture -def ep_world_of_two(monkeypatch, gloo_world_of_one): - """A two-rank expert-parallel world as far as the loader's checks see, and the checkpoint paths - it resolves; nothing else is set up.""" - resolved = [] - monkeypatch.setattr(dist, "is_initialized", lambda: True) - monkeypatch.setattr(dist, "size", lambda group=None: 2) - monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) - monkeypatch.setattr(moe, "require_transformers", lambda *args: None) - monkeypatch.setattr( - model_load_utils, "ensure_local_checkpoint", lambda p, **k: resolved.append(p) or (None, p) - ) - return resolved - - -@pytest.mark.parametrize( - ("plan", "error", "match"), - [ - (DistributionPlan(2), ValueError, "shards no routed experts"), - (DistributionPlan(2, ep_size=2, dense="fsdp"), NotImplementedError, "replicates the dense"), - ], - ids=["no-expert-parallelism", "sharded-dense"], -) -def test_the_loader_refuses_a_plan_it_cannot_load_by(plan, error, match): - """Before anything else: no process group, transformers release or checkpoint is looked at.""" - with pytest.raises(error, match=match): - load_and_prepare_ep("unused", plan) - - -def test_the_loader_refuses_a_plan_for_another_world(gloo_world_of_one): - with pytest.raises(ValueError, match="is for 2 ranks, but this run has 1"): - load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2)) - - -def test_the_loader_refuses_a_transformers_that_builds_no_model_with_torch_a2a( - monkeypatch, gloo_world_of_one -): - """transformers 5.6 checks a model's experts implementation against the ones it ships, so its - from_pretrained would refuse torch_a2a once the load had set everything up: refused first.""" - monkeypatch.setattr(dist, "is_initialized", lambda: True) - monkeypatch.setattr(dist, "size", lambda group=None: 2) - monkeypatch.setattr("transformers.__version__", "5.6.2") - - with pytest.raises(RuntimeError, match=r"requires transformers >= 5\.7"): - load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2), hf_config=SimpleNamespace()) - - -def test_a_hub_id_is_made_local_before_the_weights_load(monkeypatch, gloo_world_of_one): - """The export reads the tensors the load could not place back from the checkpoint it records, - which must be a local directory: a Hub ID there leaves the export nothing to read. The load - resolves it once the plan and the config pass, before anything reads a weight.""" - - class _ResolvedError(Exception): - """Stops the load where it resolves the checkpoint.""" - - resolved = [] - - def resolve(path, **hub_options): - resolved.append(path) - raise _ResolvedError - - monkeypatch.setattr(dist, "is_initialized", lambda: True) - monkeypatch.setattr(dist, "size", lambda group=None: 2) - monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) - monkeypatch.setattr(moe, "require_transformers", lambda *args: None) - monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", resolve) - - with pytest.raises(_ResolvedError): - load_and_prepare_ep( - "org/model", - DistributionPlan(2, ep_size=2), - hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), - ) - - assert resolved == ["org/model"] - - -def test_the_experts_are_counted_in_the_config_holding_the_plan(monkeypatch, gloo_world_of_one): - """As the distribution plan's decision counts them: a VLM whose nested config holds the MoE and - its plan has its experts counted there, not by an expert-ish field on its root config.""" - - class _CheckedError(Exception): - """Stops the load once its checks pass.""" - - def resolve(path, **hub_options): - raise _CheckedError - - monkeypatch.setattr(dist, "is_initialized", lambda: True) - monkeypatch.setattr(dist, "size", lambda group=None: 4) - monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) - monkeypatch.setattr(moe, "require_transformers", lambda *args: None) - monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", resolve) - text = SimpleNamespace(base_model_ep_plan=_EP_PLAN, num_experts=8) - - with pytest.raises(_CheckedError): - load_and_prepare_ep( - "org/model", - DistributionPlan(4, ep_size=4), - hf_config=SimpleNamespace(text_config=text, num_experts=3), - ) - - -@pytest.mark.parametrize( - ("config", "error", "match"), - [ - *( - ( - SimpleNamespace(quantization_config={"quant_method": method}), - NotImplementedError, - f"cannot dequantize a '{method}'", - ) - for method in ("gptq", "modelopt", "compressed-tensors", "mxfp4") - ), - (SimpleNamespace(), ValueError, "declare base_model_ep_plan"), - ( - SimpleNamespace(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} - ), - ValueError, - r"shards no weights of layers\.\*\.mlp\.experts through grouped_gemm", - ) - for params in ({}, {"layers.*.mlp.experts.gate_up_proj": "packed_colwise"}) - ), - ( - SimpleNamespace( - 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", - ), - ], - ids=[ - "gptq", - "modelopt", - "compressed-tensors", - "mxfp4", - "no-ep-plan", - "no-routed-experts", - "experts-container-only", - "experts-sharded-otherwise", - "experts-sharded-partly-otherwise", - ], -) -def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( - ep_world_of_two, config, error, match -): - """A format transformers' own quantizer would load as stored, for PTQ to treat as plain weights, - 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.""" - 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 == [] - - -def test_an_ep_size_that_does_not_divide_the_experts_is_refused_before_anything_loads( - ep_world_of_two, -): - """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) - - 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) - assert ep_world_of_two == [] - - -def test_a_plan_whose_expert_containers_match_no_module_is_refused(ep_world_of_two, monkeypatch): - """A plan that styles the grouped-GEMM parameters a level below the experts names a container no - module is. Loaded anyway, every module would get the dense state, and the experts' amax would - reduce over ranks holding different experts.""" - - class _Loader: - @staticmethod - def from_pretrained(path, **kwargs): - return model_with_decoder_at("model.") - - monkeypatch.setattr( - "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() - ) - 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"} - ) - - with pytest.raises(ValueError, match="that no module of"): - load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) - - -@pytest.mark.parametrize("declared", [True, False], ids=["on-the-config", "a-class-default"]) -def test_the_model_keeps_the_plan_its_checkpoint_declares( - ep_world_of_two, gloo_world_of_one, monkeypatch, declared -): - """The load shards by the torch_a2a styles, but the config the model keeps goes back to the - checkpoint's plan: an export writes it out, and nothing without ModelOpt knows those styles. A - plan the config class ships (transformers' own, e.g. Qwen3-MoE's) stays off the instance.""" - - class _PastError(Exception): - """Raised by the load's next check.""" - - model = model_with_decoder_at("model.") - - class _Loader: - @staticmethod - def from_pretrained(path, config, **kwargs): - model.config = copy.deepcopy(config) # transformers keeps a copy of the config - return model - - def next_check(*args, **kwargs): - raise _PastError - - monkeypatch.setattr( - "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() - ) - monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) - monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) - - class _Config(SimpleNamespace): - base_model_ep_plan = _EP_PLAN - - # 8 experts over 2 ranks: the stand-in model holds this rank's 4. - config = _Config(num_experts=8, **({"base_model_ep_plan": dict(_EP_PLAN)} if declared else {})) - - with pytest.raises(_PastError): - load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) - - assert model.config.base_model_ep_plan == _EP_PLAN - assert ("base_model_ep_plan" in vars(model.config)) is declared - - -def test_a_loaded_model_keeps_its_sharding_styles_out_of_its_saved_config( - ep_world_of_two, gloo_world_of_one, monkeypatch -): - """What the loader leaves on a transformers model, as transformers treats it. Built with the - torch_a2a styles, the model keeps them as ``model.tp_plan``, which calibration re-installs the - all-to-all forwards by, and the expert-parallel mesh it re-installs them over, which this stand-in - from_pretrained does not set; while its config goes back to the checkpoint's plan; and the experts - implementation its forwards dispatch on, torch_a2a, stays out of the config it serializes, - which an export writes.""" - if not Qwen3MoeConfig.base_model_ep_plan: - pytest.skip("this transformers ships no expert-parallel plan for Qwen3-MoE") - - class _PastError(Exception): - """Raised by the load's next check.""" - - built = [] - - class _Loader: - @staticmethod - def from_pretrained(path, config, **kwargs): - # As transformers' expert-parallel load builds it: this rank's 2 of the 4 experts. - config.distributed_config = DistributedConfig(enable_expert_parallel=True) - share = copy.deepcopy(config) - share.num_experts //= 2 - with torch.device("meta"): - model = Qwen3MoeForCausalLM(share) - model.config = config - built.append(model) - return model - - def next_check(*args, **kwargs): - raise _PastError - - mesh = _Mesh() - monkeypatch.setattr( - "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: mesh - ) - monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) - monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) - config = Qwen3MoeConfig( - hidden_size=16, - intermediate_size=32, - moe_intermediate_size=8, - num_hidden_layers=1, - num_attention_heads=2, - num_key_value_heads=1, - vocab_size=32, - num_experts=4, - num_experts_per_tok=1, - ) - - with pytest.raises(_PastError): - load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) - - (model,) = built - assert model.config.base_model_ep_plan == Qwen3MoeConfig.base_model_ep_plan - assert "torch_a2a_experts" in set(model.tp_plan.values()) - assert model._device_mesh is mesh - assert model.config._experts_implementation == "torch_a2a" - assert not any("experts_implementation" in key for key in model.config.to_dict()) - - -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.""" - - class _LoadedError(Exception): - """Stops the load at from_pretrained.""" - - loaded, classes = {}, [] - - class _Loader: - @staticmethod - def from_pretrained(path, **kwargs): - loaded.update(kwargs) - raise _LoadedError - - def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): - classes.append(trust_remote_code) - return _Loader - - monkeypatch.setattr( - "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) - - with pytest.raises(_LoadedError): - load_and_prepare_ep( - "org/model", DistributionPlan(2, ep_size=2), trust_remote_code=True, hf_config=config - ) - assert loaded["dtype"] is torch.float16 - assert classes == [True] - - -# ----------------------------------------------------------------------- _seed_parallel_states - - -def test_expert_containers_and_their_children_get_the_expert_state(): - """Everything else, the root and the containers' own parents included, gets the dense state; so - does a sibling whose name merely starts with a container's (Qwen2-MoE keeps - ``shared_expert_gate`` beside ``shared_expert``).""" - model = model_with_decoder_at("model.") - for layer in model.model.layers: - layer.mlp.experts.weight_quantizer = nn.Identity() # a child, as quantizer insertion adds - model.model.layers[0].mlp.experts_gate = nn.Linear(2, 1) - experts = {f"model.layers.{i}.mlp.experts": model.model.layers[i].mlp.experts for i in (0, 1)} - dense, expert = object(), object() - - _seed_parallel_states(model, experts, dense_state=dense, expert_state=expert) - - states = {name: module._parallel_state for name, module in model.named_modules()} - assert {name for name, state in states.items() if state is expert} == { - "model.layers.0.mlp.experts", - "model.layers.0.mlp.experts.weight_quantizer", - "model.layers.1.mlp.experts", - "model.layers.1.mlp.experts.weight_quantizer", - } - assert all(state is dense for state in states.values() if state is not expert) - - -# ------------------------------------------------------------------------ expert_parallel_mesh - - -@pytest.mark.parametrize( - ("plan", "is_expert_parallel"), - [(DistributionPlan(2, ep_size=2), True), (None, False), (DistributionPlan(2), False)], - ids=["expert-parallel", "no-plan", "replicated"], -) -def test_the_expert_parallel_mesh_comes_with_an_expert_parallel_plan_only(plan, is_expert_parallel): - """An expert-parallel plan: the mesh transformers kept on the model. Any other model has none, - even with a device mesh on it, as a load through from_pretrained's own TP can leave.""" - model = model_with_decoder_at("model.") - if plan is not None: - set_distribution_plan(model, plan) - model._device_mesh = mesh = object() - - assert expert_parallel_mesh(model) is (mesh if is_expert_parallel else None) - - -def test_an_expert_parallel_model_without_its_mesh_is_refused(): - """A model recorded expert-parallel that load_and_prepare_ep did not load has no mesh for - calibration to re-install the all-to-all forwards over; it is refused by name.""" - model = model_with_decoder_at("model.") - set_distribution_plan(model, DistributionPlan(2, ep_size=2)) - - with pytest.raises(RuntimeError, match="carries no device mesh"): - expert_parallel_mesh(model) - - -# --------------------------------------------------------------------------- _ckpt_quant_method - - -@pytest.mark.parametrize( - ("config", "expected"), - [ - (SimpleNamespace(quantization_config={"quant_method": "fp8"}), "fp8"), - (SimpleNamespace(quantization_config=SimpleNamespace(quant_method="mxfp4")), "mxfp4"), - ( - SimpleNamespace( - text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}) - ), - "fp8", - ), - ( - SimpleNamespace( - llm_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}) - ), - "fp8", - ), - ( - SimpleNamespace( - quantization_config={"quant_method": "mxfp4"}, - text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}), - ), - "mxfp4", - ), - (SimpleNamespace(quantization_config=FineGrainedFP8Config()), "fp8"), - (SimpleNamespace(quantization_config={"weight_block_size": [128, 128]}), None), - (SimpleNamespace(text_config=None), None), - ], - ids=[ - "dict", - "object", - "vlm-text-config", - "vlm-llm-config", - "top-level-wins", - "transformers-config-object", - "no-method", - "unquantized", - ], -) -def test_quant_method_comes_from_the_config_or_a_vlms_text_config(config, expected): - """The top-level config first, else a VLM's text or LLM config. A loaded model's config holds - transformers' config object, not the JSON dict, and callers compare against ``"fp8"`` / - ``"mxfp4"``: it reports the plain method name, not its enum.""" - assert _ckpt_quant_method(config) == expected - - -def test_from_pretrained_builds_its_mesh_from_the_ep_mesh_it_is_given(): - """Newer transformers ignores ``from_pretrained``'s ``device_mesh`` and builds a world-sized one, - which shards the experts over every rank when EP is below the world size. During the load it gets - the EP mesh instead, and its own builder is back afterwards, even when the load raises.""" - mixin = _transformers_mesh_builder() - if mixin is None: - pytest.skip("this transformers honors from_pretrained's device_mesh") - builder = mixin.initialize_distributed_mesh - ep_mesh = SimpleNamespace(device_type="cpu") - config = DistributedConfig(enable_expert_parallel=True, tp_plan="auto", tp_size=2) - with _experts_sharded_over(ep_mesh): - _, device_map, device_mesh = PreTrainedModel.prepare_distribute_model(config) - - assert device_mesh is ep_mesh - assert device_map == torch.device("cpu") - assert mixin.initialize_distributed_mesh is builder - with pytest.raises(RuntimeError), _experts_sharded_over(ep_mesh): - raise RuntimeError("the load failed") - assert mixin.initialize_distributed_mesh is builder - - -@pytest.mark.parametrize( - ("fields", "builds_its_mesh", "config", "tp_plan"), - [ - # transformers 5.6-5.9, 5.11, 5.12 - (["enable_expert_parallel"], False, {"enable_expert_parallel": True}, "auto"), - # 5.10, 5.13-5.17 - ( - ["tp_size", "tp_plan", "enable_expert_parallel"], - False, - {"enable_expert_parallel": True, "tp_plan": "auto"}, - None, - ), - # 5.18 on - ( - ["tp_size", "tp_plan", "enable_expert_parallel"], - True, - {"enable_expert_parallel": True, "tp_plan": "auto", "tp_size": 2}, - None, - ), - # no release so far - ( - ["tp_size", "enable_expert_parallel"], - True, - {"enable_expert_parallel": True, "tp_size": 2}, - "auto", - ), - ], - ids=["no-tp-plan-field", "tp-plan-field", "own-mesh-builder", "own-mesh-builder-no-tp-plan"], -) -def test_each_load_argument_goes_where_this_transformers_takes_it( - ep_world_of_two, monkeypatch, fields, builds_its_mesh, config, tp_plan -): - """``tp_plan="auto"`` goes on ``DistributedConfig`` where it has the field, else to - ``from_pretrained``, never both; a release that builds its own mesh gets it sized to the EP - degree. Each is decided by its own field, so a release with one and not the other loads too.""" - - class _LoadedError(Exception): - """Stops the load at from_pretrained.""" - - loaded = {} - - class _Loader: - @staticmethod - def from_pretrained(path, **kwargs): - loaded.update(kwargs) - raise _LoadedError - - distributed_config = dataclasses.make_dataclass( - "DistributedConfig", [(name, object, dataclasses.field(default=None)) for name in fields] - ) - monkeypatch.setattr(hf_distributed, "DistributedConfig", distributed_config) - monkeypatch.setattr( - model_load_utils, - "_transformers_mesh_builder", - lambda: SimpleNamespace(initialize_distributed_mesh=None) if builds_its_mesh else None, - ) - monkeypatch.setattr( - "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() - ) - monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - - with pytest.raises(_LoadedError): - load_and_prepare_ep( - "org/model", - DistributionPlan(2, ep_size=2), - hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), - ) - - given = {k: v for k, v in vars(loaded["distributed_config"]).items() if v is not None} - assert given == config - assert loaded.get("tp_plan") == tp_plan - - -# ----------------------------------------------------------------- one rank failing alone - - -class _PastTheMeshesError(Exception): - """Raised by the meshes' stand-in: this rank got past the load's checks.""" - - -def _meshes_in_step(*args, **kwargs): - """The meshes' subgroups, as the collective they are: a rank in it waits for every peer.""" - torch.distributed.barrier() - raise _PastTheMeshesError - - -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) - if failure != "unseen-checkpoint": - local = mock.patch.object( - model_load_utils, "ensure_local_checkpoint", side_effect=lambda p, **k: (None, p) - ) - patches.append(local) - fails_here = rank == failing_rank - # No rank has a GPU, but in the no-gpu row every rank other than the failing one has its own. - patches += _cuda(failure == "no-gpu" and not fails_here, uuid=f"GPU-{rank}") - - def fail(*args, **kwargs): - raise (MemoryError if failure == "load" else OSError)(match) - - if failure == "unreadable-config" and fails_here: - path = ckpt + "-not-on-this-node" - elif failure == "transformers" and fails_here: - patches.append(mock.patch("transformers.__version__", "5.6.2")) - elif failure == "quantized-config" and fails_here: - hf_config.quantization_config = {"quant_method": "gptq"} - elif failure == "unseen-checkpoint" and fails_here: - # Rank 0 makes the checkpoint local; a rank on a node without that storage cannot see it. - patches.append(mock.patch.object(hf_checkpoint_utils.os.path, "isdir", return_value=False)) - if failure in ("load", "record"): - loaded = model_with_decoder_at("model.") - load = fail if failure == "load" and fails_here else (lambda path, **kwargs: loaded) - loader = SimpleNamespace(__name__="Loader", from_pretrained=load) - record = fail if failure == "record" and fails_here else None - patches += [ - mock.patch("torch.distributed.device_mesh.init_device_mesh", return_value=_Mesh()), - mock.patch.object(model_load_utils, "_auto_class", return_value=loader), - mock.patch.object(model_load_utils, "record_unplaced_source_keys", side_effect=record), - ] - else: - patches.append( - mock.patch("torch.distributed.device_mesh.init_device_mesh", _meshes_in_step) - ) - refused = None - with bounded_collectives(), ExitStack() as stack: - for patch in patches: - stack.enter_context(patch) - try: # not pytest.raises: its failure would end the pooled worker instead of the job - load_and_prepare_ep(path, DistributionPlan(2, ep_size=2), hf_config=hf_config) - except Exception as error: - refused = f"{type(error).__name__}: {error}" - # The one rank's failure on every rank, before any rank went on alone into a collective. The GPU - # check decides on every rank at once, from one gather, so no rank relays another's failure. - assert refused is not None and match in refused, (rank, refused) - assert fails_here or failure == "no-gpu" or f"rank {failing_rank} failed" in refused, refused - - -@pytest.mark.skipif( - Version(transformers.__version__) < Version("5.7"), reason="the EP load needs transformers 5.7" -) -@pytest.mark.parametrize( - ("failure", "failing_rank", "match"), - [ - ("unreadable-config", 1, "not-on-this-node"), - ("transformers", 1, "requires transformers >= 5.7, but 5.6.2"), - ("quantized-config", 1, "cannot dequantize a 'gptq'"), - ("unseen-checkpoint", 1, "which this rank cannot see"), - ("no-gpu", 1, "rank(s) [1] have none while the others have theirs"), - ("load", 1, "the GPU ran out of memory"), - ("record", 1, "a weight file this rank cannot read"), - ], -) -def test_a_load_failing_on_one_rank_fails_every_rank( - cpu_dist_workers, tmp_path, failure, failing_rank, match -): - """One rank fails alone: reading the config, on its transformers, on its arguments, seeing the - checkpoint rank 0 made local, without a GPU of its own where the others have theirs, or in the - load. Every rank raises that failure, rather than leave the others waiting for it in the meshes - (stood in by a bounded barrier) or the warmup's collectives.""" - Qwen3MoeConfig(num_experts=4, num_hidden_layers=1, hidden_size=8).save_pretrained(tmp_path) - cpu_dist_workers.run( - _one_rank_fails_job, - ckpt=str(tmp_path), - failure=failure, - failing_rank=failing_rank, - match=match, - ) - - -def _gpu_layout_job(rank, size, visible, bound, shared, refused): - hf_config = SimpleNamespace(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. - for patch in _cuda(True, bound, visible, "GPU-0" if shared else f"GPU-{rank}"): - stack.enter_context(patch) - stack.enter_context( - mock.patch.object( - model_load_utils, "ensure_local_checkpoint", side_effect=lambda p, **k: (None, p) - ) - ) - stack.enter_context( - mock.patch("torch.distributed.device_mesh.init_device_mesh", _meshes_in_step) - ) - try: # not pytest.raises: its failure would end the pooled worker instead of the job - load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=hf_config) - except Exception as error: - reached = error - if refused is None: - assert isinstance(reached, _PastTheMeshesError), (rank, reached) - else: - assert isinstance(reached, RuntimeError) and refused in str(reached), (rank, reached) - - -@pytest.mark.skipif( - Version(transformers.__version__) < Version("5.7"), reason="the EP load needs transformers 5.7" -) -@pytest.mark.parametrize( - ("visible", "bound", "shared", "refused"), - [ - (1, True, False, None), - (1, False, False, "rank(s) [1] are bound to none yet"), - (2, False, False, None), - (2, True, True, "rank(s) [0, 1] are bound to the same one"), - ], - ids=[ - "one-gpu-each-bound", - "one-gpu-each-not-yet", - "every-gpu-not-yet", - "every-gpu-all-on-cuda-0", - ], -) -def test_each_rank_builds_the_meshes_on_a_gpu_of_its_own( - cpu_dist_workers, visible, bound, shared, refused -): - """Each rank builds the meshes on the GPU it is bound to, and DeviceMesh binds one not bound yet - by its local rank. Shown one GPU each, every rank bound to cuda:0 goes on; not yet bound, rank 1's - local rank names no GPU it sees. Shown every GPU and not yet bound, DeviceMesh binds each to its - own. Every rank left on cuda:0, as an import that initializes CUDA leaves a launch that binds - none, shares one GPU. A refusal comes on every rank, from one gather.""" - cpu_dist_workers.run( - _gpu_layout_job, visible=visible, bound=bound, shared=shared, refused=refused - ) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 9266e601180..5f07787ddee 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -22,10 +22,12 @@ tests of each part use a process group, two gloo ranks. """ +import copy +import dataclasses import inspect import math import os -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from types import SimpleNamespace from unittest import mock @@ -34,9 +36,11 @@ 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 +from torch import nn pytest.importorskip("accelerate") transformers = pytest.importorskip("transformers") +import transformers.distributed as hf_distributed from _test_utils.torch.transformers_models import ( _TinyQwen35Tokenizer, get_tiny_qwen3_5_moe_vl, @@ -50,29 +54,39 @@ AutoModelForImageTextToText, AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq, + FineGrainedFP8Config, LlamaConfig, LlamaForCausalLM, Phi4MultimodalConfig, + PreTrainedModel, Qwen2_5_VLConfig, Qwen3MoeConfig, + Qwen3MoeForCausalLM, T5Config, WhisperConfig, ) +from transformers.distributed import DistributedConfig import modelopt.torch.models.hf as models_hf import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist +import modelopt.torch.utils.moe as moe from modelopt.torch.quantization.plugins import accelerate as accelerate_plugin from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan -from modelopt.torch.utils.plugins import model_load_utils +from modelopt.torch.utils.plugins import hf_checkpoint_utils, model_load_utils from modelopt.torch.utils.plugins.model_load_utils import ( _auto_class, + _ckpt_quant_method, _conversion_plan, _convert_keys, + _experts_sharded_over, _resolve_target, + _seed_parallel_states, + _transformers_mesh_builder, build_meta_causal_lm, + load_and_prepare_ep, load_pretrained, record_unplaced_source_keys, resolve_distribution_plan, @@ -1830,3 +1844,717 @@ def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( with pytest.raises(error, match=match): load_pretrained(str(tmp_path), distribution=distribution) assert loaders == {} + + +# ======================================================================== the expert-parallel load +# What the expert-parallel load reads off the loaded model and its config: it takes an expert-parallel +# plan, ``_seed_parallel_states`` gives every module the parallel state its amax syncs over, and +# ``_ckpt_quant_method`` reads the checkpoint's native quantization. Everything runs in one CPU process +# but the bounded two-rank tests at the end. + +# ------------------------------------------------------------------------------- the plan + +_EP_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", +} + + +class _Mesh: + """A device mesh, as far as the loader reads one; its groups run no collective (``-1``).""" + + device_type = "cpu" + + def __getitem__(self, name): + return self + + def size(self): + return 2 + + def get_group(self): + return -1 + + +def _cuda(gpu, bound=True, count=1, uuid="GPU-0"): + """Patches pinning the CUDA state the GPU check reads: ``count`` visible GPUs and this rank on + cuda:0, the one named ``uuid``, bound to it or not yet; or, without ``gpu``, none. The check then + reads the case, not the host the test runs on, whose visible GPU no test binds.""" + state = { + "is_available": gpu, + "is_initialized": gpu and bound, + "current_device": 0, + "device_count": count if gpu else 0, + "get_device_properties": SimpleNamespace(uuid=uuid), + } + return [mock.patch(f"torch.cuda.{name}", return_value=value) for name, value in state.items()] + + +@pytest.fixture +def gloo_world_of_one(gloo_world_of_one): + """The one-rank world, with no GPU whatever the host has (:func:`_cuda`).""" + with ExitStack() as stack: + for patch in _cuda(False): + stack.enter_context(patch) + yield gloo_world_of_one + + +@pytest.fixture +def ep_world_of_two(monkeypatch, gloo_world_of_one): + """A two-rank expert-parallel world as far as the loader's checks see, and the checkpoint paths + it resolves; nothing else is set up.""" + resolved = [] + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + monkeypatch.setattr( + model_load_utils, "ensure_local_checkpoint", lambda p, **k: resolved.append(p) or (None, p) + ) + return resolved + + +@pytest.mark.parametrize( + ("plan", "error", "match"), + [ + (DistributionPlan(2), ValueError, "shards no routed experts"), + (DistributionPlan(2, ep_size=2, dense="fsdp"), NotImplementedError, "replicates the dense"), + ], + ids=["no-expert-parallelism", "sharded-dense"], +) +def test_the_loader_refuses_a_plan_it_cannot_load_by(plan, error, match): + """Before anything else: no process group, transformers release or checkpoint is looked at.""" + with pytest.raises(error, match=match): + load_and_prepare_ep("unused", plan) + + +def test_the_loader_refuses_a_plan_for_another_world(gloo_world_of_one): + with pytest.raises(ValueError, match="is for 2 ranks, but this run has 1"): + load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2)) + + +def test_the_loader_refuses_a_transformers_that_builds_no_model_with_torch_a2a( + monkeypatch, gloo_world_of_one +): + """transformers 5.6 checks a model's experts implementation against the ones it ships, so its + from_pretrained would refuse torch_a2a once the load had set everything up: refused first.""" + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr("transformers.__version__", "5.6.2") + + with pytest.raises(RuntimeError, match=r"requires transformers >= 5\.7"): + load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2), hf_config=SimpleNamespace()) + + +def test_a_hub_id_is_made_local_before_the_weights_load(monkeypatch, gloo_world_of_one): + """The export reads the tensors the load could not place back from the checkpoint it records, + which must be a local directory: a Hub ID there leaves the export nothing to read. The load + resolves it once the plan and the config pass, before anything reads a weight.""" + + class _ResolvedError(Exception): + """Stops the load where it resolves the checkpoint.""" + + resolved = [] + + def resolve(path, **hub_options): + resolved.append(path) + raise _ResolvedError + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", resolve) + + with pytest.raises(_ResolvedError): + load_and_prepare_ep( + "org/model", + DistributionPlan(2, ep_size=2), + hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), + ) + + assert resolved == ["org/model"] + + +def test_the_experts_are_counted_in_the_config_holding_the_plan(monkeypatch, gloo_world_of_one): + """As the distribution plan's decision counts them: a VLM whose nested config holds the MoE and + its plan has its experts counted there, not by an expert-ish field on its root config.""" + + class _CheckedError(Exception): + """Stops the load once its checks pass.""" + + def resolve(path, **hub_options): + raise _CheckedError + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 4) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", resolve) + text = SimpleNamespace(base_model_ep_plan=_EP_PLAN, num_experts=8) + + with pytest.raises(_CheckedError): + load_and_prepare_ep( + "org/model", + DistributionPlan(4, ep_size=4), + hf_config=SimpleNamespace(text_config=text, num_experts=3), + ) + + +@pytest.mark.parametrize( + ("config", "error", "match"), + [ + *( + ( + SimpleNamespace(quantization_config={"quant_method": method}), + NotImplementedError, + f"cannot dequantize a '{method}'", + ) + for method in ("gptq", "modelopt", "compressed-tensors", "mxfp4") + ), + (SimpleNamespace(), ValueError, "declare base_model_ep_plan"), + ( + SimpleNamespace(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} + ), + ValueError, + r"shards no weights of layers\.\*\.mlp\.experts through grouped_gemm", + ) + for params in ({}, {"layers.*.mlp.experts.gate_up_proj": "packed_colwise"}) + ), + ( + SimpleNamespace( + 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", + ), + ], + ids=[ + "gptq", + "modelopt", + "compressed-tensors", + "mxfp4", + "no-ep-plan", + "no-routed-experts", + "experts-container-only", + "experts-sharded-otherwise", + "experts-sharded-partly-otherwise", + ], +) +def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( + ep_world_of_two, config, error, match +): + """A format transformers' own quantizer would load as stored, for PTQ to treat as plain weights, + 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.""" + 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 == [] + + +def test_an_ep_size_that_does_not_divide_the_experts_is_refused_before_anything_loads( + ep_world_of_two, +): + """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) + + 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) + assert ep_world_of_two == [] + + +def test_a_plan_whose_expert_containers_match_no_module_is_refused(ep_world_of_two, monkeypatch): + """A plan that styles the grouped-GEMM parameters a level below the experts names a container no + module is. Loaded anyway, every module would get the dense state, and the experts' amax would + reduce over ranks holding different experts.""" + + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + return model_with_decoder_at("model.") + + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + 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"} + ) + + with pytest.raises(ValueError, match="that no module of"): + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) + + +@pytest.mark.parametrize("declared", [True, False], ids=["on-the-config", "a-class-default"]) +def test_the_model_keeps_the_plan_its_checkpoint_declares( + ep_world_of_two, gloo_world_of_one, monkeypatch, declared +): + """The load shards by the torch_a2a styles, but the config the model keeps goes back to the + checkpoint's plan: an export writes it out, and nothing without ModelOpt knows those styles. A + plan the config class ships (transformers' own, e.g. Qwen3-MoE's) stays off the instance.""" + + class _PastError(Exception): + """Raised by the load's next check.""" + + model = model_with_decoder_at("model.") + + class _Loader: + @staticmethod + def from_pretrained(path, config, **kwargs): + model.config = copy.deepcopy(config) # transformers keeps a copy of the config + return model + + def next_check(*args, **kwargs): + raise _PastError + + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) + + class _Config(SimpleNamespace): + base_model_ep_plan = _EP_PLAN + + # 8 experts over 2 ranks: the stand-in model holds this rank's 4. + config = _Config(num_experts=8, **({"base_model_ep_plan": dict(_EP_PLAN)} if declared else {})) + + with pytest.raises(_PastError): + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) + + assert model.config.base_model_ep_plan == _EP_PLAN + assert ("base_model_ep_plan" in vars(model.config)) is declared + + +def test_a_loaded_model_keeps_its_sharding_styles_out_of_its_saved_config( + ep_world_of_two, gloo_world_of_one, monkeypatch +): + """What the loader leaves on a transformers model, as transformers treats it. Built with the + torch_a2a styles, the model keeps them as ``model.tp_plan``, which calibration re-installs the + all-to-all forwards by, and the expert-parallel mesh it re-installs them over, which this stand-in + from_pretrained does not set; while its config goes back to the checkpoint's plan; and the experts + implementation its forwards dispatch on, torch_a2a, stays out of the config it serializes, + which an export writes.""" + if not Qwen3MoeConfig.base_model_ep_plan: + pytest.skip("this transformers ships no expert-parallel plan for Qwen3-MoE") + + class _PastError(Exception): + """Raised by the load's next check.""" + + built = [] + + class _Loader: + @staticmethod + def from_pretrained(path, config, **kwargs): + # As transformers' expert-parallel load builds it: this rank's 2 of the 4 experts. + config.distributed_config = DistributedConfig(enable_expert_parallel=True) + share = copy.deepcopy(config) + share.num_experts //= 2 + with torch.device("meta"): + model = Qwen3MoeForCausalLM(share) + model.config = config + built.append(model) + return model + + def next_check(*args, **kwargs): + raise _PastError + + mesh = _Mesh() + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: mesh + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) + config = Qwen3MoeConfig( + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + vocab_size=32, + num_experts=4, + num_experts_per_tok=1, + ) + + with pytest.raises(_PastError): + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) + + (model,) = built + assert model.config.base_model_ep_plan == Qwen3MoeConfig.base_model_ep_plan + assert "torch_a2a_experts" in set(model.tp_plan.values()) + assert model._device_mesh is mesh + assert model.config._experts_implementation == "torch_a2a" + assert not any("experts_implementation" in key for key in model.config.to_dict()) + + +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.""" + + class _LoadedError(Exception): + """Stops the load at from_pretrained.""" + + loaded, classes = {}, [] + + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + loaded.update(kwargs) + raise _LoadedError + + def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): + classes.append(trust_remote_code) + return _Loader + + monkeypatch.setattr( + "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) + + with pytest.raises(_LoadedError): + load_and_prepare_ep( + "org/model", DistributionPlan(2, ep_size=2), trust_remote_code=True, hf_config=config + ) + assert loaded["dtype"] is torch.float16 + assert classes == [True] + + +# ----------------------------------------------------------------------- _seed_parallel_states + + +def test_expert_containers_and_their_children_get_the_expert_state(): + """Everything else, the root and the containers' own parents included, gets the dense state; so + does a sibling whose name merely starts with a container's (Qwen2-MoE keeps + ``shared_expert_gate`` beside ``shared_expert``).""" + model = model_with_decoder_at("model.") + for layer in model.model.layers: + layer.mlp.experts.weight_quantizer = nn.Identity() # a child, as quantizer insertion adds + model.model.layers[0].mlp.experts_gate = nn.Linear(2, 1) + experts = {f"model.layers.{i}.mlp.experts": model.model.layers[i].mlp.experts for i in (0, 1)} + dense, expert = object(), object() + + _seed_parallel_states(model, experts, dense_state=dense, expert_state=expert) + + states = {name: module._parallel_state for name, module in model.named_modules()} + assert {name for name, state in states.items() if state is expert} == { + "model.layers.0.mlp.experts", + "model.layers.0.mlp.experts.weight_quantizer", + "model.layers.1.mlp.experts", + "model.layers.1.mlp.experts.weight_quantizer", + } + assert all(state is dense for state in states.values() if state is not expert) + + +# --------------------------------------------------------------------------- _ckpt_quant_method + + +@pytest.mark.parametrize( + ("config", "expected"), + [ + (SimpleNamespace(quantization_config={"quant_method": "fp8"}), "fp8"), + (SimpleNamespace(quantization_config=SimpleNamespace(quant_method="mxfp4")), "mxfp4"), + ( + SimpleNamespace( + text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}) + ), + "fp8", + ), + ( + SimpleNamespace( + llm_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}) + ), + "fp8", + ), + ( + SimpleNamespace( + quantization_config={"quant_method": "mxfp4"}, + text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}), + ), + "mxfp4", + ), + (SimpleNamespace(quantization_config=FineGrainedFP8Config()), "fp8"), + (SimpleNamespace(quantization_config={"weight_block_size": [128, 128]}), None), + (SimpleNamespace(text_config=None), None), + ], + ids=[ + "dict", + "object", + "vlm-text-config", + "vlm-llm-config", + "top-level-wins", + "transformers-config-object", + "no-method", + "unquantized", + ], +) +def test_quant_method_comes_from_the_config_or_a_vlms_text_config(config, expected): + """The top-level config first, else a VLM's text or LLM config. A loaded model's config holds + transformers' config object, not the JSON dict, and callers compare against ``"fp8"`` / + ``"mxfp4"``: it reports the plain method name, not its enum.""" + assert _ckpt_quant_method(config) == expected + + +def test_from_pretrained_builds_its_mesh_from_the_ep_mesh_it_is_given(): + """Newer transformers ignores ``from_pretrained``'s ``device_mesh`` and builds a world-sized one, + which shards the experts over every rank when EP is below the world size. During the load it gets + the EP mesh instead, and its own builder is back afterwards, even when the load raises.""" + mixin = _transformers_mesh_builder() + if mixin is None: + pytest.skip("this transformers honors from_pretrained's device_mesh") + builder = mixin.initialize_distributed_mesh + ep_mesh = SimpleNamespace(device_type="cpu") + config = DistributedConfig(enable_expert_parallel=True, tp_plan="auto", tp_size=2) + with _experts_sharded_over(ep_mesh): + _, device_map, device_mesh = PreTrainedModel.prepare_distribute_model(config) + + assert device_mesh is ep_mesh + assert device_map == torch.device("cpu") + assert mixin.initialize_distributed_mesh is builder + with pytest.raises(RuntimeError), _experts_sharded_over(ep_mesh): + raise RuntimeError("the load failed") + assert mixin.initialize_distributed_mesh is builder + + +@pytest.mark.parametrize( + ("fields", "builds_its_mesh", "config", "tp_plan"), + [ + # transformers 5.6-5.9, 5.11, 5.12 + (["enable_expert_parallel"], False, {"enable_expert_parallel": True}, "auto"), + # 5.10, 5.13-5.17 + ( + ["tp_size", "tp_plan", "enable_expert_parallel"], + False, + {"enable_expert_parallel": True, "tp_plan": "auto"}, + None, + ), + # 5.18 on + ( + ["tp_size", "tp_plan", "enable_expert_parallel"], + True, + {"enable_expert_parallel": True, "tp_plan": "auto", "tp_size": 2}, + None, + ), + # no release so far + ( + ["tp_size", "enable_expert_parallel"], + True, + {"enable_expert_parallel": True, "tp_size": 2}, + "auto", + ), + ], + ids=["no-tp-plan-field", "tp-plan-field", "own-mesh-builder", "own-mesh-builder-no-tp-plan"], +) +def test_each_load_argument_goes_where_this_transformers_takes_it( + ep_world_of_two, monkeypatch, fields, builds_its_mesh, config, tp_plan +): + """``tp_plan="auto"`` goes on ``DistributedConfig`` where it has the field, else to + ``from_pretrained``, never both; a release that builds its own mesh gets it sized to the EP + degree. Each is decided by its own field, so a release with one and not the other loads too.""" + + class _LoadedError(Exception): + """Stops the load at from_pretrained.""" + + loaded = {} + + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + loaded.update(kwargs) + raise _LoadedError + + distributed_config = dataclasses.make_dataclass( + "DistributedConfig", [(name, object, dataclasses.field(default=None)) for name in fields] + ) + monkeypatch.setattr(hf_distributed, "DistributedConfig", distributed_config) + monkeypatch.setattr( + model_load_utils, + "_transformers_mesh_builder", + lambda: SimpleNamespace(initialize_distributed_mesh=None) if builds_its_mesh else None, + ) + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + + with pytest.raises(_LoadedError): + load_and_prepare_ep( + "org/model", + DistributionPlan(2, ep_size=2), + hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), + ) + + given = {k: v for k, v in vars(loaded["distributed_config"]).items() if v is not None} + assert given == config + assert loaded.get("tp_plan") == tp_plan + + +# ----------------------------------------------------------------- one rank failing alone + + +class _PastTheMeshesError(Exception): + """Raised by the meshes' stand-in: this rank got past the load's checks.""" + + +def _meshes_in_step(*args, **kwargs): + """The meshes' subgroups, as the collective they are: a rank in it waits for every peer.""" + torch.distributed.barrier() + raise _PastTheMeshesError + + +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) + if failure != "unseen-checkpoint": + local = mock.patch.object( + model_load_utils, "ensure_local_checkpoint", side_effect=lambda p, **k: (None, p) + ) + patches.append(local) + fails_here = rank == failing_rank + # No rank has a GPU, but in the no-gpu row every rank other than the failing one has its own. + patches += _cuda(failure == "no-gpu" and not fails_here, uuid=f"GPU-{rank}") + + def fail(*args, **kwargs): + raise (MemoryError if failure == "load" else OSError)(match) + + if failure == "unreadable-config" and fails_here: + path = ckpt + "-not-on-this-node" + elif failure == "transformers" and fails_here: + patches.append(mock.patch("transformers.__version__", "5.6.2")) + elif failure == "quantized-config" and fails_here: + hf_config.quantization_config = {"quant_method": "gptq"} + elif failure == "unseen-checkpoint" and fails_here: + # Rank 0 makes the checkpoint local; a rank on a node without that storage cannot see it. + patches.append(mock.patch.object(hf_checkpoint_utils.os.path, "isdir", return_value=False)) + if failure in ("load", "record"): + loaded = model_with_decoder_at("model.") + load = fail if failure == "load" and fails_here else (lambda path, **kwargs: loaded) + loader = SimpleNamespace(__name__="Loader", from_pretrained=load) + record = fail if failure == "record" and fails_here else None + patches += [ + mock.patch("torch.distributed.device_mesh.init_device_mesh", return_value=_Mesh()), + mock.patch.object(model_load_utils, "_auto_class", return_value=loader), + mock.patch.object(model_load_utils, "record_unplaced_source_keys", side_effect=record), + ] + else: + patches.append( + mock.patch("torch.distributed.device_mesh.init_device_mesh", _meshes_in_step) + ) + refused = None + with bounded_collectives(), ExitStack() as stack: + for patch in patches: + stack.enter_context(patch) + try: # not pytest.raises: its failure would end the pooled worker instead of the job + load_and_prepare_ep(path, DistributionPlan(2, ep_size=2), hf_config=hf_config) + except Exception as error: + refused = f"{type(error).__name__}: {error}" + # The one rank's failure on every rank, before any rank went on alone into a collective. The GPU + # check decides on every rank at once, from one gather, so no rank relays another's failure. + assert refused is not None and match in refused, (rank, refused) + assert fails_here or failure == "no-gpu" or f"rank {failing_rank} failed" in refused, refused + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("5.7"), reason="the EP load needs transformers 5.7" +) +@pytest.mark.parametrize( + ("failure", "failing_rank", "match"), + [ + ("unreadable-config", 1, "not-on-this-node"), + ("transformers", 1, "requires transformers >= 5.7, but 5.6.2"), + ("quantized-config", 1, "cannot dequantize a 'gptq'"), + ("unseen-checkpoint", 1, "which this rank cannot see"), + ("no-gpu", 1, "rank(s) [1] have none while the others have theirs"), + ("load", 1, "the GPU ran out of memory"), + ("record", 1, "a weight file this rank cannot read"), + ], +) +def test_a_load_failing_on_one_rank_fails_every_rank( + cpu_dist_workers, tmp_path, failure, failing_rank, match +): + """One rank fails alone: reading the config, on its transformers, on its arguments, seeing the + checkpoint rank 0 made local, without a GPU of its own where the others have theirs, or in the + load. Every rank raises that failure, rather than leave the others waiting for it in the meshes + (stood in by a bounded barrier) or the warmup's collectives.""" + Qwen3MoeConfig(num_experts=4, num_hidden_layers=1, hidden_size=8).save_pretrained(tmp_path) + cpu_dist_workers.run( + _one_rank_fails_job, + ckpt=str(tmp_path), + failure=failure, + failing_rank=failing_rank, + match=match, + ) + + +def _gpu_layout_job(rank, size, visible, bound, shared, refused): + hf_config = SimpleNamespace(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. + for patch in _cuda(True, bound, visible, "GPU-0" if shared else f"GPU-{rank}"): + stack.enter_context(patch) + stack.enter_context( + mock.patch.object( + model_load_utils, "ensure_local_checkpoint", side_effect=lambda p, **k: (None, p) + ) + ) + stack.enter_context( + mock.patch("torch.distributed.device_mesh.init_device_mesh", _meshes_in_step) + ) + try: # not pytest.raises: its failure would end the pooled worker instead of the job + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=hf_config) + except Exception as error: + reached = error + if refused is None: + assert isinstance(reached, _PastTheMeshesError), (rank, reached) + else: + assert isinstance(reached, RuntimeError) and refused in str(reached), (rank, reached) + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("5.7"), reason="the EP load needs transformers 5.7" +) +@pytest.mark.parametrize( + ("visible", "bound", "shared", "refused"), + [ + (1, True, False, None), + (1, False, False, "rank(s) [1] are bound to none yet"), + (2, False, False, None), + (2, True, True, "rank(s) [0, 1] are bound to the same one"), + ], + ids=[ + "one-gpu-each-bound", + "one-gpu-each-not-yet", + "every-gpu-not-yet", + "every-gpu-all-on-cuda-0", + ], +) +def test_each_rank_builds_the_meshes_on_a_gpu_of_its_own( + cpu_dist_workers, visible, bound, shared, refused +): + """Each rank builds the meshes on the GPU it is bound to, and DeviceMesh binds one not bound yet + by its local rank. Shown one GPU each, every rank bound to cuda:0 goes on; not yet bound, rank 1's + local rank names no GPU it sees. Shown every GPU and not yet bound, DeviceMesh binds each to its + own. Every rank left on cuda:0, as an import that initializes CUDA leaves a launch that binds + none, shares one GPU. A refusal comes on every rank, from one gather.""" + cpu_dist_workers.run( + _gpu_layout_job, visible=visible, bound=bound, shared=shared, refused=refused + ) diff --git a/tests/unit/torch/utils/test_moe_a2a.py b/tests/unit/torch/utils/test_moe_a2a.py index 9dea6d86c67..090024ca24f 100644 --- a/tests/unit/torch/utils/test_moe_a2a.py +++ b/tests/unit/torch/utils/test_moe_a2a.py @@ -43,11 +43,13 @@ import modelopt.torch.quantization as mtq import modelopt.torch.utils.moe as moe +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan from modelopt.torch.utils.moe import ( _choose_ep_degree, _num_experts_from_config, _routed_expert_modules, a2a_moe, + expert_parallel_mesh, ) # Small, fully-deterministic problem. E must be divisible by every world size under test. @@ -563,3 +565,32 @@ def test_a_transformers_below_a_features_floor_is_refused(monkeypatch, installed else nullcontext() ): moe.require_transformers("5.7", "EP load") + + +# ------------------------------------------------------------------------ expert_parallel_mesh + + +@pytest.mark.parametrize( + ("plan", "is_expert_parallel"), + [(DistributionPlan(2, ep_size=2), True), (None, False), (DistributionPlan(2), False)], + ids=["expert-parallel", "no-plan", "replicated"], +) +def test_the_expert_parallel_mesh_comes_with_an_expert_parallel_plan_only(plan, is_expert_parallel): + """An expert-parallel plan: the mesh transformers kept on the model. Any other model has none, + even with a device mesh on it, as a load through from_pretrained's own TP can leave.""" + model = model_with_decoder_at("model.") + if plan is not None: + set_distribution_plan(model, plan) + model._device_mesh = mesh = object() + + assert expert_parallel_mesh(model) is (mesh if is_expert_parallel else None) + + +def test_an_expert_parallel_model_without_its_mesh_is_refused(): + """A model recorded expert-parallel that load_and_prepare_ep did not load has no mesh for + calibration to re-install the all-to-all forwards over; it is refused by name.""" + model = model_with_decoder_at("model.") + set_distribution_plan(model, DistributionPlan(2, ep_size=2)) + + with pytest.raises(RuntimeError, match="carries no device mesh"): + expert_parallel_mesh(model)