From 72b7fd0fb2bdbbc7dd32cffc4fff22ce86174594 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:22:10 +0000 Subject: [PATCH 01/21] Load the dense part sharded, and size the load warmup to each rank The expert-parallel load replicates the dense part on every rank, which does not fit when the dense part alone is bigger than a GPU. With dense="fsdp" in the plan, the dense part is now left on meta through from_pretrained, each decoder block is FSDP2-sharded around its experts over the whole world, and each rank DCP-loads only its shard. A dense part sharded within several expert-data-parallel replicas, which the export cannot write, is refused before loading, and the decision replicates it with a warning instead. The load-time allocator warmup is sized to each rank's share: expert params count 1/ep_degree, a dense part on meta nothing. This is the PR's original 37 commits in one, re-parented onto the expert-parallel loader's [2/3]: from_pretrained, with its warmup and dense-meta contexts, runs inside the loader's load guard, and the dense part's FSDP2 + DCP load runs after it, before the parallel states. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 477 +++++++++++++++++- .../torch/utils/test_expert_parallel_load.py | 210 +++++++- .../torch/utils/test_ep_dense_sharding.py | 412 +++++++++++++++ tests/unit/torch/utils/test_ep_load_warmup.py | 214 ++++++++ tests/unit/torch/utils/test_meta_dcp_load.py | 62 ++- .../unit/torch/utils/test_model_load_utils.py | 51 +- 6 files changed, 1403 insertions(+), 23 deletions(-) create mode 100644 tests/unit/torch/utils/test_ep_dense_sharding.py create mode 100644 tests/unit/torch/utils/test_ep_load_warmup.py diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 74ab4a9486d..0ebbc6620fe 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -18,6 +18,7 @@ import contextlib import copy import dataclasses +import inspect import logging import os import re @@ -62,7 +63,7 @@ is_initialized, set_distribution_plan, ) -from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap +from modelopt.torch.utils.fsdp2_utils import _off_dtype_params, fsdp2_wrap from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 from modelopt.torch.utils.moe import ( _EXPERTS_STYLES, @@ -889,6 +890,47 @@ def handed() -> DistributionPlan | None: return model +def _fsdp2_wrap_modules(model: "torch.nn.Module") -> list["torch.nn.Module"]: + """Return the transformer-block modules to FSDP2 ``fully_shard`` individually. + + Uses HuggingFace's own ``_no_split_modules`` -- the class names of the atomic blocks that must not be + split across devices (what ``device_map="auto"`` wraps at) -- falling back to modelopt's decoder-layer + auto-detection. + + For a VLM, this looks at the LANGUAGE backbone's boundaries, NOT the top-level VLM wrapper's + ``_no_split_modules`` (which for some VLMs -- e.g. Qwen2.5-VL -- ALSO lists the VISION block classes). + FSDP-sharding the vision blocks makes the calibration weight-access all-gather on those sharded vision + weights DEADLOCK; keeping the vision/audio towers REPLICATED (plain) -- the language-only calibration + never runs them and they are small enough to fit -- avoids that. This also covers VLMs whose top-level + wrapper exposes no ``_no_split_modules`` at all (e.g. Nemotron-Nano-VL). A VLM whose language + backbone shows no blocks is refused rather than wrapped by the wrapper's list. + """ + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + def _boundaries(root: "torch.nn.Module") -> "list[torch.nn.Module] | None": + ns = set(getattr(root, "_no_split_modules", None) or []) + if ns: + mods = [m for m in root.modules() if type(m).__name__ in ns] + if mods: + return mods + layers = LayerActivationCollector.get_decoder_layers(root) + return list(layers) if layers is not None else None + + # VLM -> wrap only the language backbone's blocks (vision/audio stay replicated); else the whole model. + from modelopt.torch.export.model_utils import get_language_model_from_vl + + lineage = get_language_model_from_vl(model) + root = lineage[-1] if lineage else model + mods = _boundaries(root) + if mods: + return mods + raise RuntimeError( + f"Could not determine FSDP2 wrap boundaries: {type(root).__name__} exposes no " + "`_no_split_modules` and modelopt's decoder-layer auto-detection did not match. Provide a " + "standard HF causal-LM / VLM checkpoint, or replicate the dense part (MODELOPT_FSDP_SIZE=1)." + ) + + def _map_ckpt_keys_to_model(model: "torch.nn.Module", ckpt_keys, model_keys) -> dict: """Map each checkpoint key to the model ``state_dict`` key it feeds. @@ -1490,6 +1532,18 @@ def _decide_distribution( f"{knob}={ep_size} must divide the routed experts {hf_config.model_type} shards " f"({num_experts})." ) + # Why the expert-parallel load cannot shard the dense part, if it cannot: the export cannot write + # it sharded within several expert-data-parallel replicas. + unshardable, remedy = None, "" + if ep_size > 1 and world_size // ep_size > 1: + unshardable = ( + f"within each of its {world_size // ep_size} expert-data-parallel replicas, a layout " + "the export cannot write yet" + ) + remedy = ( + " Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, or " + "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." + ) if dense == "auto": # A rank holds a whole copy of what expert parallelism does not shard, and its 1/ep_size # share of the routed experts it does. @@ -1502,7 +1556,17 @@ def _decide_distribution( ) memory = _gpu_memory() too_big = memory is not None and held > _REPLICATED_SHARE * memory - dense = "fsdp" if too_big else "replicate" + if too_big and unshardable: + warn_rank_0( + f"{hf_config.model_type}: a rank holding the dense part whole would take more than " + f"{_REPLICATED_SHARE:.0%} of a GPU, but it cannot be sharded {unshardable}; " + f"replicating it, which may run out of memory.{remedy}" + ) + dense = "fsdp" if too_big and not unshardable else "replicate" + elif dense == "fsdp" and unshardable: + raise NotImplementedError( + f"The dense part of {hf_config.model_type} cannot be sharded {unshardable}.{remedy}" + ) return DistributionPlan(world_size, ep_size, dense) @@ -1602,6 +1666,160 @@ def checkpoint_read() -> bool: return plan, hf_config +def _make_expert_param_matcher(ep_plan: dict[str, Any]): + """Build ``is_expert(param_name) -> bool`` from the model's ``base_model_ep_plan``. + + That plan is what drives EP sharding, so it replaces a fragile ``"expert" in name`` heuristic. + + A param is EP-sharded (only its ``1/ep_degree`` shard resides on this rank) iff its name, with layer + indices wildcarded exactly as HuggingFace does (``replace_layer_number_by_wildcard``) and any trailing + ``.weight`` / ``.bias`` stripped, contains an expert plan key as a run of whole dot-separated + segments. That one test covers a key equal to the name, a suffix of it (the base-relative plan keys + omit the ``model.`` / ``language_model.model.`` prefixes) and a container the name nests under, so + both per-parameter shard keys (``...experts.gate_up_proj: grouped_gemm``) and experts-container + keys (``...experts: torch_a2a_experts`` / ``nemotron_experts_ep``) match. Router entries + (``ep_router`` / ``torch_a2a_router``) are REPLICATED, not sharded, so they're excluded. If the plan is + empty the matcher is always ``False`` (-> no EP credit -> stock warmup sizing). + """ + from transformers.integrations.tensor_parallel import replace_layer_number_by_wildcard + + _router_styles = {"ep_router", "torch_a2a_router"} + keys = [k for k, v in (ep_plan or {}).items() if v not in _router_styles] + + def is_expert(pname: str) -> bool: + generic = replace_layer_number_by_wildcard(pname) + for suffix in (".weight", ".bias"): + if generic.endswith(suffix): + generic = generic[: -len(suffix)] + break + padded = "." + generic + "." + return any(("." + k + ".") in padded for k in keys) + + return is_expert + + +@contextlib.contextmanager +def _ep_aware_caching_allocator_warmup( + ep_degree: int, ep_plan: dict[str, Any], dense_on_meta: bool = False +): + """Make transformers' ``caching_allocator_warmup`` EP-aware for the duration of ``from_pretrained``. + + transformers sizes the load-time allocator warmup via ``get_total_byte_count``, which credits the + ``// world_size`` sharding division only to params matched in ``model._tp_plan``. Under the EP load + (``device_mesh`` + ``enable_expert_parallel``; sharding driven by ``base_model_ep_plan``) no expert + param is in ``_tp_plan``, so every one of them -- experts are ~95% of a large MoE -- is counted at + FULL global size. The + warmup then pre-allocates one ``torch.empty`` clamped to ``device_total - 1.2 GiB`` (~the whole GPU) + on every rank; torch keeps that as reserved cache, which later starves the NCCL a2a comm buffers. + + This temporarily replaces ``get_total_byte_count`` with a version that additionally credits EP-sharded + expert params their per-rank shard (``// ep_degree``), identifying them from ``ep_plan`` (the model's + ``base_model_ep_plan``), sizing the warmup to the real per-rank footprint (e.g. ~91 GiB instead of + ~275 GiB for Nemotron-3-Ultra-550B EP32) -- keeping the warmup's load-speed benefit without the + ~device-max transient. + + Always on (no disable knob): it only *sizes* an allocator pre-reservation, never model correctness; it + falls back to stock sizing on any incompatibility (a hook whose leading parameters differ is left in + place, and an error while sizing returns the stock number); and a mis-estimate only nudges load + speed (an over-credit under-sizes -> the allocator just grows during load; an under-credit sizes up to + the same device-max clamp as stock) -- so it can never OOM worse than, or load differently from, stock. + No-op when ``ep_degree <= 1``. + """ + if ep_degree <= 1: + yield + return + try: + import transformers.modeling_utils as _mu + except Exception: + yield + return + orig = getattr(_mu, "get_total_byte_count", None) + if orig is None: + yield + return + try: + params = list(inspect.signature(orig).parameters) + except (TypeError, ValueError): + params = [] + # A release whose hook takes other leading parameters calls it in a way this replacement cannot + # know, so it keeps the stock sizing rather than fail the load at transformers' call site. + if params[:3] != ["model", "accelerator_device_map", "hf_quantizer"]: + yield + return + + is_expert = _make_expert_param_matcher(ep_plan) + + def _ep_aware(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs): + try: + from collections import defaultdict + + try: # where current transformers keeps it (5.18); older releases, under integrations + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan + except ImportError: + try: + from transformers.integrations.tensor_parallel import _get_parameter_tp_plan + except ImportError: + _get_parameter_tp_plan = None + ws = ( + torch.distributed.get_world_size() + if torch.distributed.is_available() and torch.distributed.is_initialized() + else 1 + ) + tp_plan = (getattr(model, "_tp_plan", None) or {}) if ws > 1 else {} + tied = set(getattr(model, "all_tied_weights_keys", {}) or {}) + total: dict = defaultdict(int) + for pname, dev in accelerator_device_map.items(): + if pname in tied: + continue + try: + p = model.get_parameter_or_buffer(pname) + except Exception: + continue + dsize = ( + hf_quantizer.param_element_size(model, pname, p) + if hf_quantizer is not None + else p.element_size() + ) + nb = p.numel() * dsize + if is_expert(pname): + nb //= ep_degree # EP-sharded across the ep submesh: only the local shard resides here + elif ( + tp_plan + and _get_parameter_tp_plan is not None + and _get_parameter_tp_plan(pname, tp_plan, is_weight=True) is not None + ): + # Placed by transformers' TP branch, which never asks where the dense part goes, so + # it is on the device at its share even when the rest of the dense part is on meta. + # That branch shards over from_pretrained's device_mesh, the ep submesh. + nb //= ep_degree + elif dense_on_meta and pname.endswith((".weight", ".bias")): + # A sharded dense part: _ep_dense_meta_load forces every dense .weight/.bias to + # materialize on META during from_pretrained (they are FSDP2-sharded + DCP-loaded + # afterwards), so they occupy NO device memory in the window this warmup sizes. + # Counting them at full global size is what pushed Kimi-K3 (EP64, dense=fsdp) to a + # 185.8 GiB estimate on a 184.31 GiB card -- over the device, so the warmup fell back + # to the device-max clamp it exists to avoid. Persistent buffers are NOT forced to + # meta by that context, so they keep their full credit here. + continue + total[dev] += nb + gib = 1024**3 + sizes = ", ".join(f"{d}: {b / gib:.1f} GiB" for d, b in total.items()) + print_rank_0( + f"[MEM] EP-aware warmup sizing (experts // ep_degree={ep_degree}" + f"{', dense on meta -> uncounted' if dense_on_meta else ''}) -> {sizes} " + f"(was full-model ~device-max; EP-aware plan-driven sizing)" + ) + return total + except Exception: + return orig(model, accelerator_device_map, hf_quantizer, *args, **kwargs) + + try: + _mu.get_total_byte_count = _ep_aware + yield + finally: + _mu.get_total_byte_count = orig + + 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. @@ -1617,6 +1835,102 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st ) +def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int: + """FSDP2-shard the model's dense modules, leaving the routed experts to EP. + + FSDP2 ``fully_shard`` the model's DENSE (non-expert) modules over ``fsdp_mesh``, leaving the routed + experts (already EP-sharded on ``ep`` + the a2a transport) untouched -- the "FSDP2 dense + EP MoE" + hybrid. ``fsdp_mesh`` is the FULL world (dense shards over every rank, the same ranks the experts are + EP-sharded over). + + Wraps at the DECODER-LAYER boundary (:func:`_fsdp2_wrap_modules`: the language backbone's HF + ``_no_split_modules`` blocks, else the decoder-layer detection the pure-FSDP2 path uses) and passes + the routed-expert params as ``ignored_params`` so FSDP shards + the layer's dense params (attention, layernorms, router) while leaving the experts EP-sharded on + the ``ep`` mesh + a2a transport. As there, a param off the model's dtype (an fp32 router) stays + replicated: FSDP2 shards one dtype per group. Wrapping + the whole layer (not ``self_attn`` or individual leaves) is essential: ``fully_shard`` swaps a wrapped + module's class (``Foo`` -> ``FSDPFoo``), and modelopt inspects classes by name/source -- wrapping + ``self_attn`` breaks ``mtq``'s KV-cache AST registration (``inspect.getsource`` on the synthetic + class), and wrapping the leaf ``QuantLinear``s breaks the export's name-based ``is_quantlinear``. + Wrapping only the LAYER leaves ``self_attn`` and the q/k/v/o ``QuantLinear``s PLAIN, while + ``_get_enclosing_fsdp_module`` still finds the layer FSDP wrapper for the weight-access/writeback + contexts. The root (embed/lm_head/final norm) is intentionally NOT wrapped: those stay replicated + (and unswizzled, so their quantization-exclude still applies), matching the dense-replicated baseline. + + Returns the number of layers wrapped. + """ + from torch.distributed.fsdp import fully_shard + + # Every routed-expert param is EP-sharded on the ``ep`` mesh; FSDP must IGNORE them so it never tries + # to reshard them onto ``fsdp_mesh``. Ignored params not present in a given layer are simply no-ops. + ignored_params = _off_dtype_params(model) + for em in expert_modules: + ignored_params.update(em.parameters(recurse=True)) + + from torch.distributed.tensor import DTensor, Replicate + + ignored_param_ids = {id(p) for p in ignored_params} # id-based: Tensor.__eq__ is elementwise + + # Wrap only the OUTERMOST _no_split_modules. Some models (e.g. Kimi's MoonViT VLM) list BOTH a + # container (the vision model) and its blocks in ``_no_split_modules``; ``fully_shard`` on the parent + # already shards the children, so wrapping a child afterward hits ``DeviceMesh._concatenate`` ("Cannot + # concatenate overlapping meshes"). Drop any wrap module nested under another. + wrap_modules = _fsdp2_wrap_modules(model) + wrap_ids = {id(m) for m in wrap_modules} + nested_ids: set = set() + for m in wrap_modules: + for c in m.modules(): + if id(c) != id(m) and id(c) in wrap_ids: + nested_ids.add(id(c)) + wrap_modules = [m for m in wrap_modules if id(m) not in nested_ids] + + n = 0 + for layer in wrap_modules: + layer_ignored = set(ignored_params) # experts + any pre-sharded dense (below) + has_shardable = False + for name, p in list(layer.named_parameters(recurse=True)): + if id(p) in ignored_param_ids: + continue + if isinstance(p.data, DTensor): + if all(isinstance(pl, Replicate) for pl in p.data.placements): + # Replicate -> plain (a free view, no gather), re-registered as a plain Parameter: + # assigning .data would keep the DTensor parameter, which fully_shard then tries + # to place on both its own mesh and the FSDP one. + owner_name, _, leaf = name.rpartition(".") + owner = layer.get_submodule(owner_name) if owner_name else layer + setattr( + owner, + leaf, + nn.Parameter(p.data.to_local(), requires_grad=p.requires_grad), + ) + else: + layer_ignored.add( + p + ) # already sharded by the loader (TP) -> leave it distributed + continue + has_shardable = True + if not has_shardable: + continue # nothing left to FSDP-shard in this layer (all expert / already-sharded) + fully_shard( + layer, + mesh=fsdp_mesh, + reshard_after_forward=True, + ignored_params=layer_ignored, + ) + n += 1 + if not n: + # The dense params were forced to meta for this sharding; with nothing sharded, the load would + # materialize the whole dense part on every rank, which the plan shards it to avoid. + raise RuntimeError( + f"Sharding the dense part found {len(wrap_modules)} block(s) but no dense parameter to " + "shard in any of them (each is a routed expert's, already sharded, or off the model's " + "dominant dtype), so every rank would hold the whole dense part. Replicate the dense " + "part instead (MODELOPT_FSDP_SIZE=1)." + ) + return n + + # The checkpoint formats load_and_prepare_ep dequantizes as it loads, beside unquantized ones (None). # Not MXFP4: transformers dequantizes it into GPT-OSS's transposed expert layout, whatever the model's. _EP_SOURCE_FORMATS: frozenset[str | None] = frozenset({None}) @@ -1673,6 +1987,99 @@ def _ep_mesh(*args, **kwargs): mixin.initialize_distributed_mesh = original +@contextlib.contextmanager +def _ep_dense_meta_load(is_expert): + """Keep the dense part off the GPU during ``from_pretrained``, for the EP hybrid. + + Force every non-expert weight/bias to materialize on ``meta`` (i.e. NOT on the GPU) by making + transformers' per-param ``get_device`` return ``"meta"`` for dense keys. + + transformers 5.x's ``convert_and_load_state_dict_in_model`` places EXPERT params via the tp / + ``distributed_operation`` branch (which fuses + EP-shards them and never consults ``get_device``), and + every OTHER (dense) param via ``param_device = get_device(device_map, key)`` -> ``spawn_materialize``. + Returning ``"meta"`` there leaves dense params unmaterialized while experts still EP-shard normally. The + skipped dense params are then FSDP2-sharded + DCP-loaded (:func:`_load_ep_dense_distributed`), so dense is + never replicated per rank. Only ``.weight``/``.bias`` are forced meta -- persistent buffers (small) load + normally. Patches ``get_device`` on both ``transformers.modeling_utils`` and the module that actually + defines the loader (it's looked up as a module global there). + """ + import sys as _sys + + import transformers.modeling_utils as _mu + + loader = getattr(_mu, "convert_and_load_state_dict_in_model", None) + mods = {_mu} + loader_module = getattr(loader, "__module__", None) + dm = _sys.modules.get(loader_module) if loader_module else None + if dm is not None: + mods.add(dm) + targets = [m for m in mods if hasattr(m, "get_device")] + if not targets: + import transformers + + raise RuntimeError( + "Sharding the dense part keeps it off the device during from_pretrained through " + f"transformers' get_device, which none of {sorted(m.__name__ for m in mods)} has on " + f"transformers {transformers.__version__}. Set MODELOPT_FSDP_SIZE=1 if the dense part " + "fits one GPU." + ) + + def _dense_weight(key) -> bool: + k = str(key) + return k.endswith((".weight", ".bias")) and not is_expert(k) + + saved = [(m, m.get_device) for m in targets] + + def _make(orig): + def patched(device_map, key, *args, **kwargs): + if _dense_weight(key): + return "meta" + return orig(device_map, key, *args, **kwargs) + + return patched + + try: + for m, orig in saved: + m.get_device = _make(orig) + yield + finally: + for m, orig in saved: + m.get_device = orig + + +def _load_ep_dense_distributed(model, fsdp_mesh, local_dir, is_expert, expert_modules): + """FSDP2-shard the meta dense part and DCP-load each rank's shard, for the EP hybrid. + + Each rank's shard comes straight from the checkpoint, so the dense part is NEVER + materialized replicated. Experts are already EP-sharded (loaded by ``from_pretrained``) and left untouched. Returns + the number of dense layers sharded, like :func:`_fsdp2_shard_dense`. + + Requires the dense ``.weight``/``.bias`` params to be on ``meta`` on entry (see + :func:`_ep_dense_meta_load`). Delegates the materialize + DCP load to :func:`_meta_dcp_load`, fed the + partially-meta model (dense meta, experts already real) and the dense param names. + """ + from modelopt.torch.utils.distributed import local_rank + + device = torch.device(f"cuda:{local_rank()}") + # FSDP2-shard the dense modules. Their params are meta plain tensors -> fully_shard makes them meta Shard + # DTensors; experts are passed as ignored_params and stay EP-sharded. + n_dense = _fsdp2_shard_dense(model, fsdp_mesh, expert_modules) + # Restrict the shared loader to the dense params (experts already materialized by from_pretrained must not + # be touched). It selectively materializes only the meta dense modules (experts skipped), preserves + # buffers across to_empty, and DCP-loads each rank's dense shard. + dense_names = { + n + for n, _ in model.named_parameters() + if n.endswith((".weight", ".bias")) and not is_expert(n) + } + print_rank_0( + f"[ep] dense-distributed load: FSDP2-sharded {n_dense} dense module(s); " + f"meta+DCP-loading {len(dense_names)} dense param(s) (experts left EP-sharded)" + ) + _meta_dcp_load(model, local_dir, device, load_names=dense_names) + return n_dense + + def load_and_prepare_ep( ckpt_path: str, plan: DistributionPlan, @@ -1696,6 +2103,9 @@ def load_and_prepare_ep( the load that rests on ``from_pretrained`` running no collective, as transformers' expert-parallel load does not on 5.7 through 5.18. + With ``plan.dense == "fsdp"`` the dense part is FSDP2-sharded over the ``ep`` group instead, + each rank loading only its shard: for a model whose dense part does not fit one GPU. + 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 @@ -1704,8 +2114,14 @@ def load_and_prepare_ep( """ 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.") + if plan.is_fsdp_dense and plan.edp > 1: + # Refused before the load: the export, after the whole calibration, is what cannot write it. + raise NotImplementedError( + f"{plan!r} FSDP2-shards the dense part within each of {plan.edp} expert-data-parallel " + "replicas, a layout the export cannot write yet. Set MODELOPT_EP_DEGREE to the world " + "size if it divides the expert count, or MODELOPT_FSDP_SIZE=1 if the dense part fits " + "one GPU." + ) # Bringing the group up is the launcher's job; this load needs it live and device-bound. if not dist.is_initialized(): raise RuntimeError( @@ -1889,13 +2305,32 @@ 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 fails every rank here, before the warmup's collectives: + # Size the load-time caching-allocator warmup to the real per-rank (EP-sharded) footprint instead of + # the full model -- otherwise transformers pre-allocates ~device-max on every rank and torch holds it + # as reserved cache, starving the NCCL a2a buffers. Expert params are identified from `swapped_plan` + # (the actual EP plan). Always on: falls back to stock sizing on any issue and a mis-size only nudges + # load speed, never correctness. (See _ep_aware_caching_allocator_warmup.) + is_expert_fn = _make_expert_param_matcher(swapped_plan) + # The plan places the DENSE (non-expert) modules; the experts are always EP-sharded by from_pretrained. + # A replicated dense part loads with them, on every rank. A sharded one is left on meta during + # from_pretrained (get_device->meta via _ep_dense_meta_load), then FSDP2-sharded + DCP-loaded, so it + # is never replicated anywhere. + _dense_ctx = ( + _ep_dense_meta_load(is_expert_fn) if plan.is_fsdp_dense else contextlib.nullcontext() + ) + # A rank failing in the load fails every rank here, before the collectives below: # 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 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): + with ( + _ep_aware_caching_allocator_warmup( + ep_degree, swapped_plan, dense_on_meta=plan.is_fsdp_dense + ), + _dense_ctx, + _experts_sharded_over(ep_submesh), + ): model = loader.from_pretrained(ckpt_path, **from_pretrained_kwargs) model.eval() # The config the model keeps goes back to the checkpoint's plan, as an export writes it out @@ -1907,6 +2342,15 @@ 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 + if plan.is_fsdp_dense and dist.is_master(): + _nmeta = [n for n, p in model.named_parameters() if p.is_meta] + _dense_meta = [n for n in _nmeta if not is_expert_fn(n)] + print( + f"[ep] dense-meta load: {len(_nmeta)} meta param(s) after from_pretrained " + f"({len(_dense_meta)} dense, {len(_nmeta) - len(_dense_meta)} expert); experts materialized. " + f"e.g. dense-meta {_dense_meta[:3]}", + flush=True, + ) # 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(): @@ -1935,6 +2379,27 @@ def gpu() -> tuple[bool, bool, str | None]: "ranks holding different experts. Quantize the model without expert parallelism." ) + # "FSDP2 dense + EP MoE" hybrid: experts are EP-sharded over ``ep`` and the DENSE (non-expert) modules + # are FSDP2-sharded over the same ranks, the whole world: a dense part sharded within several + # expert-parallel groups is refused above. Dense all-gather then expert a2a run sequentially within a + # forward (same ranks), so there is no concurrent-collective deadlock. Weight amax is computed on the + # FSDP-gathered full weight (identical per rank); activation amax syncs over the whole world via the + # dense ``dp`` group registered below. + # + # A replicated dense part stays data-parallel and needs no post-load step. + if plan.is_fsdp_dense: + # Fully-distributed dense: dense was left on meta during from_pretrained; FSDP2-shard it over the + # full-world mesh and DCP-load each rank's shard directly (never replicated anywhere). + dense_fsdp_mesh = dense_mesh["dp"] + n_dense = _load_ep_dense_distributed( + model, dense_fsdp_mesh, ckpt_path, is_expert_fn, expert_modules.values() + ) + print_rank_0( + f"[ep] dense sharded: FSDP2-sharded {n_dense} dense module(s) over the whole world, " + f"mesh {dense_fsdp_mesh.mesh_dim_names}={tuple(dense_fsdp_mesh.shape)}; experts " + "EP-sharded on ep" + ) + # 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. diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 849305f4f6a..f9f50f98ea4 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -31,8 +31,11 @@ from torch.distributed.tensor import DTensor from transformers import AutoConfig, AutoModelForCausalLM, Qwen3MoeConfig +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.utils.distributed import ( DistributionPlan, + distributed_forward_loop, distribution_plan, set_distribution_plan, ) @@ -63,10 +66,20 @@ 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): +# (dense part sharded, share of the GPUs the experts span). A dense part FSDP2-sharded within several +# expert-parallel groups (sharded with half the GPUs per group) is refused by the loader: the export +# cannot write it. +_LAYOUTS = [ + pytest.param(False, 1, id="dense-replicated-ep-world"), + pytest.param(False, 2, id="dense-replicated-ep-half"), + pytest.param(True, 1, id="dense-sharded-ep-world"), +] + + +def _sharded_model_matches_the_whole_one(rank, size, ckpt_dir, fsdp_dense, 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) + plan = DistributionPlan(size, ep_size=ep_degree, dense="fsdp" if fsdp_dense else "replicate") 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 @@ -101,13 +114,198 @@ def _sharded_model_matches_the_whole_one(rank, size, ckpt_dir, ep_degree): 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"]) +@pytest.mark.parametrize(("fsdp_dense", "ep_share"), _LAYOUTS) def test_expert_parallel_model_matches_the_whole_model( - need_2_gpus, dist_workers, tmp_path, ep_share + need_2_gpus, dist_workers, tmp_path, fsdp_dense, ep_share ): - """With the experts sharded over every rank, or over half of them and replicated across the halves.""" + """With the dense part replicated, or FSDP2-sharded and loaded shard by shard through DCP; and + 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) + dist_workers.run(_sharded_model_matches_the_whole_one, str(tmp_path), fsdp_dense, ep_degree) + + +def _batches(size): + """Every rank's batch, the same on all ranks: each trains on its own, the reference on all.""" + return [ + torch.randint(0, 128, (2, 16), generator=torch.Generator().manual_seed(r)) + for r in range(size) + ] + + +def _whole_model_grads_per_batch(whole, batches, device): + """The whole model's gradient on each batch alone, as ``{parameter name: gradient}``.""" + grads = [] + for batch in batches: + whole.zero_grad(set_to_none=True) + batch = batch.to(device) + whole(batch, labels=batch).loss.backward() + grads.append({n: p.grad.detach().clone() for n, p in whole.named_parameters()}) + return grads + + +def _fail_together(problems): + """Raise on every rank if any rank found a problem, so that no rank walks on alone into a collective + the failed ones never reach (which would hang until the test timeout instead of failing).""" + gathered = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(gathered, problems) + found = [ + f"rank {r}: {problem}" + for r, rank_problems in enumerate(gathered) + for problem in rank_problems + ] + assert not found, "\n".join(found[:20]) + + +def _relative_error(actual, expected): + # bf16 gradients summed in a different order on each side: compare by relative norm. + return ( + (actual.float() - expected.float()).norm() / expected.float().norm().clamp_min(1e-12) + ).item() + + +def _assert_grads_match_the_whole_model(model, per_batch, rank, size, ep_degree): + """Routed experts collect the gradient of every token routed to them, their EP group's batches; + FSDP2-sharded dense parameters hold the gradient averaged over every rank; replicated dense + parameters see only this rank's batch.""" + ep_rank = rank % ep_degree + group = range(rank - ep_rank, rank - ep_rank + ep_degree) + per_rank = NUM_EXPERTS // ep_degree + owned = slice(ep_rank * per_rank, (ep_rank + 1) * per_rank) + problems, names = [], set() + for name, param in model.named_parameters(): + names.add(name) + if param.grad is None: + problems.append(f"{name} received no gradient") + continue + if ".mlp.experts." in name: + actual = _local(param.grad) + expected = sum(per_batch[r][name] for r in group)[owned] + elif isinstance(param.grad, DTensor): + actual = param.grad.full_tensor() + expected = sum(per_batch[r][name] for r in range(size)) / size + else: + actual = param.grad + expected = per_batch[rank][name] + error = _relative_error(actual, expected) + if not error < 3e-2: + problems.append(f"{name}: gradient off by {error:.3g} (relative norm)") + if names != set(per_batch[0]): + problems.append( + f"parameters differ from the whole model's: {sorted(names ^ set(per_batch[0]))[:5]}" + ) + _fail_together(problems) + + +def _sharded_model_backward_matches_the_whole_one(rank, size, ckpt_dir, fsdp_dense, ep_degree): + device = torch.device(f"cuda:{rank}") + config = AutoConfig.from_pretrained(ckpt_dir) + plan = DistributionPlan(size, ep_size=ep_degree, dense="fsdp" if fsdp_dense else "replicate") + model = load_and_prepare_ep(ckpt_dir, plan, hf_config=config, freeze=False) + set_distribution_plan( + model, plan + ) # loaded directly, so recorded here, as load_pretrained would + whole = AutoModelForCausalLM.from_pretrained(ckpt_dir, dtype=torch.bfloat16).to(device) + + batches = _batches(size) + mine = batches[rank].to(device) + model(mine, labels=mine).loss.backward() + _assert_grads_match_the_whole_model( + model, _whole_model_grads_per_batch(whole, batches, device), rank, size, ep_degree + ) + + +@pytest.mark.parametrize(("fsdp_dense", "ep_share"), _LAYOUTS) +def test_expert_parallel_backward_matches_the_whole_model( + need_2_gpus, dist_workers, tmp_path, fsdp_dense, ep_share +): + """Loaded trainable, the sharded model's gradients are the whole model's: tokens' gradients travel + back through the all-to-all to the experts that saw them, and to the router that chose them.""" + 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_backward_matches_the_whole_one, str(tmp_path), fsdp_dense, ep_degree + ) + + +def _calibrated_quantizers(model): + return { + name: module + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and module.is_enabled and module.amax is not None + } + + +def _quantized_sharded_model_trains_like_the_whole_one( + rank, size, ckpt_dir, fsdp_dense, fmt, ep_degree +): + cfg = {"fp8": mtq.FP8_DEFAULT_CFG, "nvfp4": mtq.NVFP4_DEFAULT_CFG}[fmt] + device = torch.device(f"cuda:{rank}") + config = AutoConfig.from_pretrained(ckpt_dir) + plan = DistributionPlan(size, ep_size=ep_degree, dense="fsdp" if fsdp_dense else "replicate") + model = load_and_prepare_ep(ckpt_dir, plan, hf_config=config, freeze=False) + set_distribution_plan( + model, plan + ) # loaded directly, so recorded here, as load_pretrained would + whole = AutoModelForCausalLM.from_pretrained(ckpt_dir, dtype=torch.bfloat16).to(device) + + # Calibrate as hf_ptq does: each rank on its own batch, the reference on all of them. + batches = _batches(size) + mine = batches[rank].to(device) + forward_loop = distributed_forward_loop( + model, [{"input_ids": mine}], device=device, ep_mesh=expert_parallel_mesh(model) + ) + mtq.quantize(model, cfg, forward_loop) + mtq.quantize(whole, cfg, lambda m: [m(b.to(device)) for b in batches]) + + # Expert-parallel calibration reaches the whole model's scales: weight scales from this rank's + # experts, activation scales synced over the ranks (equal to bf16 rounding of the activations). + per_rank = NUM_EXPERTS // ep_degree + ep_rank = rank % ep_degree + whole_quantizers = _calibrated_quantizers(whole) + sharded_quantizers = _calibrated_quantizers(model) + problems = [] if sharded_quantizers else ["no quantizer was calibrated"] + for name, quantizer in sharded_quantizers.items(): + if "_weight_quantizers." in name: + prefix, index = name.rsplit(".", 1) + name = f"{prefix}.{ep_rank * per_rank + int(index)}" + reference = whole_quantizers.get(name) + if reference is None: + problems.append(f"{name} is calibrated here but not in the whole model") + continue + error = _relative_error(quantizer.amax, reference.amax) + if not error < 2e-2: + problems.append(f"{name}: amax off by {error:.3g} (relative)") + # Train on identical scales, so that the gradients compare the transport and not calibration. + quantizer.amax = reference.amax + _fail_together(problems) + + model(mine, labels=mine).loss.backward() + _assert_grads_match_the_whole_model( + model, _whole_model_grads_per_batch(whole, batches, device), rank, size, ep_degree + ) + + +@pytest.mark.parametrize("fmt", ["fp8", "nvfp4"]) +@pytest.mark.parametrize(("fsdp_dense", "ep_share"), _LAYOUTS) +def test_quantized_expert_parallel_model_trains_like_the_whole_model( + need_2_gpus, dist_workers, tmp_path, fsdp_dense, fmt, ep_share +): + """QAT under expert parallelism: quantized and calibrated expert-parallel -- experts over every + rank or half of them, dense replicated or FSDP2-sharded -- the model's scales and its gradients + through the fake quantizers match the whole model calibrated on every rank's data.""" + 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( + _quantized_sharded_model_trains_like_the_whole_one, + str(tmp_path), + fsdp_dense, + fmt, + ep_degree, + ) diff --git a/tests/unit/torch/utils/test_ep_dense_sharding.py b/tests/unit/torch/utils/test_ep_dense_sharding.py new file mode 100644 index 00000000000..28d2af01d15 --- /dev/null +++ b/tests/unit/torch/utils/test_ep_dense_sharding.py @@ -0,0 +1,412 @@ +# 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. + +"""How the expert-parallel loader places the dense (non-expert) part of an MoE. + +``from_pretrained`` always EP-shards the experts. ``_make_expert_param_matcher`` tells their params +from the dense ones, the distribution plan says whether the dense part is replicated or sharded, +``_ep_dense_meta_load`` keeps the dense weights off the device during ``from_pretrained`` when they +are sharded, and ``_fsdp2_shard_dense`` / +``_load_ep_dense_distributed`` FSDP2-shard them around the experts. The FSDP2 cases run on a +one-rank gloo process group on CPU. +""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import FSDPModule +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor + +pytest.importorskip("transformers") + +import transformers.core_model_loading as core_model_loading +import transformers.modeling_utils as modeling_utils +from _test_utils.torch.transformers_models import get_tiny_llama +from transformers import AutoModelForCausalLM + +import modelopt.torch.utils.distributed as mto_dist +import modelopt.torch.utils.moe as moe +import modelopt.torch.utils.plugins.model_load_utils as model_load_utils +from modelopt.torch.utils.distributed import DistributionPlan +from modelopt.torch.utils.plugins.model_load_utils import ( + _ep_aware_caching_allocator_warmup, + _ep_dense_meta_load, + _fsdp2_shard_dense, + _load_ep_dense_distributed, + _make_expert_param_matcher, + load_and_prepare_ep, +) + +# ------------------------------------------------------------------ _make_expert_param_matcher + +# The plan load_and_prepare_ep shards with after the a2a swap: base-relative keys (no ``model.``). +_A2A_PLAN = { + "layers.*.mlp.gate": "torch_a2a_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "torch_a2a_experts", +} + + +# Only the per-parameter keys: a container key would match the same names on its own. +_PARAM_KEYS = {k: v for k, v in _A2A_PLAN.items() if v == "grouped_gemm"} + + +@pytest.mark.parametrize( + ("plan", "name", "expected"), + [ + pytest.param(_PARAM_KEYS, "layers.0.mlp.experts.gate_up_proj", True, id="key-unprefixed"), + pytest.param( + _PARAM_KEYS, "model.language_model.layers.3.mlp.experts.down_proj", True, id="vlm" + ), + pytest.param( + _PARAM_KEYS, + "language_model.model.layers.3.mlp.experts.gate_up_proj", + True, + id="vlm-wrapper", + ), + # NemotronH tags only the container; every param under it is an expert's. + pytest.param( + {"layers.*.mixer.experts": "torch_a2a_experts"}, + "backbone.layers.1.mixer.experts.3.down_proj.weight", + True, + id="under-a-container", + ), + pytest.param( + _A2A_PLAN, "model.layers.0.mlp.shared_experts.gate_proj.weight", False, id="shared" + ), + pytest.param( + {"layers.*.mlp.gate": "ep_router"}, "model.layers.0.mlp.gate.weight", False, id="router" + ), + # No EP credit then: the load-time warmup falls back to stock sizing. + pytest.param(None, "model.layers.0.mlp.experts.gate_up_proj", False, id="no-plan"), + ], +) +def test_expert_param_matcher(plan, name, expected): + """A param is an expert's when the plan says so, under any prefix the base-relative keys omit; + a router is replicated whatever its style. The warmup sizing tests classify a whole MoE's + experts, router, attention, embedding and head with the same matcher.""" + assert _make_expert_param_matcher(plan)(name) is expected + + +# ------------------------------------------------------------------ a block-fp8 checkpoint + + +def test_a_block_fp8_checkpoint_is_refused_as_a_format_whatever_its_dense_layout( + monkeypatch, gloo_world_of_one +): + """This load does not dequantize block-fp8 checkpoints at all, so that is the refusal a plan + sharding their dense part meets, not one advising to replicate it.""" + monkeypatch.setattr("modelopt.torch.utils.distributed.is_initialized", lambda: True) + monkeypatch.setattr("modelopt.torch.utils.distributed.size", lambda group=None: 2) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + config = SimpleNamespace(quantization_config={"quant_method": "fp8"}) + + with pytest.raises(NotImplementedError, match="cannot dequantize a 'fp8' checkpoint"): + load_and_prepare_ep("unused", DistributionPlan(2, 2, "fsdp"), hf_config=config) + + +# ------------------------------------------------------------------------- _ep_dense_meta_load + + +def _is_expert(name): + return ".experts." in name + + +def _get_device_lookups(): + """The modules whose ``get_device`` global the state-dict loaders call.""" + return modeling_utils, core_model_loading + + +def test_dense_weights_resolve_to_meta_and_the_rest_to_the_stock_device(monkeypatch): + """Only a dense ``.weight`` / ``.bias`` is sent to meta. An expert param or a buffer goes to + the stock lookup, with the caller's arguments, in both modules the loaders read it from.""" + stock_calls = [] + + def stock_get_device(device_map, key, *args, **kwargs): + stock_calls.append((key, args, kwargs)) + return device_map[""] + + lookups = _get_device_lookups() + for module in lookups: + monkeypatch.setattr(module, "get_device", stock_get_device) + device_map = {"": "cuda:0"} + expert = "model.layers.0.mlp.experts.0.gate_proj.weight" + buffer = "model.layers.0.mlp.gate.e_score_correction_bias" + + with _ep_dense_meta_load(_is_expert): + for module in lookups: + assert module.get_device(device_map, "model.layers.0.self_attn.q_proj.weight") == "meta" + assert module.get_device(device_map, "model.layers.0.self_attn.q_proj.bias") == "meta" + assert module.get_device(device_map, expert, valid_torch_device=True) == "cuda:0" + assert module.get_device(device_map, buffer) == "cuda:0" + + assert stock_calls == [(expert, (), {"valid_torch_device": True}), (buffer, (), {})] * 2 + + +def test_the_dense_load_refuses_a_transformers_without_its_get_device_hook(monkeypatch): + """Without the hook the dense part would load whole on every rank, the memory sharding it saves.""" + for module in _get_device_lookups(): + monkeypatch.delattr(module, "get_device", raising=False) + + with pytest.raises(RuntimeError, match="get_device"), _ep_dense_meta_load(_is_expert): + pass + + +def _get_device_patch(): + return _ep_dense_meta_load(_is_expert), [(m, "get_device") for m in _get_device_lookups()] + + +def _warmup_sizing_patch(): + if not hasattr(modeling_utils, "get_total_byte_count"): + pytest.skip("transformers has no get_total_byte_count to size") + sizing = _ep_aware_caching_allocator_warmup(4, _A2A_PLAN) + return sizing, [(modeling_utils, "get_total_byte_count")] + + +@pytest.mark.parametrize( + "patch", [_get_device_patch, _warmup_sizing_patch], ids=lambda f: f.__name__ +) +def test_a_load_time_patch_is_undone_even_when_the_load_raises(patch): + """``from_pretrained`` reads these transformers globals; a failed load leaves stock ones.""" + context, patched = patch() + stock = [getattr(module, attr) for module, attr in patched] + + with pytest.raises(RuntimeError, match="load failed"), context: + assert all(getattr(m, a) is not f for (m, a), f in zip(patched, stock)) + raise RuntimeError("load failed") + + assert all(getattr(m, a) is f for (m, a), f in zip(patched, stock)) + + +def test_from_pretrained_leaves_the_dense_weights_on_meta_and_loads_the_rest(tmp_path): + """Through a real CPU ``from_pretrained``: every dense weight stays on meta for the FSDP2 + + DCP load to fill, while the MLPs (standing in for experts) and the rotary buffers load.""" + src = get_tiny_llama() + src.save_pretrained(tmp_path) + + with _ep_dense_meta_load(lambda name: ".mlp." in name): + model = AutoModelForCausalLM.from_pretrained(tmp_path) + + ref_params, ref_buffers = src.state_dict(), dict(src.named_buffers()) + for name, param in model.named_parameters(): + if ".mlp." in name: + assert torch.equal(param, ref_params[name]), name + else: + assert param.is_meta, name + for name, buffer in model.named_buffers(): + assert torch.equal(buffer, ref_buffers[name]), name + + +# ------------------------------------------------------------------------- _fsdp2_shard_dense + + +@pytest.fixture +def fsdp_mesh(skip_on_windows, tmp_path_factory): + """The dense FSDP mesh of a one-rank gloo process group (FSDP2 needs a process group).""" + store = tmp_path_factory.mktemp("pg") / "store" + torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) + yield init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + torch.distributed.destroy_process_group() + + +class _Experts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.randn(4, 8, 4)) + self.down_proj = nn.Parameter(torch.randn(4, 4, 4)) + + +class _DecoderLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Linear(4, 4) + self.mlp = nn.Module() + self.mlp.gate = nn.Linear(4, 4, bias=False) # the router: dense, sharded with its layer + self.mlp.experts = _Experts() + self.input_layernorm = nn.LayerNorm(4) + + +class _MoeLM(nn.Module): + """A causal-LM-shaped MoE whose ``_no_split_modules`` names its decoder layer, as HF's do.""" + + _no_split_modules = ["_DecoderLayer"] + + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.embed_tokens = nn.Embedding(16, 4) + self.model.layers = nn.ModuleList(_DecoderLayer() for _ in range(2)) + self.model.norm = nn.LayerNorm(4) + self.lm_head = nn.Linear(4, 16, bias=False) + + def experts(self): + return [layer.mlp.experts for layer in self.model.layers] + + +def test_each_decoder_layer_is_sharded_around_its_experts(fsdp_mesh): + """A layer's dense params (attention, router, norms) become shards on the FSDP mesh; its + experts stay as loaded, for EP. The root's embedding, final norm and head stay replicated, and + a layer's submodules keep their classes (modelopt recognizes them by class).""" + model = _MoeLM() + before = {name: param.detach().clone() for name, param in model.named_parameters()} + experts = {name: param for name, param in model.named_parameters() if _is_expert(name)} + + assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2 + + assert all(isinstance(layer, FSDPModule) for layer in model.model.layers) + assert not isinstance(model, FSDPModule) + assert type(model.model.layers[0].self_attn) is nn.Linear + for name, param in model.named_parameters(): + if name in experts: + assert param is experts[name] and not isinstance(param, DTensor), name + elif name.startswith("model.layers."): + assert isinstance(param, DTensor), name + assert param.device_mesh == fsdp_mesh and param.placements == (Shard(0),), name + assert torch.equal(param.full_tensor(), before[name]), name + else: + assert not isinstance(param, DTensor), name + + +def test_a_param_off_the_models_dtype_stays_replicated(fsdp_mesh): + """FSDP2 shards one dtype per group, and HF keeps a few params in fp32, as a router: such a + param stays replicated, as the pure-FSDP2 wrap leaves it, and its layer still unshards.""" + model = _MoeLM().to(torch.bfloat16) + gate = model.model.layers[0].mlp.gate.float() + + assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2 + + assert not isinstance(gate.weight, DTensor) and gate.weight.dtype == torch.float32 + for layer in model.model.layers: + layer.unshard() # FSDP2 checks a group's dtypes here + + +def test_a_layer_holding_only_experts_is_left_unwrapped(fsdp_mesh): + model = _MoeLM() + bare = model.model.layers[1] + bare.self_attn = bare.mlp.gate = bare.input_layernorm = nn.Identity() + + assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 1 + assert not isinstance(bare, FSDPModule) + + +def test_a_dense_part_left_with_nothing_to_shard_is_refused(fsdp_mesh): + """Experts outweighing the dense part in another dtype put every dense param off the model's + dominant dtype, out of the shard groups: no layer is left to shard, and the load would hold the + whole dense part on every rank.""" + model = _MoeLM() + for experts in model.experts(): + experts.to(torch.bfloat16) + + with pytest.raises(RuntimeError, match="no dense parameter to shard"): + _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) + + +def test_a_dense_param_the_loader_already_sharded_stays_as_it_is(fsdp_mesh): + """A dense param that arrives as a ``Shard`` DTensor of the loader's own plan is not resharded; + the rest of its layer still is.""" + tp_mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("tp",)) + model = _MoeLM() + attn = model.model.layers[0].self_attn + attn.weight = nn.Parameter(distribute_tensor(attn.weight.detach(), tp_mesh, [Shard(1)])) + presharded = attn.weight + + assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2 + + assert attn.weight is presharded and attn.weight.placements == (Shard(1),) + assert model.model.layers[0].mlp.gate.weight.placements == (Shard(0),) + + +def test_a_replicated_dense_param_is_sharded_over_the_fsdp_mesh_like_a_plain_one(fsdp_mesh): + ep_mesh = init_device_mesh("cpu", (1, 1), mesh_dim_names=("edp", "ep"))["ep"] + model = _MoeLM() + gate = model.model.layers[0].mlp.gate + weight = gate.weight.detach().clone() + gate.weight = nn.Parameter(distribute_tensor(weight, ep_mesh, [Replicate()])) + + _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) + + assert gate.weight.device_mesh == fsdp_mesh and gate.weight.placements == (Shard(0),) + assert torch.equal(gate.weight.full_tensor(), weight) + + +class _VisionBlock(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + + +class _VisionTower(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList(_VisionBlock() for _ in range(2)) + + +class _MoeVlm(_MoeLM): + """Kimi's MoonViT style: ``_no_split_modules`` lists the vision tower *and* its blocks.""" + + _no_split_modules = ["_VisionTower", "_VisionBlock", "_DecoderLayer"] + + def __init__(self): + super().__init__() + self.vision_tower = _VisionTower() + + +def test_a_listed_container_is_wrapped_once_not_again_through_its_blocks(fsdp_mesh): + model = _MoeVlm() + + assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 3 + + assert isinstance(model.vision_tower, FSDPModule) + assert not any(isinstance(block, FSDPModule) for block in model.vision_tower.blocks) + + +# ------------------------------------------------------------------ _load_ep_dense_distributed + + +def test_the_dense_load_shards_first_then_asks_for_exactly_the_meta_dense_params( + fsdp_mesh, monkeypatch, tmp_path +): + """As ``_ep_dense_meta_load`` leaves the model: dense weights on meta; experts loaded, here + per-expert ``nn.Linear``s, whose ``.weight`` only the expert matcher tells from a dense one; + and a Mamba-style ``A_log`` (no ``.weight`` / ``.bias`` suffix) loaded too. The dense layers + are FSDP2-sharded before the shared meta+DCP loader runs, and it is asked for exactly the meta + dense params on this rank's GPU. That load materializes on ``cuda:``, so a + recorder stands in for it here.""" + with torch.device("meta"): + model = _MoeLM() + for layer in model.model.layers: + layer.mlp.experts = nn.ModuleList(nn.Linear(4, 4, bias=False) for _ in range(4)) + layer.A_log = nn.Parameter(torch.zeros(4)) + loads = [] + + def record_load(loaded, local_dir, device, *, load_names): + sharded = [isinstance(layer.self_attn.weight, DTensor) for layer in loaded.model.layers] + loads.append((local_dir, device, load_names, sharded)) + + monkeypatch.setattr(model_load_utils, "_meta_dcp_load", record_load) + monkeypatch.setattr(mto_dist, "local_rank", lambda: 3) + + wrapped = _load_ep_dense_distributed( + model, fsdp_mesh, str(tmp_path), _is_expert, model.experts() + ) + + assert wrapped == 2 + meta_dense = {name for name, param in model.named_parameters() if param.is_meta} + assert loads == [(str(tmp_path), torch.device("cuda", 3), meta_dense, [True, True])] + assert not any(_is_expert(name) or name.endswith("A_log") for name in meta_dense) diff --git a/tests/unit/torch/utils/test_ep_load_warmup.py b/tests/unit/torch/utils/test_ep_load_warmup.py new file mode 100644 index 00000000000..18c56d6d159 --- /dev/null +++ b/tests/unit/torch/utils/test_ep_load_warmup.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The expert-parallel sizing of transformers' load-time caching-allocator warmup. + +``from_pretrained`` reserves the bytes ``get_total_byte_count`` reports for each device in one +allocation. Under the EP load it would count every expert at full size (the experts are most of a +large MoE), so ``_ep_aware_caching_allocator_warmup`` swaps in a sizing that counts one EP shard of +each expert. The sizing is called directly here, and once through ``caching_allocator_warmup`` with +the CUDA queries and the allocation faked. +""" + +import pytest +import torch +from torch import nn + +pytest.importorskip("transformers") + +from transformers import modeling_utils + +from modelopt.torch.utils.plugins.model_load_utils import _ep_aware_caching_allocator_warmup + +if not hasattr(modeling_utils, "get_total_byte_count"): + pytest.skip("transformers has no get_total_byte_count to size", allow_module_level=True) + +_GPU = torch.device("cuda", 0) + +# The plan load_and_prepare_ep shards with after the a2a swap: base-relative keys (no ``model.``). +_A2A_PLAN = { + "layers.*.mlp.gate": "torch_a2a_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "torch_a2a_experts", +} + +# _TinyMoe's fp32 bytes: the fused experts; the dense params (embedding, attention weight + bias, +# router); the router's persistent ``e_score_correction_bias`` buffer. The tied head adds none. +_EXPERT_BYTES = 4 * (4 * 16 * 8 + 4 * 8 * 8) +_DENSE_BYTES = 4 * (16 * 8 + 8 * 8 + 8 + 4 * 8) +_BUFFER_BYTES = 4 * 4 +_ATTN_BYTES = 4 * (8 * 8 + 8) + + +class _TinyMoe(nn.Module): + """One MoE decoder layer, a head tied to the embedding, and a persistent router buffer.""" + + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.embed_tokens = nn.Embedding(16, 8) + layer = nn.Module() + layer.self_attn = nn.Linear(8, 8) + layer.mlp = nn.Module() + layer.mlp.gate = nn.Linear(8, 4, bias=False) + layer.mlp.gate.register_buffer("e_score_correction_bias", torch.zeros(4)) + layer.mlp.experts = nn.Module() + layer.mlp.experts.gate_up_proj = nn.Parameter(torch.zeros(4, 16, 8)) + layer.mlp.experts.down_proj = nn.Parameter(torch.zeros(4, 8, 8)) + self.model.layers = nn.ModuleList([layer]) + self.lm_head = nn.Linear(8, 16, bias=False) + self.lm_head.weight = self.model.embed_tokens.weight + self.all_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def get_parameter_or_buffer(self, name): + """``PreTrainedModel``'s lookup, which the sizing goes through.""" + try: + return self.get_parameter(name) + except AttributeError: + return self.get_buffer(name) + + def device_map(self): + """Every param and persistent buffer on one GPU, as ``from_pretrained`` expands it.""" + names = [name for name, _ in self.named_parameters(remove_duplicate=False)] + names += [name for name, _ in self.named_buffers()] + return dict.fromkeys(names, _GPU) + + +class _OneBytePerElement: + def param_element_size(self, model, name, param): + return 1 + + +_ALL_BYTES = _EXPERT_BYTES + _DENSE_BYTES + _BUFFER_BYTES + + +@pytest.mark.parametrize( + ("dense_on_meta", "quantizer", "stock_bytes", "per_rank_bytes"), + [ + # Experts count one EP shard; the router is replicated, so it counts in full like the rest. + pytest.param( + False, + None, + _ALL_BYTES, + _EXPERT_BYTES // 4 + _DENSE_BYTES + _BUFFER_BYTES, + id="dense-replicated", + ), + # Sharded dense: every dense .weight / .bias stays on meta during from_pretrained, while + # persistent buffers still load onto the device. + pytest.param( + True, None, _ALL_BYTES, _EXPERT_BYTES // 4 + _BUFFER_BYTES, id="dense-on-meta" + ), + # The quantizer sizes the elements: one byte each where fp32 takes four. + pytest.param( + False, + _OneBytePerElement(), + _ALL_BYTES // 4, + (_EXPERT_BYTES // 4 + _DENSE_BYTES + _BUFFER_BYTES) // 4, + id="quantizer", + ), + ], +) +def test_transformers_reserves_only_the_per_rank_footprint( + monkeypatch, dense_on_meta, quantizer, stock_bytes, per_rank_bytes +): + """End to end: the one up-front allocation ``caching_allocator_warmup`` makes shrinks from the + whole model (stock sizing, which also checks the byte counts above) to what this EP rank + holds.""" + reserved = [] + + def fake_empty(numel, *, dtype, device, **kwargs): + reserved.append((numel * torch.tensor([], dtype=dtype).element_size(), device)) + + model = _TinyMoe() + with monkeypatch.context() as m: + m.setattr(torch.cuda, "mem_get_info", lambda index: (80 * 1024**3, 80 * 1024**3)) + m.setattr(torch.cuda, "memory_reserved", lambda index: 0) + m.setattr(torch.cuda, "memory_allocated", lambda index: 0) + m.setattr(torch, "empty", fake_empty) + modeling_utils.caching_allocator_warmup(model, model.device_map(), quantizer) + with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN, dense_on_meta=dense_on_meta): + modeling_utils.caching_allocator_warmup(model, model.device_map(), quantizer) + + assert reserved == [(stock_bytes, _GPU), (per_rank_bytes, _GPU)] + + +def test_a_tp_planned_dense_param_keeps_its_tp_share_while_the_rest_stays_on_meta(monkeypatch): + """transformers places a TP-planned param itself, never asking where the dense part goes, so it + is on the device at its TP share even when the rest of the dense part stays on meta. The share + is of the ep submesh it is placed over, not of the world, which holds two of them here.""" + model = _TinyMoe() + model._tp_plan = {"model.layers.*.self_attn": "colwise"} + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 8) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + + with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN, dense_on_meta=True): + total = modeling_utils.get_total_byte_count(model, model.device_map()) + + assert total == {_GPU: _EXPERT_BYTES // 4 + _ATTN_BYTES // 4 + _BUFFER_BYTES} + + +def test_an_error_while_sizing_falls_back_to_the_stock_sizing(monkeypatch): + """The warmup only sizes an allocator reservation, so whatever the EP-aware sizing cannot + handle gets transformers' own number instead of failing the load, with whatever else a later + release passes the hook.""" + stock_calls = [] + + def stock_sizing(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs): + stock_calls.append((model, accelerator_device_map, hf_quantizer, args, kwargs)) + return {_GPU: 123} + + class _UnsupportedQuantizer: + def param_element_size(self, model, name, param): + raise TypeError("unsupported param") + + monkeypatch.setattr(modeling_utils, "get_total_byte_count", stock_sizing) + model, quantizer = _TinyMoe(), _UnsupportedQuantizer() + device_map = model.device_map() + + with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN): + sizing = modeling_utils.get_total_byte_count(model, device_map, quantizer, 1, flag=True) + assert sizing == {_GPU: 123} + + assert stock_calls == [(model, device_map, quantizer, (1,), {"flag": True})] + + +def test_a_sizing_hook_with_other_parameters_keeps_the_stock_sizing(monkeypatch): + """A release whose hook takes other leading parameters keeps transformers' own sizing: the + replacement would not know how the hook is called.""" + + def stock_sizing(model, device_map, quantizer=None): + return {_GPU: 123} + + monkeypatch.setattr(modeling_utils, "get_total_byte_count", stock_sizing) + + with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN): + assert modeling_utils.get_total_byte_count is stock_sizing + + +@pytest.mark.parametrize("ep_degree", [1, 0]) +def test_without_expert_parallelism_the_stock_sizing_stays(ep_degree): + stock = modeling_utils.get_total_byte_count + + with _ep_aware_caching_allocator_warmup(ep_degree, _A2A_PLAN): + assert modeling_utils.get_total_byte_count is stock + + +def test_a_transformers_without_the_sizing_hook_is_left_alone(monkeypatch): + monkeypatch.delattr(modeling_utils, "get_total_byte_count") + + with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN): + assert not hasattr(modeling_utils, "get_total_byte_count") diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 4c642b1644c..96558467de2 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -17,9 +17,10 @@ ``_meta_dcp_load`` materializes a model built on ``meta`` and fills it from HF safetensors through ``torch.distributed.checkpoint``; ``_map_ckpt_keys_to_model`` bridges checkpoint names to the built -model's ``state_dict`` names. The loads run for real on CPU, from tiny transformers models into -hand-written safetensors files. DCP loads in a single process without a process group; the tests that -need the distributed path start a one-rank gloo group. +model's ``state_dict`` names; ``_fsdp2_wrap_modules`` picks the blocks FSDP2 shards. The loads run +for real on CPU, from tiny transformers models into hand-written safetensors files. DCP loads in a +single process without a process group; the tests that need the distributed path start a one-rank +gloo group. """ import pytest @@ -40,6 +41,7 @@ from modelopt.torch.utils.plugins.model_load_utils import ( _copy_tied_values, + _fsdp2_wrap_modules, _map_ckpt_keys_to_model, _meta_dcp_load, ) @@ -170,6 +172,60 @@ def test_map_ckpt_keys_to_model(conversion_mapping, ckpt_keys, model_keys, expec assert _map_ckpt_keys_to_model(model, ckpt_keys, model_keys) == expected +# -------------------------------------------------------------------------- FSDP2 wrap boundaries + + +def test_wrap_modules_are_the_decoder_layers(): + """A causal LM is wrapped at its ``_no_split_modules`` blocks: the decoder layers, in order.""" + model = _meta_twin(get_tiny_llama) + + assert _fsdp2_wrap_modules(model) == list(model.model.layers) + + +def test_vlm_wraps_only_its_language_layers(): + """Qwen2.5-VL's wrapper also lists the vision block in ``_no_split_modules``, but only the language + layers may be sharded: sharded vision weights deadlock calibration, so the tower stays replicated. + """ + model = _meta_twin(_tiny_qwen2_5_vl) + assert type(model.model.visual.blocks[0]).__name__ in model._no_split_modules + + assert _fsdp2_wrap_modules(model) == list(model.model.language_model.layers) + + +@pytest.mark.parametrize("no_split_modules", [None, ["NoSuchBlock"]]) +def test_wrap_modules_fall_back_to_decoder_layer_detection(no_split_modules): + """Without usable ``_no_split_modules`` (absent, or naming no module present), modelopt's + decoder-layer auto-detection still finds the blocks. + """ + model = _meta_twin(get_tiny_llama) + model._no_split_modules = no_split_modules + + assert _fsdp2_wrap_modules(model) == list(model.model.layers) + + +def _vlm_whose_language_model_shows_no_blocks(): + """A VLM whose own ``_no_split_modules`` names its vision block too, as Qwen2.5-VL's does.""" + model = nn.Module() + model._no_split_modules = ["Linear"] + model.language_model = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2)) + model.visual = nn.Linear(2, 2) + return model + + +@pytest.mark.parametrize( + "build", + [ + pytest.param(lambda: nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2)), id="no-blocks"), + pytest.param(_vlm_whose_language_model_shows_no_blocks, id="vlm-language-model"), + ], +) +def test_wrap_modules_refuse_a_model_without_blocks(build): + """With neither source of boundaries matching, fail loudly instead of guessing a layout. For a + VLM that is its language model's: the VLM's own list would shard the vision blocks too.""" + with pytest.raises(RuntimeError, match="Could not determine FSDP2 wrap boundaries"): + _fsdp2_wrap_modules(build()) + + # --------------------------------------------------------------------------------- meta + DCP load diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 55c1dc717a2..85f3243278f 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -611,13 +611,43 @@ def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit( @needs_qwen3_moe_ep_plan -def test_the_env_knobs_override_the_automatic_plan(run_as, serve, monkeypatch, tmp_path): - monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") - monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") +def test_a_dense_part_the_expert_parallel_load_cannot_shard_is_replicated( + run_as, serve, gpu_memory, tmp_path +): + """The expert-parallel load cannot shard the dense part within several EP groups, which the + export cannot write: 2 experts over 4 ranks. Left to decide, the plan replicates it however + big, warning with the way out; asked to shard it, it refuses. Without expert parallelism the + rule does not apply.""" + run_as(4) + serve(_qwen3_moe(num_experts=2)) + gpu_memory(1) + + with pytest.warns(UserWarning, match="may run out of memory. Set MODELOPT_EP_DEGREE"): + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 2) + with pytest.raises(NotImplementedError, match="2 expert-data-parallel replicas"): + resolve_distribution_plan(str(tmp_path), dense="fsdp") + assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan( + 4, dense="fsdp" + ) + + +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize( + ("knob", "value", "plan"), + [ + ("MODELOPT_EP_DEGREE", "2", DistributionPlan(4, 2)), + ("MODELOPT_FSDP_SIZE", "world", DistributionPlan(4, 4, "fsdp")), + ], + ids=["ep-degree", "fsdp-size"], +) +def test_the_env_knobs_override_the_automatic_plan( + run_as, serve, monkeypatch, tmp_path, knob, value, plan +): + monkeypatch.setenv(knob, value) run_as(4) serve(_qwen3_moe(num_experts=8)) - assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 2, "fsdp") + assert resolve_distribution_plan(str(tmp_path))[0] == plan @needs_qwen3_moe_ep_plan @@ -1955,12 +1985,17 @@ def ep_world_of_two(monkeypatch, gloo_world_of_one): @pytest.mark.parametrize( ("plan", "error", "match"), [ - (DistributionPlan(2), ValueError, "shards no routed experts"), - (DistributionPlan(2, ep_size=2, dense="fsdp"), NotImplementedError, "replicates the dense"), + pytest.param(DistributionPlan(2), ValueError, "shards no routed experts", id="no-experts"), + # FSDP2 within each of several expert-parallel groups: a layout the export cannot write. + pytest.param( + DistributionPlan(4, 2, "fsdp"), + NotImplementedError, + "expert-data-parallel", + id="dense-hsdp", + ), ], - ids=["no-expert-parallelism", "sharded-dense"], ) -def test_the_loader_refuses_a_plan_it_cannot_load_by(plan, error, match): +def test_the_loader_refuses_a_plan_it_cannot_carry_out(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) From 4058038ed5a721f30b26e93fec123fabe0cb579b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:39:05 +0000 Subject: [PATCH 02/21] Import the dense-sharding helpers' modules at the top of the file fully_shard, DTensor and defaultdict were imported again inside the helpers that use them, though the module imports them already; Replicate, sys, transformers, transformers.modeling_utils and transformers' replace_layer_number_by_wildcard were imported inside functions with no reason to be (every supported transformers, 5.5 through 5.18, has them). They are all module-level imports now. The warmup sizing's guard for a transformers without modeling_utils goes with its local import: no 5.x release lacks it. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 40 ++++++------------- 1 file changed, 12 insertions(+), 28 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 0ebbc6620fe..7c745729b7b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -22,6 +22,7 @@ import logging import os import re +import sys import warnings from collections import defaultdict from contextlib import contextmanager @@ -31,7 +32,9 @@ import torch import torch.distributed.checkpoint as dcp import torch.nn as nn +import transformers import transformers.distributed as hf_distributed +import transformers.modeling_utils as hf_modeling_utils from huggingface_hub import snapshot_download if TYPE_CHECKING: @@ -39,8 +42,8 @@ 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 -from torch.distributed.tensor import DTensor +from torch.distributed.fsdp import FSDPModule, fully_shard +from torch.distributed.tensor import DTensor, Replicate from transformers import ( AutoConfig, AutoModel, @@ -52,6 +55,7 @@ ) from transformers.conversion_mapping import get_model_conversion_mapping from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key +from transformers.integrations.tensor_parallel import replace_layer_number_by_wildcard from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import env, moe @@ -1681,8 +1685,6 @@ def _make_expert_param_matcher(ep_plan: dict[str, Any]): (``ep_router`` / ``torch_a2a_router``) are REPLICATED, not sharded, so they're excluded. If the plan is empty the matcher is always ``False`` (-> no EP credit -> stock warmup sizing). """ - from transformers.integrations.tensor_parallel import replace_layer_number_by_wildcard - _router_styles = {"ep_router", "torch_a2a_router"} keys = [k for k, v in (ep_plan or {}).items() if v not in _router_styles] @@ -1728,12 +1730,7 @@ def _ep_aware_caching_allocator_warmup( if ep_degree <= 1: yield return - try: - import transformers.modeling_utils as _mu - except Exception: - yield - return - orig = getattr(_mu, "get_total_byte_count", None) + orig = getattr(hf_modeling_utils, "get_total_byte_count", None) if orig is None: yield return @@ -1751,8 +1748,6 @@ def _ep_aware_caching_allocator_warmup( def _ep_aware(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs): try: - from collections import defaultdict - try: # where current transformers keeps it (5.18); older releases, under integrations from transformers.distributed.tensor_parallel import _get_parameter_tp_plan except ImportError: @@ -1814,10 +1809,10 @@ def _ep_aware(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs) return orig(model, accelerator_device_map, hf_quantizer, *args, **kwargs) try: - _mu.get_total_byte_count = _ep_aware + hf_modeling_utils.get_total_byte_count = _ep_aware yield finally: - _mu.get_total_byte_count = orig + hf_modeling_utils.get_total_byte_count = orig def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_state) -> None: @@ -1860,16 +1855,11 @@ def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int: Returns the number of layers wrapped. """ - from torch.distributed.fsdp import fully_shard - # Every routed-expert param is EP-sharded on the ``ep`` mesh; FSDP must IGNORE them so it never tries # to reshard them onto ``fsdp_mesh``. Ignored params not present in a given layer are simply no-ops. ignored_params = _off_dtype_params(model) for em in expert_modules: ignored_params.update(em.parameters(recurse=True)) - - from torch.distributed.tensor import DTensor, Replicate - ignored_param_ids = {id(p) for p in ignored_params} # id-based: Tensor.__eq__ is elementwise # Wrap only the OUTERMOST _no_split_modules. Some models (e.g. Kimi's MoonViT VLM) list BOTH a @@ -2003,20 +1993,14 @@ def _ep_dense_meta_load(is_expert): normally. Patches ``get_device`` on both ``transformers.modeling_utils`` and the module that actually defines the loader (it's looked up as a module global there). """ - import sys as _sys - - import transformers.modeling_utils as _mu - - loader = getattr(_mu, "convert_and_load_state_dict_in_model", None) - mods = {_mu} + loader = getattr(hf_modeling_utils, "convert_and_load_state_dict_in_model", None) + mods = {hf_modeling_utils} loader_module = getattr(loader, "__module__", None) - dm = _sys.modules.get(loader_module) if loader_module else None + dm = sys.modules.get(loader_module) if loader_module else None if dm is not None: mods.add(dm) targets = [m for m in mods if hasattr(m, "get_device")] if not targets: - import transformers - raise RuntimeError( "Sharding the dense part keeps it off the device during from_pretrained through " f"transformers' get_device, which none of {sorted(m.__name__ for m in mods)} has on " From 9c5b36b728eacf73ecc769a206c3c4126d1f29a9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:40:34 +0000 Subject: [PATCH 03/21] Keep the dense-sharding helpers' docs to what a maintainer needs The helpers' docstrings and their call sites retold the same load mechanics at length, and the warmup sizing carried model-specific memory numbers. Each now states its contract and the one reason it is built that way: blocks are wrapped whole because fully_shard swaps a module's class, the experts and off-dtype params are left out, the root stays replicated, and the dense part stays on meta until it is sharded. The numbers move to the PR description. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 164 +++++------------- 1 file changed, 41 insertions(+), 123 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 7c745729b7b..341ccacc3c4 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -895,19 +895,11 @@ def handed() -> DistributionPlan | None: def _fsdp2_wrap_modules(model: "torch.nn.Module") -> list["torch.nn.Module"]: - """Return the transformer-block modules to FSDP2 ``fully_shard`` individually. - - Uses HuggingFace's own ``_no_split_modules`` -- the class names of the atomic blocks that must not be - split across devices (what ``device_map="auto"`` wraps at) -- falling back to modelopt's decoder-layer - auto-detection. - - For a VLM, this looks at the LANGUAGE backbone's boundaries, NOT the top-level VLM wrapper's - ``_no_split_modules`` (which for some VLMs -- e.g. Qwen2.5-VL -- ALSO lists the VISION block classes). - FSDP-sharding the vision blocks makes the calibration weight-access all-gather on those sharded vision - weights DEADLOCK; keeping the vision/audio towers REPLICATED (plain) -- the language-only calibration - never runs them and they are small enough to fit -- avoids that. This also covers VLMs whose top-level - wrapper exposes no ``_no_split_modules`` at all (e.g. Nemotron-Nano-VL). A VLM whose language - backbone shows no blocks is refused rather than wrapped by the wrapper's list. + """The blocks to FSDP2-wrap one by one: the language backbone's ``_no_split_modules`` blocks. + + Falls back to modelopt's decoder-layer detection. A VLM's own list can name vision blocks, which + the language-only calibration never runs: sharded, their weight gathers would deadlock, so they + stay replicated. A model whose language backbone shows no blocks is refused. """ from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -1671,19 +1663,11 @@ def checkpoint_read() -> bool: def _make_expert_param_matcher(ep_plan: dict[str, Any]): - """Build ``is_expert(param_name) -> bool`` from the model's ``base_model_ep_plan``. - - That plan is what drives EP sharding, so it replaces a fragile ``"expert" in name`` heuristic. - - A param is EP-sharded (only its ``1/ep_degree`` shard resides on this rank) iff its name, with layer - indices wildcarded exactly as HuggingFace does (``replace_layer_number_by_wildcard``) and any trailing - ``.weight`` / ``.bias`` stripped, contains an expert plan key as a run of whole dot-separated - segments. That one test covers a key equal to the name, a suffix of it (the base-relative plan keys - omit the ``model.`` / ``language_model.model.`` prefixes) and a container the name nests under, so - both per-parameter shard keys (``...experts.gate_up_proj: grouped_gemm``) and experts-container - keys (``...experts: torch_a2a_experts`` / ``nemotron_experts_ep``) match. Router entries - (``ep_router`` / ``torch_a2a_router``) are REPLICATED, not sharded, so they're excluded. If the plan is - empty the matcher is always ``False`` (-> no EP credit -> stock warmup sizing). + """``is_expert(param_name)``: whether ``ep_plan`` (the plan the EP load shards by) shards it. + + True when the name, layer indices wildcarded and ``.weight``/``.bias`` stripped, contains a + non-router plan key as whole dot-separated segments: the key itself, a suffix of the name (plan keys + omit the ``model.`` prefix), or a container it nests under. Router styles are replicated. """ _router_styles = {"ep_router", "torch_a2a_router"} keys = [k for k, v in (ep_plan or {}).items() if v not in _router_styles] @@ -1704,28 +1688,13 @@ def is_expert(pname: str) -> bool: def _ep_aware_caching_allocator_warmup( ep_degree: int, ep_plan: dict[str, Any], dense_on_meta: bool = False ): - """Make transformers' ``caching_allocator_warmup`` EP-aware for the duration of ``from_pretrained``. - - transformers sizes the load-time allocator warmup via ``get_total_byte_count``, which credits the - ``// world_size`` sharding division only to params matched in ``model._tp_plan``. Under the EP load - (``device_mesh`` + ``enable_expert_parallel``; sharding driven by ``base_model_ep_plan``) no expert - param is in ``_tp_plan``, so every one of them -- experts are ~95% of a large MoE -- is counted at - FULL global size. The - warmup then pre-allocates one ``torch.empty`` clamped to ``device_total - 1.2 GiB`` (~the whole GPU) - on every rank; torch keeps that as reserved cache, which later starves the NCCL a2a comm buffers. - - This temporarily replaces ``get_total_byte_count`` with a version that additionally credits EP-sharded - expert params their per-rank shard (``// ep_degree``), identifying them from ``ep_plan`` (the model's - ``base_model_ep_plan``), sizing the warmup to the real per-rank footprint (e.g. ~91 GiB instead of - ~275 GiB for Nemotron-3-Ultra-550B EP32) -- keeping the warmup's load-speed benefit without the - ~device-max transient. - - Always on (no disable knob): it only *sizes* an allocator pre-reservation, never model correctness; it - falls back to stock sizing on any incompatibility (a hook whose leading parameters differ is left in - place, and an error while sizing returns the stock number); and a mis-estimate only nudges load - speed (an over-credit under-sizes -> the allocator just grows during load; an under-credit sizes up to - the same device-max clamp as stock) -- so it can never OOM worse than, or load differently from, stock. - No-op when ``ep_degree <= 1``. + """Size transformers' load-time allocator warmup to this rank's expert-parallel share. + + transformers credits a param its ``1/world_size`` share only through ``_tp_plan``, which holds no + expert param under the EP load, so the warmup would reserve about the whole GPU on every rank and + starve the all-to-all's NCCL buffers. While the load runs, an expert param counts ``1/ep_degree`` of + itself and, with ``dense_on_meta``, a dense ``.weight``/``.bias`` nothing. It only sizes a + reservation: a hook of another signature, or an error while sizing, keeps the stock sizing. """ if ep_degree <= 1: yield @@ -1777,24 +1746,17 @@ def _ep_aware(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs) ) nb = p.numel() * dsize if is_expert(pname): - nb //= ep_degree # EP-sharded across the ep submesh: only the local shard resides here + nb //= ep_degree elif ( tp_plan and _get_parameter_tp_plan is not None and _get_parameter_tp_plan(pname, tp_plan, is_weight=True) is not None ): - # Placed by transformers' TP branch, which never asks where the dense part goes, so - # it is on the device at its share even when the rest of the dense part is on meta. - # That branch shards over from_pretrained's device_mesh, the ep submesh. + # Placed at its share over the ep submesh by transformers' TP branch, which ignores + # get_device, so it is on the device even when the rest of the dense part is on meta. nb //= ep_degree elif dense_on_meta and pname.endswith((".weight", ".bias")): - # A sharded dense part: _ep_dense_meta_load forces every dense .weight/.bias to - # materialize on META during from_pretrained (they are FSDP2-sharded + DCP-loaded - # afterwards), so they occupy NO device memory in the window this warmup sizes. - # Counting them at full global size is what pushed Kimi-K3 (EP64, dense=fsdp) to a - # 185.8 GiB estimate on a 184.31 GiB card -- over the device, so the warmup fell back - # to the device-max clamp it exists to avoid. Persistent buffers are NOT forced to - # meta by that context, so they keep their full credit here. + # On meta through the load (_ep_dense_meta_load); buffers are not, so they count. continue total[dev] += nb gib = 1024**3 @@ -1831,41 +1793,21 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int: - """FSDP2-shard the model's dense modules, leaving the routed experts to EP. - - FSDP2 ``fully_shard`` the model's DENSE (non-expert) modules over ``fsdp_mesh``, leaving the routed - experts (already EP-sharded on ``ep`` + the a2a transport) untouched -- the "FSDP2 dense + EP MoE" - hybrid. ``fsdp_mesh`` is the FULL world (dense shards over every rank, the same ranks the experts are - EP-sharded over). - - Wraps at the DECODER-LAYER boundary (:func:`_fsdp2_wrap_modules`: the language backbone's HF - ``_no_split_modules`` blocks, else the decoder-layer detection the pure-FSDP2 path uses) and passes - the routed-expert params as ``ignored_params`` so FSDP shards - the layer's dense params (attention, layernorms, router) while leaving the experts EP-sharded on - the ``ep`` mesh + a2a transport. As there, a param off the model's dtype (an fp32 router) stays - replicated: FSDP2 shards one dtype per group. Wrapping - the whole layer (not ``self_attn`` or individual leaves) is essential: ``fully_shard`` swaps a wrapped - module's class (``Foo`` -> ``FSDPFoo``), and modelopt inspects classes by name/source -- wrapping - ``self_attn`` breaks ``mtq``'s KV-cache AST registration (``inspect.getsource`` on the synthetic - class), and wrapping the leaf ``QuantLinear``s breaks the export's name-based ``is_quantlinear``. - Wrapping only the LAYER leaves ``self_attn`` and the q/k/v/o ``QuantLinear``s PLAIN, while - ``_get_enclosing_fsdp_module`` still finds the layer FSDP wrapper for the weight-access/writeback - contexts. The root (embed/lm_head/final norm) is intentionally NOT wrapped: those stay replicated - (and unswizzled, so their quantization-exclude still applies), matching the dense-replicated baseline. - - Returns the number of layers wrapped. + """FSDP2-shard the dense part over ``fsdp_mesh``, block by block; return the blocks wrapped. + + Each :func:`_fsdp2_wrap_modules` block is wrapped whole: ``fully_shard`` swaps a wrapped module's + class, and quantization and export recognize attention and linear modules by class, so those stay + plain inside the block. The routed experts (EP-sharded) and params off the model's dtype are left + out; the root (embeddings, LM head, final norm) is not wrapped and stays replicated. """ - # Every routed-expert param is EP-sharded on the ``ep`` mesh; FSDP must IGNORE them so it never tries - # to reshard them onto ``fsdp_mesh``. Ignored params not present in a given layer are simply no-ops. + # The experts are EP-sharded already; a block ignores the ones it does not hold. ignored_params = _off_dtype_params(model) for em in expert_modules: ignored_params.update(em.parameters(recurse=True)) ignored_param_ids = {id(p) for p in ignored_params} # id-based: Tensor.__eq__ is elementwise - # Wrap only the OUTERMOST _no_split_modules. Some models (e.g. Kimi's MoonViT VLM) list BOTH a - # container (the vision model) and its blocks in ``_no_split_modules``; ``fully_shard`` on the parent - # already shards the children, so wrapping a child afterward hits ``DeviceMesh._concatenate`` ("Cannot - # concatenate overlapping meshes"). Drop any wrap module nested under another. + # Only the outermost blocks: a model listing a container and its blocks (Kimi's MoonViT) would + # wrap a child of a wrapped module, which fully_shard refuses (overlapping meshes). wrap_modules = _fsdp2_wrap_modules(model) wrap_ids = {id(m) for m in wrap_modules} nested_ids: set = set() @@ -1884,9 +1826,8 @@ def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int: continue if isinstance(p.data, DTensor): if all(isinstance(pl, Replicate) for pl in p.data.placements): - # Replicate -> plain (a free view, no gather), re-registered as a plain Parameter: - # assigning .data would keep the DTensor parameter, which fully_shard then tries - # to place on both its own mesh and the FSDP one. + # Re-registered as a plain Parameter (a free view): assigning .data would keep the + # DTensor parameter, which fully_shard would place on two meshes. owner_name, _, leaf = name.rpartition(".") owner = layer.get_submodule(owner_name) if owner_name else layer setattr( @@ -1979,19 +1920,12 @@ def _ep_mesh(*args, **kwargs): @contextlib.contextmanager def _ep_dense_meta_load(is_expert): - """Keep the dense part off the GPU during ``from_pretrained``, for the EP hybrid. - - Force every non-expert weight/bias to materialize on ``meta`` (i.e. NOT on the GPU) by making - transformers' per-param ``get_device`` return ``"meta"`` for dense keys. - - transformers 5.x's ``convert_and_load_state_dict_in_model`` places EXPERT params via the tp / - ``distributed_operation`` branch (which fuses + EP-shards them and never consults ``get_device``), and - every OTHER (dense) param via ``param_device = get_device(device_map, key)`` -> ``spawn_materialize``. - Returning ``"meta"`` there leaves dense params unmaterialized while experts still EP-shard normally. The - skipped dense params are then FSDP2-sharded + DCP-loaded (:func:`_load_ep_dense_distributed`), so dense is - never replicated per rank. Only ``.weight``/``.bias`` are forced meta -- persistent buffers (small) load - normally. Patches ``get_device`` on both ``transformers.modeling_utils`` and the module that actually - defines the loader (it's looked up as a module global there). + """Keep the dense ``.weight``/``.bias`` on ``meta`` through ``from_pretrained``. + + transformers places expert params through its TP branch, which never asks ``get_device``, and every + other param where ``get_device`` says: ``"meta"`` for a dense key leaves it for + :func:`_load_ep_dense_distributed` to shard and load. Buffers load as usual. ``get_device`` is + patched in every module the loader looks it up in. """ loader = getattr(hf_modeling_utils, "convert_and_load_state_dict_in_model", None) mods = {hf_modeling_utils} @@ -2289,16 +2223,8 @@ 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") - # Size the load-time caching-allocator warmup to the real per-rank (EP-sharded) footprint instead of - # the full model -- otherwise transformers pre-allocates ~device-max on every rank and torch holds it - # as reserved cache, starving the NCCL a2a buffers. Expert params are identified from `swapped_plan` - # (the actual EP plan). Always on: falls back to stock sizing on any issue and a mis-size only nudges - # load speed, never correctness. (See _ep_aware_caching_allocator_warmup.) is_expert_fn = _make_expert_param_matcher(swapped_plan) - # The plan places the DENSE (non-expert) modules; the experts are always EP-sharded by from_pretrained. - # A replicated dense part loads with them, on every rank. A sharded one is left on meta during - # from_pretrained (get_device->meta via _ep_dense_meta_load), then FSDP2-sharded + DCP-loaded, so it - # is never replicated anywhere. + # A sharded dense part stays on meta through from_pretrained, to be sharded and loaded below. _dense_ctx = ( _ep_dense_meta_load(is_expert_fn) if plan.is_fsdp_dense else contextlib.nullcontext() ) @@ -2363,17 +2289,9 @@ def gpu() -> tuple[bool, bool, str | None]: "ranks holding different experts. Quantize the model without expert parallelism." ) - # "FSDP2 dense + EP MoE" hybrid: experts are EP-sharded over ``ep`` and the DENSE (non-expert) modules - # are FSDP2-sharded over the same ranks, the whole world: a dense part sharded within several - # expert-parallel groups is refused above. Dense all-gather then expert a2a run sequentially within a - # forward (same ranks), so there is no concurrent-collective deadlock. Weight amax is computed on the - # FSDP-gathered full weight (identical per rank); activation amax syncs over the whole world via the - # dense ``dp`` group registered below. - # - # A replicated dense part stays data-parallel and needs no post-load step. + # The dense part shards over the whole world, the experts' ranks: a forward's dense all-gathers + # and expert all-to-alls then run in turn on the same ranks, so they cannot deadlock. if plan.is_fsdp_dense: - # Fully-distributed dense: dense was left on meta during from_pretrained; FSDP2-shard it over the - # full-world mesh and DCP-load each rank's shard directly (never replicated anywhere). dense_fsdp_mesh = dense_mesh["dp"] n_dense = _load_ep_dense_distributed( model, dense_fsdp_mesh, ckpt_path, is_expert_fn, expert_modules.values() From 8780ffd811fa8c21248f0615b1a0f16858102b1d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:41:22 +0000 Subject: [PATCH 04/21] Fail every rank when one fails preparing the meta-DCP load Each rank reads the checkpoint's metadata, checks it against the model and materializes its params before DCP's collective load. A rank failing there alone (a metadata read on a node that cannot see the checkpoint, to_empty out of memory) raised while its peers went on into dcp.load and waited for it. The preparation now runs through the failure-carrying all_reduce, so a failure on any rank raises on every rank before any enters the load. Both the expert-parallel dense load and the FSDP2 load go through _meta_dcp_load. _meta_dcp_load keeps its signature: it runs _meta_dcp_prepare (the checks and the materialization, unchanged) through all_reduce and hands its state to _meta_dcp_fill (the load and the re-tie, unchanged). A single process runs it as before. A two-rank test fails rank 1's preparation, at the metadata read or at to_empty, and checks that both ranks raise rank 1's error without entering the load. Without the change rank 0 enters it. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 22 ++++++++- tests/unit/torch/utils/test_meta_dcp_load.py | 49 ++++++++++++++++++- 2 files changed, 69 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 341ccacc3c4..a6c8a00f13d 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1055,8 +1055,21 @@ def _meta_dcp_load( Every target param ends up filled -- from the checkpoint, by ``tie_weights`` with a param the checkpoint filled, or with the value it already had -- or the load fails: nothing is left as the - uninitialized memory ``to_empty`` hands out. + uninitialized memory ``to_empty`` hands out. A rank failing before ``dcp.load`` fails every rank. """ + prepared: list[tuple] = [] + # Each rank reads the checkpoint, checks it and materializes its params itself: one failing must + # fail every rank here, before the others enter dcp.load's collectives and wait there for it. + dist.all_reduce( + "the meta-DCP load's preparation", + lambda: prepared.append(_meta_dcp_prepare(model, local_dir, device, load_names)), + lambda _: None, + ) + _meta_dcp_fill(model, *prepared[0]) + + +def _meta_dcp_prepare(model: "torch.nn.Module", local_dir: str, device, load_names: set) -> tuple: + """:func:`_meta_dcp_load`'s checks and materialization, up to the state dict DCP fills.""" # Read what the checkpoint holds before changing the model, so a refusal leaves it as it was. reader = HuggingFaceStorageReader(str(local_dir)) ckpt_md = reader.read_metadata().state_dict_metadata @@ -1197,6 +1210,13 @@ def _meta_dcp_load( full_sd = model.state_dict() # the materialized tensors: to_empty replaced the meta params load_sd = {inv[mk]: full_sd[mk] for mk in full_sd if mk in inv} + return reader, load_sd, full_sd, want_params, mapped_names, emptied, kept_names + + +def _meta_dcp_fill( + model, reader, load_sd, full_sd, want_params, mapped_names, emptied, kept_names +) -> None: + """:func:`_meta_dcp_load`'s DCP load and re-tie, from :func:`_meta_dcp_prepare`'s state.""" dcp.load(load_sd, storage_reader=reader) # Taken before tying, which can swap a filled param out of its module; in state_dict order, the # same on every rank, so every rank keeps the same param of a tied group. diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 96558467de2..723c8c576d3 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -20,14 +20,18 @@ model's ``state_dict`` names; ``_fsdp2_wrap_modules`` picks the blocks FSDP2 shards. The loads run for real on CPU, from tiny transformers models into hand-written safetensors files. DCP loads in a single process without a process group; the tests that need the distributed path start a one-rank -gloo group. +gloo group, or run on the pooled two-rank one. """ +from contextlib import nullcontext +from unittest import mock + import pytest import torch import torch.distributed.checkpoint as dcp from safetensors.torch import save_file from torch import nn +from torch.distributed.checkpoint import HuggingFaceStorageReader from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import fully_shard from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor @@ -35,6 +39,7 @@ pytest.importorskip("accelerate") pytest.importorskip("transformers") +from _test_utils.torch.distributed.utils import bounded_collectives from _test_utils.torch.transformers_models import get_tiny_llama from accelerate import init_empty_weights from transformers import AutoModelForImageTextToText, Qwen2_5_VLConfig @@ -638,3 +643,45 @@ def test_a_tie_laid_out_differently_is_refused_and_left_as_it_was(gloo_world_of_ _copy_tied_values(kept, src, "lm_head.weight") assert torch.equal(local, before) + + +# ---------------------------------------------------------------- one rank failing to prepare + +# Where a rank's preparation fails, before the load's collectives: what it raises, and with what. +_FAILURES = { + "metadata": (HuggingFaceStorageReader, "read_metadata", OSError, "cannot read the checkpoint"), + "materialize": (nn.Module, "to_empty", MemoryError, "out of device memory"), +} + + +def _one_rank_fails_preparing_job(rank, size, ckpt_dir, failure): + owner, attr, error, message = _FAILURES[failure] + model = _meta_twin(get_tiny_llama) + entered, refused = [], None + + def dcp_load(state_dict, **kwargs): + # As the load's collectives do: a rank in them waits for every peer. + entered.append(rank) + torch.distributed.barrier() + + with ( + bounded_collectives(), + mock.patch.object(dcp, "load", dcp_load), + mock.patch.object(owner, attr, side_effect=error(message)) if rank == 1 else nullcontext(), + ): + try: # not pytest.raises: its failure would end the pooled worker instead of the job + _meta_dcp_load(model, ckpt_dir, "cpu", load_names=_every_name(model)) + except Exception as e: + refused = str(e) + # Rank 1's failure on both ranks, before either entered the load. + assert not entered, (rank, entered) + assert refused is not None and message in refused, (rank, refused) + assert rank == 1 or "rank 1 failed" in refused, refused + + +@pytest.mark.parametrize("failure", list(_FAILURES)) +def test_a_rank_failing_to_prepare_the_load_fails_every_rank(cpu_dist_workers, tmp_path, failure): + """Each rank reads the checkpoint and materializes its params before DCP's collective load. One + failing there, alone, fails every rank, rather than leave the others waiting in the load.""" + ckpt_dir = _save(tmp_path, get_tiny_llama().state_dict()) + cpu_dist_workers.run(_one_rank_fails_preparing_job, ckpt_dir=ckpt_dir, failure=failure) From 3f1e23bcfb42862cd92d3f1a5b340fe245275c0b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:42:16 +0000 Subject: [PATCH 05/21] Load the EP dense shards on the mesh's device, and say so briefly _load_ep_dense_distributed materialized the dense part on cuda: whatever mesh it was given, through a function-local import of local_rank. It now takes the device from the FSDP mesh, this rank's GPU on a CUDA mesh as before, so a CPU mesh loads on the CPU and the helper runs in a unit test, and it reads the rank through the module's own dist import. Its docstring and comments say what a maintainer needs: the meta dense part is FSDP2-sharded over the mesh, each rank DCP-loads its shard, the experts are left as loaded. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 21 +++++-------------- .../torch/utils/test_ep_dense_sharding.py | 7 ++----- 2 files changed, 7 insertions(+), 21 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index a6c8a00f13d..9242de94b2e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1986,25 +1986,14 @@ def patched(device_map, key, *args, **kwargs): def _load_ep_dense_distributed(model, fsdp_mesh, local_dir, is_expert, expert_modules): - """FSDP2-shard the meta dense part and DCP-load each rank's shard, for the EP hybrid. + """FSDP2-shard the meta dense part over ``fsdp_mesh`` and DCP-load each rank's shard of it. - Each rank's shard comes straight from the checkpoint, so the dense part is NEVER - materialized replicated. Experts are already EP-sharded (loaded by ``from_pretrained``) and left untouched. Returns - the number of dense layers sharded, like :func:`_fsdp2_shard_dense`. - - Requires the dense ``.weight``/``.bias`` params to be on ``meta`` on entry (see - :func:`_ep_dense_meta_load`). Delegates the materialize + DCP load to :func:`_meta_dcp_load`, fed the - partially-meta model (dense meta, experts already real) and the dense param names. + The dense ``.weight``/``.bias`` params arrive on meta (:func:`_ep_dense_meta_load`); the experts, + already loaded and EP-sharded, are left as they are. Returns the number of layers sharded. """ - from modelopt.torch.utils.distributed import local_rank - - device = torch.device(f"cuda:{local_rank()}") - # FSDP2-shard the dense modules. Their params are meta plain tensors -> fully_shard makes them meta Shard - # DTensors; experts are passed as ignored_params and stay EP-sharded. + cuda = fsdp_mesh.device_type == "cuda" + device = torch.device(fsdp_mesh.device_type, dist.local_rank() if cuda else None) n_dense = _fsdp2_shard_dense(model, fsdp_mesh, expert_modules) - # Restrict the shared loader to the dense params (experts already materialized by from_pretrained must not - # be touched). It selectively materializes only the meta dense modules (experts skipped), preserves - # buffers across to_empty, and DCP-loads each rank's dense shard. dense_names = { n for n, _ in model.named_parameters() diff --git a/tests/unit/torch/utils/test_ep_dense_sharding.py b/tests/unit/torch/utils/test_ep_dense_sharding.py index 28d2af01d15..3c6b71e114f 100644 --- a/tests/unit/torch/utils/test_ep_dense_sharding.py +++ b/tests/unit/torch/utils/test_ep_dense_sharding.py @@ -39,7 +39,6 @@ from _test_utils.torch.transformers_models import get_tiny_llama from transformers import AutoModelForCausalLM -import modelopt.torch.utils.distributed as mto_dist import modelopt.torch.utils.moe as moe import modelopt.torch.utils.plugins.model_load_utils as model_load_utils from modelopt.torch.utils.distributed import DistributionPlan @@ -386,8 +385,7 @@ def test_the_dense_load_shards_first_then_asks_for_exactly_the_meta_dense_params per-expert ``nn.Linear``s, whose ``.weight`` only the expert matcher tells from a dense one; and a Mamba-style ``A_log`` (no ``.weight`` / ``.bias`` suffix) loaded too. The dense layers are FSDP2-sharded before the shared meta+DCP loader runs, and it is asked for exactly the meta - dense params on this rank's GPU. That load materializes on ``cuda:``, so a - recorder stands in for it here.""" + dense params, on the mesh's device. A recorder stands in for that load.""" with torch.device("meta"): model = _MoeLM() for layer in model.model.layers: @@ -400,7 +398,6 @@ def record_load(loaded, local_dir, device, *, load_names): loads.append((local_dir, device, load_names, sharded)) monkeypatch.setattr(model_load_utils, "_meta_dcp_load", record_load) - monkeypatch.setattr(mto_dist, "local_rank", lambda: 3) wrapped = _load_ep_dense_distributed( model, fsdp_mesh, str(tmp_path), _is_expert, model.experts() @@ -408,5 +405,5 @@ def record_load(loaded, local_dir, device, *, load_names): assert wrapped == 2 meta_dense = {name for name, param in model.named_parameters() if param.is_meta} - assert loads == [(str(tmp_path), torch.device("cuda", 3), meta_dense, [True, True])] + assert loads == [(str(tmp_path), torch.device("cpu"), meta_dense, [True, True])] assert not any(_is_expert(name) or name.endswith("A_log") for name in meta_dense) From c4ad0c3aba1f04eaab5509ab498afa8d4a3ee990 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:42:31 +0000 Subject: [PATCH 06/21] Load both aliases of a tied param in the EP dense load The dense load named the params to fill through named_parameters(), which lists a tied embedding and LM head once, as the embedding. A tied checkpoint can store the pair as lm_head.weight alone, as from_pretrained accepts; the shared meta+DCP load then dropped that tensor as outside the names it was given, and refused the load with both aliases left uninitialized. The names now count every alias, as the shared load does for its own checks, so the stored head fills the pair. A test loads a tied tiny Qwen3-MoE whose checkpoint stores lm_head.weight alone through the dense load, as the expert-parallel load leaves it; it failed on "2 param(s) ... would stay uninitialized" before the change. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 3 ++- .../torch/utils/test_ep_dense_sharding.py | 27 ++++++++++++++++++- 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 9242de94b2e..fa25f7a0d1b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1994,9 +1994,10 @@ def _load_ep_dense_distributed(model, fsdp_mesh, local_dir, is_expert, expert_mo cuda = fsdp_mesh.device_type == "cuda" device = torch.device(fsdp_mesh.device_type, dist.local_rank() if cuda else None) n_dense = _fsdp2_shard_dense(model, fsdp_mesh, expert_modules) + # Every alias of a tied param: the checkpoint may store the one named_parameters() drops. dense_names = { n - for n, _ in model.named_parameters() + for n, _ in model.named_parameters(remove_duplicate=False) if n.endswith((".weight", ".bias")) and not is_expert(n) } print_rank_0( diff --git a/tests/unit/torch/utils/test_ep_dense_sharding.py b/tests/unit/torch/utils/test_ep_dense_sharding.py index 3c6b71e114f..741bf9ef5bb 100644 --- a/tests/unit/torch/utils/test_ep_dense_sharding.py +++ b/tests/unit/torch/utils/test_ep_dense_sharding.py @@ -36,7 +36,7 @@ import transformers.core_model_loading as core_model_loading import transformers.modeling_utils as modeling_utils -from _test_utils.torch.transformers_models import get_tiny_llama +from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3_moe from transformers import AutoModelForCausalLM import modelopt.torch.utils.moe as moe @@ -407,3 +407,28 @@ def record_load(loaded, local_dir, device, *, load_names): meta_dense = {name for name, param in model.named_parameters() if param.is_meta} assert loads == [(str(tmp_path), torch.device("cpu"), meta_dense, [True, True])] assert not any(_is_expert(name) or name.endswith("A_log") for name in meta_dense) + + +def test_a_tied_head_the_checkpoint_stores_alone_reaches_both_aliases(fsdp_mesh, tmp_path): + """A tied checkpoint can store the pair as ``lm_head.weight`` alone, which ``from_pretrained`` + loads into both aliases. With a tied MoE as the expert-parallel load leaves it -- the dense part + on meta and tied, the experts loaded -- the dense load fills both aliases with it too.""" + from safetensors.torch import load_file, save_file + + src = get_tiny_qwen3_moe(tie_word_embeddings=True) + src.save_pretrained(tmp_path) + weights = tmp_path / "model.safetensors" + stored = load_file(weights) + stored["lm_head.weight"] = stored.pop("model.embed_tokens.weight") + save_file(stored, weights, metadata={"format": "pt"}) + with _ep_dense_meta_load(_is_expert): + model = AutoModelForCausalLM.from_pretrained(tmp_path) + model.tie_weights() # as load_and_prepare_ep does once the weights load + assert model.lm_head.weight is model.model.embed_tokens.weight + assert model.lm_head.weight.is_meta + + experts = [layer.mlp.experts for layer in model.model.layers] + _load_ep_dense_distributed(model, fsdp_mesh, str(tmp_path), _is_expert, experts) + + for alias in (model.lm_head.weight, model.model.embed_tokens.weight): + assert torch.equal(alias, src.lm_head.weight) From b73b91beb49db7b5bdce0e7e8caea405b01ed999 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:44:09 +0000 Subject: [PATCH 07/21] Import safetensors at the top of the dense-sharding tests The tied-head test imported it inside the function; the repo keeps test imports at module top too. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_ep_dense_sharding.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/torch/utils/test_ep_dense_sharding.py b/tests/unit/torch/utils/test_ep_dense_sharding.py index 741bf9ef5bb..515c7ee4f4f 100644 --- a/tests/unit/torch/utils/test_ep_dense_sharding.py +++ b/tests/unit/torch/utils/test_ep_dense_sharding.py @@ -37,6 +37,7 @@ import transformers.core_model_loading as core_model_loading import transformers.modeling_utils as modeling_utils from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3_moe +from safetensors.torch import load_file, save_file from transformers import AutoModelForCausalLM import modelopt.torch.utils.moe as moe @@ -413,8 +414,6 @@ def test_a_tied_head_the_checkpoint_stores_alone_reaches_both_aliases(fsdp_mesh, """A tied checkpoint can store the pair as ``lm_head.weight`` alone, which ``from_pretrained`` loads into both aliases. With a tied MoE as the expert-parallel load leaves it -- the dense part on meta and tied, the experts loaded -- the dense load fills both aliases with it too.""" - from safetensors.torch import load_file, save_file - src = get_tiny_qwen3_moe(tie_word_embeddings=True) src.save_pretrained(tmp_path) weights = tmp_path / "model.safetensors" From 3d92a080d59c82a1ed36acaf122ff18be8baab96 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:03:57 +0000 Subject: [PATCH 08/21] Check for missing gradients before the EP test's collective comparisons The gradient comparison read a sharded gradient with full_tensor(), a collective, inside the loop that also skipped a parameter without a gradient. A rank missing one gradient the others had would skip that all-gather while they entered it: the test hung or crashed in mismatched collectives instead of reporting the missing gradient. Every rank now checks for missing gradients first and fails together before any comparison runs. Signed-off-by: Shengliang Xu --- tests/gpu/torch/utils/test_expert_parallel_load.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index f9f50f98ea4..0e14032e3cf 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -174,12 +174,14 @@ def _assert_grads_match_the_whole_model(model, per_batch, rank, size, ep_degree) group = range(rank - ep_rank, rank - ep_rank + ep_degree) per_rank = NUM_EXPERTS // ep_degree owned = slice(ep_rank * per_rank, (ep_rank + 1) * per_rank) + # Checked before the comparisons: full_tensor() below is a collective that a rank without the + # gradient would skip, leaving the others waiting in it instead of failing. + _fail_together( + [f"{name} received no gradient" for name, p in model.named_parameters() if p.grad is None] + ) problems, names = [], set() for name, param in model.named_parameters(): names.add(name) - if param.grad is None: - problems.append(f"{name} received no gradient") - continue if ".mlp.experts." in name: actual = _local(param.grad) expected = sum(per_batch[r][name] for r in group)[owned] From fa4ec4fcad25bc5753320d7b93f4380f1e63acbb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:05:15 +0000 Subject: [PATCH 09/21] Hand the meta-DCP load's prepared state over by name _meta_dcp_prepare returned a bare 7-element tuple that _meta_dcp_fill took as eight positional parameters, with nothing checking the two orders agree: swapping two of its sets, both set[str], would still run and silently change which params the unfilled-param guard covers. The state is a typed NamedTuple now, built and read by field name. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 48 ++++++++++++++----- 1 file changed, 37 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index fa25f7a0d1b..2f49fe3536a 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -27,7 +27,7 @@ from collections import defaultdict from contextlib import contextmanager from itertools import chain -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NamedTuple import torch import torch.distributed.checkpoint as dcp @@ -1057,7 +1057,7 @@ def _meta_dcp_load( checkpoint filled, or with the value it already had -- or the load fails: nothing is left as the uninitialized memory ``to_empty`` hands out. A rank failing before ``dcp.load`` fails every rank. """ - prepared: list[tuple] = [] + prepared: list[_MetaDcpPrepared] = [] # Each rank reads the checkpoint, checks it and materializes its params itself: one failing must # fail every rank here, before the others enter dcp.load's collectives and wait there for it. dist.all_reduce( @@ -1065,10 +1065,24 @@ def _meta_dcp_load( lambda: prepared.append(_meta_dcp_prepare(model, local_dir, device, load_names)), lambda _: None, ) - _meta_dcp_fill(model, *prepared[0]) + _meta_dcp_fill(model, prepared[0]) -def _meta_dcp_prepare(model: "torch.nn.Module", local_dir: str, device, load_names: set) -> tuple: +class _MetaDcpPrepared(NamedTuple): + """What :func:`_meta_dcp_prepare` hands :func:`_meta_dcp_fill`, by name.""" + + reader: HuggingFaceStorageReader + load_sd: dict[str, torch.Tensor] + full_sd: dict[str, torch.Tensor] + want_params: set[str] + mapped_names: set[str] + emptied: set[str] + kept_names: set[str] + + +def _meta_dcp_prepare( + model: "torch.nn.Module", local_dir: str, device, load_names: set +) -> _MetaDcpPrepared: """:func:`_meta_dcp_load`'s checks and materialization, up to the state dict DCP fills.""" # Read what the checkpoint holds before changing the model, so a refusal leaves it as it was. reader = HuggingFaceStorageReader(str(local_dir)) @@ -1210,17 +1224,27 @@ def _meta_dcp_prepare(model: "torch.nn.Module", local_dir: str, device, load_nam full_sd = model.state_dict() # the materialized tensors: to_empty replaced the meta params load_sd = {inv[mk]: full_sd[mk] for mk in full_sd if mk in inv} - return reader, load_sd, full_sd, want_params, mapped_names, emptied, kept_names + return _MetaDcpPrepared( + reader=reader, + load_sd=load_sd, + full_sd=full_sd, + want_params=want_params, + mapped_names=mapped_names, + emptied=emptied, + kept_names=kept_names, + ) -def _meta_dcp_fill( - model, reader, load_sd, full_sd, want_params, mapped_names, emptied, kept_names -) -> None: +def _meta_dcp_fill(model, prepared: _MetaDcpPrepared) -> None: """:func:`_meta_dcp_load`'s DCP load and re-tie, from :func:`_meta_dcp_prepare`'s state.""" - dcp.load(load_sd, storage_reader=reader) + dcp.load(prepared.load_sd, storage_reader=prepared.reader) # Taken before tying, which can swap a filled param out of its module; in state_dict order, the # same on every rank, so every rank keeps the same param of a tied group. - filled = {n: model.get_parameter(n) for n in full_sd if n in want_params and n in mapped_names} + filled = { + n: model.get_parameter(n) + for n in prepared.full_sd + if n in prepared.want_params and n in prepared.mapped_names + } if hasattr(model, "tie_weights"): model.tie_weights() # re-establish tied weights the checkpoint omitted # Tying keeps one param of each tied group, not necessarily the one the checkpoint filled: a tied @@ -1237,7 +1261,9 @@ def _meta_dcp_fill( # initialize it; this loader can't run the model's init on shards, so it fails rather than return # the model. Every rank reaches the same verdict. unfilled = sorted( - n for n in emptied - mapped_names - kept_names if id(model.get_parameter(n)) not in loaded + n + for n in prepared.emptied - prepared.mapped_names - prepared.kept_names + if id(model.get_parameter(n)) not in loaded ) if unfilled: raise RuntimeError( From ccd8b5b4e84349d08e4b7f0eef998ac7f3a3d474 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:06:44 +0000 Subject: [PATCH 10/21] Import the layer-wildcard helper from either transformers location replace_layer_number_by_wildcard, moved to module scope with the other imports, came only from transformers.integrations.tensor_parallel. On transformers 5.18 that module is a partial shim of transformers.distributed.tensor_parallel (it already lacks _get_parameter_tp_plan, which the warmup sizing looks up in both places), so a release dropping the name from it would make the whole module unimportable, not just the expert-parallel load. It is now imported from the new location first, as the warmup's lookup does, and from integrations on older releases (5.5.4 has only that one). Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 2f49fe3536a..37ffe7855db 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -55,7 +55,6 @@ ) from transformers.conversion_mapping import get_model_conversion_mapping from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key -from transformers.integrations.tensor_parallel import replace_layer_number_by_wildcard from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import env, moe @@ -85,6 +84,11 @@ read_safetensors_subset, ) +try: # where newer transformers keeps it; older releases, under integrations (now a partial shim) + from transformers.distributed.tensor_parallel import replace_layer_number_by_wildcard +except ImportError: + from transformers.integrations.tensor_parallel import replace_layer_number_by_wildcard + __all__ = [ "build_meta_causal_lm", "fsdp2_wrap", From e375bb11fec814393a95cf87d8f7411c9393063f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:47:49 +0000 Subject: [PATCH 11/21] Refuse a sharded dense part that from_pretrained did not leave on meta The dense part is sharded and loaded per rank only if from_pretrained leaves it on meta, which it does by asking get_device for each dense weight. A transformers that patched get_device but placed the dense part another way would load it whole on every rank -- the memory this layout exists to save -- and then shard and refill it without a word, the user seeing at most an out-of-memory error. The load now refuses that, naming MODELOPT_FSDP_SIZE=1, and the diagnostic it printed only on rank 0 goes through print_rank_0. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 23 +++++--- .../torch/utils/test_ep_dense_sharding.py | 52 +++++++++++++++++++ 2 files changed, 67 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 37ffe7855db..dce1247b1c5 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2292,14 +2292,21 @@ 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 - if plan.is_fsdp_dense and dist.is_master(): - _nmeta = [n for n, p in model.named_parameters() if p.is_meta] - _dense_meta = [n for n in _nmeta if not is_expert_fn(n)] - print( - f"[ep] dense-meta load: {len(_nmeta)} meta param(s) after from_pretrained " - f"({len(_dense_meta)} dense, {len(_nmeta) - len(_dense_meta)} expert); experts materialized. " - f"e.g. dense-meta {_dense_meta[:3]}", - flush=True, + if plan.is_fsdp_dense: + meta = [n for n, p in model.named_parameters() if p.is_meta] + dense_meta = [n for n in meta if not is_expert_fn(n)] + if not dense_meta: + # A transformers that places the dense part without asking get_device loads it whole + # on every rank, which is what this layout exists to avoid, and would do so silently. + raise RuntimeError( + "Sharding the dense part needs from_pretrained to leave it on meta, but no dense " + f"param came back on meta on transformers {transformers.__version__}: every rank " + "loaded the whole dense part. Set MODELOPT_FSDP_SIZE=1 if the dense part fits " + "one GPU." + ) + print_rank_0( + f"[ep] dense-meta load: {len(dense_meta)} dense param(s) on meta, " + f"{len(meta) - len(dense_meta)} expert; e.g. {dense_meta[:3]}" ) # 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. diff --git a/tests/unit/torch/utils/test_ep_dense_sharding.py b/tests/unit/torch/utils/test_ep_dense_sharding.py index 515c7ee4f4f..04abe82ccfb 100644 --- a/tests/unit/torch/utils/test_ep_dense_sharding.py +++ b/tests/unit/torch/utils/test_ep_dense_sharding.py @@ -376,6 +376,58 @@ def test_a_listed_container_is_wrapped_once_not_again_through_its_blocks(fsdp_me assert not any(isinstance(block, FSDPModule) for block in model.vision_tower.blocks) +class _ShardMesh: + """The meshes, as far as the loader reads them before it loads the dense part.""" + + device_type = "cpu" + + def __getitem__(self, name): + return self + + def size(self): + return 2 + + def get_group(self): + return -1 + + +def test_a_dense_part_from_pretrained_did_not_leave_on_meta_is_refused( + monkeypatch, gloo_world_of_one +): + """A transformers that places the dense part without asking ``get_device`` loads it whole on every + rank, the memory this layout exists to save: the load refuses that before sharding anything.""" + + class _DenseLoadedError(Exception): + """Raised by the dense load: the loader got past the check.""" + + def dense_load(*args, **kwargs): + raise _DenseLoadedError + + loader = SimpleNamespace(__name__="Loader", from_pretrained=lambda path, **kwargs: _MoeLM()) + monkeypatch.setattr("modelopt.torch.utils.distributed.is_initialized", lambda: True) + monkeypatch.setattr("modelopt.torch.utils.distributed.size", lambda group=None: 2) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *a, **k: _ShardMesh() + ) + monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", lambda p, **k: (None, p)) + 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, "_load_ep_dense_distributed", dense_load) + config = SimpleNamespace( + base_model_ep_plan={ + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + }, + num_experts=4, + ) + + with pytest.raises(RuntimeError, match="no dense param came back on meta"): + load_and_prepare_ep("unused", DistributionPlan(2, 2, "fsdp"), hf_config=config) + + # ------------------------------------------------------------------ _load_ep_dense_distributed From e9b1c9ba819b978d80d0ed87071ae61d25b5b9e2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:47:59 +0000 Subject: [PATCH 12/21] Warn when the EP-aware warmup sizing falls back to transformers' own The fallback keeps the load going, but transformers' full-model sizing can reserve most of the device and starve the all-to-all's buffers, so a bug in the EP-aware sizing showed up only as a later NCCL or memory failure. It now warns with the error it fell back from. fsdp2_wrap's docstring points at _fsdp2_wrap_modules, which picks the blocks for the expert-parallel dense load and keeps a VLM's vision tower replicated. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/fsdp2_utils.py | 5 +++++ modelopt/torch/utils/plugins/model_load_utils.py | 7 ++++++- tests/unit/torch/utils/test_ep_load_warmup.py | 5 +++-- 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index 419abb4ae56..dbc7c60f6e2 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/modelopt/torch/utils/fsdp2_utils.py @@ -91,6 +91,11 @@ def _move_to_fsdp_device(model, params: set[torch.nn.Parameter]) -> None: def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False): """Auto-detect a HF causal-LM's decoder layers and FSDP2 ``fully_shard`` each one. + Its callers load causal LMs and text-only VLMs; the expert-parallel dense load, which also takes + VLMs, picks its blocks through + :func:`~modelopt.torch.utils.plugins.model_load_utils._fsdp2_wrap_modules`, keeping a vision + tower replicated. + By default (``shard_root=True``) the root module is wrapped too, so embed/lm_head/norm are sharded instead of replicated per rank; pass ``shard_root=False`` to leave the root replicated (only decoder layers sharded). Returns the detected decoder layers so callers can reuse the diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index dce1247b1c5..39b4c588c40 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1817,7 +1817,12 @@ def _ep_aware(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs) f"(was full-model ~device-max; EP-aware plan-driven sizing)" ) return total - except Exception: + except Exception as e: + # Its fallback reserves most of the device, so say why, or a later NCCL failure has no cause. + warn_rank_0( + f"EP-aware warmup sizing failed ({type(e).__name__}: {e}); using transformers' " + "full-model sizing, which may reserve most of the device." + ) return orig(model, accelerator_device_map, hf_quantizer, *args, **kwargs) try: diff --git a/tests/unit/torch/utils/test_ep_load_warmup.py b/tests/unit/torch/utils/test_ep_load_warmup.py index 18c56d6d159..14477affbc1 100644 --- a/tests/unit/torch/utils/test_ep_load_warmup.py +++ b/tests/unit/torch/utils/test_ep_load_warmup.py @@ -164,7 +164,7 @@ def test_a_tp_planned_dense_param_keeps_its_tp_share_while_the_rest_stays_on_met def test_an_error_while_sizing_falls_back_to_the_stock_sizing(monkeypatch): """The warmup only sizes an allocator reservation, so whatever the EP-aware sizing cannot handle gets transformers' own number instead of failing the load, with whatever else a later - release passes the hook.""" + release passes the hook, and a warning: that number can reserve most of the device.""" stock_calls = [] def stock_sizing(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs): @@ -180,7 +180,8 @@ def param_element_size(self, model, name, param): device_map = model.device_map() with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN): - sizing = modeling_utils.get_total_byte_count(model, device_map, quantizer, 1, flag=True) + with pytest.warns(UserWarning, match="EP-aware warmup sizing failed"): + sizing = modeling_utils.get_total_byte_count(model, device_map, quantizer, 1, flag=True) assert sizing == {_GPU: 123} assert stock_calls == [(model, device_map, quantizer, (1,), {"flag": True})] From 0a67b30a618412995848ecab6fedc57fbf948f21 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:33:46 +0000 Subject: [PATCH 13/21] Say why _fsdp2_wrap_modules imports where it does Both imports come from packages that sit above this module, and the export package imports it, so they stay in the function, as fsdp2_wrap's does; the export one moves up beside the other. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 39b4c588c40..2119a13b46a 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -905,6 +905,9 @@ def _fsdp2_wrap_modules(model: "torch.nn.Module") -> list["torch.nn.Module"]: the language-only calibration never runs: sharded, their weight gathers would deadlock, so they stay replicated. A model whose language backbone shows no blocks is refused. """ + # Imported here, as fsdp2_wrap does: the export and quantization packages sit above this module, + # and the export package imports it. + from modelopt.torch.export.model_utils import get_language_model_from_vl from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector def _boundaries(root: "torch.nn.Module") -> "list[torch.nn.Module] | None": @@ -917,8 +920,6 @@ def _boundaries(root: "torch.nn.Module") -> "list[torch.nn.Module] | None": return list(layers) if layers is not None else None # VLM -> wrap only the language backbone's blocks (vision/audio stay replicated); else the whole model. - from modelopt.torch.export.model_utils import get_language_model_from_vl - lineage = get_language_model_from_vl(model) root = lineage[-1] if lineage else model mods = _boundaries(root) From 64c512182b2fd32b2391d5a37fb1ca5796604b38 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:43:04 +0000 Subject: [PATCH 14/21] Stand in a list for the unplaced keys in the dense-meta refusal test record_unplaced_source_keys returns a sorted list, which a later check iterates; the test's stand-in returned None. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_ep_dense_sharding.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_ep_dense_sharding.py b/tests/unit/torch/utils/test_ep_dense_sharding.py index 04abe82ccfb..73d290d5e4d 100644 --- a/tests/unit/torch/utils/test_ep_dense_sharding.py +++ b/tests/unit/torch/utils/test_ep_dense_sharding.py @@ -412,7 +412,7 @@ def dense_load(*args, **kwargs): ) monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", lambda p, **k: (None, p)) 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: []) monkeypatch.setattr(model_load_utils, "_load_ep_dense_distributed", dense_load) config = SimpleNamespace( base_model_ep_plan={ From f51aedf3f49d82daf6e5b4dca6b4f1c9d7e033c4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 07:01:41 +0000 Subject: [PATCH 15/21] Import the dense-wrap helpers at module scope LayerActivationCollector and get_language_model_from_vl were imported in _fsdp2_wrap_modules, which uses them. Imported first in a fresh interpreter, model_load_utils, modelopt.torch.export and modelopt.torch.quantization each load cleanly with both at module scope, on transformers 5.5.4, 5.9.0 and 5.18, so nothing calls for a local import there. fsdp2_wrap keeps its own, in fsdp2_utils, which imports no Hugging Face or quantization package. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 2119a13b46a..e43e7c1ebe6 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -56,6 +56,8 @@ from transformers.conversion_mapping import get_model_conversion_mapping from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key +from modelopt.torch.export.model_utils import get_language_model_from_vl +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import env, moe from modelopt.torch.utils.distributed import ( @@ -905,10 +907,6 @@ def _fsdp2_wrap_modules(model: "torch.nn.Module") -> list["torch.nn.Module"]: the language-only calibration never runs: sharded, their weight gathers would deadlock, so they stay replicated. A model whose language backbone shows no blocks is refused. """ - # Imported here, as fsdp2_wrap does: the export and quantization packages sit above this module, - # and the export package imports it. - from modelopt.torch.export.model_utils import get_language_model_from_vl - from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector def _boundaries(root: "torch.nn.Module") -> "list[torch.nn.Module] | None": ns = set(getattr(root, "_no_split_modules", None) or []) From b579ddd3e4f87d4ce04a192bd62b8d76e5c8efe5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 07:04:09 +0000 Subject: [PATCH 16/21] Warn when the warmup sizing hook has other parameters, and name EP_DEGREE=1 The EP-aware warmup sizing falls back to transformers' own when the hook's leading parameters differ, which is how a transformers upgrade would reach it; that sizing may reserve most of the device, so it now warns, as the error fallback does. The dense part's sharding refusal also names MODELOPT_EP_DEGREE=1, which keeps the dense part sharded by dropping expert parallelism. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 11 +++++++++-- tests/unit/torch/utils/test_ep_load_warmup.py | 8 ++++++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index e43e7c1ebe6..245e256096e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1586,7 +1586,8 @@ def _decide_distribution( "the export cannot write yet" ) remedy = ( - " Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, or " + " Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, " + "MODELOPT_EP_DEGREE=1 to FSDP2-shard the whole model without expert parallelism, or " "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." ) if dense == "auto": @@ -1757,8 +1758,14 @@ def _ep_aware_caching_allocator_warmup( except (TypeError, ValueError): params = [] # A release whose hook takes other leading parameters calls it in a way this replacement cannot - # know, so it keeps the stock sizing rather than fail the load at transformers' call site. + # know, so it keeps the stock sizing rather than fail the load at transformers' call site. Said, + # since that sizing may reserve most of the device and a later NCCL failure would have no cause. if params[:3] != ["model", "accelerator_device_map", "hf_quantizer"]: + warn_rank_0( + f"transformers {transformers.__version__}'s warmup sizing hook takes {params[:3]}, not " + "(model, accelerator_device_map, hf_quantizer); using its full-model sizing, which may " + "reserve most of the device." + ) yield return diff --git a/tests/unit/torch/utils/test_ep_load_warmup.py b/tests/unit/torch/utils/test_ep_load_warmup.py index 14477affbc1..17b1f714fe9 100644 --- a/tests/unit/torch/utils/test_ep_load_warmup.py +++ b/tests/unit/torch/utils/test_ep_load_warmup.py @@ -189,14 +189,18 @@ def param_element_size(self, model, name, param): def test_a_sizing_hook_with_other_parameters_keeps_the_stock_sizing(monkeypatch): """A release whose hook takes other leading parameters keeps transformers' own sizing: the - replacement would not know how the hook is called.""" + replacement would not know how the hook is called. It says so, since that sizing may reserve + most of the device.""" def stock_sizing(model, device_map, quantizer=None): return {_GPU: 123} monkeypatch.setattr(modeling_utils, "get_total_byte_count", stock_sizing) - with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN): + with ( + pytest.warns(UserWarning, match="using its full-model sizing"), + _ep_aware_caching_allocator_warmup(4, _A2A_PLAN), + ): assert modeling_utils.get_total_byte_count is stock_sizing From a79a929f110ef3a2656f63ac13c248355e81ee5d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 07:28:39 +0000 Subject: [PATCH 17/21] Say which VLMs keep their vision blocks unwrapped, and which buffers load as usual _fsdp2_wrap_modules searches a VLM from its language backbone only when get_language_model_from_vl recognizes it; any other model is searched whole, an unrecognized VLM's listed vision blocks included, as a test already pins. The dense meta load sends a buffer to meta too if it is named .weight or .bias, which no supported MoE has. The two docstrings said both unconditionally. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 245e256096e..b0e3bc52c6d 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -903,9 +903,11 @@ def handed() -> DistributionPlan | None: def _fsdp2_wrap_modules(model: "torch.nn.Module") -> list["torch.nn.Module"]: """The blocks to FSDP2-wrap one by one: the language backbone's ``_no_split_modules`` blocks. - Falls back to modelopt's decoder-layer detection. A VLM's own list can name vision blocks, which - the language-only calibration never runs: sharded, their weight gathers would deadlock, so they - stay replicated. A model whose language backbone shows no blocks is refused. + Falls back to modelopt's decoder-layer detection. A VLM :func:`get_language_model_from_vl` + recognizes is searched from its language backbone, so its vision blocks, which the language-only + calibration never runs, are not wrapped; any other model is searched whole, so an unrecognized + VLM's listed vision blocks are wrapped too. A model whose language backbone shows no blocks is + refused. """ def _boundaries(root: "torch.nn.Module") -> "list[torch.nn.Module] | None": @@ -917,7 +919,7 @@ def _boundaries(root: "torch.nn.Module") -> "list[torch.nn.Module] | None": layers = LayerActivationCollector.get_decoder_layers(root) return list(layers) if layers is not None else None - # VLM -> wrap only the language backbone's blocks (vision/audio stay replicated); else the whole model. + # A recognized VLM -> only its language backbone's blocks (vision/audio unwrapped); else the whole model. lineage = get_language_model_from_vl(model) root = lineage[-1] if lineage else model mods = _boundaries(root) @@ -1985,8 +1987,9 @@ def _ep_dense_meta_load(is_expert): transformers places expert params through its TP branch, which never asks ``get_device``, and every other param where ``get_device`` says: ``"meta"`` for a dense key leaves it for - :func:`_load_ep_dense_distributed` to shard and load. Buffers load as usual. ``get_device`` is - patched in every module the loader looks it up in. + :func:`_load_ep_dense_distributed` to shard and load. Buffers load as usual unless named + ``.weight``/``.bias``, which no supported MoE has: such a buffer would go to ``meta`` too, where + the meta+DCP load refuses it. ``get_device`` is patched in every module the loader looks it up in. """ loader = getattr(hf_modeling_utils, "convert_and_load_state_dict_in_model", None) mods = {hf_modeling_utils} From 9fe4dea09d68dd687f26d74a58a219f9308b24c5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 08:04:44 +0000 Subject: [PATCH 18/21] Move the dense-sharding and warmup tests into test_model_load_utils.py The two test files cover model_load_utils alone, so their tests join that module's test file, as CONTRIBUTING.md asks: test_ep_dense_sharding.py and test_ep_load_warmup.py become two sections at its end, with their fixtures and parametrization. The warmup file's module-level skip on a transformers without get_total_byte_count becomes a mark on its tests, and the plan constant both files defined is defined once. Signed-off-by: Shengliang Xu --- .../torch/utils/test_ep_dense_sharding.py | 485 ------------- tests/unit/torch/utils/test_ep_load_warmup.py | 219 ------ .../unit/torch/utils/test_model_load_utils.py | 648 +++++++++++++++++- 3 files changed, 647 insertions(+), 705 deletions(-) delete mode 100644 tests/unit/torch/utils/test_ep_dense_sharding.py delete mode 100644 tests/unit/torch/utils/test_ep_load_warmup.py diff --git a/tests/unit/torch/utils/test_ep_dense_sharding.py b/tests/unit/torch/utils/test_ep_dense_sharding.py deleted file mode 100644 index 73d290d5e4d..00000000000 --- a/tests/unit/torch/utils/test_ep_dense_sharding.py +++ /dev/null @@ -1,485 +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. - -"""How the expert-parallel loader places the dense (non-expert) part of an MoE. - -``from_pretrained`` always EP-shards the experts. ``_make_expert_param_matcher`` tells their params -from the dense ones, the distribution plan says whether the dense part is replicated or sharded, -``_ep_dense_meta_load`` keeps the dense weights off the device during ``from_pretrained`` when they -are sharded, and ``_fsdp2_shard_dense`` / -``_load_ep_dense_distributed`` FSDP2-shard them around the experts. The FSDP2 cases run on a -one-rank gloo process group on CPU. -""" - -from types import SimpleNamespace - -import pytest -import torch -from torch import nn -from torch.distributed.device_mesh import init_device_mesh -from torch.distributed.fsdp import FSDPModule -from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor - -pytest.importorskip("transformers") - -import transformers.core_model_loading as core_model_loading -import transformers.modeling_utils as modeling_utils -from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3_moe -from safetensors.torch import load_file, save_file -from transformers import AutoModelForCausalLM - -import modelopt.torch.utils.moe as moe -import modelopt.torch.utils.plugins.model_load_utils as model_load_utils -from modelopt.torch.utils.distributed import DistributionPlan -from modelopt.torch.utils.plugins.model_load_utils import ( - _ep_aware_caching_allocator_warmup, - _ep_dense_meta_load, - _fsdp2_shard_dense, - _load_ep_dense_distributed, - _make_expert_param_matcher, - load_and_prepare_ep, -) - -# ------------------------------------------------------------------ _make_expert_param_matcher - -# The plan load_and_prepare_ep shards with after the a2a swap: base-relative keys (no ``model.``). -_A2A_PLAN = { - "layers.*.mlp.gate": "torch_a2a_router", - "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", - "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "torch_a2a_experts", -} - - -# Only the per-parameter keys: a container key would match the same names on its own. -_PARAM_KEYS = {k: v for k, v in _A2A_PLAN.items() if v == "grouped_gemm"} - - -@pytest.mark.parametrize( - ("plan", "name", "expected"), - [ - pytest.param(_PARAM_KEYS, "layers.0.mlp.experts.gate_up_proj", True, id="key-unprefixed"), - pytest.param( - _PARAM_KEYS, "model.language_model.layers.3.mlp.experts.down_proj", True, id="vlm" - ), - pytest.param( - _PARAM_KEYS, - "language_model.model.layers.3.mlp.experts.gate_up_proj", - True, - id="vlm-wrapper", - ), - # NemotronH tags only the container; every param under it is an expert's. - pytest.param( - {"layers.*.mixer.experts": "torch_a2a_experts"}, - "backbone.layers.1.mixer.experts.3.down_proj.weight", - True, - id="under-a-container", - ), - pytest.param( - _A2A_PLAN, "model.layers.0.mlp.shared_experts.gate_proj.weight", False, id="shared" - ), - pytest.param( - {"layers.*.mlp.gate": "ep_router"}, "model.layers.0.mlp.gate.weight", False, id="router" - ), - # No EP credit then: the load-time warmup falls back to stock sizing. - pytest.param(None, "model.layers.0.mlp.experts.gate_up_proj", False, id="no-plan"), - ], -) -def test_expert_param_matcher(plan, name, expected): - """A param is an expert's when the plan says so, under any prefix the base-relative keys omit; - a router is replicated whatever its style. The warmup sizing tests classify a whole MoE's - experts, router, attention, embedding and head with the same matcher.""" - assert _make_expert_param_matcher(plan)(name) is expected - - -# ------------------------------------------------------------------ a block-fp8 checkpoint - - -def test_a_block_fp8_checkpoint_is_refused_as_a_format_whatever_its_dense_layout( - monkeypatch, gloo_world_of_one -): - """This load does not dequantize block-fp8 checkpoints at all, so that is the refusal a plan - sharding their dense part meets, not one advising to replicate it.""" - monkeypatch.setattr("modelopt.torch.utils.distributed.is_initialized", lambda: True) - monkeypatch.setattr("modelopt.torch.utils.distributed.size", lambda group=None: 2) - monkeypatch.setattr(moe, "require_transformers", lambda *args: None) - config = SimpleNamespace(quantization_config={"quant_method": "fp8"}) - - with pytest.raises(NotImplementedError, match="cannot dequantize a 'fp8' checkpoint"): - load_and_prepare_ep("unused", DistributionPlan(2, 2, "fsdp"), hf_config=config) - - -# ------------------------------------------------------------------------- _ep_dense_meta_load - - -def _is_expert(name): - return ".experts." in name - - -def _get_device_lookups(): - """The modules whose ``get_device`` global the state-dict loaders call.""" - return modeling_utils, core_model_loading - - -def test_dense_weights_resolve_to_meta_and_the_rest_to_the_stock_device(monkeypatch): - """Only a dense ``.weight`` / ``.bias`` is sent to meta. An expert param or a buffer goes to - the stock lookup, with the caller's arguments, in both modules the loaders read it from.""" - stock_calls = [] - - def stock_get_device(device_map, key, *args, **kwargs): - stock_calls.append((key, args, kwargs)) - return device_map[""] - - lookups = _get_device_lookups() - for module in lookups: - monkeypatch.setattr(module, "get_device", stock_get_device) - device_map = {"": "cuda:0"} - expert = "model.layers.0.mlp.experts.0.gate_proj.weight" - buffer = "model.layers.0.mlp.gate.e_score_correction_bias" - - with _ep_dense_meta_load(_is_expert): - for module in lookups: - assert module.get_device(device_map, "model.layers.0.self_attn.q_proj.weight") == "meta" - assert module.get_device(device_map, "model.layers.0.self_attn.q_proj.bias") == "meta" - assert module.get_device(device_map, expert, valid_torch_device=True) == "cuda:0" - assert module.get_device(device_map, buffer) == "cuda:0" - - assert stock_calls == [(expert, (), {"valid_torch_device": True}), (buffer, (), {})] * 2 - - -def test_the_dense_load_refuses_a_transformers_without_its_get_device_hook(monkeypatch): - """Without the hook the dense part would load whole on every rank, the memory sharding it saves.""" - for module in _get_device_lookups(): - monkeypatch.delattr(module, "get_device", raising=False) - - with pytest.raises(RuntimeError, match="get_device"), _ep_dense_meta_load(_is_expert): - pass - - -def _get_device_patch(): - return _ep_dense_meta_load(_is_expert), [(m, "get_device") for m in _get_device_lookups()] - - -def _warmup_sizing_patch(): - if not hasattr(modeling_utils, "get_total_byte_count"): - pytest.skip("transformers has no get_total_byte_count to size") - sizing = _ep_aware_caching_allocator_warmup(4, _A2A_PLAN) - return sizing, [(modeling_utils, "get_total_byte_count")] - - -@pytest.mark.parametrize( - "patch", [_get_device_patch, _warmup_sizing_patch], ids=lambda f: f.__name__ -) -def test_a_load_time_patch_is_undone_even_when_the_load_raises(patch): - """``from_pretrained`` reads these transformers globals; a failed load leaves stock ones.""" - context, patched = patch() - stock = [getattr(module, attr) for module, attr in patched] - - with pytest.raises(RuntimeError, match="load failed"), context: - assert all(getattr(m, a) is not f for (m, a), f in zip(patched, stock)) - raise RuntimeError("load failed") - - assert all(getattr(m, a) is f for (m, a), f in zip(patched, stock)) - - -def test_from_pretrained_leaves_the_dense_weights_on_meta_and_loads_the_rest(tmp_path): - """Through a real CPU ``from_pretrained``: every dense weight stays on meta for the FSDP2 + - DCP load to fill, while the MLPs (standing in for experts) and the rotary buffers load.""" - src = get_tiny_llama() - src.save_pretrained(tmp_path) - - with _ep_dense_meta_load(lambda name: ".mlp." in name): - model = AutoModelForCausalLM.from_pretrained(tmp_path) - - ref_params, ref_buffers = src.state_dict(), dict(src.named_buffers()) - for name, param in model.named_parameters(): - if ".mlp." in name: - assert torch.equal(param, ref_params[name]), name - else: - assert param.is_meta, name - for name, buffer in model.named_buffers(): - assert torch.equal(buffer, ref_buffers[name]), name - - -# ------------------------------------------------------------------------- _fsdp2_shard_dense - - -@pytest.fixture -def fsdp_mesh(skip_on_windows, tmp_path_factory): - """The dense FSDP mesh of a one-rank gloo process group (FSDP2 needs a process group).""" - store = tmp_path_factory.mktemp("pg") / "store" - torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) - yield init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) - torch.distributed.destroy_process_group() - - -class _Experts(nn.Module): - def __init__(self): - super().__init__() - self.gate_up_proj = nn.Parameter(torch.randn(4, 8, 4)) - self.down_proj = nn.Parameter(torch.randn(4, 4, 4)) - - -class _DecoderLayer(nn.Module): - def __init__(self): - super().__init__() - self.self_attn = nn.Linear(4, 4) - self.mlp = nn.Module() - self.mlp.gate = nn.Linear(4, 4, bias=False) # the router: dense, sharded with its layer - self.mlp.experts = _Experts() - self.input_layernorm = nn.LayerNorm(4) - - -class _MoeLM(nn.Module): - """A causal-LM-shaped MoE whose ``_no_split_modules`` names its decoder layer, as HF's do.""" - - _no_split_modules = ["_DecoderLayer"] - - def __init__(self): - super().__init__() - self.model = nn.Module() - self.model.embed_tokens = nn.Embedding(16, 4) - self.model.layers = nn.ModuleList(_DecoderLayer() for _ in range(2)) - self.model.norm = nn.LayerNorm(4) - self.lm_head = nn.Linear(4, 16, bias=False) - - def experts(self): - return [layer.mlp.experts for layer in self.model.layers] - - -def test_each_decoder_layer_is_sharded_around_its_experts(fsdp_mesh): - """A layer's dense params (attention, router, norms) become shards on the FSDP mesh; its - experts stay as loaded, for EP. The root's embedding, final norm and head stay replicated, and - a layer's submodules keep their classes (modelopt recognizes them by class).""" - model = _MoeLM() - before = {name: param.detach().clone() for name, param in model.named_parameters()} - experts = {name: param for name, param in model.named_parameters() if _is_expert(name)} - - assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2 - - assert all(isinstance(layer, FSDPModule) for layer in model.model.layers) - assert not isinstance(model, FSDPModule) - assert type(model.model.layers[0].self_attn) is nn.Linear - for name, param in model.named_parameters(): - if name in experts: - assert param is experts[name] and not isinstance(param, DTensor), name - elif name.startswith("model.layers."): - assert isinstance(param, DTensor), name - assert param.device_mesh == fsdp_mesh and param.placements == (Shard(0),), name - assert torch.equal(param.full_tensor(), before[name]), name - else: - assert not isinstance(param, DTensor), name - - -def test_a_param_off_the_models_dtype_stays_replicated(fsdp_mesh): - """FSDP2 shards one dtype per group, and HF keeps a few params in fp32, as a router: such a - param stays replicated, as the pure-FSDP2 wrap leaves it, and its layer still unshards.""" - model = _MoeLM().to(torch.bfloat16) - gate = model.model.layers[0].mlp.gate.float() - - assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2 - - assert not isinstance(gate.weight, DTensor) and gate.weight.dtype == torch.float32 - for layer in model.model.layers: - layer.unshard() # FSDP2 checks a group's dtypes here - - -def test_a_layer_holding_only_experts_is_left_unwrapped(fsdp_mesh): - model = _MoeLM() - bare = model.model.layers[1] - bare.self_attn = bare.mlp.gate = bare.input_layernorm = nn.Identity() - - assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 1 - assert not isinstance(bare, FSDPModule) - - -def test_a_dense_part_left_with_nothing_to_shard_is_refused(fsdp_mesh): - """Experts outweighing the dense part in another dtype put every dense param off the model's - dominant dtype, out of the shard groups: no layer is left to shard, and the load would hold the - whole dense part on every rank.""" - model = _MoeLM() - for experts in model.experts(): - experts.to(torch.bfloat16) - - with pytest.raises(RuntimeError, match="no dense parameter to shard"): - _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) - - -def test_a_dense_param_the_loader_already_sharded_stays_as_it_is(fsdp_mesh): - """A dense param that arrives as a ``Shard`` DTensor of the loader's own plan is not resharded; - the rest of its layer still is.""" - tp_mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("tp",)) - model = _MoeLM() - attn = model.model.layers[0].self_attn - attn.weight = nn.Parameter(distribute_tensor(attn.weight.detach(), tp_mesh, [Shard(1)])) - presharded = attn.weight - - assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2 - - assert attn.weight is presharded and attn.weight.placements == (Shard(1),) - assert model.model.layers[0].mlp.gate.weight.placements == (Shard(0),) - - -def test_a_replicated_dense_param_is_sharded_over_the_fsdp_mesh_like_a_plain_one(fsdp_mesh): - ep_mesh = init_device_mesh("cpu", (1, 1), mesh_dim_names=("edp", "ep"))["ep"] - model = _MoeLM() - gate = model.model.layers[0].mlp.gate - weight = gate.weight.detach().clone() - gate.weight = nn.Parameter(distribute_tensor(weight, ep_mesh, [Replicate()])) - - _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) - - assert gate.weight.device_mesh == fsdp_mesh and gate.weight.placements == (Shard(0),) - assert torch.equal(gate.weight.full_tensor(), weight) - - -class _VisionBlock(nn.Module): - def __init__(self): - super().__init__() - self.proj = nn.Linear(4, 4) - - -class _VisionTower(nn.Module): - def __init__(self): - super().__init__() - self.blocks = nn.ModuleList(_VisionBlock() for _ in range(2)) - - -class _MoeVlm(_MoeLM): - """Kimi's MoonViT style: ``_no_split_modules`` lists the vision tower *and* its blocks.""" - - _no_split_modules = ["_VisionTower", "_VisionBlock", "_DecoderLayer"] - - def __init__(self): - super().__init__() - self.vision_tower = _VisionTower() - - -def test_a_listed_container_is_wrapped_once_not_again_through_its_blocks(fsdp_mesh): - model = _MoeVlm() - - assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 3 - - assert isinstance(model.vision_tower, FSDPModule) - assert not any(isinstance(block, FSDPModule) for block in model.vision_tower.blocks) - - -class _ShardMesh: - """The meshes, as far as the loader reads them before it loads the dense part.""" - - device_type = "cpu" - - def __getitem__(self, name): - return self - - def size(self): - return 2 - - def get_group(self): - return -1 - - -def test_a_dense_part_from_pretrained_did_not_leave_on_meta_is_refused( - monkeypatch, gloo_world_of_one -): - """A transformers that places the dense part without asking ``get_device`` loads it whole on every - rank, the memory this layout exists to save: the load refuses that before sharding anything.""" - - class _DenseLoadedError(Exception): - """Raised by the dense load: the loader got past the check.""" - - def dense_load(*args, **kwargs): - raise _DenseLoadedError - - loader = SimpleNamespace(__name__="Loader", from_pretrained=lambda path, **kwargs: _MoeLM()) - monkeypatch.setattr("modelopt.torch.utils.distributed.is_initialized", lambda: True) - monkeypatch.setattr("modelopt.torch.utils.distributed.size", lambda group=None: 2) - monkeypatch.setattr(moe, "require_transformers", lambda *args: None) - monkeypatch.setattr( - "torch.distributed.device_mesh.init_device_mesh", lambda *a, **k: _ShardMesh() - ) - monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", lambda p, **k: (None, p)) - monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) - monkeypatch.setattr(model_load_utils, "_load_ep_dense_distributed", dense_load) - config = SimpleNamespace( - base_model_ep_plan={ - "layers.*.mlp.gate": "ep_router", - "layers.*.mlp.experts": "moe_tp_experts", - "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", - "layers.*.mlp.experts.down_proj": "grouped_gemm", - }, - num_experts=4, - ) - - with pytest.raises(RuntimeError, match="no dense param came back on meta"): - load_and_prepare_ep("unused", DistributionPlan(2, 2, "fsdp"), hf_config=config) - - -# ------------------------------------------------------------------ _load_ep_dense_distributed - - -def test_the_dense_load_shards_first_then_asks_for_exactly_the_meta_dense_params( - fsdp_mesh, monkeypatch, tmp_path -): - """As ``_ep_dense_meta_load`` leaves the model: dense weights on meta; experts loaded, here - per-expert ``nn.Linear``s, whose ``.weight`` only the expert matcher tells from a dense one; - and a Mamba-style ``A_log`` (no ``.weight`` / ``.bias`` suffix) loaded too. The dense layers - are FSDP2-sharded before the shared meta+DCP loader runs, and it is asked for exactly the meta - dense params, on the mesh's device. A recorder stands in for that load.""" - with torch.device("meta"): - model = _MoeLM() - for layer in model.model.layers: - layer.mlp.experts = nn.ModuleList(nn.Linear(4, 4, bias=False) for _ in range(4)) - layer.A_log = nn.Parameter(torch.zeros(4)) - loads = [] - - def record_load(loaded, local_dir, device, *, load_names): - sharded = [isinstance(layer.self_attn.weight, DTensor) for layer in loaded.model.layers] - loads.append((local_dir, device, load_names, sharded)) - - monkeypatch.setattr(model_load_utils, "_meta_dcp_load", record_load) - - wrapped = _load_ep_dense_distributed( - model, fsdp_mesh, str(tmp_path), _is_expert, model.experts() - ) - - assert wrapped == 2 - meta_dense = {name for name, param in model.named_parameters() if param.is_meta} - assert loads == [(str(tmp_path), torch.device("cpu"), meta_dense, [True, True])] - assert not any(_is_expert(name) or name.endswith("A_log") for name in meta_dense) - - -def test_a_tied_head_the_checkpoint_stores_alone_reaches_both_aliases(fsdp_mesh, tmp_path): - """A tied checkpoint can store the pair as ``lm_head.weight`` alone, which ``from_pretrained`` - loads into both aliases. With a tied MoE as the expert-parallel load leaves it -- the dense part - on meta and tied, the experts loaded -- the dense load fills both aliases with it too.""" - src = get_tiny_qwen3_moe(tie_word_embeddings=True) - src.save_pretrained(tmp_path) - weights = tmp_path / "model.safetensors" - stored = load_file(weights) - stored["lm_head.weight"] = stored.pop("model.embed_tokens.weight") - save_file(stored, weights, metadata={"format": "pt"}) - with _ep_dense_meta_load(_is_expert): - model = AutoModelForCausalLM.from_pretrained(tmp_path) - model.tie_weights() # as load_and_prepare_ep does once the weights load - assert model.lm_head.weight is model.model.embed_tokens.weight - assert model.lm_head.weight.is_meta - - experts = [layer.mlp.experts for layer in model.model.layers] - _load_ep_dense_distributed(model, fsdp_mesh, str(tmp_path), _is_expert, experts) - - for alias in (model.lm_head.weight, model.model.embed_tokens.weight): - assert torch.equal(alias, src.lm_head.weight) diff --git a/tests/unit/torch/utils/test_ep_load_warmup.py b/tests/unit/torch/utils/test_ep_load_warmup.py deleted file mode 100644 index 17b1f714fe9..00000000000 --- a/tests/unit/torch/utils/test_ep_load_warmup.py +++ /dev/null @@ -1,219 +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. - -"""The expert-parallel sizing of transformers' load-time caching-allocator warmup. - -``from_pretrained`` reserves the bytes ``get_total_byte_count`` reports for each device in one -allocation. Under the EP load it would count every expert at full size (the experts are most of a -large MoE), so ``_ep_aware_caching_allocator_warmup`` swaps in a sizing that counts one EP shard of -each expert. The sizing is called directly here, and once through ``caching_allocator_warmup`` with -the CUDA queries and the allocation faked. -""" - -import pytest -import torch -from torch import nn - -pytest.importorskip("transformers") - -from transformers import modeling_utils - -from modelopt.torch.utils.plugins.model_load_utils import _ep_aware_caching_allocator_warmup - -if not hasattr(modeling_utils, "get_total_byte_count"): - pytest.skip("transformers has no get_total_byte_count to size", allow_module_level=True) - -_GPU = torch.device("cuda", 0) - -# The plan load_and_prepare_ep shards with after the a2a swap: base-relative keys (no ``model.``). -_A2A_PLAN = { - "layers.*.mlp.gate": "torch_a2a_router", - "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", - "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "torch_a2a_experts", -} - -# _TinyMoe's fp32 bytes: the fused experts; the dense params (embedding, attention weight + bias, -# router); the router's persistent ``e_score_correction_bias`` buffer. The tied head adds none. -_EXPERT_BYTES = 4 * (4 * 16 * 8 + 4 * 8 * 8) -_DENSE_BYTES = 4 * (16 * 8 + 8 * 8 + 8 + 4 * 8) -_BUFFER_BYTES = 4 * 4 -_ATTN_BYTES = 4 * (8 * 8 + 8) - - -class _TinyMoe(nn.Module): - """One MoE decoder layer, a head tied to the embedding, and a persistent router buffer.""" - - def __init__(self): - super().__init__() - self.model = nn.Module() - self.model.embed_tokens = nn.Embedding(16, 8) - layer = nn.Module() - layer.self_attn = nn.Linear(8, 8) - layer.mlp = nn.Module() - layer.mlp.gate = nn.Linear(8, 4, bias=False) - layer.mlp.gate.register_buffer("e_score_correction_bias", torch.zeros(4)) - layer.mlp.experts = nn.Module() - layer.mlp.experts.gate_up_proj = nn.Parameter(torch.zeros(4, 16, 8)) - layer.mlp.experts.down_proj = nn.Parameter(torch.zeros(4, 8, 8)) - self.model.layers = nn.ModuleList([layer]) - self.lm_head = nn.Linear(8, 16, bias=False) - self.lm_head.weight = self.model.embed_tokens.weight - self.all_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - - def get_parameter_or_buffer(self, name): - """``PreTrainedModel``'s lookup, which the sizing goes through.""" - try: - return self.get_parameter(name) - except AttributeError: - return self.get_buffer(name) - - def device_map(self): - """Every param and persistent buffer on one GPU, as ``from_pretrained`` expands it.""" - names = [name for name, _ in self.named_parameters(remove_duplicate=False)] - names += [name for name, _ in self.named_buffers()] - return dict.fromkeys(names, _GPU) - - -class _OneBytePerElement: - def param_element_size(self, model, name, param): - return 1 - - -_ALL_BYTES = _EXPERT_BYTES + _DENSE_BYTES + _BUFFER_BYTES - - -@pytest.mark.parametrize( - ("dense_on_meta", "quantizer", "stock_bytes", "per_rank_bytes"), - [ - # Experts count one EP shard; the router is replicated, so it counts in full like the rest. - pytest.param( - False, - None, - _ALL_BYTES, - _EXPERT_BYTES // 4 + _DENSE_BYTES + _BUFFER_BYTES, - id="dense-replicated", - ), - # Sharded dense: every dense .weight / .bias stays on meta during from_pretrained, while - # persistent buffers still load onto the device. - pytest.param( - True, None, _ALL_BYTES, _EXPERT_BYTES // 4 + _BUFFER_BYTES, id="dense-on-meta" - ), - # The quantizer sizes the elements: one byte each where fp32 takes four. - pytest.param( - False, - _OneBytePerElement(), - _ALL_BYTES // 4, - (_EXPERT_BYTES // 4 + _DENSE_BYTES + _BUFFER_BYTES) // 4, - id="quantizer", - ), - ], -) -def test_transformers_reserves_only_the_per_rank_footprint( - monkeypatch, dense_on_meta, quantizer, stock_bytes, per_rank_bytes -): - """End to end: the one up-front allocation ``caching_allocator_warmup`` makes shrinks from the - whole model (stock sizing, which also checks the byte counts above) to what this EP rank - holds.""" - reserved = [] - - def fake_empty(numel, *, dtype, device, **kwargs): - reserved.append((numel * torch.tensor([], dtype=dtype).element_size(), device)) - - model = _TinyMoe() - with monkeypatch.context() as m: - m.setattr(torch.cuda, "mem_get_info", lambda index: (80 * 1024**3, 80 * 1024**3)) - m.setattr(torch.cuda, "memory_reserved", lambda index: 0) - m.setattr(torch.cuda, "memory_allocated", lambda index: 0) - m.setattr(torch, "empty", fake_empty) - modeling_utils.caching_allocator_warmup(model, model.device_map(), quantizer) - with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN, dense_on_meta=dense_on_meta): - modeling_utils.caching_allocator_warmup(model, model.device_map(), quantizer) - - assert reserved == [(stock_bytes, _GPU), (per_rank_bytes, _GPU)] - - -def test_a_tp_planned_dense_param_keeps_its_tp_share_while_the_rest_stays_on_meta(monkeypatch): - """transformers places a TP-planned param itself, never asking where the dense part goes, so it - is on the device at its TP share even when the rest of the dense part stays on meta. The share - is of the ep submesh it is placed over, not of the world, which holds two of them here.""" - model = _TinyMoe() - model._tp_plan = {"model.layers.*.self_attn": "colwise"} - monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) - monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 8) - monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) - - with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN, dense_on_meta=True): - total = modeling_utils.get_total_byte_count(model, model.device_map()) - - assert total == {_GPU: _EXPERT_BYTES // 4 + _ATTN_BYTES // 4 + _BUFFER_BYTES} - - -def test_an_error_while_sizing_falls_back_to_the_stock_sizing(monkeypatch): - """The warmup only sizes an allocator reservation, so whatever the EP-aware sizing cannot - handle gets transformers' own number instead of failing the load, with whatever else a later - release passes the hook, and a warning: that number can reserve most of the device.""" - stock_calls = [] - - def stock_sizing(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs): - stock_calls.append((model, accelerator_device_map, hf_quantizer, args, kwargs)) - return {_GPU: 123} - - class _UnsupportedQuantizer: - def param_element_size(self, model, name, param): - raise TypeError("unsupported param") - - monkeypatch.setattr(modeling_utils, "get_total_byte_count", stock_sizing) - model, quantizer = _TinyMoe(), _UnsupportedQuantizer() - device_map = model.device_map() - - with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN): - with pytest.warns(UserWarning, match="EP-aware warmup sizing failed"): - sizing = modeling_utils.get_total_byte_count(model, device_map, quantizer, 1, flag=True) - assert sizing == {_GPU: 123} - - assert stock_calls == [(model, device_map, quantizer, (1,), {"flag": True})] - - -def test_a_sizing_hook_with_other_parameters_keeps_the_stock_sizing(monkeypatch): - """A release whose hook takes other leading parameters keeps transformers' own sizing: the - replacement would not know how the hook is called. It says so, since that sizing may reserve - most of the device.""" - - def stock_sizing(model, device_map, quantizer=None): - return {_GPU: 123} - - monkeypatch.setattr(modeling_utils, "get_total_byte_count", stock_sizing) - - with ( - pytest.warns(UserWarning, match="using its full-model sizing"), - _ep_aware_caching_allocator_warmup(4, _A2A_PLAN), - ): - assert modeling_utils.get_total_byte_count is stock_sizing - - -@pytest.mark.parametrize("ep_degree", [1, 0]) -def test_without_expert_parallelism_the_stock_sizing_stays(ep_degree): - stock = modeling_utils.get_total_byte_count - - with _ep_aware_caching_allocator_warmup(ep_degree, _A2A_PLAN): - assert modeling_utils.get_total_byte_count is stock - - -def test_a_transformers_without_the_sizing_hook_is_left_alone(monkeypatch): - monkeypatch.delattr(modeling_utils, "get_total_byte_count") - - with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN): - assert not hasattr(modeling_utils, "get_total_byte_count") diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 85f3243278f..3f13c2fa35f 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -36,16 +36,23 @@ import torch from _test_utils.torch.distributed.utils import bounded_collectives from _test_utils.torch.moe_models import model_with_decoder_at -from safetensors.torch import save_file +from safetensors.torch import load_file, save_file from torch import nn +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import FSDPModule +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor pytest.importorskip("accelerate") transformers = pytest.importorskip("transformers") +import transformers.core_model_loading as core_model_loading import transformers.distributed as hf_distributed +import transformers.modeling_utils as modeling_utils from _test_utils.torch.transformers_models import ( _TinyQwen35Tokenizer, + get_tiny_llama, get_tiny_qwen3_5_moe_vl, get_tiny_qwen3_5_vl, + get_tiny_qwen3_moe, ) from packaging.version import Version from transformers import ( @@ -82,7 +89,12 @@ _ckpt_quant_method, _conversion_plan, _convert_keys, + _ep_aware_caching_allocator_warmup, + _ep_dense_meta_load, _experts_sharded_over, + _fsdp2_shard_dense, + _load_ep_dense_distributed, + _make_expert_param_matcher, _resolve_target, _seed_parallel_states, _transformers_mesh_builder, @@ -2630,3 +2642,637 @@ def test_each_rank_builds_the_meshes_on_a_gpu_of_its_own( cpu_dist_workers.run( _gpu_layout_job, visible=visible, bound=bound, shared=shared, refused=refused ) + + +# ------------------------------------------------------- the expert-parallel load's dense part +# ``from_pretrained`` always EP-shards the experts. ``_make_expert_param_matcher`` tells their params +# from the dense ones, the distribution plan says whether the dense part is replicated or sharded, +# ``_ep_dense_meta_load`` keeps the dense weights off the device during ``from_pretrained`` when they +# are sharded, and ``_fsdp2_shard_dense`` / ``_load_ep_dense_distributed`` FSDP2-shard them around +# the experts. The FSDP2 cases run on a one-rank gloo process group on CPU. + +# ------------------------------------------------------------------ _make_expert_param_matcher + +# The plan load_and_prepare_ep shards with after the a2a swap: base-relative keys (no ``model.``). +_A2A_PLAN = { + "layers.*.mlp.gate": "torch_a2a_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "torch_a2a_experts", +} + + +# Only the per-parameter keys: a container key would match the same names on its own. +_PARAM_KEYS = {k: v for k, v in _A2A_PLAN.items() if v == "grouped_gemm"} + + +@pytest.mark.parametrize( + ("plan", "name", "expected"), + [ + pytest.param(_PARAM_KEYS, "layers.0.mlp.experts.gate_up_proj", True, id="key-unprefixed"), + pytest.param( + _PARAM_KEYS, "model.language_model.layers.3.mlp.experts.down_proj", True, id="vlm" + ), + pytest.param( + _PARAM_KEYS, + "language_model.model.layers.3.mlp.experts.gate_up_proj", + True, + id="vlm-wrapper", + ), + # NemotronH tags only the container; every param under it is an expert's. + pytest.param( + {"layers.*.mixer.experts": "torch_a2a_experts"}, + "backbone.layers.1.mixer.experts.3.down_proj.weight", + True, + id="under-a-container", + ), + pytest.param( + _A2A_PLAN, "model.layers.0.mlp.shared_experts.gate_proj.weight", False, id="shared" + ), + pytest.param( + {"layers.*.mlp.gate": "ep_router"}, "model.layers.0.mlp.gate.weight", False, id="router" + ), + # No EP credit then: the load-time warmup falls back to stock sizing. + pytest.param(None, "model.layers.0.mlp.experts.gate_up_proj", False, id="no-plan"), + ], +) +def test_expert_param_matcher(plan, name, expected): + """A param is an expert's when the plan says so, under any prefix the base-relative keys omit; + a router is replicated whatever its style. The warmup sizing tests classify a whole MoE's + experts, router, attention, embedding and head with the same matcher.""" + assert _make_expert_param_matcher(plan)(name) is expected + + +# ------------------------------------------------------------------ a block-fp8 checkpoint + + +def test_a_block_fp8_checkpoint_is_refused_as_a_format_whatever_its_dense_layout( + monkeypatch, gloo_world_of_one +): + """This load does not dequantize block-fp8 checkpoints at all, so that is the refusal a plan + sharding their dense part meets, not one advising to replicate it.""" + monkeypatch.setattr("modelopt.torch.utils.distributed.is_initialized", lambda: True) + monkeypatch.setattr("modelopt.torch.utils.distributed.size", lambda group=None: 2) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + config = SimpleNamespace(quantization_config={"quant_method": "fp8"}) + + with pytest.raises(NotImplementedError, match="cannot dequantize a 'fp8' checkpoint"): + load_and_prepare_ep("unused", DistributionPlan(2, 2, "fsdp"), hf_config=config) + + +# ------------------------------------------------------------------------- _ep_dense_meta_load + + +def _is_expert(name): + return ".experts." in name + + +def _get_device_lookups(): + """The modules whose ``get_device`` global the state-dict loaders call.""" + return modeling_utils, core_model_loading + + +def test_dense_weights_resolve_to_meta_and_the_rest_to_the_stock_device(monkeypatch): + """Only a dense ``.weight`` / ``.bias`` is sent to meta. An expert param or a buffer goes to + the stock lookup, with the caller's arguments, in both modules the loaders read it from.""" + stock_calls = [] + + def stock_get_device(device_map, key, *args, **kwargs): + stock_calls.append((key, args, kwargs)) + return device_map[""] + + lookups = _get_device_lookups() + for module in lookups: + monkeypatch.setattr(module, "get_device", stock_get_device) + device_map = {"": "cuda:0"} + expert = "model.layers.0.mlp.experts.0.gate_proj.weight" + buffer = "model.layers.0.mlp.gate.e_score_correction_bias" + + with _ep_dense_meta_load(_is_expert): + for module in lookups: + assert module.get_device(device_map, "model.layers.0.self_attn.q_proj.weight") == "meta" + assert module.get_device(device_map, "model.layers.0.self_attn.q_proj.bias") == "meta" + assert module.get_device(device_map, expert, valid_torch_device=True) == "cuda:0" + assert module.get_device(device_map, buffer) == "cuda:0" + + assert stock_calls == [(expert, (), {"valid_torch_device": True}), (buffer, (), {})] * 2 + + +def test_the_dense_load_refuses_a_transformers_without_its_get_device_hook(monkeypatch): + """Without the hook the dense part would load whole on every rank, the memory sharding it saves.""" + for module in _get_device_lookups(): + monkeypatch.delattr(module, "get_device", raising=False) + + with pytest.raises(RuntimeError, match="get_device"), _ep_dense_meta_load(_is_expert): + pass + + +def _get_device_patch(): + return _ep_dense_meta_load(_is_expert), [(m, "get_device") for m in _get_device_lookups()] + + +def _warmup_sizing_patch(): + if not hasattr(modeling_utils, "get_total_byte_count"): + pytest.skip("transformers has no get_total_byte_count to size") + sizing = _ep_aware_caching_allocator_warmup(4, _A2A_PLAN) + return sizing, [(modeling_utils, "get_total_byte_count")] + + +@pytest.mark.parametrize( + "patch", [_get_device_patch, _warmup_sizing_patch], ids=lambda f: f.__name__ +) +def test_a_load_time_patch_is_undone_even_when_the_load_raises(patch): + """``from_pretrained`` reads these transformers globals; a failed load leaves stock ones.""" + context, patched = patch() + stock = [getattr(module, attr) for module, attr in patched] + + with pytest.raises(RuntimeError, match="load failed"), context: + assert all(getattr(m, a) is not f for (m, a), f in zip(patched, stock)) + raise RuntimeError("load failed") + + assert all(getattr(m, a) is f for (m, a), f in zip(patched, stock)) + + +def test_from_pretrained_leaves_the_dense_weights_on_meta_and_loads_the_rest(tmp_path): + """Through a real CPU ``from_pretrained``: every dense weight stays on meta for the FSDP2 + + DCP load to fill, while the MLPs (standing in for experts) and the rotary buffers load.""" + src = get_tiny_llama() + src.save_pretrained(tmp_path) + + with _ep_dense_meta_load(lambda name: ".mlp." in name): + model = AutoModelForCausalLM.from_pretrained(tmp_path) + + ref_params, ref_buffers = src.state_dict(), dict(src.named_buffers()) + for name, param in model.named_parameters(): + if ".mlp." in name: + assert torch.equal(param, ref_params[name]), name + else: + assert param.is_meta, name + for name, buffer in model.named_buffers(): + assert torch.equal(buffer, ref_buffers[name]), name + + +# ------------------------------------------------------------------------- _fsdp2_shard_dense + + +@pytest.fixture +def fsdp_mesh(skip_on_windows, tmp_path_factory): + """The dense FSDP mesh of a one-rank gloo process group (FSDP2 needs a process group).""" + store = tmp_path_factory.mktemp("pg") / "store" + torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) + yield init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + torch.distributed.destroy_process_group() + + +class _LayerExperts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.randn(4, 8, 4)) + self.down_proj = nn.Parameter(torch.randn(4, 4, 4)) + + +class _DecoderLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Linear(4, 4) + self.mlp = nn.Module() + self.mlp.gate = nn.Linear(4, 4, bias=False) # the router: dense, sharded with its layer + self.mlp.experts = _LayerExperts() + self.input_layernorm = nn.LayerNorm(4) + + +class _MoeLM(nn.Module): + """A causal-LM-shaped MoE whose ``_no_split_modules`` names its decoder layer, as HF's do.""" + + _no_split_modules = ["_DecoderLayer"] + + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.embed_tokens = nn.Embedding(16, 4) + self.model.layers = nn.ModuleList(_DecoderLayer() for _ in range(2)) + self.model.norm = nn.LayerNorm(4) + self.lm_head = nn.Linear(4, 16, bias=False) + + def experts(self): + return [layer.mlp.experts for layer in self.model.layers] + + +def test_each_decoder_layer_is_sharded_around_its_experts(fsdp_mesh): + """A layer's dense params (attention, router, norms) become shards on the FSDP mesh; its + experts stay as loaded, for EP. The root's embedding, final norm and head stay replicated, and + a layer's submodules keep their classes (modelopt recognizes them by class).""" + model = _MoeLM() + before = {name: param.detach().clone() for name, param in model.named_parameters()} + experts = {name: param for name, param in model.named_parameters() if _is_expert(name)} + + assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2 + + assert all(isinstance(layer, FSDPModule) for layer in model.model.layers) + assert not isinstance(model, FSDPModule) + assert type(model.model.layers[0].self_attn) is nn.Linear + for name, param in model.named_parameters(): + if name in experts: + assert param is experts[name] and not isinstance(param, DTensor), name + elif name.startswith("model.layers."): + assert isinstance(param, DTensor), name + assert param.device_mesh == fsdp_mesh and param.placements == (Shard(0),), name + assert torch.equal(param.full_tensor(), before[name]), name + else: + assert not isinstance(param, DTensor), name + + +def test_a_param_off_the_models_dtype_stays_replicated(fsdp_mesh): + """FSDP2 shards one dtype per group, and HF keeps a few params in fp32, as a router: such a + param stays replicated, as the pure-FSDP2 wrap leaves it, and its layer still unshards.""" + model = _MoeLM().to(torch.bfloat16) + gate = model.model.layers[0].mlp.gate.float() + + assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2 + + assert not isinstance(gate.weight, DTensor) and gate.weight.dtype == torch.float32 + for layer in model.model.layers: + layer.unshard() # FSDP2 checks a group's dtypes here + + +def test_a_layer_holding_only_experts_is_left_unwrapped(fsdp_mesh): + model = _MoeLM() + bare = model.model.layers[1] + bare.self_attn = bare.mlp.gate = bare.input_layernorm = nn.Identity() + + assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 1 + assert not isinstance(bare, FSDPModule) + + +def test_a_dense_part_left_with_nothing_to_shard_is_refused(fsdp_mesh): + """Experts outweighing the dense part in another dtype put every dense param off the model's + dominant dtype, out of the shard groups: no layer is left to shard, and the load would hold the + whole dense part on every rank.""" + model = _MoeLM() + for experts in model.experts(): + experts.to(torch.bfloat16) + + with pytest.raises(RuntimeError, match="no dense parameter to shard"): + _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) + + +def test_a_dense_param_the_loader_already_sharded_stays_as_it_is(fsdp_mesh): + """A dense param that arrives as a ``Shard`` DTensor of the loader's own plan is not resharded; + the rest of its layer still is.""" + tp_mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("tp",)) + model = _MoeLM() + attn = model.model.layers[0].self_attn + attn.weight = nn.Parameter(distribute_tensor(attn.weight.detach(), tp_mesh, [Shard(1)])) + presharded = attn.weight + + assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2 + + assert attn.weight is presharded and attn.weight.placements == (Shard(1),) + assert model.model.layers[0].mlp.gate.weight.placements == (Shard(0),) + + +def test_a_replicated_dense_param_is_sharded_over_the_fsdp_mesh_like_a_plain_one(fsdp_mesh): + ep_mesh = init_device_mesh("cpu", (1, 1), mesh_dim_names=("edp", "ep"))["ep"] + model = _MoeLM() + gate = model.model.layers[0].mlp.gate + weight = gate.weight.detach().clone() + gate.weight = nn.Parameter(distribute_tensor(weight, ep_mesh, [Replicate()])) + + _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) + + assert gate.weight.device_mesh == fsdp_mesh and gate.weight.placements == (Shard(0),) + assert torch.equal(gate.weight.full_tensor(), weight) + + +class _VisionBlock(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + + +class _VisionTower(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList(_VisionBlock() for _ in range(2)) + + +class _MoeVlm(_MoeLM): + """Kimi's MoonViT style: ``_no_split_modules`` lists the vision tower *and* its blocks.""" + + _no_split_modules = ["_VisionTower", "_VisionBlock", "_DecoderLayer"] + + def __init__(self): + super().__init__() + self.vision_tower = _VisionTower() + + +def test_a_listed_container_is_wrapped_once_not_again_through_its_blocks(fsdp_mesh): + model = _MoeVlm() + + assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 3 + + assert isinstance(model.vision_tower, FSDPModule) + assert not any(isinstance(block, FSDPModule) for block in model.vision_tower.blocks) + + +class _ShardMesh: + """The meshes, as far as the loader reads them before it loads the dense part.""" + + device_type = "cpu" + + def __getitem__(self, name): + return self + + def size(self): + return 2 + + def get_group(self): + return -1 + + +def test_a_dense_part_from_pretrained_did_not_leave_on_meta_is_refused( + monkeypatch, gloo_world_of_one +): + """A transformers that places the dense part without asking ``get_device`` loads it whole on every + rank, the memory this layout exists to save: the load refuses that before sharding anything.""" + + class _DenseLoadedError(Exception): + """Raised by the dense load: the loader got past the check.""" + + def dense_load(*args, **kwargs): + raise _DenseLoadedError + + loader = SimpleNamespace(__name__="Loader", from_pretrained=lambda path, **kwargs: _MoeLM()) + monkeypatch.setattr("modelopt.torch.utils.distributed.is_initialized", lambda: True) + monkeypatch.setattr("modelopt.torch.utils.distributed.size", lambda group=None: 2) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *a, **k: _ShardMesh() + ) + monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", lambda p, **k: (None, p)) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: loader) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) + monkeypatch.setattr(model_load_utils, "_load_ep_dense_distributed", dense_load) + config = SimpleNamespace( + base_model_ep_plan={ + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + }, + num_experts=4, + ) + + with pytest.raises(RuntimeError, match="no dense param came back on meta"): + load_and_prepare_ep("unused", DistributionPlan(2, 2, "fsdp"), hf_config=config) + + +# ------------------------------------------------------------------ _load_ep_dense_distributed + + +def test_the_dense_load_shards_first_then_asks_for_exactly_the_meta_dense_params( + fsdp_mesh, monkeypatch, tmp_path +): + """As ``_ep_dense_meta_load`` leaves the model: dense weights on meta; experts loaded, here + per-expert ``nn.Linear``s, whose ``.weight`` only the expert matcher tells from a dense one; + and a Mamba-style ``A_log`` (no ``.weight`` / ``.bias`` suffix) loaded too. The dense layers + are FSDP2-sharded before the shared meta+DCP loader runs, and it is asked for exactly the meta + dense params, on the mesh's device. A recorder stands in for that load.""" + with torch.device("meta"): + model = _MoeLM() + for layer in model.model.layers: + layer.mlp.experts = nn.ModuleList(nn.Linear(4, 4, bias=False) for _ in range(4)) + layer.A_log = nn.Parameter(torch.zeros(4)) + loads = [] + + def record_load(loaded, local_dir, device, *, load_names): + sharded = [isinstance(layer.self_attn.weight, DTensor) for layer in loaded.model.layers] + loads.append((local_dir, device, load_names, sharded)) + + monkeypatch.setattr(model_load_utils, "_meta_dcp_load", record_load) + + wrapped = _load_ep_dense_distributed( + model, fsdp_mesh, str(tmp_path), _is_expert, model.experts() + ) + + assert wrapped == 2 + meta_dense = {name for name, param in model.named_parameters() if param.is_meta} + assert loads == [(str(tmp_path), torch.device("cpu"), meta_dense, [True, True])] + assert not any(_is_expert(name) or name.endswith("A_log") for name in meta_dense) + + +def test_a_tied_head_the_checkpoint_stores_alone_reaches_both_aliases(fsdp_mesh, tmp_path): + """A tied checkpoint can store the pair as ``lm_head.weight`` alone, which ``from_pretrained`` + loads into both aliases. With a tied MoE as the expert-parallel load leaves it -- the dense part + on meta and tied, the experts loaded -- the dense load fills both aliases with it too.""" + src = get_tiny_qwen3_moe(tie_word_embeddings=True) + src.save_pretrained(tmp_path) + weights = tmp_path / "model.safetensors" + stored = load_file(weights) + stored["lm_head.weight"] = stored.pop("model.embed_tokens.weight") + save_file(stored, weights, metadata={"format": "pt"}) + with _ep_dense_meta_load(_is_expert): + model = AutoModelForCausalLM.from_pretrained(tmp_path) + model.tie_weights() # as load_and_prepare_ep does once the weights load + assert model.lm_head.weight is model.model.embed_tokens.weight + assert model.lm_head.weight.is_meta + + experts = [layer.mlp.experts for layer in model.model.layers] + _load_ep_dense_distributed(model, fsdp_mesh, str(tmp_path), _is_expert, experts) + + for alias in (model.lm_head.weight, model.model.embed_tokens.weight): + assert torch.equal(alias, src.lm_head.weight) + + +# ------------------------------------------------- the expert-parallel load's allocator warmup +# ``from_pretrained`` reserves the bytes ``get_total_byte_count`` reports for each device in one +# allocation. Under the EP load it would count every expert at full size (the experts are most of a +# large MoE), so ``_ep_aware_caching_allocator_warmup`` swaps in a sizing that counts one EP shard of +# each expert. The sizing is called directly here, and once through ``caching_allocator_warmup`` with +# the CUDA queries and the allocation faked. + +needs_sizing_hook = pytest.mark.skipif( + not hasattr(modeling_utils, "get_total_byte_count"), + reason="transformers has no get_total_byte_count to size", +) + +_GPU = torch.device("cuda", 0) + +# _TinyMoe's fp32 bytes: the fused experts; the dense params (embedding, attention weight + bias, +# router); the router's persistent ``e_score_correction_bias`` buffer. The tied head adds none. +_EXPERT_BYTES = 4 * (4 * 16 * 8 + 4 * 8 * 8) +_DENSE_BYTES = 4 * (16 * 8 + 8 * 8 + 8 + 4 * 8) +_BUFFER_BYTES = 4 * 4 +_ATTN_BYTES = 4 * (8 * 8 + 8) + + +class _TinyMoe(nn.Module): + """One MoE decoder layer, a head tied to the embedding, and a persistent router buffer.""" + + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.embed_tokens = nn.Embedding(16, 8) + layer = nn.Module() + layer.self_attn = nn.Linear(8, 8) + layer.mlp = nn.Module() + layer.mlp.gate = nn.Linear(8, 4, bias=False) + layer.mlp.gate.register_buffer("e_score_correction_bias", torch.zeros(4)) + layer.mlp.experts = nn.Module() + layer.mlp.experts.gate_up_proj = nn.Parameter(torch.zeros(4, 16, 8)) + layer.mlp.experts.down_proj = nn.Parameter(torch.zeros(4, 8, 8)) + self.model.layers = nn.ModuleList([layer]) + self.lm_head = nn.Linear(8, 16, bias=False) + self.lm_head.weight = self.model.embed_tokens.weight + self.all_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def get_parameter_or_buffer(self, name): + """``PreTrainedModel``'s lookup, which the sizing goes through.""" + try: + return self.get_parameter(name) + except AttributeError: + return self.get_buffer(name) + + def device_map(self): + """Every param and persistent buffer on one GPU, as ``from_pretrained`` expands it.""" + names = [name for name, _ in self.named_parameters(remove_duplicate=False)] + names += [name for name, _ in self.named_buffers()] + return dict.fromkeys(names, _GPU) + + +class _OneBytePerElement: + def param_element_size(self, model, name, param): + return 1 + + +_ALL_BYTES = _EXPERT_BYTES + _DENSE_BYTES + _BUFFER_BYTES + + +@needs_sizing_hook +@pytest.mark.parametrize( + ("dense_on_meta", "quantizer", "stock_bytes", "per_rank_bytes"), + [ + # Experts count one EP shard; the router is replicated, so it counts in full like the rest. + pytest.param( + False, + None, + _ALL_BYTES, + _EXPERT_BYTES // 4 + _DENSE_BYTES + _BUFFER_BYTES, + id="dense-replicated", + ), + # Sharded dense: every dense .weight / .bias stays on meta during from_pretrained, while + # persistent buffers still load onto the device. + pytest.param( + True, None, _ALL_BYTES, _EXPERT_BYTES // 4 + _BUFFER_BYTES, id="dense-on-meta" + ), + # The quantizer sizes the elements: one byte each where fp32 takes four. + pytest.param( + False, + _OneBytePerElement(), + _ALL_BYTES // 4, + (_EXPERT_BYTES // 4 + _DENSE_BYTES + _BUFFER_BYTES) // 4, + id="quantizer", + ), + ], +) +def test_transformers_reserves_only_the_per_rank_footprint( + monkeypatch, dense_on_meta, quantizer, stock_bytes, per_rank_bytes +): + """End to end: the one up-front allocation ``caching_allocator_warmup`` makes shrinks from the + whole model (stock sizing, which also checks the byte counts above) to what this EP rank + holds.""" + reserved = [] + + def fake_empty(numel, *, dtype, device, **kwargs): + reserved.append((numel * torch.tensor([], dtype=dtype).element_size(), device)) + + model = _TinyMoe() + with monkeypatch.context() as m: + m.setattr(torch.cuda, "mem_get_info", lambda index: (80 * 1024**3, 80 * 1024**3)) + m.setattr(torch.cuda, "memory_reserved", lambda index: 0) + m.setattr(torch.cuda, "memory_allocated", lambda index: 0) + m.setattr(torch, "empty", fake_empty) + modeling_utils.caching_allocator_warmup(model, model.device_map(), quantizer) + with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN, dense_on_meta=dense_on_meta): + modeling_utils.caching_allocator_warmup(model, model.device_map(), quantizer) + + assert reserved == [(stock_bytes, _GPU), (per_rank_bytes, _GPU)] + + +@needs_sizing_hook +def test_a_tp_planned_dense_param_keeps_its_tp_share_while_the_rest_stays_on_meta(monkeypatch): + """transformers places a TP-planned param itself, never asking where the dense part goes, so it + is on the device at its TP share even when the rest of the dense part stays on meta. The share + is of the ep submesh it is placed over, not of the world, which holds two of them here.""" + model = _TinyMoe() + model._tp_plan = {"model.layers.*.self_attn": "colwise"} + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 8) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + + with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN, dense_on_meta=True): + total = modeling_utils.get_total_byte_count(model, model.device_map()) + + assert total == {_GPU: _EXPERT_BYTES // 4 + _ATTN_BYTES // 4 + _BUFFER_BYTES} + + +@needs_sizing_hook +def test_an_error_while_sizing_falls_back_to_the_stock_sizing(monkeypatch): + """The warmup only sizes an allocator reservation, so whatever the EP-aware sizing cannot + handle gets transformers' own number instead of failing the load, with whatever else a later + release passes the hook, and a warning: that number can reserve most of the device.""" + stock_calls = [] + + def stock_sizing(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs): + stock_calls.append((model, accelerator_device_map, hf_quantizer, args, kwargs)) + return {_GPU: 123} + + class _UnsupportedQuantizer: + def param_element_size(self, model, name, param): + raise TypeError("unsupported param") + + monkeypatch.setattr(modeling_utils, "get_total_byte_count", stock_sizing) + model, quantizer = _TinyMoe(), _UnsupportedQuantizer() + device_map = model.device_map() + + with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN): + with pytest.warns(UserWarning, match="EP-aware warmup sizing failed"): + sizing = modeling_utils.get_total_byte_count(model, device_map, quantizer, 1, flag=True) + assert sizing == {_GPU: 123} + + assert stock_calls == [(model, device_map, quantizer, (1,), {"flag": True})] + + +@needs_sizing_hook +def test_a_sizing_hook_with_other_parameters_keeps_the_stock_sizing(monkeypatch): + """A release whose hook takes other leading parameters keeps transformers' own sizing: the + replacement would not know how the hook is called. It says so, since that sizing may reserve + most of the device.""" + + def stock_sizing(model, device_map, quantizer=None): + return {_GPU: 123} + + monkeypatch.setattr(modeling_utils, "get_total_byte_count", stock_sizing) + + with ( + pytest.warns(UserWarning, match="using its full-model sizing"), + _ep_aware_caching_allocator_warmup(4, _A2A_PLAN), + ): + assert modeling_utils.get_total_byte_count is stock_sizing + + +@needs_sizing_hook +@pytest.mark.parametrize("ep_degree", [1, 0]) +def test_without_expert_parallelism_the_stock_sizing_stays(ep_degree): + stock = modeling_utils.get_total_byte_count + + with _ep_aware_caching_allocator_warmup(ep_degree, _A2A_PLAN): + assert modeling_utils.get_total_byte_count is stock + + +@needs_sizing_hook +def test_a_transformers_without_the_sizing_hook_is_left_alone(monkeypatch): + monkeypatch.delattr(modeling_utils, "get_total_byte_count") + + with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN): + assert not hasattr(modeling_utils, "get_total_byte_count") From 8a213756d875c8bc89cd4db1be17ab200f3c63e4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:09:47 +0000 Subject: [PATCH 19/21] Materialize the sharded dense part on the GPU each rank is bound to The dense load put each rank's shard on the GPU its local rank numbers. The experts and FSDP2's compute are on the GPU the rank is bound to, and a launch that shows each process one GPU binds every rank to cuda:0 while the local ranks run 0, 1, ...: rank 1 asked for cuda:1, which it does not see. The dense load now takes the bound GPU, as the load routing places a whole or sharded load. The new test, local rank 1 bound to cuda:0, fails without the change. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 4 ++- .../unit/torch/utils/test_model_load_utils.py | 27 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index b0e3bc52c6d..7e8632c195e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2035,8 +2035,10 @@ def _load_ep_dense_distributed(model, fsdp_mesh, local_dir, is_expert, expert_mo The dense ``.weight``/``.bias`` params arrive on meta (:func:`_ep_dense_meta_load`); the experts, already loaded and EP-sharded, are left as they are. Returns the number of layers sharded. """ + # The GPU the rank is bound to, as the experts and FSDP2's compute are: not the one its local + # rank numbers, which a launch showing each process one GPU does not bind. cuda = fsdp_mesh.device_type == "cuda" - device = torch.device(fsdp_mesh.device_type, dist.local_rank() if cuda else None) + device = torch.device(fsdp_mesh.device_type, _own_gpu_index() if cuda else None) n_dense = _fsdp2_shard_dense(model, fsdp_mesh, expert_modules) # Every alias of a tied param: the checkpoint may store the one named_parameters() drops. dense_names = { diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 3f13c2fa35f..7089c33901c 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -3061,6 +3061,33 @@ def record_load(loaded, local_dir, device, *, load_names): assert not any(_is_expert(name) or name.endswith("A_log") for name in meta_dense) +def test_the_dense_load_materializes_on_the_gpu_the_rank_is_bound_to(fsdp_mesh, monkeypatch): + """Not on the one its local rank numbers: a launch that shows each process one GPU binds every + rank to cuda:0 while the local ranks run 0, 1, ..., and the experts and FSDP2's compute are on + the bound one.""" + loads = [] + monkeypatch.setenv("LOCAL_RANK", "1") + for name, value in [ + ("is_available", True), + ("is_initialized", True), + ("current_device", 0), + ("device_count", 1), + ]: + monkeypatch.setattr(torch.cuda, name, lambda value=value: value) + monkeypatch.setattr(model_load_utils, "_fsdp2_shard_dense", lambda *args: 0) + + def record_load(model, local_dir, device, *, load_names): + loads.append(device) + + monkeypatch.setattr(model_load_utils, "_meta_dcp_load", record_load) + + _load_ep_dense_distributed( + nn.Linear(2, 2), SimpleNamespace(device_type="cuda"), "unused", lambda name: False, [] + ) + + assert loads == [torch.device("cuda", 0)] + + def test_a_tied_head_the_checkpoint_stores_alone_reaches_both_aliases(fsdp_mesh, tmp_path): """A tied checkpoint can store the pair as ``lm_head.weight`` alone, which ``from_pretrained`` loads into both aliases. With a tied MoE as the expert-parallel load leaves it -- the dense part From 855600dd7c60d7f03b078f4b082cce3130c16a25 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:09:57 +0000 Subject: [PATCH 20/21] Fail every rank when one rank fails to wrap the sharded dense part Wrapping the dense part in FSDP2 allocates each rank's shards on its own, so it can fail on one rank alone (out of memory, say), and the ranks that wrapped went on into DCP's collectives to wait for it there. The wrapping, and the dense names it collects, now run in the load's failure-carrying guard: every rank raises that failure before DCP. The bounded two-rank test fails the wrapping on one rank and fails without the guard. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 16 ++++---- .../unit/torch/utils/test_model_load_utils.py | 40 +++++++++++++++++++ 2 files changed, 49 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 7e8632c195e..20216a80185 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2039,13 +2039,15 @@ def _load_ep_dense_distributed(model, fsdp_mesh, local_dir, is_expert, expert_mo # rank numbers, which a launch showing each process one GPU does not bind. cuda = fsdp_mesh.device_type == "cuda" device = torch.device(fsdp_mesh.device_type, _own_gpu_index() if cuda else None) - n_dense = _fsdp2_shard_dense(model, fsdp_mesh, expert_modules) - # Every alias of a tied param: the checkpoint may store the one named_parameters() drops. - dense_names = { - n - for n, _ in model.named_parameters(remove_duplicate=False) - if n.endswith((".weight", ".bias")) and not is_expert(n) - } + # A rank failing to wrap (out of memory, say) fails every rank here, before DCP's collectives. + with _refused_together("the dense part's FSDP2 wrapping", alone=False): + n_dense = _fsdp2_shard_dense(model, fsdp_mesh, expert_modules) + # Every alias of a tied param: the checkpoint may store the one named_parameters() drops. + dense_names = { + n + for n, _ in model.named_parameters(remove_duplicate=False) + if n.endswith((".weight", ".bias")) and not is_expert(n) + } print_rank_0( f"[ep] dense-distributed load: FSDP2-sharded {n_dense} dense module(s); " f"meta+DCP-loading {len(dense_names)} dense param(s) (experts left EP-sharded)" diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 7089c33901c..536360dfc35 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -3088,6 +3088,46 @@ def record_load(model, local_dir, device, *, load_names): assert loads == [torch.device("cuda", 0)] +class _PastTheWrappingError(Exception): + """Raised in DCP's place: the rank went on past the wrapping.""" + + +def _one_rank_fails_the_wrapping_job(rank, size): + def wrap(*args): + if rank == 1: + raise MemoryError("out of memory sharding a dense layer") + return 0 + + def dcp_load(*args, **kwargs): + torch.distributed.barrier() # DCP's collectives: a rank in them waits for every peer + raise _PastTheWrappingError + + refused = None + with ( + bounded_collectives(), + mock.patch.object(model_load_utils, "_fsdp2_shard_dense", wrap), + mock.patch.object(model_load_utils, "_meta_dcp_load", dcp_load), + ): + try: # not pytest.raises: its failure would end the pooled worker instead of the job + _load_ep_dense_distributed( + nn.Linear(2, 2), SimpleNamespace(device_type="cpu"), "unused", lambda n: False, [] + ) + except Exception as error: + refused = f"{type(error).__name__}: {error}" + assert refused is not None and "out of memory sharding a dense layer" in refused, ( + rank, + refused, + ) + assert rank == 1 or "rank 1 failed" in refused, refused + + +def test_a_rank_failing_to_wrap_the_dense_part_fails_every_rank(cpu_dist_workers): + """Wrapping allocates each rank's shards on its own, so it can fail on one rank alone; every + rank then raises that failure, rather than leave the others waiting for it in DCP's + collectives (stood in by a bounded barrier).""" + cpu_dist_workers.run(_one_rank_fails_the_wrapping_job) + + def test_a_tied_head_the_checkpoint_stores_alone_reaches_both_aliases(fsdp_mesh, tmp_path): """A tied checkpoint can store the pair as ``lm_head.weight`` alone, which ``from_pretrained`` loads into both aliases. With a tied MoE as the expert-parallel load leaves it -- the dense part From dd04431612beeb45a636f424b1c388a3b0e02d2a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:31:20 +0000 Subject: [PATCH 21/21] Say so when transformers has no warmup sizing hook to replace Claude's review: of the three ways the load warmup's EP-aware sizing falls back to transformers' own, only a release without get_total_byte_count did so silently. That sizing may reserve most of the device, and a later NCCL or out-of-memory failure would then have no stated cause, so it now warns, naming the transformers version, as its siblings do. The test of that release now expects the warning, and fails without it. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 8 +++++++- tests/unit/torch/utils/test_model_load_utils.py | 5 ++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 20216a80185..f67748281ef 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1746,13 +1746,19 @@ def _ep_aware_caching_allocator_warmup( expert param under the EP load, so the warmup would reserve about the whole GPU on every rank and starve the all-to-all's NCCL buffers. While the load runs, an expert param counts ``1/ep_degree`` of itself and, with ``dense_on_meta``, a dense ``.weight``/``.bias`` nothing. It only sizes a - reservation: a hook of another signature, or an error while sizing, keeps the stock sizing. + reservation: a release without the hook, a hook of another signature, or an error while sizing, + keeps the stock sizing, and says so. """ if ep_degree <= 1: yield return orig = getattr(hf_modeling_utils, "get_total_byte_count", None) if orig is None: + warn_rank_0( + f"transformers {transformers.__version__} has no get_total_byte_count in " + "modeling_utils to size the load warmup by; using its own sizing, which may reserve " + "most of the device." + ) yield return try: diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 536360dfc35..15c04ddfadc 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -3341,5 +3341,8 @@ def test_without_expert_parallelism_the_stock_sizing_stays(ep_degree): def test_a_transformers_without_the_sizing_hook_is_left_alone(monkeypatch): monkeypatch.delattr(modeling_utils, "get_total_byte_count") - with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN): + with ( + pytest.warns(UserWarning, match="has no get_total_byte_count"), + _ep_aware_caching_allocator_warmup(4, _A2A_PLAN), + ): assert not hasattr(modeling_utils, "get_total_byte_count")