diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index a57b65d38fc..f398fb82bfc 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -21,6 +21,7 @@ import re from itertools import chain from typing import TYPE_CHECKING, Any +from warnings import warn import torch import torch.nn as nn @@ -43,6 +44,7 @@ barrier, broadcast_state_dict, is_initialized, + rank, set_distribution_plan, ) from modelopt.torch.utils.logging import print_rank_0 @@ -666,6 +668,254 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False return decoder_layers +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 _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, + device, + *, + load_names: set, +) -> None: + """Materialize a meta-built model's ``load_names`` params and fill them through torch DCP. + + 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``. 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`` 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 ``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. + """ + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint import HuggingFaceStorageReader + + # 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)) + + # ``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(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 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. + 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 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"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 " + "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]}" + ) + 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. + 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( 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 new file mode 100644 index 00000000000..c05f7357d25 --- /dev/null +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -0,0 +1,507 @@ +# 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. 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 +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, Replicate, Shard, distribute_tensor + +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.plugins.model_load_utils import ( + _copy_tied_values, + _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 _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( + 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 + + +@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() + if conversion_mapping is not None: + model._checkpoint_conversion_mapping = conversion_mapping + + assert _map_ckpt_keys_to_model(model, ckpt_keys, model_keys) == expected + + +# --------------------------------------------------------------------------------- meta + DCP load + + +def test_load_fills_a_meta_model_keeps_its_buffers_and_reties_it(tmp_path): + """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", load_names=_every_name(model)) + + 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_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. (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) + # 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", load_names=_every_name(model)) + + assert model.lm_head.weight is model.model.embed_tokens.weight + assert torch.equal(model.lm_head.weight, ckpt["lm_head.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", load_names=_every_name(model)) + + 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 _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_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.""" + + 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", load_names=_every_name(model)) + + 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).""" + + 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, even one that requires grad. + """ + ckpt = { + "proj.weight": torch.randn(4, 4), + "proj.bias": torch.randn(4), + "scale": torch.full((4,), 3.0), + } + 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)) + + 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"]) + + +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", load_names=_every_name(model)) + + +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, 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", 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", load_names=_every_name(model)) + assert loads == [1] + + +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", load_names=_every_name(model)) + + 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_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 + 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 layer in model.model.layers: + fully_shard(layer, mesh=mesh) + + _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.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 _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)