diff --git a/docs/source/features/checkpoint-loading.md b/docs/source/features/checkpoint-loading.md index 6caa3a0fcd95..018a7bad6967 100644 --- a/docs/source/features/checkpoint-loading.md +++ b/docs/source/features/checkpoint-loading.md @@ -1,3 +1,8 @@ + + # Checkpoint Loading The PyTorch backend provides a flexible and extensible infrastructure for loading model checkpoints from different formats, such as HuggingFace (HF). This system allows you to load models from various sources (e.g., HuggingFace or custom formats) by implementing the required components, such as the checkpoint’s weight loader, mapper, and configuration parser. @@ -83,6 +88,39 @@ Currently, HF checkpoint loader is the primary built-in format, supporting: - **Configuration parser** - Parsing HF stored configuration information to TRTLLM `ModelConfig` object - **Weights Mapping** - Converting HF weights into TRTLLM compatible representation +### Layer-wise Safetensors Loading + +The PyTorch backend can load supported Hugging Face safetensors checkpoints one +semantic model layer at a time. This bounds the live raw checkpoint tensors to +top-level weights or one model layer instead of materializing the complete +checkpoint in host memory before model-specific transformations finish. + +Enable the prototype option in the LLM configuration: + +```yaml +enable_hf_layerwise_loading: true +``` + +For example, save the option as `config.yaml` and start the server with: + +```bash +trtllm-serve --model --config config.yaml +``` + +The loader groups tensors by semantic layer independently of safetensors shard +boundaries, so inputs needed by same-layer projection and expert fusions remain +available together. The default eager loading path is unchanged when the option +is `false`. + +Current limitations: + +- Only DeepSeek V4 implements the required model-specific layer scoping. +- Only safetensors checkpoints are supported; `.bin` and `.pth` checkpoints + continue to use eager loading. +- A separate speculative draft checkpoint is not supported in layer-wise mode. +- Closing each bucket releases its mmap-backed tensor storage, but filesystem + data may remain in the reclaimable Linux page cache. + ## Using Checkpoint Loaders ### Basic Usage diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py index d00a24114745..04f6dce612b4 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py @@ -1,11 +1,14 @@ -from typing import Optional +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Iterator, Optional from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \ BaseCheckpointLoader from tensorrt_llm._torch.models.checkpoints.base_config_loader import \ BaseConfigLoader -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ - BaseWeightLoader +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ( + BaseWeightLoader, ConsumableWeightsDict) from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ BaseWeightMapper from tensorrt_llm._torch.models.checkpoints.hf.config_loader import \ @@ -73,3 +76,19 @@ def config_loader(self) -> Optional[BaseConfigLoader]: @property def checkpoint_format(self) -> str: return self._checkpoint_format + + def iter_layer_weight_buckets( + self, checkpoint_dir: str, + **kwargs: Any) -> Iterator[ConsumableWeightsDict]: + """Yield semantic-layer weight buckets from the underlying HF loader. + + Args: + checkpoint_dir: Directory containing Hugging Face checkpoint files. + **kwargs: Arguments forwarded to the concrete weight loader. + + Yields: + A consumable dictionary containing top-level tensors or tensors + for exactly one base/MTP checkpoint layer. + """ + return self._weight_loader.iter_layer_weight_buckets( + checkpoint_dir, **kwargs) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index a4e8f898548b..04055d7e2e81 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -13,12 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. import glob +import json import multiprocessing import os +import re import threading from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor -from typing import Any, List +from contextlib import ExitStack +from typing import Any, Iterator, List, Tuple import psutil import safetensors @@ -54,6 +57,116 @@ class HfWeightLoader(BaseWeightLoader): Loads weights from SafeTensors/bin/pth files. """ + @staticmethod + def _layer_bucket_name(weight_name: str) -> Tuple[str, int]: + """Return a stable, layer-atomic bucket for a checkpoint tensor.""" + match = re.match(r"^(?:model\.)?layers\.(\d+)\.", weight_name) + if match is not None: + return ("layer", int(match.group(1))) + match = re.match(r"^mtp\.(\d+)\.", weight_name) + if match is not None: + return ("mtp", int(match.group(1))) + return ("top", 0) + + def iter_layer_weight_buckets(self, + checkpoint_dir: str, + mapping: Mapping, + use_consolidated: bool = False, + **kwargs) -> Iterator[ConsumableWeightsDict]: + """Yield real CPU tensors one semantic layer at a time. + + The safetensors index is used only as metadata. Tensor storage remains + mmap-backed and the file handles for a bucket stay alive until the + caller advances the iterator. This keeps cross-shard, same-layer + fusions atomic without materializing the complete checkpoint. + + Args: + checkpoint_dir: Directory containing safetensors files and an + optional Hugging Face safetensors index. + mapping: Distributed mapping accepted for parity with eager + checkpoint loading. Bucketing itself is rank-independent. + use_consolidated: Select consolidated files instead of ordinary + Hugging Face shards. + **kwargs: Extra eager-loader arguments that do not affect bucket + planning. + + Yields: + A consumable dictionary for all top-level tensors or one numeric + base/MTP layer. File handles remain open until the next iteration. + + Raises: + RuntimeError: If the selected safetensors files are missing or + their index/key metadata is incomplete or ambiguous. + """ + del mapping, kwargs + weight_files = glob.glob(f"{checkpoint_dir}/*.safetensors") + weight_files = [ + path for path in weight_files + if ("consolidated" in os.path.basename(path)) == use_consolidated + ] + if not weight_files: + checkpoint_kind = "consolidated " if use_consolidated else "" + raise RuntimeError( + f"Layer-wise loading requires {checkpoint_kind}safetensors " + f"weights in {checkpoint_dir}.") + + files_by_name = {os.path.basename(path): path for path in weight_files} + index_files = sorted( + glob.glob(f"{checkpoint_dir}/*.safetensors.index.json")) + entries: List[Tuple[str, str]] = [] + use_index = bool(index_files) and not use_consolidated and all( + "consolidated" not in name for name in files_by_name) + if use_index: + with open(index_files[0], encoding="utf-8") as index_file: + weight_map = json.load(index_file).get("weight_map", {}) + missing_files = set(weight_map.values()) - set(files_by_name) + if missing_files: + raise RuntimeError( + f"Safetensors index {index_files[0]} references missing " + f"checkpoint files: {sorted(missing_files)}") + for weight_name, file_name in weight_map.items(): + entries.append((weight_name, files_by_name[file_name])) + if not entries: + raise RuntimeError( + f"Safetensors index {index_files[0]} has an empty weight_map." + ) + else: + seen_keys = set() + for path in sorted(weight_files): + with safetensors.safe_open(path, framework="pt", + device="cpu") as handle: + for key in handle.keys(): + if key in seen_keys: + raise RuntimeError( + "Duplicate tensor key found across safetensors " + f"files without an index: {key}") + seen_keys.add(key) + entries.append((key, path)) + + buckets = {} + for weight_name, path in entries: + bucket = self._layer_bucket_name(weight_name) + buckets.setdefault(bucket, []).append((weight_name, path)) + + def bucket_sort_key(bucket): + kind, index = bucket + return ({"top": 0, "layer": 1, "mtp": 2}[kind], index) + + for bucket in sorted(buckets, key=bucket_sort_key): + tensors = {} + with ExitStack() as stack: + handles = {} + for weight_name, path in buckets[bucket]: + if path not in handles: + handles[path] = stack.enter_context( + safetensors.safe_open(path, + framework="pt", + device="cpu")) + tensors[weight_name] = handles[path].get_tensor(weight_name) + logger.info("Loading HF weight bucket %s with %d tensors", + bucket, len(tensors)) + yield ConsumableWeightsDict(tensors) + @staticmethod def _is_weight_cache_enabled() -> bool: return os.environ.get(_WEIGHT_CACHE_ENV, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 2fb611ecb999..e07be988d661 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -28,6 +28,7 @@ import copy import math import os +import re from typing import TYPE_CHECKING, Dict, List, Optional if TYPE_CHECKING: @@ -286,6 +287,29 @@ def _copy_deepseek_v4_fused_a_weight_scale( module.weight_scale.data.copy_(fused_a_scale) +def _is_deepseek_v4_semantic_layer_target( + module_layer: int, + active_layer: int, + num_hidden_layers: int, + num_nextn_predict_layers: int, +) -> bool: + """Return whether a runtime layer belongs to a checkpoint layer bucket. + + Runtime MTP replicas map back to their checkpoint MTP layer by modulo. + Reject base layers before that operation: Python maps negative values into + the non-negative residue range, which would make base layers look like MTP + replicas, especially when the divisor is one. + """ + if active_layer < num_hidden_layers: + return module_layer == active_layer + if module_layer < num_hidden_layers: + return False + + checkpoint_mtp_idx = active_layer - num_hidden_layers + runtime_mtp_idx = module_layer - num_hidden_layers + return runtime_mtp_idx % num_nextn_predict_layers == checkpoint_mtp_idx + + def _remap_deepseek_v4_checkpoint_keys( weights: Dict, num_hidden_layers: int, kv_lora_rank: int = 448 ) -> Dict: @@ -314,11 +338,10 @@ def _remap_deepseek_v4_checkpoint_keys( ``lm_head`` via ``shared_head``. Flash omits this key entirely; Flash-Base carries it but matches the main head, so we let the main head win. """ - mtp_layer_prefix = f"model.layers.{num_hidden_layers}" routed_moe_scale_name = "weight_scale_inv" for key, value in weights.items(): if ( - key.startswith("layers.") + (key.startswith("layers.") or re.match(r"^mtp\.\d+\.", key)) and ".ffn.experts." in key and key.endswith(".weight") and getattr(value, "ndim", 0) == 2 @@ -444,14 +467,17 @@ def _emit_or_collect(model_key: str, tensor: torch.Tensor): out[f"model.{k}"] = v continue - # mtp.0.head.weight is intentionally dropped (Flash-Base only); see + # mtp..head.weight is intentionally dropped (Flash-Base only); see # docstring caveats. - if k == "mtp.0.head.weight": + if re.match(r"^mtp\.\d+\.head\.weight$", k): continue - # mtp.0.* — route to model.layers.{num_hidden_layers}.* - if k.startswith("mtp.0."): - rest = k[len("mtp.0.") :] + # mtp..* — route to model.layers.{num_hidden_layers + n}.* + mtp_match = re.match(r"^mtp\.(\d+)\.(.*)$", k) + if mtp_match is not None: + mtp_idx = int(mtp_match.group(1)) + mtp_layer_prefix = f"model.layers.{num_hidden_layers + mtp_idx}" + rest = mtp_match.group(2) # MTP-only keys: enorm, hnorm map directly; norm maps to # shared_head.norm; hc_head_* maps under shared_head; e_proj / # h_proj are loaded as two separate Linear modules (no fused @@ -524,13 +550,58 @@ def __init__(self, model, is_draft_model: bool = False): self.model_config = model.model_config self.is_draft_model = is_draft_model - def load_weights(self, weights: Dict, skip_modules: List[str] = []): + def load_weights( + self, + weights: Dict, + skip_modules: Optional[List[str]] = None, + initial_bucket_loading: bool = False, + ): + """Remap and load eager or semantic-layer DeepSeek V4 weights. + + Args: + weights: Raw checkpoint tensors or model-named tensors. + skip_modules: Module-name fragments to omit while loading. + initial_bucket_loading: Restrict traversal and synthesized defaults + to the single semantic layer represented by ``weights``. A + top-level bucket instead skips all decoder layers. + + Raises: + ValueError: If an initial bucket mixes top-level tensors with a + layer or contains more than one semantic layer. + """ + skip_modules = skip_modules or [] # If the checkpoint uses raw DS-V4 keys (layers.X.attn.wkv.weight, # mtp.0.*, embed.weight, head.weight), rewrite them to the model's # named-parameter keys before iterating modules. The detection is by # presence of any top-level "layers." key; HF-style checkpoints use # "model.layers." and skip this branch. - if any(k == "embed.weight" or k.startswith("layers.") for k in weights): + raw_checkpoint = any( + k in ("embed.weight", "head.weight", "norm.weight") + or k.startswith(("layers.", "mtp.", "hc_head_")) + for k in weights + ) + active_layer = None + if initial_bucket_loading: + layer_indices = set() + has_top_level = False + for key in weights: + if key.startswith("layers."): + layer_indices.add(int(key.split(".", 2)[1])) + elif (mtp_match := re.match(r"^mtp\.(\d+)\.", key)) is not None: + layer_indices.add(self.config.num_hidden_layers + int(mtp_match.group(1))) + elif key.startswith("model.layers."): + layer_indices.add(int(key.split(".", 3)[2])) + else: + has_top_level = True + if len(layer_indices) > 1 or (layer_indices and has_top_level): + raise ValueError( + "Initial weight bucket must contain exactly one model " + f"layer or top-level weights, got layers={sorted(layer_indices)}, " + f"has_top_level={has_top_level}." + ) + active_layer = next(iter(layer_indices), None) + + if raw_checkpoint: weights = _remap_deepseek_v4_checkpoint_keys( weights, num_hidden_layers=self.config.num_hidden_layers, @@ -551,6 +622,11 @@ def load_weights(self, weights: Dict, skip_modules: List[str] = []): ) model_params = dict(self.model.named_parameters()) for pname, p in model_params.items(): + if initial_bucket_loading: + if active_layer is None or not pname.startswith( + f"model.layers.{active_layer}." + ): + continue if pname in weights: continue if any(pname.endswith(s) for s in _ones_suffixes): @@ -816,6 +892,22 @@ def load_flat_hc_weights(module, names: List[str]) -> bool: for name, module in tqdm(all_named_modules.items(), desc="Loading weights"): if name.startswith("draft_model"): continue + if initial_bucket_loading: + if active_layer is None: + if name.startswith("model.layers."): + continue + else: + is_target_module = False + if name.startswith("model.layers."): + module_layer = int(name.split(".", 3)[2]) + is_target_module = _is_deepseek_v4_semantic_layer_target( + module_layer=module_layer, + active_layer=active_layer, + num_hidden_layers=self.config.num_hidden_layers, + num_nextn_predict_layers=self.config.num_nextn_predict_layers, + ) + if not is_target_module: + continue names = name.split(".") parent_module_name = ".".join(names[:-1]) @@ -2564,9 +2656,9 @@ def forward( **kwargs, ) - def load_weights(self, weights: Dict): + def load_weights(self, weights: Dict, initial_bucket_loading: bool = False): weight_loader = DeepseekV4WeightLoader(self) - weight_loader.load_weights(weights) + weight_loader.load_weights(weights, initial_bucket_loading=initial_bucket_loading) def post_load_weights(self): layers = self.model.layers[: self.config.num_hidden_layers] diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 57cd50b01cc3..af2cbcb82971 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -594,23 +594,55 @@ def init_meta_tensor(t: torch.Tensor): load_weights_kwargs[ "prepare_post_transform_receiver"] = self._setup_aliases - if hasattr(model, 'llm_checkpoint_dir'): - weights = checkpoint_loader.load_weights( - model.llm_checkpoint_dir, **load_weights_kwargs) + model_checkpoint_dir = (model.llm_checkpoint_dir if hasattr( + model, 'llm_checkpoint_dir') else checkpoint_dir) + layerwise_loading = (checkpoint_loader.checkpoint_format == "HF" + and getattr(self.llm_args, + 'enable_hf_layerwise_loading', + False)) + if layerwise_loading: + if loads_draft_weights: + raise RuntimeError( + "HF layer-wise loading does not yet support a " + "separate speculative draft checkpoint.") + # Validate the model contract before constructing the + # generator: advancing it would open files and may start + # mutating the model with the first bucket. + load_parameters = inspect.signature( + model.load_weights).parameters + if "initial_bucket_loading" not in load_parameters: + raise RuntimeError( + "enable_hf_layerwise_loading is enabled, but " + f"{type(model).__name__}.load_weights does not support " + "initial_bucket_loading.") + self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( + model, config) + logger.info( + "Loading HF checkpoint with one semantic layer per host-memory bucket." + ) + for weight_bucket in checkpoint_loader.iter_layer_weight_buckets( + model_checkpoint_dir, **load_weights_kwargs): + self._call_load_weights(model.load_weights, + weight_bucket, + self.weight_mapper, + initial_bucket_loading=True) + # copy_(..., non_blocking=True) and backend transforms + # may still consume mmap-backed source storage. + torch.cuda.synchronize() + del weight_bucket else: weights = checkpoint_loader.load_weights( - checkpoint_dir, **load_weights_kwargs) - - # When MX P2P succeeds, weights are already in model params. - # A non-empty dict contains size-mismatched tensors that - # should be merged via the standard disk pipeline. - weights_preloaded = checkpoint_loader.is_weights_preloaded() - self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( - model, config) - - if weights: - self._call_load_weights(model.load_weights, weights, - self.weight_mapper) + model_checkpoint_dir, **load_weights_kwargs) + + # When MX P2P succeeds, weights are already in model params. + # A non-empty dict contains size-mismatched tensors that + # should be merged via the standard disk pipeline. + weights_preloaded = checkpoint_loader.is_weights_preloaded() + self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( + model, config) + if weights: + self._call_load_weights(model.load_weights, weights, + self.weight_mapper) if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( ): @@ -1396,9 +1428,23 @@ def _call_load_weights(self, load_method: Callable, weights, weight_mapper, - allow_partial_loading: bool = False): - """Calls the model's weight loading method with the correct arguments.""" - args = inspect.getfullargspec(load_method).args + allow_partial_loading: bool = False, + initial_bucket_loading: bool = False): + """Call a model weight loader with only the options it declares. + + Args: + load_method: Bound model method that consumes checkpoint weights. + weights: Weight dictionary passed to ``load_method``. + weight_mapper: Initialized checkpoint-to-model name mapper. + allow_partial_loading: Whether a reload may omit model weights. + initial_bucket_loading: Whether ``weights`` is one bucket from the + initial semantic-layer checkpoint load. + + Raises: + AssertionError: If a requested option is unsupported by the model + method. + """ + args = inspect.signature(load_method).parameters kargs = {} if "weight_mapper" in args: kargs["weight_mapper"] = weight_mapper @@ -1406,4 +1452,9 @@ def _call_load_weights(self, kargs["allow_partial_loading"] = allow_partial_loading else: assert allow_partial_loading is False, "allow_partial_loading is not supported for this model" + if "initial_bucket_loading" in args: + kargs["initial_bucket_loading"] = initial_bucket_loading + else: + assert initial_bucket_loading is False, \ + "initial_bucket_loading is not supported for this model" load_method(weights, **kargs) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a7a07926a4e1..1ee33e0b6476 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4576,6 +4576,15 @@ def set_default_capture_num_tokens(self) -> 'TorchCompileConfig': class TorchLlmArgs(BaseLlmArgs): # PyTorch backend specific configurations + enable_hf_layerwise_loading: bool = Field( + default=False, + description= + "Load supported Hugging Face safetensors checkpoints one semantic " + "model layer at a time to reduce peak host memory during model " + "startup. Currently supported only for DeepSeek V4 and incompatible " + "with a separate speculative draft checkpoint.", + status="prototype") + garbage_collection_gen0_threshold: int = Field( default=20000, description= diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index ca6e762490cf..f970271c8146 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -404,6 +404,13 @@ "kind": "value", "path": "enable_energy_metrics" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "enable_hf_layerwise_loading" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/unittest/_torch/executor/test_model_loader_layerwise.py b/tests/unittest/_torch/executor/test_model_loader_layerwise.py new file mode 100644 index 000000000000..1c0438078896 --- /dev/null +++ b/tests/unittest/_torch/executor/test_model_loader_layerwise.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for semantic-layer checkpoint orchestration in ModelLoader.""" + +from contextlib import contextmanager, nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_mod +from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader +from tensorrt_llm.llmapi.llm_args import LoadFormat + + +class _LayerwiseModel(nn.Module): + def __init__(self, events, *, fail_on_bucket=None): + super().__init__() + self._events = events + self._fail_on_bucket = fail_on_bucket + self.llm_checkpoint_dir = "/model-checkpoint" + + def _apply(self, _fn): + return self + + def to(self, *_args, **_kwargs): + return self + + def load_weights(self, weights, weight_mapper, *, initial_bucket_loading=False): + bucket_name = next(iter(weights)) + self._events.append(("load", bucket_name, weight_mapper, initial_bucket_loading)) + if bucket_name == self._fail_on_bucket: + raise RuntimeError("bucket load failed") + + def post_load_weights(self): + self._events.append(("post_load",)) + + +class _UnsupportedModel(_LayerwiseModel): + def load_weights(self, weights, weight_mapper): + self._events.append(("load", next(iter(weights)), weight_mapper)) + + +@contextmanager +def _moe_context(_config, _mapping): + yield None + + +def _make_loader(monkeypatch, model, *, spec_config=None, layerwise_loading=True): + loader = ModelLoader( + llm_args=SimpleNamespace( + load_format=LoadFormat.AUTO, + enable_hf_layerwise_loading=layerwise_loading, + ), + mapping=MagicMock(name="mapping"), + spec_config=spec_config, + sparse_attention_config=None, + max_num_tokens=128, + max_seq_len=128, + ) + loader._load_and_validate_config = MagicMock( + return_value=SimpleNamespace(name="config", mapping=SimpleNamespace()) + ) + monkeypatch.setattr(model_loader_mod, "timing", lambda *_args, **_kwargs: nullcontext()) + monkeypatch.setattr(model_loader_mod, "maybe_create_moe_load_balancer", _moe_context) + monkeypatch.setattr(model_loader_mod, "MetaInitMode", lambda: nullcontext()) + monkeypatch.setattr( + model_loader_mod.AutoModelForCausalLM, "from_config", MagicMock(return_value=model) + ) + monkeypatch.setattr(model_loader_mod, "get_rank_model_storage", lambda _model: 0) + monkeypatch.setattr( + torch.cuda, + "current_stream", + lambda: SimpleNamespace(synchronize=lambda: None), + ) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + return loader + + +def _make_checkpoint_loader(events, buckets): + checkpoint_loader = MagicMock(name="checkpoint_loader") + checkpoint_loader.checkpoint_format = "HF" + mapper = MagicMock(name="weight_mapper") + + def initialize_mapper(_model, _config): + events.append(("mapper", mapper)) + return mapper + + def iter_buckets(checkpoint_dir, **kwargs): + events.append(("iterate", checkpoint_dir, kwargs)) + try: + yield from buckets + finally: + events.append(("iterator_closed",)) + + checkpoint_loader.get_initialized_weight_mapper.side_effect = initialize_mapper + checkpoint_loader.iter_layer_weight_buckets.side_effect = iter_buckets + return checkpoint_loader, mapper + + +def test_layerwise_model_loader_orders_mapper_buckets_sync_and_post_load(monkeypatch): + events = [] + model = _LayerwiseModel(events) + loader = _make_loader(monkeypatch, model) + checkpoint_loader, mapper = _make_checkpoint_loader( + events, + [{"top.weight": object()}, {"layers.0.weight": object()}], + ) + synchronize = MagicMock(side_effect=lambda: events.append(("synchronize",))) + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + + loaded_model, _ = loader.load("/checkpoint", checkpoint_loader) + + assert loaded_model is model + assert events[:7] == [ + ("mapper", mapper), + ( + "iterate", + "/model-checkpoint", + { + "mapping": loader.mapping, + "model": model, + "source_identity": None, + }, + ), + ("load", "top.weight", mapper, True), + ("synchronize",), + ("load", "layers.0.weight", mapper, True), + ("synchronize",), + ("iterator_closed",), + ] + assert events[-1] == ("post_load",) + assert synchronize.call_count == 2 + checkpoint_loader.load_weights.assert_not_called() + + +def test_layerwise_model_loader_is_disabled_by_default_config(monkeypatch): + events = [] + model = _LayerwiseModel(events) + loader = _make_loader(monkeypatch, model, layerwise_loading=False) + checkpoint_loader, mapper = _make_checkpoint_loader(events, [{"layers.0.weight": object()}]) + checkpoint_loader.load_weights.return_value = {"eager.weight": object()} + checkpoint_loader.is_weights_preloaded.return_value = False + checkpoint_loader.get_initialized_weight_mapper.return_value = mapper + + loaded_model, _ = loader.load("/checkpoint", checkpoint_loader) + + assert loaded_model is model + checkpoint_loader.load_weights.assert_called_once() + checkpoint_loader.iter_layer_weight_buckets.assert_not_called() + assert ("load", "eager.weight", mapper, False) in events + + +def test_layerwise_model_loader_rejects_unsupported_model_before_iteration(monkeypatch): + events = [] + loader = _make_loader(monkeypatch, _UnsupportedModel(events)) + checkpoint_loader, _mapper = _make_checkpoint_loader(events, [{"top.weight": object()}]) + + with pytest.raises(RuntimeError, match="does not support initial_bucket_loading"): + loader.load("/checkpoint", checkpoint_loader) + + checkpoint_loader.get_initialized_weight_mapper.assert_not_called() + checkpoint_loader.iter_layer_weight_buckets.assert_not_called() + assert events == [] + + +def test_layerwise_model_loader_rejects_draft_checkpoint_before_iteration(monkeypatch): + events = [] + spec_config = MagicMock() + spec_config.spec_dec_mode.need_load_draft_weights.return_value = True + loader = _make_loader(monkeypatch, _LayerwiseModel(events), spec_config=spec_config) + checkpoint_loader, _mapper = _make_checkpoint_loader(events, [{"top.weight": object()}]) + + with pytest.raises(RuntimeError, match="separate speculative draft checkpoint"): + loader.load("/checkpoint", checkpoint_loader) + + checkpoint_loader.get_initialized_weight_mapper.assert_not_called() + checkpoint_loader.iter_layer_weight_buckets.assert_not_called() + assert events == [] + + +def test_layerwise_model_loader_closes_iterator_and_stops_after_consumer_error(monkeypatch): + events = [] + model = _LayerwiseModel(events, fail_on_bucket="layers.1.weight") + loader = _make_loader(monkeypatch, model) + checkpoint_loader, mapper = _make_checkpoint_loader( + events, + [ + {"layers.0.weight": object()}, + {"layers.1.weight": object()}, + {"layers.2.weight": object()}, + ], + ) + synchronize = MagicMock(side_effect=lambda: events.append(("synchronize",))) + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + + with pytest.raises(RuntimeError, match="bucket load failed"): + loader.load("/checkpoint", checkpoint_loader) + + assert ("load", "layers.0.weight", mapper, True) in events + assert ("load", "layers.1.weight", mapper, True) in events + assert not any(event[:2] == ("load", "layers.2.weight") for event in events) + assert events[-1] == ("iterator_closed",) + assert synchronize.call_count == 1 + checkpoint_loader.post_load_apply.assert_not_called() + checkpoint_loader.post_load_publish.assert_not_called() diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 106ca0e5c59d..f81c245611a0 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -5,6 +5,7 @@ import textwrap import weakref from copy import deepcopy +from types import SimpleNamespace import pytest import torch @@ -31,8 +32,10 @@ DeepseekV4ForCausalLM, DeepseekV4Gate, DeepseekV4MTP, + DeepseekV4WeightLoader, _copy_deepseek_v4_fused_a_weight_scale, _deepseek_v4_pos_embd_params, + _is_deepseek_v4_semantic_layer_target, _remap_deepseek_v4_checkpoint_keys, _resolve_enable_fused_hc, ) @@ -187,6 +190,98 @@ def test_deepseek_v4_weight_remap_for_fp8_routed_experts(): assert "model.layers.0.mlp.experts.0.w1.weight_scale" not in remapped +def test_deepseek_v4_weight_remap_for_mtp_mxfp4_routed_experts(): + weights = { + "mtp.0.ffn.experts.0.w1.weight": torch.tensor([[-1, 2], [3, -4]], dtype=torch.int8), + "mtp.0.ffn.experts.0.w1.scale": torch.tensor([1, 2], dtype=torch.int8), + } + + remapped = _remap_deepseek_v4_checkpoint_keys(weights, num_hidden_layers=61, kv_lora_rank=448) + + assert remapped["model.layers.61.mlp.experts.0.w1.weight"].dtype == torch.uint8 + assert remapped["model.layers.61.mlp.experts.0.w1.weight_scale"].dtype == torch.uint8 + + +def test_deepseek_v4_weight_remap_supports_multiple_mtp_layers(): + weights = { + "mtp.0.enorm.weight": torch.tensor([1.0]), + "mtp.1.enorm.weight": torch.tensor([2.0]), + "mtp.1.head.weight": torch.tensor([3.0]), + } + + remapped = _remap_deepseek_v4_checkpoint_keys(weights, num_hidden_layers=61, kv_lora_rank=448) + + assert torch.equal(remapped["model.layers.61.enorm.weight"], weights["mtp.0.enorm.weight"]) + assert torch.equal(remapped["model.layers.62.enorm.weight"], weights["mtp.1.enorm.weight"]) + assert all(not key.endswith("head.weight") for key in remapped) + + +@pytest.mark.parametrize( + "weights", + [ + { + "layers.0.attn.weight": torch.tensor([1.0]), + "layers.1.attn.weight": torch.tensor([2.0]), + }, + { + "embed.weight": torch.tensor([1.0]), + "layers.0.attn.weight": torch.tensor([2.0]), + }, + { + "mtp.0.enorm.weight": torch.tensor([1.0]), + "mtp.1.enorm.weight": torch.tensor([2.0]), + }, + ], +) +def test_deepseek_v4_layerwise_loading_rejects_non_atomic_bucket(weights): + loader = DeepseekV4WeightLoader.__new__(DeepseekV4WeightLoader) + loader.config = SimpleNamespace(num_hidden_layers=61) + + with pytest.raises(ValueError, match="exactly one model layer or top-level weights"): + loader.load_weights(weights, initial_bucket_loading=True) + + +@pytest.mark.parametrize( + "module_layer,active_layer,num_hidden_layers,num_nextn_predict_layers,expected", + [ + # Base-layer buckets select only the exact base layer. + (3, 3, 61, 1, True), + (4, 3, 61, 1, False), + # Regression: negative modulo must not classify base layers as MTP. + (0, 61, 61, 1, False), + (60, 61, 61, 1, False), + (0, 62, 61, 2, False), + # One checkpoint MTP layer may back multiple runtime replicas. + (61, 61, 61, 1, True), + (62, 61, 61, 1, True), + (63, 61, 61, 1, True), + # Multiple checkpoint MTP layers retain their modulo mapping. + (61, 61, 61, 2, True), + (63, 61, 61, 2, True), + (62, 61, 61, 2, False), + (62, 62, 61, 2, True), + (64, 62, 61, 2, True), + (63, 62, 61, 2, False), + ], +) +def test_deepseek_v4_semantic_layer_target( + module_layer, + active_layer, + num_hidden_layers, + num_nextn_predict_layers, + expected, +): + assert ( + _is_deepseek_v4_semantic_layer_target( + module_layer, + active_layer, + num_hidden_layers, + num_nextn_predict_layers, + ) + is expected + ) + + def test_deepseek_v4_fused_a_weight_scale_rebuilds_fp8_shape(): module = torch.nn.Module() module.weight = torch.nn.Parameter(torch.empty((2048, 7168), dtype=torch.float8_e4m3fn)) diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py index 08ebaa09c2f8..30d5deed18b1 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py @@ -1,6 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json from unittest import mock import pytest +import safetensors +import torch +from safetensors.torch import save_file from tensorrt_llm._torch.models.checkpoints import HfWeightLoader from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict @@ -11,6 +18,247 @@ class MyError(Exception): pass +def test_layerwise_safetensors_keeps_cross_shard_layer_atomic(tmp_path): + shard1 = tmp_path / "model-00001-of-00002.safetensors" + shard2 = tmp_path / "model-00002-of-00002.safetensors" + save_file( + { + "embed.weight": torch.tensor([1.0]), + "layers.0.attn.compressor.wkv.weight": torch.tensor([2.0]), + }, + str(shard1), + ) + save_file( + { + "layers.0.attn.compressor.wgate.weight": torch.tensor([3.0]), + "layers.1.attn.wq_b.weight": torch.tensor([4.0]), + "mtp.0.enorm.weight": torch.tensor([5.0]), + }, + str(shard2), + ) + weight_map = { + "embed.weight": shard1.name, + "layers.0.attn.compressor.wkv.weight": shard1.name, + "layers.0.attn.compressor.wgate.weight": shard2.name, + "layers.1.attn.wq_b.weight": shard2.name, + "mtp.0.enorm.weight": shard2.name, + } + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}), encoding="utf-8" + ) + + loader = HfWeightLoader() + buckets = [] + for bucket in loader.iter_layer_weight_buckets(str(tmp_path), mapping=Mapping()): + buckets.append(set(bucket.keys())) + + assert buckets == [ + {"embed.weight"}, + { + "layers.0.attn.compressor.wkv.weight", + "layers.0.attn.compressor.wgate.weight", + }, + {"layers.1.attn.wq_b.weight"}, + {"mtp.0.enorm.weight"}, + ] + + +def test_layerwise_safetensors_without_index_discovers_keys(tmp_path): + save_file( + { + "model.layers.2.foo.weight": torch.tensor([2.0]), + "lm_head.weight": torch.tensor([1.0]), + }, + str(tmp_path / "model.safetensors"), + ) + + loader = HfWeightLoader() + buckets = [ + set(bucket.keys()) + for bucket in loader.iter_layer_weight_buckets(str(tmp_path), mapping=Mapping()) + ] + + assert buckets == [{"lm_head.weight"}, {"model.layers.2.foo.weight"}] + + +def test_layerwise_safetensors_rejects_missing_index_shard(tmp_path): + shard1 = tmp_path / "model-00001-of-00002.safetensors" + save_file({"layers.0.foo.weight": torch.tensor([1.0])}, str(shard1)) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "layers.0.foo.weight": shard1.name, + "layers.1.foo.weight": "model-00002-of-00002.safetensors", + } + } + ), + encoding="utf-8", + ) + + loader = HfWeightLoader() + with pytest.raises(RuntimeError, match="missing checkpoint files"): + next(loader.iter_layer_weight_buckets(str(tmp_path), mapping=Mapping())) + + +def test_layerwise_safetensors_uses_natural_layer_order_and_preserves_values(tmp_path): + save_file( + { + "layers.10.weight": torch.tensor([10.0]), + "layers.2.weight": torch.tensor([2.0]), + "mtp.1.weight": torch.tensor([101.0]), + "mtp.0.weight": torch.tensor([100.0]), + }, + str(tmp_path / "model.safetensors"), + ) + + buckets = list(HfWeightLoader().iter_layer_weight_buckets(str(tmp_path), mapping=Mapping())) + + assert [next(iter(bucket)) for bucket in buckets] == [ + "layers.2.weight", + "layers.10.weight", + "mtp.0.weight", + "mtp.1.weight", + ] + assert [next(iter(bucket.values())).item() for bucket in buckets] == [2.0, 10.0, 100.0, 101.0] + + +def test_layerwise_safetensors_rejects_empty_index(tmp_path): + save_file({"layers.0.weight": torch.tensor([1.0])}, str(tmp_path / "model.safetensors")) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {}}), encoding="utf-8" + ) + + with pytest.raises(RuntimeError, match="empty weight_map"): + next(HfWeightLoader().iter_layer_weight_buckets(str(tmp_path), mapping=Mapping())) + + +def test_layerwise_safetensors_rejects_duplicate_keys_without_index(tmp_path): + for shard_idx in range(2): + save_file( + {"layers.0.weight": torch.tensor([float(shard_idx)])}, + str(tmp_path / f"model-{shard_idx}.safetensors"), + ) + + with pytest.raises(RuntimeError, match="Duplicate tensor key"): + next(HfWeightLoader().iter_layer_weight_buckets(str(tmp_path), mapping=Mapping())) + + +def test_layerwise_safetensors_strictly_selects_consolidated_files(tmp_path): + save_file({"layers.0.weight": torch.tensor([1.0])}, str(tmp_path / "model.safetensors")) + save_file( + {"layers.1.weight": torch.tensor([2.0])}, + str(tmp_path / "model-consolidated.safetensors"), + ) + + ordinary = list(HfWeightLoader().iter_layer_weight_buckets(str(tmp_path), mapping=Mapping())) + consolidated = list( + HfWeightLoader().iter_layer_weight_buckets( + str(tmp_path), mapping=Mapping(), use_consolidated=True + ) + ) + + assert [set(bucket) for bucket in ordinary] == [{"layers.0.weight"}] + assert [set(bucket) for bucket in consolidated] == [{"layers.1.weight"}] + + +@pytest.mark.parametrize("use_consolidated", [False, True]) +def test_layerwise_safetensors_does_not_fallback_to_wrong_file_kind(tmp_path, use_consolidated): + filename = "model.safetensors" if use_consolidated else "model-consolidated.safetensors" + save_file({"layers.0.weight": torch.tensor([1.0])}, str(tmp_path / filename)) + + with pytest.raises(RuntimeError, match="requires .*safetensors weights"): + next( + HfWeightLoader().iter_layer_weight_buckets( + str(tmp_path), mapping=Mapping(), use_consolidated=use_consolidated + ) + ) + + +def test_layerwise_safetensors_handles_follow_bucket_lifetime(tmp_path, monkeypatch): + shard = tmp_path / "model.safetensors" + save_file( + { + "layers.0.weight": torch.tensor([1.0]), + "layers.1.weight": torch.tensor([2.0]), + }, + str(shard), + ) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "layers.0.weight": shard.name, + "layers.1.weight": shard.name, + } + } + ), + encoding="utf-8", + ) + opened = [] + closed = [] + real_safe_open = safetensors.safe_open + + class RecordingHandle: + def __init__(self, path, **kwargs): + self._context = real_safe_open(path, **kwargs) + self._path = path + + def __enter__(self): + self._handle = self._context.__enter__() + opened.append(self._path) + return self._handle + + def __exit__(self, *args): + closed.append(self._path) + return self._context.__exit__(*args) + + monkeypatch.setattr(safetensors, "safe_open", RecordingHandle) + iterator = HfWeightLoader().iter_layer_weight_buckets(str(tmp_path), mapping=Mapping()) + + first = next(iterator) + assert first["layers.0.weight"].item() == 1.0 + assert len(opened) == 1 + assert not closed + + second = next(iterator) + assert second["layers.1.weight"].item() == 2.0 + assert len(opened) == 2 + assert len(closed) == 1 + + iterator.close() + assert len(closed) == 2 + + +def test_layerwise_safetensors_closes_handle_on_consumer_exception(tmp_path, monkeypatch): + shard = tmp_path / "model.safetensors" + save_file({"layers.0.weight": torch.tensor([1.0])}, str(shard)) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"layers.0.weight": shard.name}}), encoding="utf-8" + ) + closed = [] + real_safe_open = safetensors.safe_open + + class RecordingHandle: + def __init__(self, path, **kwargs): + self._context = real_safe_open(path, **kwargs) + + def __enter__(self): + return self._context.__enter__() + + def __exit__(self, *args): + closed.append(True) + return self._context.__exit__(*args) + + monkeypatch.setattr(safetensors, "safe_open", RecordingHandle) + + with pytest.raises(MyError): + for _bucket in HfWeightLoader().iter_layer_weight_buckets(str(tmp_path), mapping=Mapping()): + raise MyError + + assert closed == [True] + + @pytest.fixture(autouse=True) def clean_weight_cache(): HfWeightLoader._clear_weight_cache() diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 7ff208c41ce3..de5d4e571ab7 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -74,6 +74,10 @@ methods: default: null status: prototype # Runtime behavior + enable_hf_layerwise_loading: + annotation: bool + default: False + status: prototype garbage_collection_gen0_threshold: annotation: int default: 20000 diff --git a/tests/unittest/llmapi/test_layerwise_loading_args.py b/tests/unittest/llmapi/test_layerwise_loading_args.py new file mode 100644 index 000000000000..f43c0b5228d6 --- /dev/null +++ b/tests/unittest/llmapi/test_layerwise_loading_args.py @@ -0,0 +1,32 @@ +# 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. + +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + +def test_layerwise_loading_defaults_to_false_and_is_prototype(): + args = TorchLlmArgs(model="gpt2") + + assert args.enable_hf_layerwise_loading is False + field = TorchLlmArgs.model_fields["enable_hf_layerwise_loading"] + assert field.json_schema_extra["status"] == "prototype" + + +def test_layerwise_loading_accepts_true_and_round_trips(): + args = TorchLlmArgs(model="gpt2", enable_hf_layerwise_loading=True) + + data = args.model_dump() + assert data["enable_hf_layerwise_loading"] is True + assert TorchLlmArgs(**data).enable_hf_layerwise_loading is True