From 904828d4ba2908c5c4e888c2d3b5761ce4bdd4fc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 21:22:25 +0000 Subject: [PATCH 01/12] Load each rank's checkpoint shard straight into a meta-built model A distributed load that materializes the whole model on every rank and then throws most of it away costs world x the memory and I/O. These helpers let a loader build the model on meta and fill only what each rank keeps. _meta_dcp_load materializes each rank's params on its device and fills them from the HF safetensors through torch DCP's HuggingFaceStorageReader. Placement is decided entirely by the caller: dcp.load reads a full tensor into a plain (replicated) param and only the local slice into a Shard DTensor, so leaving the model as is, fully_shard-ing all of it, or fully_shard-ing only the dense modules gives replicated, sharded, or the FSDP-dense + EP-expert hybrid -- with no full-model copy for the sharded cases. _map_ckpt_keys_to_model maps each checkpoint key to the state_dict key it feeds, reproducing what from_pretrained resolves: exact match, then the model's _checkpoint_conversion_mapping, then the longest unique trailing-component match for the transformers 5 VLM prefix nesting. _fsdp2_wrap_modules picks the blocks to fully_shard individually from HF's own _no_split_modules, looking at the language backbone rather than a VLM wrapper that would also list its vision blocks. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 218 ++++++++++++++++++++++++++++ 1 file changed, 218 insertions(+) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index e42d8015269..c48f621b7ad 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -18,6 +18,7 @@ import functools import io import os +import re import sys import time import traceback @@ -531,6 +532,223 @@ def distributed_run(timeout: timedelta | None = None, device: torch.device | str cleanup() +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). + """ + 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) + if lineage and lineage[-1] is not model: + lm_mods = _boundaries(lineage[-1]) + if lm_mods: + return lm_mods + + mods = _boundaries(model) + if mods: + return mods + raise RuntimeError( + "Could not determine FSDP2 wrap boundaries: model exposes no `_no_split_modules` and modelopt's " + "decoder-layer auto-detection did not match. Provide a standard HF causal-LM / VLM checkpoint." + ) + + +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. + + Map each CHECKPOINT key to the built model's ``state_dict`` key it feeds, reproducing what + ``from_pretrained`` resolves on load -- so a DCP by-name load fills the right tensors. + + Three passes per checkpoint key (first hit wins): (1) exact match (the common case); (2) the model's + ``_checkpoint_conversion_mapping`` regexes if it carries any (first pattern that fires); (3) longest + UNIQUE trailing-component match. Pass 3 covers the transformers>=5 VLM prefix-nesting + (``model.layers.*`` checkpoint -> ``model.language_model.layers.*`` model, ``visual.*`` -> ``model.visual.*``) + WITHOUT hard-coding it: the leaf path is preserved, so the ckpt key's component suffix uniquely locates + the model key. Ambiguous (>1 candidate) or absent suffixes are left unmapped -> caught by the caller's + coverage guard rather than mis-assigned. + """ + key_map: dict = {} + remaining = set(ckpt_keys) + # Pass 1: exact + for ck in list(remaining): + if ck in model_keys: + key_map[ck] = ck + remaining.discard(ck) + # Pass 2: model-declared regex conversion mapping (rare; e.g. leaf renames) + conv = getattr(model, "_checkpoint_conversion_mapping", None) or {} + if conv and remaining: + compiled = [(re.compile(p), r) for p, r in conv.items()] + for ck in list(remaining): + mk = ck + for rx, repl in compiled: + mk, n = rx.subn(repl, mk) + if n: + break + if mk in model_keys: + key_map[ck] = mk + remaining.discard(ck) + # Pass 3: longest unique component-suffix match + if remaining: + from collections import defaultdict + + suffix_idx: dict = defaultdict(list) + for mk in model_keys: + comps = mk.split(".") + for k in range(1, len(comps) + 1): + suffix_idx[tuple(comps[-k:])].append(mk) + for ck in remaining: + comps = ck.split(".") + for k in range(len(comps), 0, -1): + cands = suffix_idx.get(tuple(comps[-k:])) + if cands and len(cands) == 1: + key_map[ck] = cands[0] + break + return key_map + + +def _meta_dcp_load( + model: "torch.nn.Module", + local_dir: str, + device, + *, + load_names: "set | None" = None, + skip_names: "set | None" = None, + max_unloaded_frac: float = 0.5, +) -> None: + """The single meta + torch-DCP loader shared by every distributed dense path. + + materialize each rank's params on ``device`` and fill them from the HF safetensors via DCP's + ``HuggingFaceStorageReader``. + + PLACEMENT is decided entirely by the CALLER, not here: ``dcp.load`` reads a FULL tensor into a plain + (replicated) param and only the LOCAL SLICE into a ``Shard`` ``DTensor``. So the caller picks the layout + by what it does before calling this -- nothing (-> Replicate, full per rank), + ``fully_shard`` the whole model (``load_and_prepare_fsdp2`` -> Shard), or ``fully_shard`` only the dense + modules (the EP hybrid -> dense Shard, experts left as loaded). No full-model copy for the sharded cases. + + ``load_names`` (optional) restricts the checkpoint load to a subset of model param names -- the EP-dense + case, where the experts are ALREADY materialized by ``from_pretrained`` and must not be touched. ``None`` + loads the whole model (DP / FSDP2-full). + + Materialization is SELECTIVE (module-by-module, only where params are meta) so already-real params (EP + experts) are never re-emptied. ``to_empty`` also re-empties BUFFERS (rotary ``inv_freq`` etc., + any the + caller already loaded), so they are snapshotted and restored -- a params-only checkpoint load would + otherwise leave them garbage and the calibration forward would drift (weights load fine but activation + amax is wrong). Persistent buffers the checkpoint DOES carry are (re)loaded via DCP below. + """ + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint import HuggingFaceStorageReader + + # Snapshot real buffers (to_empty re-empties them), materialize params, restore buffers. Whole-model + # meta (DP / FSDP2-full, ``load_names is None``) uses ``model.to_empty`` -- the FSDP2-aware root + # materialization; a per-module selective loop corrupts the ROOT-sharded FSDP state (its state_dict then + # returns wrong keys). The PARTIAL case (EP-dense: dense meta, experts already real) MUST be selective so + # the real experts are not re-emptied -- there no root is fully_shard'd, so per-module to_empty is safe. + saved_buffers = {n: b.detach().clone() for n, b in model.named_buffers() if not b.is_meta} + if load_names is None: + model.to_empty(device=device) + else: + for m in model.modules(): + if any(p.is_meta for p in m.parameters(recurse=False)): + m.to_empty(device=device, recurse=False) + for n, b in model.named_buffers(): + src = saved_buffers.get(n) + if src is not None: + b.copy_(src.to(device=b.device, dtype=b.dtype)) + + # ``dcp.load`` matches state_dict entries to checkpoint entries BY NAME, but a checkpoint's stored key + # names can differ from the built model's ``state_dict`` names, and from_pretrained bridges the gap on + # load. E.g. Qwen2.5-VL's checkpoint stores ``model.layers.*`` / ``visual.*`` while transformers>=5 nests + # them as ``model.language_model.layers.*`` / ``model.visual.*`` -- a pure prefix-insertion with the leaf + # path preserved, and there is NO ``_checkpoint_conversion_mapping`` on the model to read. A raw by-name + # match would SILENTLY skip every renamed weight (they stay ``to_empty`` garbage -> a wrecked model that + # only surfaces as a confusing downstream crash). Reproduce from_pretrained's result generically via + # :func:`_map_ckpt_keys_to_model`: exact -> conversion-mapping regexes -> longest UNIQUE suffix. + reader = HuggingFaceStorageReader(str(local_dir)) + ckpt_keys = set(reader.read_metadata().state_dict_metadata.keys()) + if not ckpt_keys: + raise RuntimeError( + f"meta+DCP load: the checkpoint at {local_dir} has NO tensors -- the resolved snapshot carries no " + "safetensors (a config-only / incomplete HF revision; e.g. refs/main pointing at a weightless " + "snapshot). Re-download the checkpoint or pass a complete local snapshot dir / pinned revision." + ) + full_sd = ( + model.state_dict() + ) # (sharded ``DTensor`` or plain) params + replicated persistent buffers + target_names = set(full_sd) if load_names is None else (set(load_names) & set(full_sd)) + key_map = _map_ckpt_keys_to_model(model, ckpt_keys, target_names) # {ckpt_key: model_key} + # Build load_sd in the model's state_dict ORDER (identical on every rank) -- NOT key_map/set iteration + # order, which is hash-randomized per process. dcp.load's collective load planner deadlocks if ranks + # present the same keys in different orders, so a deterministic, rank-consistent order is required. + inv = {mk: ck for ck, mk in key_map.items()} # model_key -> ckpt_key (1:1 on model keys) + load_sd = {inv[mk]: full_sd[mk] for mk in full_sd if mk in inv} + mapped_names = set(key_map.values()) + # Fail LOUD if the mapping left too much of the TARGET WEIGHT unloaded -- a key scheme this path did not + # resolve; a silent partial load produces a garbage-weight model. Weight by NUMEL, not param COUNT: a + # fused-expert MoE (Qwen3-MoE: per-expert checkpoint -> grouped ``experts.gate_up_proj``/``down_proj`` + # model params DCP can't fuse) leaves only ~2 params/layer unloaded but they are ~95% of the weight, so a + # count-based guard misses it. A small tail is normal (tied weights + a VLM's not-in-checkpoint embeds). + # ``skip_names`` = params loaded by a SEPARATE path (e.g. Qwen3-MoE fused grouped experts, filled by + # :func:`_load_fused_experts_sharded` which fuses per-expert checkpoint tensors DCP can't concat) -- + # exclude them from the coverage guard so it doesn't (correctly) flag them as unloaded here. + skip_names = skip_names or set() + numel = { + n: p.numel() + for n, p in model.named_parameters() + if n in target_names and n not in skip_names + } + want_params = set(numel) + unloaded = want_params - mapped_names + total_nel = max(1, sum(numel.values())) + unloaded_nel = sum(numel[n] for n in unloaded) + if unloaded_nel > max_unloaded_frac * total_nel: + raise RuntimeError( + f"meta+DCP load: {len(unloaded)}/{len(want_params)} params ({100 * unloaded_nel / total_nel:.0f}% " + f"of the target weight) matched no checkpoint key -- the checkpoint's key naming was not resolved " + f"to this model's state_dict (exact + conversion-mapping + suffix match all missed), or it needs a " + f"weight conversion DCP can't do (FUSED EXPERTS, split qkv, on-the-fly dequant). Unloaded e.g.: " + f"{sorted(unloaded)[:5]}. Use the from_pretrained load (MODELOPT_FSDP2_SHARDED_LOAD=0 / " + f"MODELOPT_EP_DENSE_DISTRIBUTED_LOAD=0)." + ) + if is_initialized() and rank() == 0: + loaded_bufs = mapped_names - {n for n, _ in model.named_parameters()} + if unloaded: + warn( + f"meta+DCP load: {len(unloaded)} param(s) not loaded from checkpoint (tied? re-tied): " + f"{sorted(unloaded)[:5]}" + ) + if loaded_bufs: + warn( + f"meta+DCP load: loaded {len(loaded_bufs)} persistent buffer(s) from checkpoint: " + f"{sorted(loaded_bufs)[:5]}" + ) + dcp.load(load_sd, storage_reader=reader) + if hasattr(model, "tie_weights"): + model.tie_weights() # re-establish tied weights the checkpoint omitted + + def cleanup(): """Cleans up the distributed environment. From 8e65bc50338e230e39dd1d4a04df724ab2f2189c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:07:10 +0000 Subject: [PATCH 02/12] Test the meta-DCP checkpoint loader Cover how _map_ckpt_keys_to_model bridges checkpoint names to the built model's -- exact, declared conversion, unique suffix, which resolves a VLM's nested prefixes -- which blocks _fsdp2_wrap_modules shards, and what _meta_dcp_load does: fill a meta model from safetensors through DCP while keeping its buffers and re-tying its head, load only load_names, guard coverage by numel, and fill FSDP2-sharded params in place. The loads run for real on CPU. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_meta_dcp_load.py | 440 +++++++++++++++++++ 1 file changed, 440 insertions(+) create mode 100644 tests/unit/torch/utils/test_meta_dcp_load.py diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py new file mode 100644 index 00000000000..f000b86abe2 --- /dev/null +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -0,0 +1,440 @@ +# 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 meta-device + torch-DCP checkpoint loader in ``modelopt.torch.utils.distributed``. + +``_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; ``_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 two tests that need the distributed path start a +one-rank gloo group. +""" + +import pytest +import torch +import torch.distributed.checkpoint as dcp +from safetensors.torch import save_file +from torch import nn +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard +from torch.distributed.tensor import DTensor + +pytest.importorskip("accelerate") +pytest.importorskip("transformers") + +from _test_utils.torch.transformers_models import get_tiny_llama +from accelerate import init_empty_weights +from transformers import AutoModelForImageTextToText, Qwen2_5_VLConfig + +from modelopt.torch.utils.distributed import ( + _fsdp2_wrap_modules, + _map_ckpt_keys_to_model, + _meta_dcp_load, +) + + +def _save(ckpt_dir, tensors): + """Write ``tensors`` as a one-file HF safetensors checkpoint in ``ckpt_dir``; return its path.""" + ckpt_dir.mkdir(parents=True, exist_ok=True) + tensors = {k: v.detach().contiguous().clone() for k, v in tensors.items()} + save_file(tensors, str(ckpt_dir / "model.safetensors")) + return str(ckpt_dir) + + +def _meta_twin(build, **kwargs): + """``build(**kwargs)`` with params on ``meta`` and buffers real, as the loader's callers hold it.""" + with init_empty_weights(): + return build(**kwargs) + + +def _tiny_qwen2_5_vl(): + """A tiny Qwen2.5-VL: the LM nests at ``model.language_model``, the vision tower at ``model.visual``.""" + cfg = Qwen2_5_VLConfig( + text_config={ + "hidden_size": 16, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "vocab_size": 32, + "max_position_embeddings": 16, + }, + vision_config={ + "depth": 2, + "hidden_size": 16, + "intermediate_size": 16, + "num_heads": 2, + "out_hidden_size": 16, + "patch_size": 2, + "spatial_merge_size": 1, + "temporal_patch_size": 1, + "fullatt_block_indexes": [1], + }, + tie_word_embeddings=False, + ) + return AutoModelForImageTextToText.from_config(cfg) + + +def _flat_vl_name(key): + """The name a Qwen2.5-VL checkpoint stores ``key`` under: ``model.layers.*`` and ``visual.*``.""" + if key.startswith("model.language_model."): + return "model." + key.removeprefix("model.language_model.") + return key.removeprefix("model.") if key.startswith("model.visual.") else key + + +# ---------------------------------------------------------------- checkpoint -> model key mapping + + +def test_exact_name_match_takes_precedence(): + """A checkpoint name the model has verbatim maps to itself before any conversion pattern runs.""" + model = nn.Module() + model._checkpoint_conversion_mapping = {r"^model\.": "model.language_model."} + model_keys = {"model.norm.weight", "model.language_model.norm.weight"} + + key_map = _map_ckpt_keys_to_model(model, {"model.norm.weight"}, model_keys) + + assert key_map == {"model.norm.weight": "model.norm.weight"} + + +def test_conversion_mapping_applies_the_first_pattern_that_fires(): + """Patterns are tried in order and only the first that fires rewrites a name (no chaining); the + rewrite counts only if the model has the result, otherwise the name falls through to suffix match. + """ + model = nn.Module() + model._checkpoint_conversion_mapping = { + r"^model\.(?!language_model\.)": "model.language_model.", + r"^model\.language_model\.": "chained.", + r"^visual\.": "model.vision_tower.", # this model calls it ``model.visual`` + } + model_keys = { + "model.language_model.layers.0.mlp.down_proj.weight", + "chained.layers.0.mlp.down_proj.weight", + "model.visual.proj.weight", + } + ckpt_keys = {"model.layers.0.mlp.down_proj.weight", "visual.proj.weight"} + + assert _map_ckpt_keys_to_model(model, ckpt_keys, model_keys) == { + "model.layers.0.mlp.down_proj.weight": "model.language_model.layers.0.mlp.down_proj.weight", + "visual.proj.weight": "model.visual.proj.weight", + } + + +def test_suffix_match_resolves_the_vlm_prefix_nesting(): + """transformers>=5 nests a VLM's LM and vision tower one level below its checkpoint names, and the + model carries no ``_checkpoint_conversion_mapping``. The longest component suffix naming exactly + one model key bridges that: ``layers.0`` tells LM layer 0 from vision block 0, which shares the rest. + """ + model_keys = { + "model.language_model.embed_tokens.weight", + "model.language_model.layers.0.mlp.down_proj.weight", + "model.language_model.layers.1.mlp.down_proj.weight", + "model.language_model.norm.weight", + "model.visual.blocks.0.mlp.down_proj.weight", + "lm_head.weight", + } + expected = { + "model.embed_tokens.weight": "model.language_model.embed_tokens.weight", + "model.layers.0.mlp.down_proj.weight": "model.language_model.layers.0.mlp.down_proj.weight", + "model.layers.1.mlp.down_proj.weight": "model.language_model.layers.1.mlp.down_proj.weight", + "model.norm.weight": "model.language_model.norm.weight", + "visual.blocks.0.mlp.down_proj.weight": "model.visual.blocks.0.mlp.down_proj.weight", + "lm_head.weight": "lm_head.weight", + } + + assert _map_ckpt_keys_to_model(nn.Module(), expected.keys(), model_keys) == expected + + +def test_ambiguous_or_unknown_suffixes_stay_unmapped(): + """A name whose suffixes match several model keys (two towers both have ``layers.0``) or none (a + quantization scale with no slot) is left for the coverage guard, never assigned to a guess. + """ + model_keys = { + "model.language_model.layers.0.mlp.down_proj.weight", + "model.audio_tower.layers.0.mlp.down_proj.weight", + } + ckpt_keys = { + "model.layers.0.mlp.down_proj.weight", + "model.layers.0.mlp.down_proj.weight_scale", + } + + assert _map_ckpt_keys_to_model(nn.Module(), ckpt_keys, model_keys) == {} + + +# -------------------------------------------------------------------------- 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 test_wrap_modules_refuse_a_model_without_blocks(): + """With neither source of boundaries matching, fail loudly instead of guessing a layout.""" + with pytest.raises(RuntimeError, match="Could not determine FSDP2 wrap boundaries"): + _fsdp2_wrap_modules(nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))) + + +# --------------------------------------------------------------------------------- meta + DCP load + + +def test_load_fills_a_meta_model_keeps_its_buffers_and_reties_it(tmp_path): + """Whole-model load: every checkpoint tensor lands, the real rotary buffers survive ``to_empty``, + and the tied ``lm_head`` the checkpoint omits is re-tied instead of left as empty memory. + """ + src = get_tiny_llama(tie_word_embeddings=True) + ckpt = {k: v for k, v in src.state_dict().items() if k != "lm_head.weight"} + ckpt_dir = _save(tmp_path, ckpt) + model = _meta_twin(get_tiny_llama, tie_word_embeddings=True) + buffers = {n: b.clone() for n, b in model.named_buffers()} + assert buffers and all(p.is_meta for p in model.parameters()) + + _meta_dcp_load(model, ckpt_dir, "cpu") + + loaded = model.state_dict() + for name, value in ckpt.items(): + assert torch.equal(loaded[name], value), name + for name, buf in model.named_buffers(): + assert torch.equal(buf, buffers[name]), name + assert model.lm_head.weight is model.model.embed_tokens.weight + + +def test_flat_vlm_checkpoint_names_reach_the_nested_params(tmp_path): + """Qwen2.5-VL's checkpoint says ``model.layers.*`` / ``visual.*`` where the built model has + ``model.language_model.*`` / ``model.visual.*``; a by-name load would fill none of them. + """ + src = _tiny_qwen2_5_vl() + ref = src.state_dict() + ckpt = {_flat_vl_name(k): v for k, v in ref.items()} + assert ckpt.keys() & ref.keys() == {"lm_head.weight"} # every other name differs + model = _meta_twin(_tiny_qwen2_5_vl) + + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu") + + loaded = model.state_dict() + assert loaded.keys() == ref.keys() + for name, value in ref.items(): + assert torch.equal(loaded[name], value), name + + +def test_load_names_leave_already_real_params_untouched(tmp_path): + """EP-dense: with ``load_names`` only those meta params are materialized and loaded. The MLPs stand + in for experts ``from_pretrained`` already placed; the checkpoint holds other values for them, and + they must keep their own (same tensors, neither re-emptied nor reloaded). + """ + src = get_tiny_llama() + ckpt_dir = _save(tmp_path, src.state_dict()) + model = _meta_twin(get_tiny_llama) + experts = {} + for i, layer in enumerate(model.model.layers): + layer.mlp.to_empty(device="cpu") + for name, param in layer.mlp.named_parameters(): + nn.init.constant_(param, -7.0) + experts[f"model.layers.{i}.mlp.{name}"] = param + dense = {n for n, _ in model.named_parameters()} - experts.keys() + + _meta_dcp_load(model, ckpt_dir, "cpu", load_names=dense) + + params = dict(model.named_parameters()) + for name, param in experts.items(): + assert params[name] is param and bool((param == -7.0).all()), name + ref = src.state_dict() + for name in dense: + assert torch.equal(params[name], ref[name]), name + + +class _WithBuffers(nn.Module): + """A linear, a persistent buffer (checkpoints carry it) and a non-persistent one (built at init).""" + + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + self.register_buffer("scale", torch.ones(4)) + self.register_buffer("table", torch.arange(4.0), persistent=False) + + +def test_persistent_buffers_load_from_the_checkpoint_and_others_survive(tmp_path): + """``to_empty`` wipes buffers too: one the checkpoint carries is loaded from it, and one no + checkpoint can carry (non-persistent) gets its pre-load value back. + """ + ckpt = { + "proj.weight": torch.randn(4, 4), + "proj.bias": torch.randn(4), + "scale": torch.full((4,), 3.0), + } + with init_empty_weights(): + model = _WithBuffers() + + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu") + + assert torch.equal(model.scale, ckpt["scale"]) + assert torch.equal(model.table, torch.arange(4.0)) + assert torch.equal(model.proj.weight, ckpt["proj.weight"]) + + +class _OneBigManySmall(nn.Module): + """One param holding ~99% of the weight beside twelve tiny ones.""" + + def __init__(self): + super().__init__() + self.big = nn.Linear(64, 64, bias=False) + self.small = nn.ModuleList(nn.Linear(2, 2) for _ in range(6)) + + +_SMALL = {f"small.{i}.{kind}" for i in range(6) for kind in ("weight", "bias")} + + +def _guard_case(ckpt_dir, missing): + """A meta ``_OneBigManySmall`` and a checkpoint with all of its tensors except ``missing``.""" + with init_empty_weights(): + model = _OneBigManySmall() + ckpt = {k: torch.randn(v.shape) for k, v in model.state_dict().items() if k not in missing} + return model, ckpt, _save(ckpt_dir, ckpt) + + +def test_coverage_guard_weighs_unloaded_params_by_numel(tmp_path): + """Missing the one ~99% param fails loudly while missing the twelve tiny ones (~1%) loads. A + count-based guard gets both backwards, and the fused-expert miss it exists for is a few huge params. + """ + model, _, ckpt_dir = _guard_case(tmp_path / "big_missing", missing={"big.weight"}) + with pytest.raises(RuntimeError, match="matched no checkpoint key"): + _meta_dcp_load(model, ckpt_dir, "cpu") + + model, ckpt, ckpt_dir = _guard_case(tmp_path / "small_missing", missing=_SMALL) + _meta_dcp_load(model, ckpt_dir, "cpu") + assert torch.equal(model.big.weight, ckpt["big.weight"]) + + +def test_skip_names_exempt_params_filled_by_another_path(tmp_path): + """Params a separate loader fills (fused grouped experts) are left out of the guard's tally.""" + model, ckpt, ckpt_dir = _guard_case(tmp_path, missing={"big.weight"}) + + _meta_dcp_load(model, ckpt_dir, "cpu", skip_names={"big.weight"}) + + assert torch.equal(model.small[0].weight, ckpt["small.0.weight"]) + + +def test_max_unloaded_frac_sets_the_guard_threshold(tmp_path): + """The ~1% shortfall the default tolerates fails once the tolerance is zero.""" + model, _, ckpt_dir = _guard_case(tmp_path, missing=_SMALL) + + with pytest.raises(RuntimeError, match="matched no checkpoint key"): + _meta_dcp_load(model, ckpt_dir, "cpu", max_unloaded_frac=0.0) + + +def test_checkpoint_without_tensors_is_rejected(tmp_path): + """A weightless (config-only) snapshot fails with its own message, not as an all-unloaded model.""" + (tmp_path / "config.json").write_text("{}") + with init_empty_weights(): + model = _OneBigManySmall() + + with pytest.raises(RuntimeError, match="has NO tensors"): + _meta_dcp_load(model, str(tmp_path), "cpu") + + +def test_load_request_follows_the_model_state_dict_order(tmp_path, monkeypatch): + """``dcp.load`` gets the checkpoint names in the model's ``state_dict`` order, which every rank + shares. Key-map or set order is hash-seeded per process, and differing orders across ranks + deadlock the collective load planner. + """ + src = _tiny_qwen2_5_vl() + ckpt_dir = _save(tmp_path, {_flat_vl_name(k): v for k, v in src.state_dict().items()}) + model = _meta_twin(_tiny_qwen2_5_vl) + requests = [] + real_load = dcp.load + + def recording_load(state_dict, **kwargs): + requests.append(list(state_dict)) + return real_load(state_dict, **kwargs) + + monkeypatch.setattr(dcp, "load", recording_load) + _meta_dcp_load(model, ckpt_dir, "cpu") + + assert requests == [[_flat_vl_name(k) for k in model.state_dict()]] + + +# ------------------------------------------------------------------- under a process group + + +@pytest.fixture +def gloo_world_of_one(skip_on_windows, tmp_path_factory): + """A one-rank gloo process group, so the loader takes its distributed path.""" + store = tmp_path_factory.mktemp("pg") / "store" + torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) + yield + torch.distributed.destroy_process_group() + + +def test_rank0_reports_unloaded_params_and_loaded_buffers(tmp_path, gloo_world_of_one): + """In a distributed run rank 0 names the params the checkpoint left unfilled (below the guard's + bar) and the persistent buffers it filled. + """ + ckpt = {"proj.weight": torch.randn(4, 4), "scale": torch.full((4,), 3.0)} + with init_empty_weights(): + model = _WithBuffers() + + with pytest.warns(UserWarning) as record: + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu") + + messages = [str(w.message) for w in record] + assert any("1 param(s) not loaded" in m and "'proj.bias'" in m for m in messages), messages + assert any("1 persistent buffer(s)" in m and "'scale'" in m for m in messages), messages + + +def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one): + """FSDP2-full: the blocks and the root are ``fully_shard``-ed on meta before the load. The params + must stay sharded ``DTensor``s holding the checkpoint values, and the buffers must survive the + root ``to_empty``. + """ + src = get_tiny_llama() + ckpt_dir = _save(tmp_path, src.state_dict()) + model = _meta_twin(get_tiny_llama) + buffers = {n: b.clone() for n, b in model.named_buffers()} + mesh = init_device_mesh("cpu", (1,)) + for block in _fsdp2_wrap_modules(model): + fully_shard(block, mesh=mesh) + fully_shard(model, mesh=mesh) + + _meta_dcp_load(model, ckpt_dir, "cpu") + + ref = src.state_dict() + for name, param in model.named_parameters(): + assert isinstance(param, DTensor), name + assert torch.equal(param.full_tensor(), ref[name]), name + for name, buf in model.named_buffers(): + assert torch.equal(buf, buffers[name]), name From 9d4fd1efb909019ea013ffc4cbc5ce578392d1fa Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:07:34 +0000 Subject: [PATCH 03/12] Map each model key from at most one checkpoint key The suffix pass of _map_ckpt_keys_to_model mapped a checkpoint key the model has no slot for onto any key its suffix named uniquely -- including one another checkpoint key already named exactly -- and the later mapping won. DeepSeek-V3's MTP layer does this: model.layers.61.shared_head.norm.weight lands on model.norm.weight, so a meta-DCP load would fill the final norm with the MTP head's. A truncated model loading a deeper checkpoint got its last layer's weights in layer 0. With load_names, uniqueness was judged among the targets alone, so a per-expert key could claim a dense param; and when two keys collided, set order, which differs per process, picked the tensor. A model key now takes one checkpoint key: no pass maps onto a key an earlier pass claimed, and two claimants in one pass are both left to the coverage guard. Uniqueness is judged against the whole model, and the load keeps only the mappings it fills. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 51 ++++++++++++---- tests/unit/torch/utils/test_meta_dcp_load.py | 63 ++++++++++++++++++++ 2 files changed, 101 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index c48f621b7ad..695b471dbbd 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -588,43 +588,61 @@ def _map_ckpt_keys_to_model(model: "torch.nn.Module", ckpt_keys, model_keys) -> WITHOUT hard-coding it: the leaf path is preserved, so the ckpt key's component suffix uniquely locates the model key. Ambiguous (>1 candidate) or absent suffixes are left unmapped -> caught by the caller's coverage guard rather than mis-assigned. + + ``model_keys`` must be the WHOLE model's keys, even when only some of them will be loaded: a suffix is + unique only if no other model key shares it, and the caller drops the mappings it does not load. Each + model key takes at most one checkpoint key -- a later pass never maps onto a key an earlier pass + claimed, and two checkpoint keys claiming one key in the same pass are both left unmapped. Otherwise a + key the model has no place for (an MTP layer's ``model.layers.61.shared_head.norm.weight``, a layer a + truncated model lacks) lands on ``model.norm.weight`` by suffix alone and overwrites it, with set order + picking the winner. """ + from collections import defaultdict + key_map: dict = {} remaining = set(ckpt_keys) + + def claim(hits: dict) -> None: + claimed = set(key_map.values()) + for mk, cks in hits.items(): + if len(cks) == 1 and mk not in claimed: + key_map[cks[0]] = mk + remaining.difference_update(cks) + # Pass 1: exact - for ck in list(remaining): - if ck in model_keys: - key_map[ck] = ck - remaining.discard(ck) + claim({ck: [ck] for ck in remaining if ck in model_keys}) # Pass 2: model-declared regex conversion mapping (rare; e.g. leaf renames) conv = getattr(model, "_checkpoint_conversion_mapping", None) or {} if conv and remaining: compiled = [(re.compile(p), r) for p, r in conv.items()] - for ck in list(remaining): + hits: dict = defaultdict(list) + for ck in remaining: mk = ck for rx, repl in compiled: mk, n = rx.subn(repl, mk) if n: break if mk in model_keys: - key_map[ck] = mk - remaining.discard(ck) - # Pass 3: longest unique component-suffix match + hits[mk].append(ck) + claim(hits) + # Pass 3: longest unique component-suffix match. Candidates only grow as the suffix shortens, so the + # longest suffix with any candidate decides. if remaining: - from collections import defaultdict - suffix_idx: dict = defaultdict(list) for mk in model_keys: comps = mk.split(".") for k in range(1, len(comps) + 1): suffix_idx[tuple(comps[-k:])].append(mk) + hits = defaultdict(list) for ck in remaining: comps = ck.split(".") for k in range(len(comps), 0, -1): cands = suffix_idx.get(tuple(comps[-k:])) - if cands and len(cands) == 1: - key_map[ck] = cands[0] + if cands: + if len(cands) == 1: + hits[cands[0]].append(ck) break + claim(hits) return key_map @@ -698,7 +716,14 @@ def _meta_dcp_load( model.state_dict() ) # (sharded ``DTensor`` or plain) params + replicated persistent buffers target_names = set(full_sd) if load_names is None else (set(load_names) & set(full_sd)) - key_map = _map_ckpt_keys_to_model(model, ckpt_keys, target_names) # {ckpt_key: model_key} + # {ckpt_key: model_key}, mapped against the WHOLE model and then cut to what this load fills: a suffix + # unique among the targets alone can still belong to a key outside them (an expert's, when only the + # dense part loads here). + key_map = { + ck: mk + for ck, mk in _map_ckpt_keys_to_model(model, ckpt_keys, set(full_sd)).items() + if mk in target_names + } # Build load_sd in the model's state_dict ORDER (identical on every rank) -- NOT key_map/set iteration # order, which is hash-randomized per process. dcp.load's collective load planner deadlocks if ranks # present the same keys in different orders, so a deterministic, rank-consistent order is required. diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index f000b86abe2..b7e11ca47f6 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -173,6 +173,33 @@ def test_ambiguous_or_unknown_suffixes_stay_unmapped(): assert _map_ckpt_keys_to_model(nn.Module(), ckpt_keys, model_keys) == {} +def test_a_key_the_model_has_no_place_for_never_displaces_a_match(): + """DeepSeek-V3's MTP layer and the layers a truncated model lacks have no slot in the model, yet + their suffixes name a key uniquely -- one another checkpoint key already names exactly. Mapped, + they would overwrite that key's tensor; they must stay unmapped. + """ + model_keys = {"model.norm.weight", "model.layers.0.input_layernorm.weight"} + ckpt_keys = model_keys | { + "model.layers.61.shared_head.norm.weight", + "model.layers.1.input_layernorm.weight", + } + + assert _map_ckpt_keys_to_model(nn.Module(), ckpt_keys, model_keys) == { + key: key for key in model_keys + } + + +def test_two_keys_claiming_one_model_key_map_neither(): + """Two checkpoint names reaching the same model key by suffix are equally plausible; loading + either would let set order pick the tensor, so neither is mapped.""" + model_keys = {"model.language_model.norm.weight"} + + assert ( + _map_ckpt_keys_to_model(nn.Module(), {"model.norm.weight", "mtp.norm.weight"}, model_keys) + == {} + ) + + # -------------------------------------------------------------------------- FSDP2 wrap boundaries @@ -278,6 +305,42 @@ def test_load_names_leave_already_real_params_untouched(tmp_path): assert torch.equal(params[name], ref[name]), name +class _DenseAndExpert(nn.Module): + """A dense param the EP-dense load fills and an expert one ``from_pretrained`` already placed.""" + + def __init__(self): + super().__init__() + self.dense = nn.Linear(4, 4, bias=False) + self.expert = nn.Linear(4, 4, bias=False) + + +def test_load_names_judge_suffixes_against_the_whole_model(tmp_path): + """Among the dense names alone ``weight`` names only ``dense.weight``, so a checkpoint key with no + slot in the model (a per-expert ``experts.0.weight``) would land on it. Across the whole model + ``weight`` also names the expert, so that key stays unmapped and the dense param gets its own. + """ + ckpt = {"dense.weight": torch.randn(4, 4), "experts.0.weight": torch.randn(4, 4)} + model = _meta_twin(_DenseAndExpert) + model.expert.to_empty(device="cpu") + + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names={"dense.weight"}) + + assert torch.equal(model.dense.weight, ckpt["dense.weight"]) + + +def test_a_truncated_model_keeps_its_own_layers(tmp_path): + """A 1-layer model loading a 2-layer checkpoint gets the checkpoint's layer 0. Layer 1's names + reach layer 0's keys by suffix alone and must not overwrite them.""" + src = get_tiny_llama() + model = _meta_twin(get_tiny_llama, num_hidden_layers=1) + + _meta_dcp_load(model, _save(tmp_path, src.state_dict()), "cpu") + + ref = src.state_dict() + for name, value in model.state_dict().items(): + assert torch.equal(value, ref[name]), name + + class _WithBuffers(nn.Module): """A linear, a persistent buffer (checkpoints carry it) and a non-persistent one (built at init).""" From 6ced73fa84f2742198306ee7bc240180460d4d92 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:34:54 +0000 Subject: [PATCH 04/12] Keep a real param that shares a module with a meta one With load_names, _meta_dcp_load materializes each module that has a meta param with to_empty(recurse=False), which re-empties every direct param of that module, and it restored only buffers. A param already loaded beside a meta one then stayed uninitialized memory: in the expert-parallel dense load, each MoE router's e_score_correction_bias (a parameter in Kimi's and DeepSeek-V3's remote code) sits next to the gate weight forced to meta, and nothing outside load_names refills it. Snapshot such params and restore them after materializing, as for buffers. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 17 ++++++++++++++-- tests/unit/torch/utils/test_meta_dcp_load.py | 21 ++++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 695b471dbbd..f5fbadfca73 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -685,16 +685,29 @@ def _meta_dcp_load( # returns wrong keys). The PARTIAL case (EP-dense: dense meta, experts already real) MUST be selective so # the real experts are not re-emptied -- there no root is fully_shard'd, so per-module to_empty is safe. saved_buffers = {n: b.detach().clone() for n, b in model.named_buffers() if not b.is_meta} + # A module's ``to_empty`` re-empties every direct param, so a param already real beside a meta one + # (a router's ``e_score_correction_bias`` next to its meta ``weight``) is kept here too: nothing + # outside ``load_names`` would refill it. + saved_params: dict[str, torch.Tensor] = {} if load_names is None: model.to_empty(device=device) else: - for m in model.modules(): - if any(p.is_meta for p in m.parameters(recurse=False)): + for name, m in model.named_modules(): + direct = dict(m.named_parameters(recurse=False)) + if any(p.is_meta for p in direct.values()): + prefix = f"{name}." if name else "" + saved_params.update( + (prefix + n, p.detach().clone()) for n, p in direct.items() if not p.is_meta + ) m.to_empty(device=device, recurse=False) for n, b in model.named_buffers(): src = saved_buffers.get(n) if src is not None: b.copy_(src.to(device=b.device, dtype=b.dtype)) + with torch.no_grad(): + for n, src in saved_params.items(): + p = model.get_parameter(n) + p.copy_(src.to(device=p.device, dtype=p.dtype)) # ``dcp.load`` matches state_dict entries to checkpoint entries BY NAME, but a checkpoint's stored key # names can differ from the built model's ``state_dict`` names, and from_pretrained bridges the gap on diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index b7e11ca47f6..7e7ffac5b06 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -305,6 +305,27 @@ def test_load_names_leave_already_real_params_untouched(tmp_path): assert torch.equal(params[name], ref[name]), name +class _Router(nn.Module): + """A router whose bias correction is a param ``from_pretrained`` loads, beside a meta weight.""" + + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.empty(4, 4, device="meta")) + self.e_score_correction_bias = nn.Parameter(torch.arange(4.0)) + + +def test_a_real_param_beside_a_meta_one_keeps_its_value(tmp_path): + """Materializing the router re-empties all of its params, but only ``weight`` is reloaded.""" + model = nn.Module() + model.gate = _Router() + ckpt = {"gate.weight": torch.randn(4, 4), "gate.e_score_correction_bias": torch.zeros(4)} + + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names={"gate.weight"}) + + assert torch.equal(model.gate.weight, ckpt["gate.weight"]) + assert torch.equal(model.gate.e_score_correction_bias, torch.arange(4.0)) + + class _DenseAndExpert(nn.Module): """A dense param the EP-dense load fills and an expert one ``from_pretrained`` already placed.""" From 4360fb13f41d79b54c66299c898d0e6f3a927719 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 23:24:04 +0000 Subject: [PATCH 05/12] Move the meta-DCP checkpoint load to the model loaders Loading each rank's checkpoint shard into a meta-built model is a model load; it moves from distributed.py to model_load_utils with the other loaders, together with the key mapping and the FSDP2 wrap boundaries it works with. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 256 ----------------- .../torch/utils/plugins/model_load_utils.py | 257 ++++++++++++++++++ tests/unit/torch/utils/test_meta_dcp_load.py | 2 +- 3 files changed, 258 insertions(+), 257 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index f5fbadfca73..e42d8015269 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -18,7 +18,6 @@ import functools import io import os -import re import sys import time import traceback @@ -532,261 +531,6 @@ def distributed_run(timeout: timedelta | None = None, device: torch.device | str cleanup() -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). - """ - 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) - if lineage and lineage[-1] is not model: - lm_mods = _boundaries(lineage[-1]) - if lm_mods: - return lm_mods - - mods = _boundaries(model) - if mods: - return mods - raise RuntimeError( - "Could not determine FSDP2 wrap boundaries: model exposes no `_no_split_modules` and modelopt's " - "decoder-layer auto-detection did not match. Provide a standard HF causal-LM / VLM checkpoint." - ) - - -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. - - Map each CHECKPOINT key to the built model's ``state_dict`` key it feeds, reproducing what - ``from_pretrained`` resolves on load -- so a DCP by-name load fills the right tensors. - - Three passes per checkpoint key (first hit wins): (1) exact match (the common case); (2) the model's - ``_checkpoint_conversion_mapping`` regexes if it carries any (first pattern that fires); (3) longest - UNIQUE trailing-component match. Pass 3 covers the transformers>=5 VLM prefix-nesting - (``model.layers.*`` checkpoint -> ``model.language_model.layers.*`` model, ``visual.*`` -> ``model.visual.*``) - WITHOUT hard-coding it: the leaf path is preserved, so the ckpt key's component suffix uniquely locates - the model key. Ambiguous (>1 candidate) or absent suffixes are left unmapped -> caught by the caller's - coverage guard rather than mis-assigned. - - ``model_keys`` must be the WHOLE model's keys, even when only some of them will be loaded: a suffix is - unique only if no other model key shares it, and the caller drops the mappings it does not load. Each - model key takes at most one checkpoint key -- a later pass never maps onto a key an earlier pass - claimed, and two checkpoint keys claiming one key in the same pass are both left unmapped. Otherwise a - key the model has no place for (an MTP layer's ``model.layers.61.shared_head.norm.weight``, a layer a - truncated model lacks) lands on ``model.norm.weight`` by suffix alone and overwrites it, with set order - picking the winner. - """ - from collections import defaultdict - - key_map: dict = {} - remaining = set(ckpt_keys) - - def claim(hits: dict) -> None: - claimed = set(key_map.values()) - for mk, cks in hits.items(): - if len(cks) == 1 and mk not in claimed: - key_map[cks[0]] = mk - remaining.difference_update(cks) - - # Pass 1: exact - claim({ck: [ck] for ck in remaining if ck in model_keys}) - # Pass 2: model-declared regex conversion mapping (rare; e.g. leaf renames) - conv = getattr(model, "_checkpoint_conversion_mapping", None) or {} - if conv and remaining: - compiled = [(re.compile(p), r) for p, r in conv.items()] - hits: dict = defaultdict(list) - for ck in remaining: - mk = ck - for rx, repl in compiled: - mk, n = rx.subn(repl, mk) - if n: - break - if mk in model_keys: - hits[mk].append(ck) - claim(hits) - # Pass 3: longest unique component-suffix match. Candidates only grow as the suffix shortens, so the - # longest suffix with any candidate decides. - if remaining: - suffix_idx: dict = defaultdict(list) - for mk in model_keys: - comps = mk.split(".") - for k in range(1, len(comps) + 1): - suffix_idx[tuple(comps[-k:])].append(mk) - hits = defaultdict(list) - for ck in remaining: - comps = ck.split(".") - for k in range(len(comps), 0, -1): - cands = suffix_idx.get(tuple(comps[-k:])) - if cands: - if len(cands) == 1: - hits[cands[0]].append(ck) - break - claim(hits) - return key_map - - -def _meta_dcp_load( - model: "torch.nn.Module", - local_dir: str, - device, - *, - load_names: "set | None" = None, - skip_names: "set | None" = None, - max_unloaded_frac: float = 0.5, -) -> None: - """The single meta + torch-DCP loader shared by every distributed dense path. - - materialize each rank's params on ``device`` and fill them from the HF safetensors via DCP's - ``HuggingFaceStorageReader``. - - PLACEMENT is decided entirely by the CALLER, not here: ``dcp.load`` reads a FULL tensor into a plain - (replicated) param and only the LOCAL SLICE into a ``Shard`` ``DTensor``. So the caller picks the layout - by what it does before calling this -- nothing (-> Replicate, full per rank), - ``fully_shard`` the whole model (``load_and_prepare_fsdp2`` -> Shard), or ``fully_shard`` only the dense - modules (the EP hybrid -> dense Shard, experts left as loaded). No full-model copy for the sharded cases. - - ``load_names`` (optional) restricts the checkpoint load to a subset of model param names -- the EP-dense - case, where the experts are ALREADY materialized by ``from_pretrained`` and must not be touched. ``None`` - loads the whole model (DP / FSDP2-full). - - Materialization is SELECTIVE (module-by-module, only where params are meta) so already-real params (EP - experts) are never re-emptied. ``to_empty`` also re-empties BUFFERS (rotary ``inv_freq`` etc., + any the - caller already loaded), so they are snapshotted and restored -- a params-only checkpoint load would - otherwise leave them garbage and the calibration forward would drift (weights load fine but activation - amax is wrong). Persistent buffers the checkpoint DOES carry are (re)loaded via DCP below. - """ - import torch.distributed.checkpoint as dcp - from torch.distributed.checkpoint import HuggingFaceStorageReader - - # Snapshot real buffers (to_empty re-empties them), materialize params, restore buffers. Whole-model - # meta (DP / FSDP2-full, ``load_names is None``) uses ``model.to_empty`` -- the FSDP2-aware root - # materialization; a per-module selective loop corrupts the ROOT-sharded FSDP state (its state_dict then - # returns wrong keys). The PARTIAL case (EP-dense: dense meta, experts already real) MUST be selective so - # the real experts are not re-emptied -- there no root is fully_shard'd, so per-module to_empty is safe. - saved_buffers = {n: b.detach().clone() for n, b in model.named_buffers() if not b.is_meta} - # A module's ``to_empty`` re-empties every direct param, so a param already real beside a meta one - # (a router's ``e_score_correction_bias`` next to its meta ``weight``) is kept here too: nothing - # outside ``load_names`` would refill it. - saved_params: dict[str, torch.Tensor] = {} - if load_names is None: - model.to_empty(device=device) - else: - for name, m in model.named_modules(): - direct = dict(m.named_parameters(recurse=False)) - if any(p.is_meta for p in direct.values()): - prefix = f"{name}." if name else "" - saved_params.update( - (prefix + n, p.detach().clone()) for n, p in direct.items() if not p.is_meta - ) - m.to_empty(device=device, recurse=False) - for n, b in model.named_buffers(): - src = saved_buffers.get(n) - if src is not None: - b.copy_(src.to(device=b.device, dtype=b.dtype)) - with torch.no_grad(): - for n, src in saved_params.items(): - p = model.get_parameter(n) - p.copy_(src.to(device=p.device, dtype=p.dtype)) - - # ``dcp.load`` matches state_dict entries to checkpoint entries BY NAME, but a checkpoint's stored key - # names can differ from the built model's ``state_dict`` names, and from_pretrained bridges the gap on - # load. E.g. Qwen2.5-VL's checkpoint stores ``model.layers.*`` / ``visual.*`` while transformers>=5 nests - # them as ``model.language_model.layers.*`` / ``model.visual.*`` -- a pure prefix-insertion with the leaf - # path preserved, and there is NO ``_checkpoint_conversion_mapping`` on the model to read. A raw by-name - # match would SILENTLY skip every renamed weight (they stay ``to_empty`` garbage -> a wrecked model that - # only surfaces as a confusing downstream crash). Reproduce from_pretrained's result generically via - # :func:`_map_ckpt_keys_to_model`: exact -> conversion-mapping regexes -> longest UNIQUE suffix. - reader = HuggingFaceStorageReader(str(local_dir)) - ckpt_keys = set(reader.read_metadata().state_dict_metadata.keys()) - if not ckpt_keys: - raise RuntimeError( - f"meta+DCP load: the checkpoint at {local_dir} has NO tensors -- the resolved snapshot carries no " - "safetensors (a config-only / incomplete HF revision; e.g. refs/main pointing at a weightless " - "snapshot). Re-download the checkpoint or pass a complete local snapshot dir / pinned revision." - ) - full_sd = ( - model.state_dict() - ) # (sharded ``DTensor`` or plain) params + replicated persistent buffers - target_names = set(full_sd) if load_names is None else (set(load_names) & set(full_sd)) - # {ckpt_key: model_key}, mapped against the WHOLE model and then cut to what this load fills: a suffix - # unique among the targets alone can still belong to a key outside them (an expert's, when only the - # dense part loads here). - key_map = { - ck: mk - for ck, mk in _map_ckpt_keys_to_model(model, ckpt_keys, set(full_sd)).items() - if mk in target_names - } - # Build load_sd in the model's state_dict ORDER (identical on every rank) -- NOT key_map/set iteration - # order, which is hash-randomized per process. dcp.load's collective load planner deadlocks if ranks - # present the same keys in different orders, so a deterministic, rank-consistent order is required. - inv = {mk: ck for ck, mk in key_map.items()} # model_key -> ckpt_key (1:1 on model keys) - load_sd = {inv[mk]: full_sd[mk] for mk in full_sd if mk in inv} - mapped_names = set(key_map.values()) - # Fail LOUD if the mapping left too much of the TARGET WEIGHT unloaded -- a key scheme this path did not - # resolve; a silent partial load produces a garbage-weight model. Weight by NUMEL, not param COUNT: a - # fused-expert MoE (Qwen3-MoE: per-expert checkpoint -> grouped ``experts.gate_up_proj``/``down_proj`` - # model params DCP can't fuse) leaves only ~2 params/layer unloaded but they are ~95% of the weight, so a - # count-based guard misses it. A small tail is normal (tied weights + a VLM's not-in-checkpoint embeds). - # ``skip_names`` = params loaded by a SEPARATE path (e.g. Qwen3-MoE fused grouped experts, filled by - # :func:`_load_fused_experts_sharded` which fuses per-expert checkpoint tensors DCP can't concat) -- - # exclude them from the coverage guard so it doesn't (correctly) flag them as unloaded here. - skip_names = skip_names or set() - numel = { - n: p.numel() - for n, p in model.named_parameters() - if n in target_names and n not in skip_names - } - want_params = set(numel) - unloaded = want_params - mapped_names - total_nel = max(1, sum(numel.values())) - unloaded_nel = sum(numel[n] for n in unloaded) - if unloaded_nel > max_unloaded_frac * total_nel: - raise RuntimeError( - f"meta+DCP load: {len(unloaded)}/{len(want_params)} params ({100 * unloaded_nel / total_nel:.0f}% " - f"of the target weight) matched no checkpoint key -- the checkpoint's key naming was not resolved " - f"to this model's state_dict (exact + conversion-mapping + suffix match all missed), or it needs a " - f"weight conversion DCP can't do (FUSED EXPERTS, split qkv, on-the-fly dequant). Unloaded e.g.: " - f"{sorted(unloaded)[:5]}. Use the from_pretrained load (MODELOPT_FSDP2_SHARDED_LOAD=0 / " - f"MODELOPT_EP_DENSE_DISTRIBUTED_LOAD=0)." - ) - if is_initialized() and rank() == 0: - loaded_bufs = mapped_names - {n for n, _ in model.named_parameters()} - if unloaded: - warn( - f"meta+DCP load: {len(unloaded)} param(s) not loaded from checkpoint (tied? re-tied): " - f"{sorted(unloaded)[:5]}" - ) - if loaded_bufs: - warn( - f"meta+DCP load: loaded {len(loaded_bufs)} persistent buffer(s) from checkpoint: " - f"{sorted(loaded_bufs)[:5]}" - ) - dcp.load(load_sd, storage_reader=reader) - if hasattr(model, "tie_weights"): - model.tie_weights() # re-establish tied weights the checkpoint omitted - - def cleanup(): """Cleans up the distributed environment. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index a57b65d38fc..db30b75054d 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -21,6 +21,7 @@ import re from itertools import chain from typing import TYPE_CHECKING, Any +from warnings import warn import torch import torch.nn as nn @@ -43,6 +44,7 @@ barrier, broadcast_state_dict, is_initialized, + rank, set_distribution_plan, ) from modelopt.torch.utils.logging import print_rank_0 @@ -666,6 +668,261 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False return decoder_layers +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). + """ + 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) + if lineage and lineage[-1] is not model: + lm_mods = _boundaries(lineage[-1]) + if lm_mods: + return lm_mods + + mods = _boundaries(model) + if mods: + return mods + raise RuntimeError( + "Could not determine FSDP2 wrap boundaries: model exposes no `_no_split_modules` and modelopt's " + "decoder-layer auto-detection did not match. Provide a standard HF causal-LM / VLM checkpoint." + ) + + +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. + + Map each CHECKPOINT key to the built model's ``state_dict`` key it feeds, reproducing what + ``from_pretrained`` resolves on load -- so a DCP by-name load fills the right tensors. + + Three passes per checkpoint key (first hit wins): (1) exact match (the common case); (2) the model's + ``_checkpoint_conversion_mapping`` regexes if it carries any (first pattern that fires); (3) longest + UNIQUE trailing-component match. Pass 3 covers the transformers>=5 VLM prefix-nesting + (``model.layers.*`` checkpoint -> ``model.language_model.layers.*`` model, ``visual.*`` -> ``model.visual.*``) + WITHOUT hard-coding it: the leaf path is preserved, so the ckpt key's component suffix uniquely locates + the model key. Ambiguous (>1 candidate) or absent suffixes are left unmapped -> caught by the caller's + coverage guard rather than mis-assigned. + + ``model_keys`` must be the WHOLE model's keys, even when only some of them will be loaded: a suffix is + unique only if no other model key shares it, and the caller drops the mappings it does not load. Each + model key takes at most one checkpoint key -- a later pass never maps onto a key an earlier pass + claimed, and two checkpoint keys claiming one key in the same pass are both left unmapped. Otherwise a + key the model has no place for (an MTP layer's ``model.layers.61.shared_head.norm.weight``, a layer a + truncated model lacks) lands on ``model.norm.weight`` by suffix alone and overwrites it, with set order + picking the winner. + """ + from collections import defaultdict + + key_map: dict = {} + remaining = set(ckpt_keys) + + def claim(hits: dict) -> None: + claimed = set(key_map.values()) + for mk, cks in hits.items(): + if len(cks) == 1 and mk not in claimed: + key_map[cks[0]] = mk + remaining.difference_update(cks) + + # Pass 1: exact + claim({ck: [ck] for ck in remaining if ck in model_keys}) + # Pass 2: model-declared regex conversion mapping (rare; e.g. leaf renames) + conv = getattr(model, "_checkpoint_conversion_mapping", None) or {} + if conv and remaining: + compiled = [(re.compile(p), r) for p, r in conv.items()] + hits: dict = defaultdict(list) + for ck in remaining: + mk = ck + for rx, repl in compiled: + mk, n = rx.subn(repl, mk) + if n: + break + if mk in model_keys: + hits[mk].append(ck) + claim(hits) + # Pass 3: longest unique component-suffix match. Candidates only grow as the suffix shortens, so the + # longest suffix with any candidate decides. + if remaining: + suffix_idx: dict = defaultdict(list) + for mk in model_keys: + comps = mk.split(".") + for k in range(1, len(comps) + 1): + suffix_idx[tuple(comps[-k:])].append(mk) + hits = defaultdict(list) + for ck in remaining: + comps = ck.split(".") + for k in range(len(comps), 0, -1): + cands = suffix_idx.get(tuple(comps[-k:])) + if cands: + if len(cands) == 1: + hits[cands[0]].append(ck) + break + claim(hits) + return key_map + + +def _meta_dcp_load( + model: "torch.nn.Module", + local_dir: str, + device, + *, + load_names: "set | None" = None, + skip_names: "set | None" = None, + max_unloaded_frac: float = 0.5, +) -> None: + """The single meta + torch-DCP loader shared by every distributed dense path. + + materialize each rank's params on ``device`` and fill them from the HF safetensors via DCP's + ``HuggingFaceStorageReader``. + + PLACEMENT is decided entirely by the CALLER, not here: ``dcp.load`` reads a FULL tensor into a plain + (replicated) param and only the LOCAL SLICE into a ``Shard`` ``DTensor``. So the caller picks the layout + by what it does before calling this -- nothing (-> Replicate, full per rank), + ``fully_shard`` the whole model (``load_and_prepare_fsdp2`` -> Shard), or ``fully_shard`` only the dense + modules (the EP hybrid -> dense Shard, experts left as loaded). No full-model copy for the sharded cases. + + ``load_names`` (optional) restricts the checkpoint load to a subset of model param names -- the EP-dense + case, where the experts are ALREADY materialized by ``from_pretrained`` and must not be touched. ``None`` + loads the whole model (DP / FSDP2-full). + + Materialization is SELECTIVE (module-by-module, only where params are meta) so already-real params (EP + experts) are never re-emptied. ``to_empty`` also re-empties BUFFERS (rotary ``inv_freq`` etc., + any the + caller already loaded), so they are snapshotted and restored -- a params-only checkpoint load would + otherwise leave them garbage and the calibration forward would drift (weights load fine but activation + amax is wrong). Persistent buffers the checkpoint DOES carry are (re)loaded via DCP below. + """ + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint import HuggingFaceStorageReader + + # Snapshot real buffers (to_empty re-empties them), materialize params, restore buffers. Whole-model + # meta (DP / FSDP2-full, ``load_names is None``) uses ``model.to_empty`` -- the FSDP2-aware root + # materialization; a per-module selective loop corrupts the ROOT-sharded FSDP state (its state_dict then + # returns wrong keys). The PARTIAL case (EP-dense: dense meta, experts already real) MUST be selective so + # the real experts are not re-emptied -- there no root is fully_shard'd, so per-module to_empty is safe. + saved_buffers = {n: b.detach().clone() for n, b in model.named_buffers() if not b.is_meta} + # A module's ``to_empty`` re-empties every direct param, so a param already real beside a meta one + # (a router's ``e_score_correction_bias`` next to its meta ``weight``) is kept here too: nothing + # outside ``load_names`` would refill it. + saved_params: dict[str, torch.Tensor] = {} + if load_names is None: + model.to_empty(device=device) + else: + for name, m in model.named_modules(): + direct = dict(m.named_parameters(recurse=False)) + if any(p.is_meta for p in direct.values()): + prefix = f"{name}." if name else "" + saved_params.update( + (prefix + n, p.detach().clone()) for n, p in direct.items() if not p.is_meta + ) + m.to_empty(device=device, recurse=False) + for n, b in model.named_buffers(): + src = saved_buffers.get(n) + if src is not None: + b.copy_(src.to(device=b.device, dtype=b.dtype)) + with torch.no_grad(): + for n, src in saved_params.items(): + p = model.get_parameter(n) + p.copy_(src.to(device=p.device, dtype=p.dtype)) + + # ``dcp.load`` matches state_dict entries to checkpoint entries BY NAME, but a checkpoint's stored key + # names can differ from the built model's ``state_dict`` names, and from_pretrained bridges the gap on + # load. E.g. Qwen2.5-VL's checkpoint stores ``model.layers.*`` / ``visual.*`` while transformers>=5 nests + # them as ``model.language_model.layers.*`` / ``model.visual.*`` -- a pure prefix-insertion with the leaf + # path preserved, and there is NO ``_checkpoint_conversion_mapping`` on the model to read. A raw by-name + # match would SILENTLY skip every renamed weight (they stay ``to_empty`` garbage -> a wrecked model that + # only surfaces as a confusing downstream crash). Reproduce from_pretrained's result generically via + # :func:`_map_ckpt_keys_to_model`: exact -> conversion-mapping regexes -> longest UNIQUE suffix. + reader = HuggingFaceStorageReader(str(local_dir)) + ckpt_keys = set(reader.read_metadata().state_dict_metadata.keys()) + if not ckpt_keys: + raise RuntimeError( + f"meta+DCP load: the checkpoint at {local_dir} has NO tensors -- the resolved snapshot carries no " + "safetensors (a config-only / incomplete HF revision; e.g. refs/main pointing at a weightless " + "snapshot). Re-download the checkpoint or pass a complete local snapshot dir / pinned revision." + ) + full_sd = ( + model.state_dict() + ) # (sharded ``DTensor`` or plain) params + replicated persistent buffers + target_names = set(full_sd) if load_names is None else (set(load_names) & set(full_sd)) + # {ckpt_key: model_key}, mapped against the WHOLE model and then cut to what this load fills: a suffix + # unique among the targets alone can still belong to a key outside them (an expert's, when only the + # dense part loads here). + key_map = { + ck: mk + for ck, mk in _map_ckpt_keys_to_model(model, ckpt_keys, set(full_sd)).items() + if mk in target_names + } + # Build load_sd in the model's state_dict ORDER (identical on every rank) -- NOT key_map/set iteration + # order, which is hash-randomized per process. dcp.load's collective load planner deadlocks if ranks + # present the same keys in different orders, so a deterministic, rank-consistent order is required. + inv = {mk: ck for ck, mk in key_map.items()} # model_key -> ckpt_key (1:1 on model keys) + load_sd = {inv[mk]: full_sd[mk] for mk in full_sd if mk in inv} + mapped_names = set(key_map.values()) + # Fail LOUD if the mapping left too much of the TARGET WEIGHT unloaded -- a key scheme this path did not + # resolve; a silent partial load produces a garbage-weight model. Weight by NUMEL, not param COUNT: a + # fused-expert MoE (Qwen3-MoE: per-expert checkpoint -> grouped ``experts.gate_up_proj``/``down_proj`` + # model params DCP can't fuse) leaves only ~2 params/layer unloaded but they are ~95% of the weight, so a + # count-based guard misses it. A small tail is normal (tied weights + a VLM's not-in-checkpoint embeds). + # ``skip_names`` = params loaded by a SEPARATE path (e.g. Qwen3-MoE fused grouped experts, filled by + # :func:`_load_fused_experts_sharded` which fuses per-expert checkpoint tensors DCP can't concat) -- + # exclude them from the coverage guard so it doesn't (correctly) flag them as unloaded here. + skip_names = skip_names or set() + numel = { + n: p.numel() + for n, p in model.named_parameters() + if n in target_names and n not in skip_names + } + want_params = set(numel) + unloaded = want_params - mapped_names + total_nel = max(1, sum(numel.values())) + unloaded_nel = sum(numel[n] for n in unloaded) + if unloaded_nel > max_unloaded_frac * total_nel: + raise RuntimeError( + f"meta+DCP load: {len(unloaded)}/{len(want_params)} params ({100 * unloaded_nel / total_nel:.0f}% " + f"of the target weight) matched no checkpoint key -- the checkpoint's key naming was not resolved " + f"to this model's state_dict (exact + conversion-mapping + suffix match all missed), or it needs a " + f"weight conversion DCP can't do (FUSED EXPERTS, split qkv, on-the-fly dequant). Unloaded e.g.: " + f"{sorted(unloaded)[:5]}. Use the from_pretrained load (MODELOPT_FSDP2_SHARDED_LOAD=0 / " + f"MODELOPT_EP_DENSE_DISTRIBUTED_LOAD=0)." + ) + if is_initialized() and rank() == 0: + loaded_bufs = mapped_names - {n for n, _ in model.named_parameters()} + if unloaded: + warn( + f"meta+DCP load: {len(unloaded)} param(s) not loaded from checkpoint (tied? re-tied): " + f"{sorted(unloaded)[:5]}" + ) + if loaded_bufs: + warn( + f"meta+DCP load: loaded {len(loaded_bufs)} persistent buffer(s) from checkpoint: " + f"{sorted(loaded_bufs)[:5]}" + ) + dcp.load(load_sd, storage_reader=reader) + if hasattr(model, "tie_weights"): + model.tie_weights() # re-establish tied weights the checkpoint omitted + + def parallel_load_and_prepare_fsdp2( ckpt_path: str, device: torch.device, diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 7e7ffac5b06..e83cb24496b 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -39,7 +39,7 @@ from accelerate import init_empty_weights from transformers import AutoModelForImageTextToText, Qwen2_5_VLConfig -from modelopt.torch.utils.distributed import ( +from modelopt.torch.utils.plugins.model_load_utils import ( _fsdp2_wrap_modules, _map_ckpt_keys_to_model, _meta_dcp_load, From e8c5b035d1b3144f269b4dd2e1a171ed1e424915 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 20:48:29 +0000 Subject: [PATCH 06/12] Fail the meta+DCP load on params the checkpoint does not fill The load tolerated up to half the target weight missing and returned the model with whatever it lacked still holding the memory to_empty handed out. Only tied weights are refilled after the load; an untied param with no checkpoint tensor stayed garbage while the load reported success. The load now fails, naming them, for every re-emptied target param that no checkpoint tensor filled and tie_weights did not tie to a loaded one. Params that already had values -- restored beside a meta one, or in a module the load leaves alone -- keep them. The numel guard stays as the check that fails before the load, and its advice names the knob that exists (MODELOPT_FSDP_SIZE=1) instead of two removed ones. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 40 +++++++--- tests/unit/torch/utils/test_meta_dcp_load.py | 74 ++++++++++++++----- 2 files changed, 85 insertions(+), 29 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index db30b75054d..1c1c58c7df6 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -811,6 +811,9 @@ def _meta_dcp_load( caller already loaded), so they are snapshotted and restored -- a params-only checkpoint load would otherwise leave them garbage and the calibration forward would drift (weights load fine but activation amax is wrong). Persistent buffers the checkpoint DOES carry are (re)loaded via DCP below. + + Every target param ends up filled -- from the checkpoint, by ``tie_weights``, or with the value it + already had -- or the load fails: nothing is left as the uninitialized memory ``to_empty`` hands out. """ import torch.distributed.checkpoint as dcp from torch.distributed.checkpoint import HuggingFaceStorageReader @@ -826,8 +829,10 @@ def _meta_dcp_load( # outside ``load_names`` would refill it. saved_params: dict[str, torch.Tensor] = {} if load_names is None: + emptied = {n for n, _ in model.named_parameters(remove_duplicate=False)} model.to_empty(device=device) else: + emptied = set() for name, m in model.named_modules(): direct = dict(m.named_parameters(recurse=False)) if any(p.is_meta for p in direct.values()): @@ -835,6 +840,7 @@ def _meta_dcp_load( saved_params.update( (prefix + n, p.detach().clone()) for n, p in direct.items() if not p.is_meta ) + emptied.update(prefix + n for n in direct) m.to_empty(device=device, recurse=False) for n, b in model.named_buffers(): src = saved_buffers.get(n) @@ -879,11 +885,12 @@ def _meta_dcp_load( inv = {mk: ck for ck, mk in key_map.items()} # model_key -> ckpt_key (1:1 on model keys) load_sd = {inv[mk]: full_sd[mk] for mk in full_sd if mk in inv} mapped_names = set(key_map.values()) - # Fail LOUD if the mapping left too much of the TARGET WEIGHT unloaded -- a key scheme this path did not - # resolve; a silent partial load produces a garbage-weight model. Weight by NUMEL, not param COUNT: a - # fused-expert MoE (Qwen3-MoE: per-expert checkpoint -> grouped ``experts.gate_up_proj``/``down_proj`` - # model params DCP can't fuse) leaves only ~2 params/layer unloaded but they are ~95% of the weight, so a - # count-based guard misses it. A small tail is normal (tied weights + a VLM's not-in-checkpoint embeds). + # Fail fast, BEFORE the load, if the mapping left too much of the TARGET WEIGHT unloaded -- a key scheme + # this path did not resolve. Weight by NUMEL, not param COUNT: a fused-expert MoE (Qwen3-MoE: per-expert + # checkpoint -> grouped ``experts.gate_up_proj``/``down_proj`` model params DCP can't fuse) leaves only + # ~2 params/layer unloaded but they are ~95% of the weight, so a count-based guard misses it. A smaller + # shortfall can be tied weights, which ``tie_weights`` fills after the load; whatever it does not fill + # fails there. # ``skip_names`` = params loaded by a SEPARATE path (e.g. Qwen3-MoE fused grouped experts, filled by # :func:`_load_fused_experts_sharded` which fuses per-expert checkpoint tensors DCP can't concat) -- # exclude them from the coverage guard so it doesn't (correctly) flag them as unloaded here. @@ -903,16 +910,11 @@ def _meta_dcp_load( f"of the target weight) matched no checkpoint key -- the checkpoint's key naming was not resolved " f"to this model's state_dict (exact + conversion-mapping + suffix match all missed), or it needs a " f"weight conversion DCP can't do (FUSED EXPERTS, split qkv, on-the-fly dequant). Unloaded e.g.: " - f"{sorted(unloaded)[:5]}. Use the from_pretrained load (MODELOPT_FSDP2_SHARDED_LOAD=0 / " - f"MODELOPT_EP_DENSE_DISTRIBUTED_LOAD=0)." + f"{sorted(unloaded)[:5]}. Replicating the dense part (MODELOPT_FSDP_SIZE=1) loads it through " + "from_pretrained instead." ) if is_initialized() and rank() == 0: loaded_bufs = mapped_names - {n for n, _ in model.named_parameters()} - if unloaded: - warn( - f"meta+DCP load: {len(unloaded)} param(s) not loaded from checkpoint (tied? re-tied): " - f"{sorted(unloaded)[:5]}" - ) if loaded_bufs: warn( f"meta+DCP load: loaded {len(loaded_bufs)} persistent buffer(s) from checkpoint: " @@ -921,6 +923,20 @@ def _meta_dcp_load( dcp.load(load_sd, storage_reader=reader) if hasattr(model, "tie_weights"): model.tie_weights() # re-establish tied weights the checkpoint omitted + # A re-emptied target param no checkpoint tensor filled, and that tying did not make a loaded param, + # still holds ``to_empty``'s memory. from_pretrained would 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. + loaded = {id(model.get_parameter(n)) for n in want_params & mapped_names} + unfilled = sorted( + n + for n in (unloaded & emptied) - saved_params.keys() + if id(model.get_parameter(n)) not in loaded + ) + if unfilled: + raise RuntimeError( + f"meta+DCP load: {len(unfilled)} param(s) have no checkpoint tensor and are tied to no loaded " + f"param, so they would stay uninitialized: {unfilled[:5]}" + ) def parallel_load_and_prepare_fsdp2( diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index e83cb24496b..b8b766dc0c2 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -326,6 +326,28 @@ def test_a_real_param_beside_a_meta_one_keeps_its_value(tmp_path): assert torch.equal(model.gate.e_score_correction_bias, torch.arange(4.0)) +def test_target_params_that_already_had_values_need_no_checkpoint_tensor(tmp_path): + """A target param that was real before the load -- beside a meta one, so restored after + ``to_empty``, or in a module the load never re-empties -- keeps its value when the checkpoint + lacks it, instead of failing as unfilled.""" + model = nn.Module() + model.gate = _Router() + model.expert = nn.Linear(2, 2, bias=False) + expert = model.expert.weight.detach().clone() + ckpt = {"gate.weight": torch.randn(4, 4)} + + _meta_dcp_load( + model, + _save(tmp_path, ckpt), + "cpu", + load_names={"gate.weight", "gate.e_score_correction_bias", "expert.weight"}, + ) + + assert torch.equal(model.gate.weight, ckpt["gate.weight"]) + assert torch.equal(model.gate.e_score_correction_bias, torch.arange(4.0)) + assert torch.equal(model.expert.weight, expert) + + class _DenseAndExpert(nn.Module): """A dense param the EP-dense load fills and an expert one ``from_pretrained`` already placed.""" @@ -391,6 +413,19 @@ def test_persistent_buffers_load_from_the_checkpoint_and_others_survive(tmp_path assert torch.equal(model.proj.weight, ckpt["proj.weight"]) +def test_an_untied_param_the_checkpoint_lacks_fails_the_load(tmp_path): + """``proj.bias`` has no checkpoint tensor and tying fills nothing: it still holds the memory + ``to_empty`` gave it, so the load fails naming it instead of returning the model.""" + ckpt = {"proj.weight": torch.randn(4, 4), "scale": torch.full((4,), 3.0)} + with init_empty_weights(): + model = _WithBuffers() + + with pytest.raises( + RuntimeError, match=r"1 param\(s\) have no checkpoint tensor.*\['proj.bias'\]" + ): + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu") + + class _OneBigManySmall(nn.Module): """One param holding ~99% of the weight beside twelve tiny ones.""" @@ -411,17 +446,28 @@ def _guard_case(ckpt_dir, missing): return model, ckpt, _save(ckpt_dir, ckpt) -def test_coverage_guard_weighs_unloaded_params_by_numel(tmp_path): - """Missing the one ~99% param fails loudly while missing the twelve tiny ones (~1%) loads. A - count-based guard gets both backwards, and the fused-expert miss it exists for is a few huge params. +def test_coverage_guard_weighs_unloaded_params_by_numel(tmp_path, monkeypatch): + """Missing the one ~99% param fails before the load; missing the twelve tiny ones (~1%) gets past + that guard and fails after the load instead. A count-based guard gets both backwards, and the + fused-expert miss it exists for is a few huge params. """ + loads = [] + real_load = dcp.load + + def counting_load(state_dict, **kwargs): + loads.append(len(state_dict)) + return real_load(state_dict, **kwargs) + + monkeypatch.setattr(dcp, "load", counting_load) model, _, ckpt_dir = _guard_case(tmp_path / "big_missing", missing={"big.weight"}) with pytest.raises(RuntimeError, match="matched no checkpoint key"): _meta_dcp_load(model, ckpt_dir, "cpu") + assert not loads - model, ckpt, ckpt_dir = _guard_case(tmp_path / "small_missing", missing=_SMALL) - _meta_dcp_load(model, ckpt_dir, "cpu") - assert torch.equal(model.big.weight, ckpt["big.weight"]) + model, _, ckpt_dir = _guard_case(tmp_path / "small_missing", missing=_SMALL) + with pytest.raises(RuntimeError, match="12 param\\(s\\) have no checkpoint tensor"): + _meta_dcp_load(model, ckpt_dir, "cpu") + assert loads == [1] def test_skip_names_exempt_params_filled_by_another_path(tmp_path): @@ -434,7 +480,7 @@ def test_skip_names_exempt_params_filled_by_another_path(tmp_path): def test_max_unloaded_frac_sets_the_guard_threshold(tmp_path): - """The ~1% shortfall the default tolerates fails once the tolerance is zero.""" + """The ~1% shortfall the default lets past the early guard fails it once the tolerance is zero.""" model, _, ckpt_dir = _guard_case(tmp_path, missing=_SMALL) with pytest.raises(RuntimeError, match="matched no checkpoint key"): @@ -484,21 +530,15 @@ def gloo_world_of_one(skip_on_windows, tmp_path_factory): torch.distributed.destroy_process_group() -def test_rank0_reports_unloaded_params_and_loaded_buffers(tmp_path, gloo_world_of_one): - """In a distributed run rank 0 names the params the checkpoint left unfilled (below the guard's - bar) and the persistent buffers it filled. - """ - ckpt = {"proj.weight": torch.randn(4, 4), "scale": torch.full((4,), 3.0)} +def test_rank0_reports_the_persistent_buffers_it_loaded(tmp_path, gloo_world_of_one): + """In a distributed run rank 0 names the persistent buffers the checkpoint filled.""" + ckpt = {"proj.weight": torch.randn(4, 4), "proj.bias": torch.randn(4), "scale": torch.ones(4)} with init_empty_weights(): model = _WithBuffers() - with pytest.warns(UserWarning) as record: + with pytest.warns(UserWarning, match=r"1 persistent buffer\(s\).*'scale'"): _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu") - messages = [str(w.message) for w in record] - assert any("1 param(s) not loaded" in m and "'proj.bias'" in m for m in messages), messages - assert any("1 persistent buffer(s)" in m and "'scale'" in m for m in messages), messages - def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one): """FSDP2-full: the blocks and the root are ``fully_shard``-ed on meta before the load. The params From d9ac551ecc08912dc5b166c2728bced392ad0f2c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 23:24:28 +0000 Subject: [PATCH 07/12] Keep a tied checkpoint's values when tying keeps the alias it left out A tied checkpoint may store only lm_head.weight. The load filled the head, then tie_weights replaced it with the embedding's empty param, and the strict check, reading the params after the tie, counted that param as loaded: both tied weights came back uninitialized. The params the checkpoint filled are now taken before tying, and a kept param tying left unfilled takes its filled alias's values, as from_pretrained leaves it. The copy is local, for plain tensors and alike-sharded DTensors; any other layout is refused. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 40 +++++++++++++++- tests/unit/torch/utils/test_meta_dcp_load.py | 48 +++++++++++++++++++ 2 files changed, 87 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 1c1c58c7df6..8ea05af288f 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -782,6 +782,34 @@ def claim(hits: dict) -> None: return key_map +def _copy_tied_values(kept: torch.Tensor, src: torch.Tensor, name: str) -> None: + """Copy ``name``'s checkpoint values from ``src`` into ``kept``, the tied param that replaced it. + + The copy is local: plain tensors, or DTensors laid out alike (one mesh, the same placements), as + FSDP2 shards tied params. Any other layout would need a redistribution and is refused instead. + """ + from torch.distributed.tensor import DTensor + + alike = kept.shape == src.shape and ( + not (isinstance(kept, DTensor) or isinstance(src, DTensor)) + or ( + isinstance(kept, DTensor) + and isinstance(src, DTensor) + and kept.device_mesh == src.device_mesh + and kept.placements == src.placements + ) + ) + if not alike: + raise RuntimeError( + f"meta+DCP load: tie_weights replaced {name}, which the checkpoint filled, with a param " + "laid out differently, so its values cannot be carried over" + ) + with torch.no_grad(): + if isinstance(kept, DTensor): + kept, src = kept.to_local(), src.to_local() + kept.copy_(src) + + def _meta_dcp_load( model: "torch.nn.Module", local_dir: str, @@ -921,12 +949,22 @@ def _meta_dcp_load( f"{sorted(loaded_bufs)[:5]}" ) dcp.load(load_sd, storage_reader=reader) + # Taken before tying, which can swap a filled param out of its module. + filled = {n: model.get_parameter(n) for n in want_params & 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 + # checkpoint storing only ``lm_head.weight`` loses its values to the embedding's empty param. The kept + # param takes the filled one's values, as from_pretrained leaves it. + loaded = {id(p) for p in filled.values()} + for name, src in filled.items(): + kept = model.get_parameter(name) + if id(kept) not in loaded: + _copy_tied_values(kept, src, name) + loaded.add(id(kept)) # A re-emptied target param no checkpoint tensor filled, and that tying did not make a loaded param, # still holds ``to_empty``'s memory. from_pretrained would 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. - loaded = {id(model.get_parameter(n)) for n in want_params & mapped_names} unfilled = sorted( n for n in (unloaded & emptied) - saved_params.keys() diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index b8b766dc0c2..ad2c172ce89 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -261,6 +261,30 @@ def test_load_fills_a_meta_model_keeps_its_buffers_and_reties_it(tmp_path): assert model.lm_head.weight is model.model.embed_tokens.weight +@pytest.mark.parametrize("stored", ["model.embed_tokens.weight", "lm_head.weight"]) +def test_a_tie_keeps_the_checkpoint_values_whichever_alias_holds_them( + tmp_path, monkeypatch, stored +): + """A tied checkpoint stores one alias; tying keeps the embedding's parameter. When the checkpoint + stored only ``lm_head.weight``, the kept parameter must take its values rather than keep the memory + ``to_empty`` gave it, as ``from_pretrained`` leaves it.""" + src = get_tiny_llama(tie_word_embeddings=True) + tied = {"model.embed_tokens.weight", "lm_head.weight"} + ckpt = {k: v for k, v in src.state_dict().items() if k not in tied - {stored}} + ckpt_dir = _save(tmp_path, ckpt) + model = _meta_twin(get_tiny_llama, tie_word_embeddings=True) + # Fresh storage holds NaN, so a parameter nothing wrote to shows. + empty_like = torch.empty_like + monkeypatch.setattr( + torch, "empty_like", lambda t, **kw: empty_like(t, **kw).fill_(float("nan")) + ) + + _meta_dcp_load(model, ckpt_dir, "cpu") + + assert model.lm_head.weight is model.model.embed_tokens.weight + assert torch.equal(model.lm_head.weight, ckpt[stored]) + + def test_flat_vlm_checkpoint_names_reach_the_nested_params(tmp_path): """Qwen2.5-VL's checkpoint says ``model.layers.*`` / ``visual.*`` where the built model has ``model.language_model.*`` / ``model.visual.*``; a by-name load would fill none of them. @@ -562,3 +586,27 @@ def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one): assert torch.equal(param.full_tensor(), ref[name]), name for name, buf in model.named_buffers(): assert torch.equal(buf, buffers[name]), name + + +def test_an_fsdp2_tie_keeps_the_values_of_the_alias_the_checkpoint_stored( + tmp_path, gloo_world_of_one, monkeypatch +): + """The same tie under FSDP2, where the kept param and the filled one are sharded alike.""" + src = get_tiny_llama(tie_word_embeddings=True) + ckpt = {k: v for k, v in src.state_dict().items() if k != "model.embed_tokens.weight"} + ckpt_dir = _save(tmp_path, ckpt) + model = _meta_twin(get_tiny_llama, tie_word_embeddings=True) + mesh = init_device_mesh("cpu", (1,)) + for block in _fsdp2_wrap_modules(model): + fully_shard(block, mesh=mesh) + fully_shard(model, mesh=mesh) + empty_like = torch.empty_like + monkeypatch.setattr( + torch, "empty_like", lambda t, **kw: empty_like(t, **kw).fill_(float("nan")) + ) + + _meta_dcp_load(model, ckpt_dir, "cpu") + + assert isinstance(model.lm_head.weight, DTensor) + assert torch.equal(model.lm_head.weight.full_tensor(), ckpt["lm_head.weight"]) + assert torch.equal(model.model.embed_tokens.weight.full_tensor(), ckpt["lm_head.weight"]) From ba053d2b98a58db58aa766bae53fefade6bb8112 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 00:28:03 +0000 Subject: [PATCH 08/12] Test that a tie laid out differently is refused and left untouched _copy_tied_values refuses a tied pair it would have to redistribute: a DTensor against a plain tensor, DTensors placed differently or on different meshes, or mismatched shapes. No test reached that branch, so a version accepting every pair passed. Each layout now raises, and the kept param keeps what it held. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_meta_dcp_load.py | 40 +++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index ad2c172ce89..21758ffec91 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -30,7 +30,7 @@ from torch import nn from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import fully_shard -from torch.distributed.tensor import DTensor +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor pytest.importorskip("accelerate") pytest.importorskip("transformers") @@ -40,6 +40,7 @@ from transformers import AutoModelForImageTextToText, Qwen2_5_VLConfig from modelopt.torch.utils.plugins.model_load_utils import ( + _copy_tied_values, _fsdp2_wrap_modules, _map_ckpt_keys_to_model, _meta_dcp_load, @@ -610,3 +611,40 @@ def test_an_fsdp2_tie_keeps_the_values_of_the_alias_the_checkpoint_stored( assert isinstance(model.lm_head.weight, DTensor) assert torch.equal(model.lm_head.weight.full_tensor(), ckpt["lm_head.weight"]) assert torch.equal(model.model.embed_tokens.weight.full_tensor(), ckpt["lm_head.weight"]) + + +def _tie_pair(layout): + """A kept param and the filled one it replaced, laid out as ``layout`` names.""" + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + kept, values = torch.zeros(4, 4), torch.arange(16.0).reshape(4, 4) + if layout == "dtensor-vs-plain": + return distribute_tensor(kept, mesh, [Shard(0)]), values + if layout == "plain-vs-dtensor": + return kept, distribute_tensor(values, mesh, [Shard(0)]) + if layout == "other-placements": + return distribute_tensor(kept, mesh, [Shard(0)]), distribute_tensor( + values, mesh, [Replicate()] + ) + if layout == "other-mesh": + other = init_device_mesh("cpu", (1,), mesh_dim_names=("ep",)) + return distribute_tensor(kept, mesh, [Shard(0)]), distribute_tensor( + values, other, [Shard(0)] + ) + return kept, values.reshape(2, 8) # other-shape + + +@pytest.mark.parametrize( + "layout", + ["dtensor-vs-plain", "plain-vs-dtensor", "other-placements", "other-mesh", "other-shape"], +) +def test_a_tie_laid_out_differently_is_refused_and_left_as_it_was(gloo_world_of_one, layout): + """Carrying a tied param's values over is a local copy. A pair it would have to redistribute is + refused rather than counted as loaded, and the kept param keeps what it held.""" + kept, src = _tie_pair(layout) + local = kept.to_local() if isinstance(kept, DTensor) else kept + before = local.clone() + + with pytest.raises(RuntimeError, match="laid out differently"): + _copy_tied_values(kept, src, "lm_head.weight") + + assert torch.equal(local, before) From 4381e72ab58c34dd02843dac1b0c3659b61e304a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 06:59:49 +0000 Subject: [PATCH 09/12] Drop the meta-DCP load options no caller passes The only caller this series gives _meta_dcp_load, the EP hybrid's dense load, names the dense params it loads and shards only the decoder layers. Nothing passes skip_names or max_unloaded_frac, and nothing leaves load_names out to load the whole model, the mode a root-sharded FSDP2 layout needed. load_names is now required, and the loader keeps only the module-by-module materialization that load takes. The tests name what each load fills, and the FSDP2 fill test shards the layers alone, as the EP hybrid does. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 70 +++++--------- tests/unit/torch/utils/test_meta_dcp_load.py | 94 ++++++------------- 2 files changed, 55 insertions(+), 109 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 8ea05af288f..3ecd2dca00d 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -815,30 +815,25 @@ def _meta_dcp_load( local_dir: str, device, *, - load_names: "set | None" = None, - skip_names: "set | None" = None, - max_unloaded_frac: float = 0.5, + load_names: set, ) -> None: - """The single meta + torch-DCP loader shared by every distributed dense path. + """Materialize a meta-built model's ``load_names`` params and fill them through torch DCP. - materialize each rank's params on ``device`` and fill them from the HF safetensors via DCP's + Each rank materializes the params on ``device`` and fills them from the HF safetensors via DCP's ``HuggingFaceStorageReader``. PLACEMENT is decided entirely by the CALLER, not here: ``dcp.load`` reads a FULL tensor into a plain - (replicated) param and only the LOCAL SLICE into a ``Shard`` ``DTensor``. So the caller picks the layout - by what it does before calling this -- nothing (-> Replicate, full per rank), - ``fully_shard`` the whole model (``load_and_prepare_fsdp2`` -> Shard), or ``fully_shard`` only the dense - modules (the EP hybrid -> dense Shard, experts left as loaded). No full-model copy for the sharded cases. + (replicated) param and only the LOCAL SLICE into a ``Shard`` ``DTensor``. The EP hybrid ``fully_shard``-s + its dense modules first, so each rank loads only its dense shard, and no full copy is made. - ``load_names`` (optional) restricts the checkpoint load to a subset of model param names -- the EP-dense - case, where the experts are ALREADY materialized by ``from_pretrained`` and must not be touched. ``None`` - loads the whole model (DP / FSDP2-full). + ``load_names`` restricts the checkpoint load to those ``state_dict`` names -- the dense params in the EP + hybrid, whose experts are ALREADY materialized by ``from_pretrained`` and must not be touched. Materialization is SELECTIVE (module-by-module, only where params are meta) so already-real params (EP experts) are never re-emptied. ``to_empty`` also re-empties BUFFERS (rotary ``inv_freq`` etc., + any the caller already loaded), so they are snapshotted and restored -- a params-only checkpoint load would otherwise leave them garbage and the calibration forward would drift (weights load fine but activation - amax is wrong). Persistent buffers the checkpoint DOES carry are (re)loaded via DCP below. + amax is wrong). Persistent buffers ``load_names`` names are (re)loaded from the checkpoint via DCP below. Every target param ends up filled -- from the checkpoint, by ``tie_weights``, or with the value it already had -- or the load fails: nothing is left as the uninitialized memory ``to_empty`` hands out. @@ -846,30 +841,25 @@ def _meta_dcp_load( import torch.distributed.checkpoint as dcp from torch.distributed.checkpoint import HuggingFaceStorageReader - # Snapshot real buffers (to_empty re-empties them), materialize params, restore buffers. Whole-model - # meta (DP / FSDP2-full, ``load_names is None``) uses ``model.to_empty`` -- the FSDP2-aware root - # materialization; a per-module selective loop corrupts the ROOT-sharded FSDP state (its state_dict then - # returns wrong keys). The PARTIAL case (EP-dense: dense meta, experts already real) MUST be selective so - # the real experts are not re-emptied -- there no root is fully_shard'd, so per-module to_empty is safe. + # Snapshot real buffers (to_empty re-empties them), materialize params, restore buffers. Materialize + # selectively, module by module, so the real experts are not re-emptied. That is safe only because no + # root is fully_shard-ed: a per-module loop corrupts a ROOT-sharded FSDP state (its state_dict then + # returns wrong keys). saved_buffers = {n: b.detach().clone() for n, b in model.named_buffers() if not b.is_meta} # A module's ``to_empty`` re-empties every direct param, so a param already real beside a meta one # (a router's ``e_score_correction_bias`` next to its meta ``weight``) is kept here too: nothing # outside ``load_names`` would refill it. saved_params: dict[str, torch.Tensor] = {} - if load_names is None: - emptied = {n for n, _ in model.named_parameters(remove_duplicate=False)} - model.to_empty(device=device) - else: - emptied = set() - for name, m in model.named_modules(): - direct = dict(m.named_parameters(recurse=False)) - if any(p.is_meta for p in direct.values()): - prefix = f"{name}." if name else "" - saved_params.update( - (prefix + n, p.detach().clone()) for n, p in direct.items() if not p.is_meta - ) - emptied.update(prefix + n for n in direct) - m.to_empty(device=device, recurse=False) + emptied: set[str] = set() + for name, m in model.named_modules(): + direct = dict(m.named_parameters(recurse=False)) + if any(p.is_meta for p in direct.values()): + prefix = f"{name}." if name else "" + saved_params.update( + (prefix + n, p.detach().clone()) for n, p in direct.items() if not p.is_meta + ) + emptied.update(prefix + n for n in direct) + m.to_empty(device=device, recurse=False) for n, b in model.named_buffers(): src = saved_buffers.get(n) if src is not None: @@ -898,7 +888,7 @@ def _meta_dcp_load( full_sd = ( model.state_dict() ) # (sharded ``DTensor`` or plain) params + replicated persistent buffers - target_names = set(full_sd) if load_names is None else (set(load_names) & set(full_sd)) + target_names = set(load_names) & set(full_sd) # {ckpt_key: model_key}, mapped against the WHOLE model and then cut to what this load fills: a suffix # unique among the targets alone can still belong to a key outside them (an expert's, when only the # dense part loads here). @@ -913,26 +903,18 @@ def _meta_dcp_load( inv = {mk: ck for ck, mk in key_map.items()} # model_key -> ckpt_key (1:1 on model keys) load_sd = {inv[mk]: full_sd[mk] for mk in full_sd if mk in inv} mapped_names = set(key_map.values()) - # Fail fast, BEFORE the load, if the mapping left too much of the TARGET WEIGHT unloaded -- a key scheme + # Fail fast, BEFORE the load, if the mapping left over half of the TARGET WEIGHT unloaded -- a key scheme # this path did not resolve. Weight by NUMEL, not param COUNT: a fused-expert MoE (Qwen3-MoE: per-expert # checkpoint -> grouped ``experts.gate_up_proj``/``down_proj`` model params DCP can't fuse) leaves only # ~2 params/layer unloaded but they are ~95% of the weight, so a count-based guard misses it. A smaller # shortfall can be tied weights, which ``tie_weights`` fills after the load; whatever it does not fill # fails there. - # ``skip_names`` = params loaded by a SEPARATE path (e.g. Qwen3-MoE fused grouped experts, filled by - # :func:`_load_fused_experts_sharded` which fuses per-expert checkpoint tensors DCP can't concat) -- - # exclude them from the coverage guard so it doesn't (correctly) flag them as unloaded here. - skip_names = skip_names or set() - numel = { - n: p.numel() - for n, p in model.named_parameters() - if n in target_names and n not in skip_names - } + numel = {n: p.numel() for n, p in model.named_parameters() if n in target_names} want_params = set(numel) unloaded = want_params - mapped_names total_nel = max(1, sum(numel.values())) unloaded_nel = sum(numel[n] for n in unloaded) - if unloaded_nel > max_unloaded_frac * total_nel: + if 2 * unloaded_nel > total_nel: raise RuntimeError( f"meta+DCP load: {len(unloaded)}/{len(want_params)} params ({100 * unloaded_nel / total_nel:.0f}% " f"of the target weight) matched no checkpoint key -- the checkpoint's key naming was not resolved " diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 21758ffec91..7e05935b10d 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -19,8 +19,8 @@ ``torch.distributed.checkpoint``; ``_map_ckpt_keys_to_model`` bridges checkpoint names to the built 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 two tests that need the distributed path start a -one-rank gloo group. +single process without a process group; the tests that need the distributed path start a one-rank +gloo group. """ import pytest @@ -61,6 +61,11 @@ def _meta_twin(build, **kwargs): return build(**kwargs) +def _every_name(model): + """Every ``state_dict`` name of ``model``: what a load into a model with nothing placed fills.""" + return set(model.state_dict()) + + def _tiny_qwen2_5_vl(): """A tiny Qwen2.5-VL: the LM nests at ``model.language_model``, the vision tower at ``model.visual``.""" cfg = Qwen2_5_VLConfig( @@ -242,8 +247,8 @@ def test_wrap_modules_refuse_a_model_without_blocks(): def test_load_fills_a_meta_model_keeps_its_buffers_and_reties_it(tmp_path): - """Whole-model load: every checkpoint tensor lands, the real rotary buffers survive ``to_empty``, - and the tied ``lm_head`` the checkpoint omits is re-tied instead of left as empty memory. + """Every checkpoint tensor lands, the real rotary buffers survive ``to_empty``, and the tied + ``lm_head`` the checkpoint omits is re-tied instead of left as empty memory. """ src = get_tiny_llama(tie_word_embeddings=True) ckpt = {k: v for k, v in src.state_dict().items() if k != "lm_head.weight"} @@ -252,7 +257,7 @@ def test_load_fills_a_meta_model_keeps_its_buffers_and_reties_it(tmp_path): buffers = {n: b.clone() for n, b in model.named_buffers()} assert buffers and all(p.is_meta for p in model.parameters()) - _meta_dcp_load(model, ckpt_dir, "cpu") + _meta_dcp_load(model, ckpt_dir, "cpu", load_names=_every_name(model)) loaded = model.state_dict() for name, value in ckpt.items(): @@ -280,7 +285,7 @@ def test_a_tie_keeps_the_checkpoint_values_whichever_alias_holds_them( torch, "empty_like", lambda t, **kw: empty_like(t, **kw).fill_(float("nan")) ) - _meta_dcp_load(model, ckpt_dir, "cpu") + _meta_dcp_load(model, ckpt_dir, "cpu", load_names=_every_name(model)) assert model.lm_head.weight is model.model.embed_tokens.weight assert torch.equal(model.lm_head.weight, ckpt[stored]) @@ -296,7 +301,7 @@ def test_flat_vlm_checkpoint_names_reach_the_nested_params(tmp_path): assert ckpt.keys() & ref.keys() == {"lm_head.weight"} # every other name differs model = _meta_twin(_tiny_qwen2_5_vl) - _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu") + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names=_every_name(model)) loaded = model.state_dict() assert loaded.keys() == ref.keys() @@ -402,7 +407,7 @@ def test_a_truncated_model_keeps_its_own_layers(tmp_path): src = get_tiny_llama() model = _meta_twin(get_tiny_llama, num_hidden_layers=1) - _meta_dcp_load(model, _save(tmp_path, src.state_dict()), "cpu") + _meta_dcp_load(model, _save(tmp_path, src.state_dict()), "cpu", load_names=_every_name(model)) ref = src.state_dict() for name, value in model.state_dict().items(): @@ -431,7 +436,7 @@ def test_persistent_buffers_load_from_the_checkpoint_and_others_survive(tmp_path with init_empty_weights(): model = _WithBuffers() - _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu") + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names=_every_name(model)) assert torch.equal(model.scale, ckpt["scale"]) assert torch.equal(model.table, torch.arange(4.0)) @@ -448,7 +453,7 @@ def test_an_untied_param_the_checkpoint_lacks_fails_the_load(tmp_path): with pytest.raises( RuntimeError, match=r"1 param\(s\) have no checkpoint tensor.*\['proj.bias'\]" ): - _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu") + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names=_every_name(model)) class _OneBigManySmall(nn.Module): @@ -486,32 +491,15 @@ def counting_load(state_dict, **kwargs): monkeypatch.setattr(dcp, "load", counting_load) model, _, ckpt_dir = _guard_case(tmp_path / "big_missing", missing={"big.weight"}) with pytest.raises(RuntimeError, match="matched no checkpoint key"): - _meta_dcp_load(model, ckpt_dir, "cpu") + _meta_dcp_load(model, ckpt_dir, "cpu", load_names=_every_name(model)) assert not loads model, _, ckpt_dir = _guard_case(tmp_path / "small_missing", missing=_SMALL) with pytest.raises(RuntimeError, match="12 param\\(s\\) have no checkpoint tensor"): - _meta_dcp_load(model, ckpt_dir, "cpu") + _meta_dcp_load(model, ckpt_dir, "cpu", load_names=_every_name(model)) assert loads == [1] -def test_skip_names_exempt_params_filled_by_another_path(tmp_path): - """Params a separate loader fills (fused grouped experts) are left out of the guard's tally.""" - model, ckpt, ckpt_dir = _guard_case(tmp_path, missing={"big.weight"}) - - _meta_dcp_load(model, ckpt_dir, "cpu", skip_names={"big.weight"}) - - assert torch.equal(model.small[0].weight, ckpt["small.0.weight"]) - - -def test_max_unloaded_frac_sets_the_guard_threshold(tmp_path): - """The ~1% shortfall the default lets past the early guard fails it once the tolerance is zero.""" - model, _, ckpt_dir = _guard_case(tmp_path, missing=_SMALL) - - with pytest.raises(RuntimeError, match="matched no checkpoint key"): - _meta_dcp_load(model, ckpt_dir, "cpu", max_unloaded_frac=0.0) - - def test_checkpoint_without_tensors_is_rejected(tmp_path): """A weightless (config-only) snapshot fails with its own message, not as an all-unloaded model.""" (tmp_path / "config.json").write_text("{}") @@ -519,7 +507,7 @@ def test_checkpoint_without_tensors_is_rejected(tmp_path): model = _OneBigManySmall() with pytest.raises(RuntimeError, match="has NO tensors"): - _meta_dcp_load(model, str(tmp_path), "cpu") + _meta_dcp_load(model, str(tmp_path), "cpu", load_names=_every_name(model)) def test_load_request_follows_the_model_state_dict_order(tmp_path, monkeypatch): @@ -538,7 +526,7 @@ def recording_load(state_dict, **kwargs): return real_load(state_dict, **kwargs) monkeypatch.setattr(dcp, "load", recording_load) - _meta_dcp_load(model, ckpt_dir, "cpu") + _meta_dcp_load(model, ckpt_dir, "cpu", load_names=_every_name(model)) assert requests == [[_flat_vl_name(k) for k in model.state_dict()]] @@ -562,57 +550,33 @@ def test_rank0_reports_the_persistent_buffers_it_loaded(tmp_path, gloo_world_of_ model = _WithBuffers() with pytest.warns(UserWarning, match=r"1 persistent buffer\(s\).*'scale'"): - _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu") + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names=_every_name(model)) def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one): - """FSDP2-full: the blocks and the root are ``fully_shard``-ed on meta before the load. The params - must stay sharded ``DTensor``s holding the checkpoint values, and the buffers must survive the - root ``to_empty``. + """The EP hybrid's layout: the decoder layers are ``fully_shard``-ed on meta before the load, the + root is not. The layers' params must stay sharded ``DTensor``s holding the checkpoint values, the + root's plain params must hold theirs, and the buffers must survive ``to_empty``. """ src = get_tiny_llama() ckpt_dir = _save(tmp_path, src.state_dict()) model = _meta_twin(get_tiny_llama) buffers = {n: b.clone() for n, b in model.named_buffers()} mesh = init_device_mesh("cpu", (1,)) - for block in _fsdp2_wrap_modules(model): - fully_shard(block, mesh=mesh) - fully_shard(model, mesh=mesh) + for layer in model.model.layers: + fully_shard(layer, mesh=mesh) - _meta_dcp_load(model, ckpt_dir, "cpu") + _meta_dcp_load(model, ckpt_dir, "cpu", load_names=_every_name(model)) ref = src.state_dict() for name, param in model.named_parameters(): - assert isinstance(param, DTensor), name - assert torch.equal(param.full_tensor(), ref[name]), name + assert isinstance(param, DTensor) == name.startswith("model.layers."), name + loaded = param.full_tensor() if isinstance(param, DTensor) else param + assert torch.equal(loaded, ref[name]), name for name, buf in model.named_buffers(): assert torch.equal(buf, buffers[name]), name -def test_an_fsdp2_tie_keeps_the_values_of_the_alias_the_checkpoint_stored( - tmp_path, gloo_world_of_one, monkeypatch -): - """The same tie under FSDP2, where the kept param and the filled one are sharded alike.""" - src = get_tiny_llama(tie_word_embeddings=True) - ckpt = {k: v for k, v in src.state_dict().items() if k != "model.embed_tokens.weight"} - ckpt_dir = _save(tmp_path, ckpt) - model = _meta_twin(get_tiny_llama, tie_word_embeddings=True) - mesh = init_device_mesh("cpu", (1,)) - for block in _fsdp2_wrap_modules(model): - fully_shard(block, mesh=mesh) - fully_shard(model, mesh=mesh) - empty_like = torch.empty_like - monkeypatch.setattr( - torch, "empty_like", lambda t, **kw: empty_like(t, **kw).fill_(float("nan")) - ) - - _meta_dcp_load(model, ckpt_dir, "cpu") - - assert isinstance(model.lm_head.weight, DTensor) - assert torch.equal(model.lm_head.weight.full_tensor(), ckpt["lm_head.weight"]) - assert torch.equal(model.model.embed_tokens.weight.full_tensor(), ckpt["lm_head.weight"]) - - def _tie_pair(layout): """A kept param and the filled one it replaced, laid out as ``layout`` names.""" mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) From d7077d5b737895f70acd34c69614ca89312f2a55 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:00:26 +0000 Subject: [PATCH 10/12] Leave the FSDP2 wrap boundaries to the dense sharding that uses them _fsdp2_wrap_modules picks the blocks FSDP2 shards, and the EP hybrid's dense sharding, later in this series, is its only caller. Nothing here calls it, so it and its tests move to where that sharding lands. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 43 ------------------ tests/unit/torch/utils/test_meta_dcp_load.py | 45 ++----------------- 2 files changed, 3 insertions(+), 85 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 3ecd2dca00d..92c6e8a1c20 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -668,49 +668,6 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False return decoder_layers -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). - """ - 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) - if lineage and lineage[-1] is not model: - lm_mods = _boundaries(lineage[-1]) - if lm_mods: - return lm_mods - - mods = _boundaries(model) - if mods: - return mods - raise RuntimeError( - "Could not determine FSDP2 wrap boundaries: model exposes no `_no_split_modules` and modelopt's " - "decoder-layer auto-detection did not match. Provide a standard HF causal-LM / VLM checkpoint." - ) - - 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. diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 7e05935b10d..c8c4fa31f55 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -17,10 +17,9 @@ ``_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; ``_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. +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. """ import pytest @@ -41,7 +40,6 @@ from modelopt.torch.utils.plugins.model_load_utils import ( _copy_tied_values, - _fsdp2_wrap_modules, _map_ckpt_keys_to_model, _meta_dcp_load, ) @@ -206,43 +204,6 @@ def test_two_keys_claiming_one_model_key_map_neither(): ) -# -------------------------------------------------------------------------- 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 test_wrap_modules_refuse_a_model_without_blocks(): - """With neither source of boundaries matching, fail loudly instead of guessing a layout.""" - with pytest.raises(RuntimeError, match="Could not determine FSDP2 wrap boundaries"): - _fsdp2_wrap_modules(nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))) - - # --------------------------------------------------------------------------------- meta + DCP load From a940fe246015a2e9f73853f30be9a11af1919e9e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 16:30:16 +0000 Subject: [PATCH 11/12] Restore a meta-DCP load's saved buffers under no_grad A buffer can require grad, and copying its saved value back in place outside no_grad raises for a leaf tensor, or records the copy in autograd otherwise. The buffers now go back in the same no_grad block as the params beside them. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 8 ++++---- tests/unit/torch/utils/test_meta_dcp_load.py | 3 ++- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 92c6e8a1c20..f398fb82bfc 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -817,11 +817,11 @@ def _meta_dcp_load( ) emptied.update(prefix + n for n in direct) m.to_empty(device=device, recurse=False) - for n, b in model.named_buffers(): - src = saved_buffers.get(n) - if src is not None: - b.copy_(src.to(device=b.device, dtype=b.dtype)) with torch.no_grad(): + for n, b in model.named_buffers(): + src = saved_buffers.get(n) + if src is not None: + b.copy_(src.to(device=b.device, dtype=b.dtype)) for n, src in saved_params.items(): p = model.get_parameter(n) p.copy_(src.to(device=p.device, dtype=p.dtype)) diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index c8c4fa31f55..86236079bae 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -387,7 +387,7 @@ def __init__(self): def test_persistent_buffers_load_from_the_checkpoint_and_others_survive(tmp_path): """``to_empty`` wipes buffers too: one the checkpoint carries is loaded from it, and one no - checkpoint can carry (non-persistent) gets its pre-load value back. + checkpoint can carry (non-persistent) gets its pre-load value back, even one that requires grad. """ ckpt = { "proj.weight": torch.randn(4, 4), @@ -396,6 +396,7 @@ def test_persistent_buffers_load_from_the_checkpoint_and_others_survive(tmp_path } with init_empty_weights(): model = _WithBuffers() + model.table.requires_grad_(True) # restoring it in place must not be recorded by autograd _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names=_every_name(model)) From 8fc1baeccac380f45d519ad3658e2da24d16e559 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 18:00:27 +0000 Subject: [PATCH 12/12] Trim the meta-DCP load's tests: one key-mapping table, no duplicates Merge the four _map_ckpt_keys_to_model tests that differ only in their inputs (an exact name first, the first conversion pattern that fires, an ambiguous or unknown suffix, two names for one slot) into one parametrized table. Drop the tests a real load already covers, each checked by breaking the code it guards: - the VLM suffix-nesting test: the flat-VLM load and the request-order test run the same suffix pass on a real Qwen2.5-VL; - the never-displace-a-match test: the truncated-model load fails when a later pass may take a claimed key; - the embedding-stored case of the tie test: the load-and-retie test loads that layout; - the real-param-beside-a-meta-one test: the already-real-target and load_names tests fail when the restore or the load_names filter goes. Drop the two tests that only pinned message wording: the no-tensors message (without its check the coverage guard still refuses the empty checkpoint) and rank 0's loaded-buffers warning. The messages stay. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_meta_dcp_load.py | 207 +++++++------------ 1 file changed, 69 insertions(+), 138 deletions(-) diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 86236079bae..c05f7357d25 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -102,106 +102,72 @@ def _flat_vl_name(key): # ---------------------------------------------------------------- checkpoint -> model key mapping -def test_exact_name_match_takes_precedence(): - """A checkpoint name the model has verbatim maps to itself before any conversion pattern runs.""" - model = nn.Module() - model._checkpoint_conversion_mapping = {r"^model\.": "model.language_model."} - model_keys = {"model.norm.weight", "model.language_model.norm.weight"} - - key_map = _map_ckpt_keys_to_model(model, {"model.norm.weight"}, model_keys) - - assert key_map == {"model.norm.weight": "model.norm.weight"} - - -def test_conversion_mapping_applies_the_first_pattern_that_fires(): - """Patterns are tried in order and only the first that fires rewrites a name (no chaining); the - rewrite counts only if the model has the result, otherwise the name falls through to suffix match. - """ +@pytest.mark.parametrize( + ("conversion_mapping", "ckpt_keys", "model_keys", "expected"), + [ + pytest.param( + # A name the model has verbatim maps to itself before any conversion pattern runs. + {r"^model\.": "model.language_model."}, + {"model.norm.weight"}, + {"model.norm.weight", "model.language_model.norm.weight"}, + {"model.norm.weight": "model.norm.weight"}, + id="exact-name-first", + ), + pytest.param( + # Patterns are tried in order and only the first that fires rewrites a name (no + # chaining); a rewrite counts only if the model has the result, otherwise the name falls + # through to the suffix match (this model calls its vision tower ``model.visual``). + { + r"^model\.(?!language_model\.)": "model.language_model.", + r"^model\.language_model\.": "chained.", + r"^visual\.": "model.vision_tower.", + }, + {"model.layers.0.mlp.down_proj.weight", "visual.proj.weight"}, + { + "model.language_model.layers.0.mlp.down_proj.weight", + "chained.layers.0.mlp.down_proj.weight", + "model.visual.proj.weight", + }, + { + "model.layers.0.mlp.down_proj.weight": ( + "model.language_model.layers.0.mlp.down_proj.weight" + ), + "visual.proj.weight": "model.visual.proj.weight", + }, + id="first-conversion-pattern", + ), + pytest.param( + # Suffixes naming several model keys (two towers both have ``layers.0``) or none (a + # quantization scale with no slot) are left for the coverage guard, never guessed. + None, + {"model.layers.0.mlp.down_proj.weight", "model.layers.0.mlp.down_proj.weight_scale"}, + { + "model.language_model.layers.0.mlp.down_proj.weight", + "model.audio_tower.layers.0.mlp.down_proj.weight", + }, + {}, + id="ambiguous-or-unknown-suffix", + ), + pytest.param( + # Two names reaching one model key by suffix are equally plausible: loading either + # would let set order pick the tensor. + None, + {"model.norm.weight", "mtp.norm.weight"}, + {"model.language_model.norm.weight"}, + {}, + id="two-names-one-slot", + ), + ], +) +def test_map_ckpt_keys_to_model(conversion_mapping, ckpt_keys, model_keys, expected): + """Exact names first, then the model's conversion mapping, then the longest unique component + suffix; anything ambiguous stays unmapped. The suffix pass and the one-key-per-slot rule run for + real in the flat-VLM and truncated-model loads below.""" model = nn.Module() - model._checkpoint_conversion_mapping = { - r"^model\.(?!language_model\.)": "model.language_model.", - r"^model\.language_model\.": "chained.", - r"^visual\.": "model.vision_tower.", # this model calls it ``model.visual`` - } - model_keys = { - "model.language_model.layers.0.mlp.down_proj.weight", - "chained.layers.0.mlp.down_proj.weight", - "model.visual.proj.weight", - } - ckpt_keys = {"model.layers.0.mlp.down_proj.weight", "visual.proj.weight"} - - assert _map_ckpt_keys_to_model(model, ckpt_keys, model_keys) == { - "model.layers.0.mlp.down_proj.weight": "model.language_model.layers.0.mlp.down_proj.weight", - "visual.proj.weight": "model.visual.proj.weight", - } - - -def test_suffix_match_resolves_the_vlm_prefix_nesting(): - """transformers>=5 nests a VLM's LM and vision tower one level below its checkpoint names, and the - model carries no ``_checkpoint_conversion_mapping``. The longest component suffix naming exactly - one model key bridges that: ``layers.0`` tells LM layer 0 from vision block 0, which shares the rest. - """ - model_keys = { - "model.language_model.embed_tokens.weight", - "model.language_model.layers.0.mlp.down_proj.weight", - "model.language_model.layers.1.mlp.down_proj.weight", - "model.language_model.norm.weight", - "model.visual.blocks.0.mlp.down_proj.weight", - "lm_head.weight", - } - expected = { - "model.embed_tokens.weight": "model.language_model.embed_tokens.weight", - "model.layers.0.mlp.down_proj.weight": "model.language_model.layers.0.mlp.down_proj.weight", - "model.layers.1.mlp.down_proj.weight": "model.language_model.layers.1.mlp.down_proj.weight", - "model.norm.weight": "model.language_model.norm.weight", - "visual.blocks.0.mlp.down_proj.weight": "model.visual.blocks.0.mlp.down_proj.weight", - "lm_head.weight": "lm_head.weight", - } - - assert _map_ckpt_keys_to_model(nn.Module(), expected.keys(), model_keys) == expected - - -def test_ambiguous_or_unknown_suffixes_stay_unmapped(): - """A name whose suffixes match several model keys (two towers both have ``layers.0``) or none (a - quantization scale with no slot) is left for the coverage guard, never assigned to a guess. - """ - model_keys = { - "model.language_model.layers.0.mlp.down_proj.weight", - "model.audio_tower.layers.0.mlp.down_proj.weight", - } - ckpt_keys = { - "model.layers.0.mlp.down_proj.weight", - "model.layers.0.mlp.down_proj.weight_scale", - } - - assert _map_ckpt_keys_to_model(nn.Module(), ckpt_keys, model_keys) == {} - - -def test_a_key_the_model_has_no_place_for_never_displaces_a_match(): - """DeepSeek-V3's MTP layer and the layers a truncated model lacks have no slot in the model, yet - their suffixes name a key uniquely -- one another checkpoint key already names exactly. Mapped, - they would overwrite that key's tensor; they must stay unmapped. - """ - model_keys = {"model.norm.weight", "model.layers.0.input_layernorm.weight"} - ckpt_keys = model_keys | { - "model.layers.61.shared_head.norm.weight", - "model.layers.1.input_layernorm.weight", - } - - assert _map_ckpt_keys_to_model(nn.Module(), ckpt_keys, model_keys) == { - key: key for key in model_keys - } + if conversion_mapping is not None: + model._checkpoint_conversion_mapping = conversion_mapping - -def test_two_keys_claiming_one_model_key_map_neither(): - """Two checkpoint names reaching the same model key by suffix are equally plausible; loading - either would let set order pick the tensor, so neither is mapped.""" - model_keys = {"model.language_model.norm.weight"} - - assert ( - _map_ckpt_keys_to_model(nn.Module(), {"model.norm.weight", "mtp.norm.weight"}, model_keys) - == {} - ) + assert _map_ckpt_keys_to_model(model, ckpt_keys, model_keys) == expected # --------------------------------------------------------------------------------- meta + DCP load @@ -228,16 +194,13 @@ def test_load_fills_a_meta_model_keeps_its_buffers_and_reties_it(tmp_path): assert model.lm_head.weight is model.model.embed_tokens.weight -@pytest.mark.parametrize("stored", ["model.embed_tokens.weight", "lm_head.weight"]) -def test_a_tie_keeps_the_checkpoint_values_whichever_alias_holds_them( - tmp_path, monkeypatch, stored -): +def test_a_tie_keeps_the_values_of_the_alias_the_checkpoint_stored(tmp_path, monkeypatch): """A tied checkpoint stores one alias; tying keeps the embedding's parameter. When the checkpoint stored only ``lm_head.weight``, the kept parameter must take its values rather than keep the memory - ``to_empty`` gave it, as ``from_pretrained`` leaves it.""" + ``to_empty`` gave it, as ``from_pretrained`` leaves it. (A checkpoint storing the embedding is the + load-and-retie test above.)""" src = get_tiny_llama(tie_word_embeddings=True) - tied = {"model.embed_tokens.weight", "lm_head.weight"} - ckpt = {k: v for k, v in src.state_dict().items() if k not in tied - {stored}} + ckpt = {k: v for k, v in src.state_dict().items() if k != "model.embed_tokens.weight"} ckpt_dir = _save(tmp_path, ckpt) model = _meta_twin(get_tiny_llama, tie_word_embeddings=True) # Fresh storage holds NaN, so a parameter nothing wrote to shows. @@ -249,7 +212,7 @@ def test_a_tie_keeps_the_checkpoint_values_whichever_alias_holds_them( _meta_dcp_load(model, ckpt_dir, "cpu", load_names=_every_name(model)) assert model.lm_head.weight is model.model.embed_tokens.weight - assert torch.equal(model.lm_head.weight, ckpt[stored]) + assert torch.equal(model.lm_head.weight, ckpt["lm_head.weight"]) def test_flat_vlm_checkpoint_names_reach_the_nested_params(tmp_path): @@ -305,18 +268,6 @@ def __init__(self): self.e_score_correction_bias = nn.Parameter(torch.arange(4.0)) -def test_a_real_param_beside_a_meta_one_keeps_its_value(tmp_path): - """Materializing the router re-empties all of its params, but only ``weight`` is reloaded.""" - model = nn.Module() - model.gate = _Router() - ckpt = {"gate.weight": torch.randn(4, 4), "gate.e_score_correction_bias": torch.zeros(4)} - - _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names={"gate.weight"}) - - assert torch.equal(model.gate.weight, ckpt["gate.weight"]) - assert torch.equal(model.gate.e_score_correction_bias, torch.arange(4.0)) - - def test_target_params_that_already_had_values_need_no_checkpoint_tensor(tmp_path): """A target param that was real before the load -- beside a meta one, so restored after ``to_empty``, or in a module the load never re-empties -- keeps its value when the checkpoint @@ -462,16 +413,6 @@ def counting_load(state_dict, **kwargs): assert loads == [1] -def test_checkpoint_without_tensors_is_rejected(tmp_path): - """A weightless (config-only) snapshot fails with its own message, not as an all-unloaded model.""" - (tmp_path / "config.json").write_text("{}") - with init_empty_weights(): - model = _OneBigManySmall() - - with pytest.raises(RuntimeError, match="has NO tensors"): - _meta_dcp_load(model, str(tmp_path), "cpu", load_names=_every_name(model)) - - def test_load_request_follows_the_model_state_dict_order(tmp_path, monkeypatch): """``dcp.load`` gets the checkpoint names in the model's ``state_dict`` order, which every rank shares. Key-map or set order is hash-seeded per process, and differing orders across ranks @@ -505,16 +446,6 @@ def gloo_world_of_one(skip_on_windows, tmp_path_factory): torch.distributed.destroy_process_group() -def test_rank0_reports_the_persistent_buffers_it_loaded(tmp_path, gloo_world_of_one): - """In a distributed run rank 0 names the persistent buffers the checkpoint filled.""" - ckpt = {"proj.weight": torch.randn(4, 4), "proj.bias": torch.randn(4), "scale": torch.ones(4)} - with init_empty_weights(): - model = _WithBuffers() - - with pytest.warns(UserWarning, match=r"1 persistent buffer\(s\).*'scale'"): - _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names=_every_name(model)) - - def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one): """The EP hybrid's layout: the decoder layers are ``fully_shard``-ed on meta before the load, the root is not. The layers' params must stay sharded ``DTensor``s holding the checkpoint values, the