From 1d0c605bc76d152078b2098423b841fee96a4931 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 21:22:25 +0000 Subject: [PATCH 01/17] 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 181bcff632d..086663b77dd 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 @@ -269,6 +270,223 @@ def setup(timeout: timedelta | None = None): torch.distributed.init_process_group("cpu:gloo,cuda:nccl", timeout=timeout) +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 7c0f3c2e71d343fc5ff29d1ca60e43464a49e597 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:07:10 +0000 Subject: [PATCH 02/17] 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 a9b116f5b4e47976a0fca16af66e8fa0d2c237ec Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:07:34 +0000 Subject: [PATCH 03/17] 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 086663b77dd..f1927edf8bb 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -326,43 +326,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 @@ -436,7 +454,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 beb6a31f27ed1cfda441d04eb97bd99d24dc5417 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:34:54 +0000 Subject: [PATCH 04/17] 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 f1927edf8bb..b08d1d0315a 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -423,16 +423,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 1b41bed62f5f51f776c451147e2a03651d7529d7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 23:24:04 +0000 Subject: [PATCH 05/17] 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 b08d1d0315a..181bcff632d 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 @@ -270,261 +269,6 @@ def setup(timeout: timedelta | None = None): torch.distributed.init_process_group("cpu:gloo,cuda:nccl", timeout=timeout) -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 fdf5b1fe798..d03e074e8a3 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -20,6 +20,7 @@ import re from itertools import chain from typing import TYPE_CHECKING, Any +from warnings import warn import torch import torch.nn as nn @@ -38,6 +39,7 @@ broadcast_state_dict, fsdp2_wrap, is_initialized, + rank, ) from modelopt.torch.utils.logging import print_rank_0 from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -449,6 +451,261 @@ def load_pretrained( return model +def _fsdp2_wrap_modules(model: "torch.nn.Module") -> list["torch.nn.Module"]: + """Return the transformer-block modules to FSDP2 ``fully_shard`` individually. + + Uses HuggingFace's own ``_no_split_modules`` -- the class names of the atomic blocks that must not be + split across devices (what ``device_map="auto"`` wraps at) -- falling back to modelopt's decoder-layer + auto-detection. + + For a VLM, this looks at the LANGUAGE backbone's boundaries, NOT the top-level VLM wrapper's + ``_no_split_modules`` (which for some VLMs -- e.g. Qwen2.5-VL -- ALSO lists the VISION block classes). + FSDP-sharding the vision blocks makes the calibration weight-access all-gather on those sharded vision + weights DEADLOCK; keeping the vision/audio towers REPLICATED (plain) -- the language-only calibration + never runs them and they are small enough to fit -- avoids that. This also covers VLMs whose top-level + wrapper exposes no ``_no_split_modules`` at all (e.g. Nemotron-Nano-VL). + """ + 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 0eb6137701e2ac3ae935b0b2c6898b17aecf3e90 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 20:48:29 +0000 Subject: [PATCH 06/17] 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 d03e074e8a3..21134b7b46f 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -594,6 +594,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 @@ -609,8 +612,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()): @@ -618,6 +623,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) @@ -662,11 +668,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. @@ -686,16 +693,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: " @@ -704,6 +706,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 6bc3c2b550a51ea57dd4a34f54b74870eaea3a8f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 23:24:28 +0000 Subject: [PATCH 07/17] 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 21134b7b46f..8e06d8dcaf3 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -565,6 +565,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, @@ -704,12 +732,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 05fb3c9d0f9c5e5da13e21c59b4f4b446b308b79 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 00:28:03 +0000 Subject: [PATCH 08/17] 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 c13910f7a08d7eb85960c14b5d6550661fee920f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 06:59:49 +0000 Subject: [PATCH 09/17] 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 8e06d8dcaf3..957ac30b686 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -598,30 +598,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. @@ -629,30 +624,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: @@ -681,7 +671,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). @@ -696,26 +686,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 203d9c76362cbc3b25182c198e96ddb3c4bad590 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:00:26 +0000 Subject: [PATCH 10/17] 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 957ac30b686..733b81d3fc3 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -451,49 +451,6 @@ def load_pretrained( return model -def _fsdp2_wrap_modules(model: "torch.nn.Module") -> list["torch.nn.Module"]: - """Return the transformer-block modules to FSDP2 ``fully_shard`` individually. - - Uses HuggingFace's own ``_no_split_modules`` -- the class names of the atomic blocks that must not be - split across devices (what ``device_map="auto"`` wraps at) -- falling back to modelopt's decoder-layer - auto-detection. - - For a VLM, this looks at the LANGUAGE backbone's boundaries, NOT the top-level VLM wrapper's - ``_no_split_modules`` (which for some VLMs -- e.g. Qwen2.5-VL -- ALSO lists the VISION block classes). - FSDP-sharding the vision blocks makes the calibration weight-access all-gather on those sharded vision - weights DEADLOCK; keeping the vision/audio towers REPLICATED (plain) -- the language-only calibration - never runs them and they are small enough to fit -- avoids that. This also covers VLMs whose top-level - wrapper exposes no ``_no_split_modules`` at all (e.g. Nemotron-Nano-VL). - """ - 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 0b3c65dc1e171eadf5a8763930e8d7569a23c577 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 16:30:16 +0000 Subject: [PATCH 11/17] 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 733b81d3fc3..c62e30609ed 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -600,11 +600,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 c6e61adc57385024f8d18525434f973a30a1424c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 18:00:27 +0000 Subject: [PATCH 12/17] 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 From b9229be179f315223533829093e2aaee7d90f878 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 20:05:18 +0000 Subject: [PATCH 13/17] Hold the meta-DCP load's unfilled check to every param it re-emptied The check that no re-emptied param is left holding to_empty's memory looked only at params load_names named, so a meta param it left out was materialized, skipped by DCP and returned uninitialized; it now covers every param the load re-emptied. The load also reads the checkpoint's metadata before it changes the model, frees its buffer and param snapshots once they are restored rather than through the DCP load, warns through warn_rank_0, and takes the filled params in state_dict order, the same on every rank. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 54 ++++++++++--------- tests/unit/torch/utils/test_meta_dcp_load.py | 9 ++-- 2 files changed, 34 insertions(+), 29 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c62e30609ed..ae228490b6b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -20,7 +20,6 @@ import re from itertools import chain from typing import TYPE_CHECKING, Any -from warnings import warn import torch import torch.nn as nn @@ -39,9 +38,8 @@ broadcast_state_dict, fsdp2_wrap, is_initialized, - rank, ) -from modelopt.torch.utils.logging import print_rank_0 +from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( checkpoint_tensor_files, ensure_local_checkpoint, @@ -581,6 +579,16 @@ def _meta_dcp_load( import torch.distributed.checkpoint as dcp from torch.distributed.checkpoint import HuggingFaceStorageReader + # Read what the checkpoint holds before changing the model, so a refusal leaves it as it was. + 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." + ) + # 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 @@ -608,6 +616,9 @@ def _meta_dcp_load( for n, src in saved_params.items(): p = model.get_parameter(n) p.copy_(src.to(device=p.device, dtype=p.dtype)) + # The clones are restored; only the names of the params they kept are needed from here on. + kept_names = set(saved_params) + del saved_buffers, saved_params # ``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 @@ -617,14 +628,6 @@ def _meta_dcp_load( # 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 @@ -663,16 +666,16 @@ def _meta_dcp_load( 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 loaded_bufs: - warn( - f"meta+DCP load: loaded {len(loaded_bufs)} persistent buffer(s) from checkpoint: " - f"{sorted(loaded_bufs)[:5]}" - ) + loaded_bufs = mapped_names - {n for n, _ in model.named_parameters()} + if loaded_bufs: + warn_rank_0( + 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) - # 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} + # Taken before tying, which can swap a filled param out of its module; in state_dict order, the + # same on every rank, so every rank keeps the same param of a tied group. + filled = {n: model.get_parameter(n) for n in full_sd if n in want_params and n in mapped_names} 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 @@ -684,13 +687,12 @@ def _meta_dcp_load( 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. + # A re-emptied param no checkpoint tensor filled, and that tying did not make a loaded param, still + # holds ``to_empty``'s memory, whether or not ``load_names`` named it. 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. unfilled = sorted( - n - for n in (unloaded & emptied) - saved_params.keys() - if id(model.get_parameter(n)) not in loaded + n for n in emptied - mapped_names - kept_names if id(model.get_parameter(n)) not in loaded ) if unfilled: raise RuntimeError( diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index c05f7357d25..e07d6be9015 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -356,17 +356,20 @@ 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): +@pytest.mark.parametrize("named", [True, False], ids=["named", "left-out-of-load-names"]) +def test_an_untied_param_the_checkpoint_lacks_fails_the_load(tmp_path, named): """``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.""" + ``to_empty`` gave it, so the load fails naming it instead of returning the model, also when + ``load_names`` leaves it out.""" ckpt = {"proj.weight": torch.randn(4, 4), "scale": torch.full((4,), 3.0)} with init_empty_weights(): model = _WithBuffers() + load_names = _every_name(model) - ({"proj.bias"} if not named else set()) with pytest.raises( RuntimeError, match=r"1 param\(s\) have no checkpoint tensor.*\['proj.bias'\]" ): - _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names=_every_name(model)) + _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names=load_names) class _OneBigManySmall(nn.Module): From 2b7b04a9d837066719fbf254bded40c9c0c9bdaa Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:10:49 +0000 Subject: [PATCH 14/17] Decide the meta-DCP load's refusals before it materializes, and refuse three more The checkpoint's names are mapped onto the model, and the coverage guard run, before any module is re-emptied, so every refusal leaves the model as it was; the DCP read still fills the materialized tensors. Three cases it would get wrong are now refused there: - a root-sharded model, whose FSDP2 state the module-by-module materialization does not keep; - a buffer on meta, which a checkpoint of params cannot fill (a non-persistent one is not in it); - a checkpoint tensor the mapping fills that the model's load recorded as unplaced, for the export to carry over raw: it would be written twice. The mapper's docstring says why it does not resolve keys through _target_param. The FSDP2 test now also runs a forward after the load and matches the source model's logits: FSDP2 all-gathers the shards it tracks, which full_tensor() does not read. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 104 +++++++++++------- tests/unit/torch/utils/test_meta_dcp_load.py | 41 +++++++ 2 files changed, 108 insertions(+), 37 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ae228490b6b..0ca9f5b070e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -470,6 +470,10 @@ def _map_ckpt_keys_to_model(model: "torch.nn.Module", ckpt_keys, model_keys) -> 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. + + It does not resolve keys through :func:`_target_param`: a model built on meta has no record of what + ``from_pretrained`` ran, and the stock mapping lacks a remote-code model's renames. The loader instead + refuses a mapping that fills a tensor the model's load recorded as unplaced, which the export carries. """ from collections import defaultdict @@ -588,63 +592,56 @@ def _meta_dcp_load( "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." ) + from torch.distributed.fsdp import FSDPModule - # 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] = {} - 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) - 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)) - # The clones are restored; only the names of the params they kept are needed from here on. - kept_names = set(saved_params) - del saved_buffers, saved_params + if isinstance(model, FSDPModule): + raise RuntimeError( + "meta+DCP load: the root module is fully_shard-ed. The load materializes module by module, " + "so already-real params survive, which a root-sharded FSDP2 state does not take: shard the " + "decoder layers only." + ) + meta_buffers = sorted(n for n, b in model.named_buffers() if b.is_meta) + if meta_buffers: + raise RuntimeError( + f"meta+DCP load: {len(meta_buffers)} buffer(s) are on meta, which a checkpoint of params " + f"cannot fill (a non-persistent buffer is not in it): {meta_buffers[:3]}. Build the model " + "with its buffers materialized, as accelerate's init_empty_weights does." + ) # ``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 + # path preserved, which the model carries no ``_checkpoint_conversion_mapping`` attribute for. 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. - full_sd = ( - model.state_dict() - ) # (sharded ``DTensor`` or plain) params + replicated persistent buffers - target_names = set(load_names) & set(full_sd) + model_keys = list(model.state_dict()) # (sharded or plain) params + persistent buffers + target_names = set(load_names) & set(model_keys) # {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() + for ck, mk in _map_ckpt_keys_to_model(model, ckpt_keys, set(model_keys)).items() if mk in target_names } + # The export carries the checkpoint tensors the model has no parameter for over raw, as the load + # that built the model recorded them (record_unplaced_source_keys): filling one here too would + # write it twice, once quantized. + carried = sorted( + set(key_map) & set(getattr(model, "_modelopt_unplaced_source_keys", None) or ()) + ) + if carried: + raise RuntimeError( + f"meta+DCP load: {len(carried)} checkpoint tensor(s) map onto this model's params by name, " + f"but its load recorded them as having no param, for the export to carry over raw: " + f"{carried[:3]}. The two would write them twice." + ) # 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 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 @@ -672,6 +669,39 @@ def _meta_dcp_load( f"meta+DCP load: loaded {len(loaded_bufs)} persistent buffer(s) from checkpoint: " f"{sorted(loaded_bufs)[:5]}" ) + # 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] = {} + 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) + 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)) + # The clones are restored; only the names of the params they kept are needed from here on. + kept_names = set(saved_params) + del saved_buffers, saved_params + + full_sd = model.state_dict() # the materialized tensors: to_empty replaced the meta params + load_sd = {inv[mk]: full_sd[mk] for mk in full_sd if mk in inv} dcp.load(load_sd, storage_reader=reader) # Taken before tying, which can swap a filled param out of its module; in state_dict order, the # same on every rank, so every rank keeps the same param of a tied group. diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index e07d6be9015..6a654adad76 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -471,6 +471,47 @@ def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one): assert torch.equal(loaded, ref[name]), name for name, buf in model.named_buffers(): assert torch.equal(buf, buffers[name]), name + # FSDP2 all-gathers the shards it tracks, which full_tensor() does not read: the forward shows + # the loaded values reach the computation. + ids = torch.randint(0, src.config.vocab_size, (2, 8)) + with torch.no_grad(): + assert torch.equal(model(ids).logits, src(ids).logits) + + +def _root_sharded(model): + fully_shard(model, mesh=init_device_mesh("cpu", (1,))) + + +def _with_a_meta_buffer(model): + model.model.register_buffer("scratch", torch.empty(2, device="meta"), persistent=False) + + +def _recorded_as_carried_over(model): + model._modelopt_unplaced_source_keys = ["model.norm.weight"] + + +@pytest.mark.parametrize( + ("change", "match"), + [ + (_root_sharded, "root module is fully_shard-ed"), + (_with_a_meta_buffer, r"buffer\(s\) are on meta"), + (_recorded_as_carried_over, "recorded them as having no param"), + ], + ids=["root-sharded", "meta-buffer", "carried-over-raw"], +) +def test_a_load_it_cannot_do_rightly_is_refused_before_anything_changes( + tmp_path, gloo_world_of_one, change, match +): + """A root-sharded model, which the module-by-module materialization would corrupt; a buffer on + meta, which a checkpoint of params cannot fill; a tensor the export carries over raw, which + filling would write twice. Each is refused with every param still on meta.""" + ckpt_dir = _save(tmp_path, get_tiny_llama().state_dict()) + model = _meta_twin(get_tiny_llama) + change(model) + + with pytest.raises(RuntimeError, match=match): + _meta_dcp_load(model, ckpt_dir, "cpu", load_names=_every_name(model)) + assert all(p.is_meta for p in model.parameters()) def _tie_pair(layout): From 2d8f8adb12faef3d98ae8416410b128ec9ee325d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:04:39 +0000 Subject: [PATCH 15/17] Count both names of a shared tied param in the meta-DCP load named_parameters() lists a parameter that two modules share once. With the real tie that from_pretrained's tie_weights() leaves (lm_head.weight and the embedding as one Parameter), lm_head.weight fell out of the params the load expects to fill: a checkpoint storing only lm_head.weight was refused, naming the embedding as unfilled, and a loaded lm_head.weight was reported as a persistent buffer. Both reads now take every alias (remove_duplicate=False), as unplaced_source_keys does, and the tie test also runs on a model holding one shared Parameter. Also: the function-local imports move to the module top (DTensor was imported twice); the suffix index keeps at most two candidates per suffix, since only "exactly one" is read; and the docstring says a param counts as filled by tying only when tied to one the checkpoint filled. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 31 ++++++++++--------- tests/unit/torch/utils/test_meta_dcp_load.py | 15 +++++++-- 2 files changed, 28 insertions(+), 18 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 0ca9f5b070e..0f33661e829 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -18,16 +18,20 @@ import logging import os import re +from collections import defaultdict from itertools import chain from typing import TYPE_CHECKING, Any import torch +import torch.distributed.checkpoint as dcp import torch.nn as nn from huggingface_hub import snapshot_download if TYPE_CHECKING: from collections.abc import Iterable +from torch.distributed.checkpoint import HuggingFaceStorageReader from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict +from torch.distributed.fsdp import FSDPModule from torch.distributed.tensor import DTensor from transformers import AutoConfig, AutoModel, AutoModelForCausalLM from transformers.conversion_mapping import get_model_conversion_mapping @@ -475,8 +479,6 @@ def _map_ckpt_keys_to_model(model: "torch.nn.Module", ckpt_keys, model_keys) -> ``from_pretrained`` ran, and the stock mapping lacks a remote-code model's renames. The loader instead refuses a mapping that fills a tensor the model's load recorded as unplaced, which the export carries. """ - from collections import defaultdict - key_map: dict = {} remaining = set(ckpt_keys) @@ -510,7 +512,9 @@ def claim(hits: dict) -> None: for mk in model_keys: comps = mk.split(".") for k in range(1, len(comps) + 1): - suffix_idx[tuple(comps[-k:])].append(mk) + cands = suffix_idx[tuple(comps[-k:])] + if len(cands) < 2: # only "exactly one candidate" is read below + cands.append(mk) hits = defaultdict(list) for ck in remaining: comps = ck.split(".") @@ -530,8 +534,6 @@ def _copy_tied_values(kept: torch.Tensor, src: torch.Tensor, name: str) -> None: 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 ( @@ -577,12 +579,10 @@ def _meta_dcp_load( otherwise leave them garbage and the calibration forward would drift (weights load fine but activation 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. + Every target param ends up filled -- from the checkpoint, by ``tie_weights`` with a param the + checkpoint filled, or with the value it already had -- or the load fails: nothing is left as the + uninitialized memory ``to_empty`` hands out. """ - import torch.distributed.checkpoint as dcp - from torch.distributed.checkpoint import HuggingFaceStorageReader - # Read what the checkpoint holds before changing the model, so a refusal leaves it as it was. reader = HuggingFaceStorageReader(str(local_dir)) ckpt_keys = set(reader.read_metadata().state_dict_metadata.keys()) @@ -592,8 +592,6 @@ def _meta_dcp_load( "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." ) - from torch.distributed.fsdp import FSDPModule - if isinstance(model, FSDPModule): raise RuntimeError( "meta+DCP load: the root module is fully_shard-ed. The load materializes module by module, " @@ -648,8 +646,11 @@ def _meta_dcp_load( # 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. - numel = {n: p.numel() for n, p in model.named_parameters() if n in target_names} + # fails there. Every alias of a tied param counts, as in unplaced_source_keys: the checkpoint can + # store the one named_parameters() drops by default. + numel = { + n: p.numel() for n, p in model.named_parameters(remove_duplicate=False) if n in target_names + } want_params = set(numel) unloaded = want_params - mapped_names total_nel = max(1, sum(numel.values())) @@ -663,7 +664,7 @@ def _meta_dcp_load( f"{sorted(unloaded)[:5]}. Replicating the dense part (MODELOPT_FSDP_SIZE=1) loads it through " "from_pretrained instead." ) - loaded_bufs = mapped_names - {n for n, _ in model.named_parameters()} + loaded_bufs = mapped_names - {n for n, _ in model.named_parameters(remove_duplicate=False)} if loaded_bufs: warn_rank_0( f"meta+DCP load: loaded {len(loaded_bufs)} persistent buffer(s) from checkpoint: " diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 6a654adad76..4a77cc5d425 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -194,15 +194,23 @@ 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 -def test_a_tie_keeps_the_values_of_the_alias_the_checkpoint_stored(tmp_path, monkeypatch): +@pytest.mark.parametrize("shared", [False, True], ids=["aliases-apart", "one-shared-param"]) +def test_a_tie_keeps_the_values_of_the_alias_the_checkpoint_stored( + tmp_path, monkeypatch, recwarn, shared +): """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. (A checkpoint storing the embedding is the - load-and-retie test above.)""" + ``to_empty`` gave it, as ``from_pretrained`` leaves it. That holds whether the meta model holds the + aliases apart, as ``init_empty_weights`` builds it, or as one parameter, as ``from_pretrained``'s + tying leaves it, which ``named_parameters()`` lists once; neither alias is reported as a buffer. + (A checkpoint storing the embedding is the load-and-retie test above.)""" 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) + if shared: + model.tie_weights() # outside init_empty_weights: one Parameter under both names + assert model.lm_head.weight is model.model.embed_tokens.weight # Fresh storage holds NaN, so a parameter nothing wrote to shows. empty_like = torch.empty_like monkeypatch.setattr( @@ -213,6 +221,7 @@ def test_a_tie_keeps_the_values_of_the_alias_the_checkpoint_stored(tmp_path, mon assert model.lm_head.weight is model.model.embed_tokens.weight assert torch.equal(model.lm_head.weight, ckpt["lm_head.weight"]) + assert not [w for w in recwarn if "buffer" in str(w.message)] def test_flat_vlm_checkpoint_names_reach_the_nested_params(tmp_path): From 5ed2c3ab91f46bd33feabb80d5a5aa29a277e3df Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 01:59:20 +0000 Subject: [PATCH 16/17] Refuse mismatched shapes before the meta-DCP load re-empties the model The meta-DCP load matched checkpoint names to the model's without comparing shapes. dcp.load's planner refuses a tensor whose stored shape is not its param's, as when the model was built from a config with another vocab_size, but only after to_empty has re-emptied the model, and as a CheckpointException, which derives from BaseException. The load now compares every mapped pair with the checkpoint metadata it already reads, beside its other refusals, so a mismatch is refused by name with the model still on meta. The coverage guard weighed every target param, so a large one that was already real, which keeps its value without a checkpoint tensor, counted as weight the key mapping missed and could refuse a load that needed nothing for it. Only meta params weigh in now. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 43 +++++++++++++------ tests/unit/torch/utils/test_meta_dcp_load.py | 17 ++++++-- 2 files changed, 44 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 0f33661e829..128577014c6 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -585,7 +585,8 @@ def _meta_dcp_load( """ # Read what the checkpoint holds before changing the model, so a refusal leaves it as it was. reader = HuggingFaceStorageReader(str(local_dir)) - ckpt_keys = set(reader.read_metadata().state_dict_metadata.keys()) + ckpt_md = reader.read_metadata().state_dict_metadata + ckpt_keys = set(ckpt_md) if not ckpt_keys: raise RuntimeError( f"meta+DCP load: the checkpoint at {local_dir} has NO tensors -- the resolved snapshot carries no " @@ -614,14 +615,15 @@ def _meta_dcp_load( # 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. - model_keys = list(model.state_dict()) # (sharded or plain) params + persistent buffers - target_names = set(load_names) & set(model_keys) + # (sharded or plain) params + persistent buffers; a DTensor's shape is its global one. + model_shapes = {k: t.shape for k, t in model.state_dict().items()} + target_names = set(load_names) & set(model_shapes) # {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(model_keys)).items() + for ck, mk in _map_ckpt_keys_to_model(model, ckpt_keys, set(model_shapes)).items() if mk in target_names } # The export carries the checkpoint tensors the model has no parameter for over raw, as the load @@ -636,6 +638,19 @@ def _meta_dcp_load( f"but its load recorded them as having no param, for the export to carry over raw: " f"{carried[:3]}. The two would write them twice." ) + # dcp.load refuses a checkpoint tensor whose shape is not its model tensor's only once to_empty has + # re-emptied the model; refuse it while the model is untouched. + mismatched = sorted( + f"{ck} {tuple(ckpt_md[ck].size)} -> {mk} {tuple(model_shapes[mk])}" + for ck, mk in key_map.items() + if ckpt_md[ck].size != model_shapes[mk] + ) + if mismatched: + raise RuntimeError( + f"meta+DCP load: {len(mismatched)} checkpoint tensor(s) differ in shape from the model " + f"tensor they map onto, as when the model's config does not match the checkpoint: " + f"{mismatched[:3]}" + ) # 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. @@ -647,18 +662,22 @@ def _meta_dcp_load( # ~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. Every alias of a tied param counts, as in unplaced_source_keys: the checkpoint can - # store the one named_parameters() drops by default. - numel = { - n: p.numel() for n, p in model.named_parameters(remove_duplicate=False) if n in target_names - } - want_params = set(numel) - unloaded = want_params - mapped_names + # store the one named_parameters() drops by default. Only meta params weigh in: one already real + # keeps its value without a checkpoint tensor. + want_params: set[str] = set() + numel: dict[str, int] = {} + for n, p in model.named_parameters(remove_duplicate=False): + if n in target_names: + want_params.add(n) + if p.is_meta: + numel[n] = p.numel() + unloaded = numel.keys() - mapped_names total_nel = max(1, sum(numel.values())) unloaded_nel = sum(numel[n] for n in unloaded) 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 " + f"meta+DCP load: {len(unloaded)}/{len(numel)} meta params ({100 * unloaded_nel / total_nel:.0f}% " + f"of their 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]}. Replicating the dense part (MODELOPT_FSDP_SIZE=1) loads it through " diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 4a77cc5d425..18d913ea84c 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -280,10 +280,11 @@ def __init__(self): 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.""" + lacks it, instead of failing as unfilled or, however large (the expert outweighs the one meta + param), as weight the key mapping missed.""" model = nn.Module() model.gate = _Router() - model.expert = nn.Linear(2, 2, bias=False) + model.expert = nn.Linear(64, 64, bias=False) expert = model.expert.weight.detach().clone() ckpt = {"gate.weight": torch.randn(4, 4)} @@ -499,21 +500,29 @@ def _recorded_as_carried_over(model): model._modelopt_unplaced_source_keys = ["model.norm.weight"] +def _with_another_vocab_size(model): + head = model.lm_head + model.lm_head = nn.Linear(head.in_features, 2 * head.out_features, bias=False, device="meta") + + @pytest.mark.parametrize( ("change", "match"), [ (_root_sharded, "root module is fully_shard-ed"), (_with_a_meta_buffer, r"buffer\(s\) are on meta"), (_recorded_as_carried_over, "recorded them as having no param"), + (_with_another_vocab_size, r"lm_head.weight \(32, 32\) -> lm_head.weight \(64, 32\)"), ], - ids=["root-sharded", "meta-buffer", "carried-over-raw"], + ids=["root-sharded", "meta-buffer", "carried-over-raw", "other-shape"], ) def test_a_load_it_cannot_do_rightly_is_refused_before_anything_changes( tmp_path, gloo_world_of_one, change, match ): """A root-sharded model, which the module-by-module materialization would corrupt; a buffer on meta, which a checkpoint of params cannot fill; a tensor the export carries over raw, which - filling would write twice. Each is refused with every param still on meta.""" + filling would write twice; a checkpoint tensor of another shape than its param, as a head built + for another vocab_size has, which ``dcp.load`` refuses only once the model is re-emptied. Each + is refused with every param still on meta.""" ckpt_dir = _save(tmp_path, get_tiny_llama().state_dict()) model = _meta_twin(get_tiny_llama) change(model) From 93322217d0a76414f796f5a23b3af07c8536cdd6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:26:58 +0000 Subject: [PATCH 17/17] Count every name of a param the meta-DCP load re-empties to_empty gives a param a module registers under two names a tensor under each, so the unfilled guard must know both: listing a module's direct params without their duplicates hid the second name, which kept to_empty's memory with nothing to report it. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 5 ++-- tests/unit/torch/utils/test_meta_dcp_load.py | 24 +++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 128577014c6..f1a7fd4fc17 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -696,11 +696,12 @@ def _meta_dcp_load( 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. + # outside ``load_names`` would refill it. Every name counts: ``to_empty`` gives a param the module + # registers twice a tensor under each name. saved_params: dict[str, torch.Tensor] = {} emptied: set[str] = set() for name, m in model.named_modules(): - direct = dict(m.named_parameters(recurse=False)) + direct = dict(m.named_parameters(recurse=False, remove_duplicate=False)) if any(p.is_meta for p in direct.values()): prefix = f"{name}." if name else "" saved_params.update( diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 18d913ea84c..764bce545c6 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -382,6 +382,30 @@ def test_an_untied_param_the_checkpoint_lacks_fails_the_load(tmp_path, named): _meta_dcp_load(model, _save(tmp_path, ckpt), "cpu", load_names=load_names) +class _Aliased(nn.Module): + """One param a module registers under two names.""" + + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.empty(4, 4, device="meta")) + self.alias = self.weight + + +def test_a_second_name_of_a_param_the_checkpoint_lacks_fails_the_load(tmp_path): + """``to_empty`` gives each name of a param the module registers twice a tensor of its own, so + with the checkpoint storing only the first, the second still holds empty memory: the load + fails naming it.""" + model = _Aliased() + + with pytest.raises(RuntimeError, match=r"\['alias'\]"): + _meta_dcp_load( + model, + _save(tmp_path, {"weight": torch.randn(4, 4)}), + "cpu", + load_names=_every_name(model), + ) + + class _OneBigManySmall(nn.Module): """One param holding ~99% of the weight beside twelve tiny ones."""