From 3e74fe2ff67c006248716e1a25beff2fdafc9357 Mon Sep 17 00:00:00 2001 From: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com> Date: Wed, 15 Jul 2026 01:26:43 -0700 Subject: [PATCH 1/4] [None][feat] add layer-wise safetensors loading Signed-off-by: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com> --- ...layerwise_hf_safetensors_loading_design.md | 358 ++++++++++++++++++ .../checkpoints/hf/checkpoint_loader.py | 12 + .../models/checkpoints/hf/weight_loader.py | 104 ++++- .../_torch/models/modeling_deepseekv4.py | 82 +++- .../_torch/pyexecutor/model_loader.py | 69 +++- .../modeling/test_modeling_deepseekv4.py | 17 + .../checkpoints/hf/test_weight_loader.py | 82 ++++ 7 files changed, 695 insertions(+), 29 deletions(-) create mode 100644 docs/source/torch/features/layerwise_hf_safetensors_loading_design.md diff --git a/docs/source/torch/features/layerwise_hf_safetensors_loading_design.md b/docs/source/torch/features/layerwise_hf_safetensors_loading_design.md new file mode 100644 index 000000000000..4041bfd1bd7f --- /dev/null +++ b/docs/source/torch/features/layerwise_hf_safetensors_loading_design.md @@ -0,0 +1,358 @@ + + +# Layer-wise Hugging Face Safetensors Loading for DeepSeek V4 + +## Status + +This document describes the experimental implementation on branch +`fix/dsv4-mixed-precision`. The implementation is opt-in, supports the PyTorch +backend, and currently enables semantic-layer loading only for DeepSeek V4. + +The immediate goal is to avoid host OOM while loading the +DeepSeek-V4-Pro-NVFP4 checkpoint on a 952 GiB GB300 node. It is not yet intended +to be a generic streaming-loader API for every model or checkpoint format. + +## Problem statement + +The normal Hugging Face loader materializes the complete checkpoint as one +host-side weight dictionary before the model-specific loader remaps, splits, +concatenates, transforms, and copies weights to the device. For a large NVFP4 +checkpoint, all four tensor-parallel ranks can therefore retain both raw +checkpoint tensors and transformation intermediates at the same time. The +observed aggregate demand exceeded node memory. + +DeepSeek V4 transformations are not all tensor-local. Examples include fused +attention projections, compressor `wkv` plus `wgate`, fused MoE projections, +scale conversion, TP slicing, and MTP name remapping. Streaming one tensor at a +time would break these fusion contracts. The smallest practical unit is a +semantic model layer. + +## Goals + +- Bound live raw checkpoint storage to top-level tensors or one semantic layer. +- Preserve the existing DeepSeek V4 remap, fusion, quantization, TP/EP split, + and module-copy logic. +- Keep all tensors needed by a same-layer fusion in the same bucket, even when + the tensors reside in different safetensors shards. +- Release a bucket only after asynchronous device consumers have completed. +- Leave the default eager loading path unchanged unless explicitly enabled. +- Fail before silently producing a partially or incorrectly loaded model. + +## Non-goals + +- Streaming `.bin` or `.pth` checkpoints. +- Supporting a separate speculative draft checkpoint in the first version. +- Eliminating the Linux page cache populated by checkpoint reads. +- Eliminating every host-side transformation temporary inside one layer. +- Providing a public LLM configuration field. The current switch is an + environment variable to minimize API surface while the design is reviewed. +- Making arbitrary models layer-wise loadable without model-specific support. + +## User-facing activation and compatibility boundary + +Set: + +```text +TRTLLM_HF_LAYERWISE_SAFETENSORS=1 +``` + +Accepted true values are `1`, `true`, `yes`, and `on`, case-insensitively. The +feature remains disabled by default. + +When enabled, `ModelLoader` verifies that `model.load_weights` explicitly +accepts `initial_bucket_loading` before opening the first bucket. Consequently, +an unsupported model fails early instead of receiving partial weights. The +DeepSeek V4 implementation is currently the only intended consumer. + +## Architecture + +```text +HF safetensors index + | + v +HfWeightLoader: build metadata-only bucket plan + | top -> layer 0 -> ... -> layer N-1 -> MTP + v +HfCheckpointLoader: expose capability and bucket iterator + | + v +ModelLoader: initialize mapper once, then iterate buckets + | + v +DeepseekV4WeightLoader: remap and load only the active semantic layer + | + v +CUDA synchronize -> drop bucket -> close safetensors handles + | + v +Run the existing global post-load path once +``` + +The implementation changes four runtime layers: + +1. `HfWeightLoader` plans and materializes semantic buckets. +2. `HfCheckpointLoader` exposes the opt-in capability and iterator. +3. `ModelLoader` owns bucket lifetime and synchronization. +4. `DeepseekV4WeightLoader` scopes existing loading logic to the active bucket. + +## Bucket planning and lifetime + +### Metadata discovery + +`HfWeightLoader.iter_layer_weight_buckets` first discovers `.safetensors` +files. For an ordinary sharded checkpoint, it reads the first +`*.safetensors.index.json` as metadata and builds `(tensor name, shard path)` +entries without loading tensor storage. + +Before model mutation, it verifies that every shard referenced by the index is +present and that the index has a non-empty `weight_map`. This prevents an +incomplete checkpoint directory from being interpreted as a valid partial +model. + +For checkpoints without an index, the loader enumerates keys from each +safetensors file and rejects duplicate keys across files. Consolidated files +are intentionally handled independently of an ordinary sharded index when +`use_consolidated=True`. + +### Semantic classification + +Names are assigned to buckets as follows: + +| Checkpoint name | Bucket | +|---|---| +| `layers..*` | `("layer", n)` | +| `model.layers..*` | `("layer", n)` | +| `mtp..*` | `("mtp", n)` | +| Everything else | `("top", 0)` | + +Buckets are yielded in deterministic order: top-level weights, numeric base +layers, then numeric MTP layers. Classification is independent of shard +boundaries, so all same-layer fusion inputs remain atomic. + +### Storage lifetime + +Only the current bucket is materialized. An `ExitStack` keeps every +`safe_open` handle used by that bucket alive across the `yield`. Tensor storage +therefore remains valid while the model-specific loader consumes it. Handles +close when the caller advances the iterator or unwinds it after an exception. + +After `model.load_weights` returns, `ModelLoader` calls +`torch.cuda.synchronize()`. This is required because non-blocking copies or +backend transforms may still reference mmap-backed source storage. The bucket +reference is deleted only after synchronization. + +The Linux page cache may remain large after handles close. It is reclaimable +file-backed memory and is distinct from Python-owned anonymous tensor storage. + +## Model-loader orchestration + +The eager path remains the default. In layer-wise mode, `ModelLoader`: + +1. Resolves `model.llm_checkpoint_dir` when present, otherwise the ordinary + checkpoint directory. +2. Rejects a separate speculative draft checkpoint. +3. Checks the model loading capability before consuming the iterator. +4. Initializes the existing checkpoint weight mapper once. +5. Calls `model.load_weights(bucket, initial_bucket_loading=True)` for each + bucket. +6. Synchronizes CUDA and releases the bucket. +7. Continues through the existing post-load path after all buckets complete. + +`post_load_weights` is deliberately not run per bucket. Global alias setup, +derived state, and other finalization still occur once after the complete model +has been populated. + +An error after earlier buckets have loaded leaves the in-process model +partially mutated, but startup fails and the model is never published. The +implementation does not attempt transactional rollback. + +## DeepSeek V4 model behavior + +`DeepseekV4WeightLoader.load_weights` accepts +`initial_bucket_loading=False` by default. When true, it derives one +`active_layer` from checkpoint keys and validates that the bucket is either: + +- top-level weights only, or +- exactly one base/MTP semantic layer with no top-level weights mixed in. + +The existing raw-key remap still performs the same fusion and quantization +logic. Named-parameter default synthesis and named-module traversal are scoped +to the active layer: + +- A top bucket skips every `model.layers.*` module. +- A base-layer bucket visits only the matching runtime layer. +- An MTP bucket visits runtime MTP replicas and reuses the existing modulo + mapping back to checkpoint MTP layers. +- Synthesized indexer defaults are created only for the active base layer. + +No known DeepSeek V4 checkpoint fusion crosses a semantic layer boundary. The +attention, compressor, MoE, scale, and TP transformations used by this loader +consume tensors from one layer or from the top-level bucket. + +The same patch also fixes routed-expert scale-layout detection for an isolated +`mtp.0.*` bucket. Without recognizing the MTP prefix, MTP-only loading could +select the wrong MXFP4/NVFP4 interpretation. This is a correctness fix rather +than an intentional numerical change. + +## Numerical-equivalence argument + +Layer-wise loading changes scheduling and lifetime, not the mathematical +operations. It uses the same checkpoint tensors, mapper, remap functions, +fusion order, TP/EP slicing, scale transforms, destination dtypes, and module +copy routines as eager loading. CUDA synchronization adds an ordering barrier +but does not change values. + +The equivalence argument depends on these invariants: + +- Every fusion input is in the same bucket. +- The model-specific loader never reads a different base layer while loading + the active layer. +- Global post-load operations run once after all buckets. +- Missing or malformed buckets fail loudly rather than being silently skipped. +- MTP runtime replicas map to the same checkpoint MTP layer as eager loading. + +These invariants require unit tests and an end-to-end accuracy gate; successful +model startup alone is not sufficient proof of numerical equivalence. + +## Unit-test design + +### Tests already implemented + +`tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py` contains: + +- `test_layerwise_safetensors_keeps_cross_shard_layer_atomic`: verifies that + same-layer tensors split across shards are yielded together and that order is + top, base layers, MTP. +- `test_layerwise_safetensors_without_index_discovers_keys`: verifies fallback + key discovery and support for `model.layers.` names. +- `test_layerwise_safetensors_rejects_missing_index_shard`: verifies that all + index-referenced shards are validated before the first bucket is yielded. + +`tests/unittest/_torch/modeling/test_modeling_deepseekv4.py` contains: + +- `test_deepseek_v4_weight_remap_for_mtp_mxfp4_routed_experts`: verifies MTP + routed-expert scale-layout detection and remapped dtypes. + +These tests are CPU-only in intent. Full test-file collection still requires a +Linux TensorRT-LLM development environment because module imports depend on the +TensorRT-LLM runtime stack. + +### Required tests before upstreaming + +#### HF bucket iterator + +- Environment parsing: false by default and all documented true spellings. +- Natural numeric order: layer 2 precedes layer 10 regardless of index order. +- Empty index rejection. +- Duplicate-key rejection for no-index multi-file checkpoints. +- Consolidated selection does not accidentally consume an ordinary sharded + index or non-consolidated files. +- Non-safetensors input fails with the documented error. +- Handle lifetime: handles remain open during bucket consumption and close when + advancing, on normal completion, and on consumer exception. +- Tensor contents exactly match the source files for cross-shard buckets. + +#### ModelLoader orchestration + +Use fake checkpoint loaders and a small fake model to avoid GPU/model weights: + +- Mapper initialization occurs once and before the first bucket load. +- One model load call occurs per bucket in iterator order. +- `initial_bucket_loading=True` is passed only on the layer-wise path. +- CUDA synchronization occurs once after every successful bucket load and + before the iterator advances. +- An unsupported model fails before the iterator is consumed or model state is + changed. +- A separate draft checkpoint fails before bucket consumption. +- `llm_checkpoint_dir` is honored. +- Consumer exceptions close current handles and prevent later buckets and + global post-load publication. +- The ordinary eager, preloaded, MX, GMS, and draft-loading paths retain their + existing call ordering when the environment variable is unset. + +#### DeepSeek V4 bucket scoping + +- Reject a bucket containing two base layers. +- Reject a bucket mixing top-level and layer weights. +- A top bucket never visits layer modules. +- A base-layer bucket visits only its matching runtime layer. +- Synthesized indexer defaults are limited to the active layer. +- An MTP bucket reaches every intended runtime MTP replica and preserves modulo + checkpoint mapping. +- Missing members of each fused group fail with the expected key/error rather + than silently loading a partial fused tensor. +- A small synthetic model loaded eagerly and bucket-by-bucket produces an + identical `state_dict`. Device-only transforms may be mocked for a CPU test, + followed by a real-GPU equivalence test. + +### Accuracy and integration gates + +The following tests complement, but do not replace, unit tests: + +1. Load the real checkpoint with TEP4 in CTX-only and GEN-only configurations. +2. Confirm all ranks traverse top, layers 0 through 60, and MTP without OOM. +3. Compare eager and layer-wise loading on a smaller compatible fixture using + exact state-dict equality after post-load. +4. With fixed prompt, seed, and greedy decoding, compare generated token IDs. +5. If logits are available, compare first-token logits or top-k logits under an + agreed tolerance and check for NaN/Inf. +6. Record per-task `MaxRSS`, job cgroup `memory.peak`, current anonymous memory, + and file cache separately. +7. Run CTX/GEN throughput smoke tests to catch runtime-only fusion, shape, or + communication failures. + +## Current validation evidence + +As of 2026-07-14: + +- Static checkpoint-index validation found 285,660 tensors in 63 semantic + buckets: one top bucket, 61 base layers, and one MTP bucket. +- `py_compile` passed for the modified source and test files. +- `git diff --check` passed. +- Real TEP4 CTX-only and disaggregated CTX/GEN workers loaded every bucket and + reached server/benchmark startup without host OOM. +- The observed node cgroup peak was approximately 874-878 GiB on a 952 GiB + node. After load, anonymous memory dropped substantially; most remaining + cgroup memory was reclaimable checkpoint page cache. +- The focused pytest cases have not yet run in a complete development image; + login and benchmark images used during bring-up did not include pytest and + compute nodes could not fetch it from PyPI. + +## Known limitations and review focus + +- All ranks still read the full checkpoint namespace; this is layer-bounded, + not rank-presharded I/O. +- Peak host memory is now below node capacity but remains high. Per-rank + historical `MaxRSS` was approximately 227-235 GiB in the real checkpoint + test. Review should separate anonymous transformation storage from mmap/file + cache accounting. +- Bucket classification is name-pattern based and currently understands top, + base layers, and MTP. New DeepSeek V4 checkpoint namespaces must be audited. +- Capability discovery uses Python signature inspection. Decorated loaders or + future keyword-only signatures should be considered during review. +- The iterator API currently lives on the HF concrete loader rather than a + generic base interface. +- There is no fused-group preflight manifest. Completeness currently relies on + index validation plus strict failures in existing model-specific transforms. +- The opt-in is process-global. It intentionally fails early for unsupported + models, but a future public configuration may need per-model scope. +- Page cache remains charged to the job cgroup until reclaimed by the kernel. + This is expected and must not be confused with live Python-owned weights. + +## Reviewer checklist + +- Verify eager behavior is byte-for-byte unchanged when the environment + variable is unset. +- Audit every DeepSeek V4 fusion for cross-layer dependencies. +- Audit bucket release ordering around asynchronous CUDA work. +- Verify top, base-layer, and MTP name classification against all supported + checkpoint naming variants. +- Check MTP replica routing for `num_nextn_predict_layers > 1`. +- Check exception behavior for missing shards, missing fused members, and + mid-stream consumer failure. +- Confirm no post-load transform is incorrectly executed once per bucket. +- Require eager-versus-layer-wise state/output equivalence before generalizing + the feature to additional models. diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py index d00a24114745..3fc9360bd2a2 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from typing import Optional from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \ @@ -73,3 +76,12 @@ def config_loader(self) -> Optional[BaseConfigLoader]: @property def checkpoint_format(self) -> str: return self._checkpoint_format + + def is_layerwise_loading_enabled(self) -> bool: + return (hasattr(self._weight_loader, + "is_layerwise_safetensors_enabled") + and self._weight_loader.is_layerwise_safetensors_enabled()) + + def iter_layer_weight_buckets(self, checkpoint_dir: str, **kwargs): + 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..abecace3dcca 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 contextlib import ExitStack from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor -from typing import Any, List +from typing import Any, Iterator, List, Tuple import psutil import safetensors @@ -36,6 +39,7 @@ from tensorrt_llm.mapping import Mapping _WEIGHT_CACHE_ENV = "TRTLLM_HF_WEIGHT_CACHE" +_LAYERWISE_SAFETENSORS_ENV = "TRTLLM_HF_LAYERWISE_SAFETENSORS" _WEIGHT_CACHE_MAX_ENTRIES_ENV = "TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES" # Default to a single cached checkpoint: each entry pins a full copy of the # raw weights in CPU RAM, so callers wanting cross-model caching must opt in @@ -54,6 +58,104 @@ class HfWeightLoader(BaseWeightLoader): Loads weights from SafeTensors/bin/pth files. """ + @staticmethod + def is_layerwise_safetensors_enabled() -> bool: + return os.environ.get(_LAYERWISE_SAFETENSORS_ENV, + "0").lower() in ("1", "true", "yes", "on") + + @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. + """ + 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..105227b66937 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: @@ -314,11 +315,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 +444,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 +527,40 @@ 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: List[str] = [], + initial_bucket_loading: bool = False): # 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 +581,10 @@ 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 +850,29 @@ 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]) + if active_layer >= self.config.num_hidden_layers: + # Runtime may instantiate multiple MTP replicas. + # The existing name rewrite below maps each replica + # back to the checkpoint MTP layer modulo the number + # of checkpoint next-token layers. + checkpoint_mtp_idx = active_layer - self.config.num_hidden_layers + runtime_mtp_idx = module_layer - self.config.num_hidden_layers + is_target_module = ( + runtime_mtp_idx % self.config.num_nextn_predict_layers + == checkpoint_mtp_idx + ) + else: + is_target_module = module_layer == active_layer + if not is_target_module: + continue names = name.split(".") parent_module_name = ".".join(names[:-1]) @@ -2564,9 +2621,12 @@ 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 7596c2a68291..5312359619e9 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -558,23 +558,52 @@ def init_meta_tensor(t: torch.Tensor): self._is_mx_staged_receiver_allowlisted(model) and not loads_draft_weights) - 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 = ( + hasattr(checkpoint_loader, 'is_layerwise_loading_enabled') + and checkpoint_loader.is_layerwise_loading_enabled()) + if layerwise_loading: + if loads_draft_weights: + raise RuntimeError( + "HF layer-wise loading does not yet support a " + "separate speculative draft checkpoint.") + load_parameters = inspect.signature( + model.load_weights).parameters + if "initial_bucket_loading" not in load_parameters: + raise RuntimeError( + "TRTLLM_HF_LAYERWISE_SAFETENSORS 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( ): @@ -1300,9 +1329,10 @@ def _call_load_weights(self, load_method: Callable, weights, weight_mapper, - allow_partial_loading: bool = False): + allow_partial_loading: bool = False, + initial_bucket_loading: bool = False): """Calls the model's weight loading method with the correct arguments.""" - args = inspect.getfullargspec(load_method).args + args = inspect.signature(load_method).parameters kargs = {} if "weight_mapper" in args: kargs["weight_mapper"] = weight_mapper @@ -1310,4 +1340,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/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 106ca0e5c59d..874a2d155253 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -187,6 +187,23 @@ 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_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..52d56057db77 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,12 @@ +# 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 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 +17,82 @@ 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())) + + @pytest.fixture(autouse=True) def clean_weight_cache(): HfWeightLoader._clear_weight_cache() From 2217f86151693a1774e73339fd533825c0720fcb Mon Sep 17 00:00:00 2001 From: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com> Date: Wed, 15 Jul 2026 01:26:44 -0700 Subject: [PATCH 2/4] [None][test] document and validate layer-wise loading Signed-off-by: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com> --- ...layerwise_hf_safetensors_loading_design.md | 53 ++--- .../checkpoints/hf/checkpoint_loader.py | 24 +- .../models/checkpoints/hf/weight_loader.py | 32 ++- .../_torch/models/modeling_deepseekv4.py | 40 +++- .../_torch/pyexecutor/model_loader.py | 30 ++- .../executor/test_model_loader_layerwise.py | 189 +++++++++++++++ .../modeling/test_modeling_deepseekv4.py | 56 ++++- .../checkpoints/hf/test_weight_loader.py | 225 ++++++++++++++++-- 8 files changed, 549 insertions(+), 100 deletions(-) create mode 100644 tests/unittest/_torch/executor/test_model_loader_layerwise.py diff --git a/docs/source/torch/features/layerwise_hf_safetensors_loading_design.md b/docs/source/torch/features/layerwise_hf_safetensors_loading_design.md index 4041bfd1bd7f..318a397d71d7 100644 --- a/docs/source/torch/features/layerwise_hf_safetensors_loading_design.md +++ b/docs/source/torch/features/layerwise_hf_safetensors_loading_design.md @@ -230,53 +230,31 @@ model startup alone is not sufficient proof of numerical equivalence. key discovery and support for `model.layers.` names. - `test_layerwise_safetensors_rejects_missing_index_shard`: verifies that all index-referenced shards are validated before the first bucket is yielded. +- Parameterized environment tests cover the disabled default and all accepted + true spellings. +- Numeric-order, empty-index, duplicate-key, consolidated-selection, tensor + value, and handle-lifetime tests cover metadata and resource boundaries. + +`tests/unittest/_torch/executor/test_model_loader_layerwise.py` contains fake +model and checkpoint-loader orchestration tests. They verify mapper ordering, +one call and one CUDA barrier per successful bucket, keyword-only capability +detection, `llm_checkpoint_dir`, iterator cleanup after consumer failure, and +early rejection of unsupported models and separate draft checkpoints. `tests/unittest/_torch/modeling/test_modeling_deepseekv4.py` contains: - `test_deepseek_v4_weight_remap_for_mtp_mxfp4_routed_experts`: verifies MTP routed-expert scale-layout detection and remapped dtypes. +- Multi-MTP remapping and malformed semantic-bucket rejection tests. These tests are CPU-only in intent. Full test-file collection still requires a Linux TensorRT-LLM development environment because module imports depend on the TensorRT-LLM runtime stack. -### Required tests before upstreaming - -#### HF bucket iterator - -- Environment parsing: false by default and all documented true spellings. -- Natural numeric order: layer 2 precedes layer 10 regardless of index order. -- Empty index rejection. -- Duplicate-key rejection for no-index multi-file checkpoints. -- Consolidated selection does not accidentally consume an ordinary sharded - index or non-consolidated files. -- Non-safetensors input fails with the documented error. -- Handle lifetime: handles remain open during bucket consumption and close when - advancing, on normal completion, and on consumer exception. -- Tensor contents exactly match the source files for cross-shard buckets. - -#### ModelLoader orchestration - -Use fake checkpoint loaders and a small fake model to avoid GPU/model weights: - -- Mapper initialization occurs once and before the first bucket load. -- One model load call occurs per bucket in iterator order. -- `initial_bucket_loading=True` is passed only on the layer-wise path. -- CUDA synchronization occurs once after every successful bucket load and - before the iterator advances. -- An unsupported model fails before the iterator is consumed or model state is - changed. -- A separate draft checkpoint fails before bucket consumption. -- `llm_checkpoint_dir` is honored. -- Consumer exceptions close current handles and prevent later buckets and - global post-load publication. -- The ordinary eager, preloaded, MX, GMS, and draft-loading paths retain their - existing call ordering when the environment variable is unset. +### Remaining tests before upstreaming #### DeepSeek V4 bucket scoping -- Reject a bucket containing two base layers. -- Reject a bucket mixing top-level and layer weights. - A top bucket never visits layer modules. - A base-layer bucket visits only its matching runtime layer. - Synthesized indexer defaults are limited to the active layer. @@ -317,9 +295,10 @@ As of 2026-07-14: - The observed node cgroup peak was approximately 874-878 GiB on a 952 GiB node. After load, anonymous memory dropped substantially; most remaining cgroup memory was reclaimable checkpoint page cache. -- The focused pytest cases have not yet run in a complete development image; - login and benchmark images used during bring-up did not include pytest and - compute nodes could not fetch it from PyPI. +- In the branch-matched development image on a Slurm GPU node, the complete + DeepSeek V4 modeling file passed (41 tests). The complete HF weight-loader, + layer-wise ModelLoader, and existing MX/GMS ModelLoader files passed together + (62 tests). ## Known limitations and review focus diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py index 3fc9360bd2a2..c488c6e794de 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py @@ -1,14 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from typing import Optional +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 \ @@ -78,10 +78,22 @@ def checkpoint_format(self) -> str: return self._checkpoint_format def is_layerwise_loading_enabled(self) -> bool: - return (hasattr(self._weight_loader, - "is_layerwise_safetensors_enabled") + """Return whether opt-in HF safetensors layer-wise loading is enabled.""" + return (hasattr(self._weight_loader, "is_layerwise_safetensors_enabled") and self._weight_loader.is_layerwise_safetensors_enabled()) - def iter_layer_weight_buckets(self, checkpoint_dir: str, **kwargs): + 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 abecace3dcca..97976cc27bcb 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -18,9 +18,9 @@ import os import re import threading -from contextlib import ExitStack from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack from typing import Any, Iterator, List, Tuple import psutil @@ -60,6 +60,7 @@ class HfWeightLoader(BaseWeightLoader): @staticmethod def is_layerwise_safetensors_enabled() -> bool: + """Return whether semantic-layer safetensors loading is enabled.""" return os.environ.get(_LAYERWISE_SAFETENSORS_ENV, "0").lower() in ("1", "true", "yes", "on") @@ -74,18 +75,35 @@ def _layer_bucket_name(weight_name: str) -> Tuple[str, int]: 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]: + 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") diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 105227b66937..740e6e7ac567 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -527,10 +527,26 @@ 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] = [], - initial_bucket_loading: bool = False): + 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 @@ -539,7 +555,8 @@ def load_weights(self, raw_checkpoint = any( k in ("embed.weight", "head.weight", "norm.weight") or k.startswith(("layers.", "mtp.", "hc_head_")) - for k in weights) + for k in weights + ) active_layer = None if initial_bucket_loading: layer_indices = set() @@ -557,7 +574,8 @@ def load_weights(self, 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}.") + f"has_top_level={has_top_level}." + ) active_layer = next(iter(layer_indices), None) if raw_checkpoint: @@ -583,7 +601,8 @@ def load_weights(self, 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}."): + f"model.layers.{active_layer}." + ): continue if pname in weights: continue @@ -2621,12 +2640,9 @@ def forward( **kwargs, ) - def load_weights(self, - weights: Dict, - initial_bucket_loading: bool = False): + def load_weights(self, weights: Dict, initial_bucket_loading: bool = False): weight_loader = DeepseekV4WeightLoader(self) - weight_loader.load_weights( - weights, initial_bucket_loading=initial_bucket_loading) + 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 5312359619e9..4d95bdde7f0d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -561,13 +561,17 @@ def init_meta_tensor(t: torch.Tensor): model_checkpoint_dir = (model.llm_checkpoint_dir if hasattr( model, 'llm_checkpoint_dir') else checkpoint_dir) layerwise_loading = ( - hasattr(checkpoint_loader, 'is_layerwise_loading_enabled') + checkpoint_loader.checkpoint_format == "HF" and hasattr( + checkpoint_loader, 'is_layerwise_loading_enabled') and checkpoint_loader.is_layerwise_loading_enabled()) 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: @@ -582,11 +586,10 @@ def init_meta_tensor(t: torch.Tensor): ) 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) + 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() @@ -1331,7 +1334,20 @@ def _call_load_weights(self, weight_mapper, allow_partial_loading: bool = False, initial_bucket_loading: bool = False): - """Calls the model's weight loading method with the correct arguments.""" + """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: 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..ff42881e52c7 --- /dev/null +++ b/tests/unittest/_torch/executor/test_model_loader_layerwise.py @@ -0,0 +1,189 @@ +# 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): + loader = ModelLoader( + llm_args=SimpleNamespace(load_format=LoadFormat.AUTO), + 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" + checkpoint_loader.is_layerwise_loading_enabled.return_value = True + 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_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 874a2d155253..14f893072466 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,6 +32,7 @@ DeepseekV4ForCausalLM, DeepseekV4Gate, DeepseekV4MTP, + DeepseekV4WeightLoader, _copy_deepseek_v4_fused_a_weight_scale, _deepseek_v4_pos_embd_params, _remap_deepseek_v4_checkpoint_keys, @@ -189,19 +191,53 @@ def test_deepseek_v4_weight_remap_for_fp8_routed_experts(): 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), + "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) + 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 + 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) def test_deepseek_v4_fused_a_weight_scale_rebuilds_fp8_shape(): 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 52d56057db77..82e2a39add9f 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py @@ -5,6 +5,7 @@ from unittest import mock import pytest +import safetensors import torch from safetensors.torch import save_file @@ -24,13 +25,17 @@ def test_layerwise_safetensors_keeps_cross_shard_layer_atomic(tmp_path): { "embed.weight": torch.tensor([1.0]), "layers.0.attn.compressor.wkv.weight": torch.tensor([2.0]), - }, str(shard1)) + }, + 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)) + }, + str(shard2), + ) weight_map = { "embed.weight": shard1.name, "layers.0.attn.compressor.wkv.weight": shard1.name, @@ -39,12 +44,12 @@ def test_layerwise_safetensors_keeps_cross_shard_layer_atomic(tmp_path): "mtp.0.enorm.weight": shard2.name, } (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"weight_map": weight_map}), encoding="utf-8") + 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()): + for bucket in loader.iter_layer_weight_buckets(str(tmp_path), mapping=Mapping()): buckets.append(set(bucket.keys())) assert buckets == [ @@ -59,38 +64,216 @@ def test_layerwise_safetensors_keeps_cross_shard_layer_atomic(tmp_path): 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")) + 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()) + 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"}] + 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", + json.dumps( + { + "weight_map": { + "layers.0.foo.weight": shard1.name, + "layers.1.foo.weight": "model-00002-of-00002.safetensors", + } } - }), - encoding="utf-8") + ), + encoding="utf-8", + ) loader = HfWeightLoader() with pytest.raises(RuntimeError, match="missing checkpoint files"): + next(loader.iter_layer_weight_buckets(str(tmp_path), mapping=Mapping())) + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "Yes", "on", "ON"]) +def test_layerwise_safetensors_environment_true_values(monkeypatch, value): + monkeypatch.setenv("TRTLLM_HF_LAYERWISE_SAFETENSORS", value) + + assert HfWeightLoader.is_layerwise_safetensors_enabled() + + +@pytest.mark.parametrize("value", [None, "0", "false", "no", "off", "unexpected"]) +def test_layerwise_safetensors_environment_false_values(monkeypatch, value): + if value is None: + monkeypatch.delenv("TRTLLM_HF_LAYERWISE_SAFETENSORS", raising=False) + else: + monkeypatch.setenv("TRTLLM_HF_LAYERWISE_SAFETENSORS", value) + + assert not HfWeightLoader.is_layerwise_safetensors_enabled() + + +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( - loader.iter_layer_weight_buckets(str(tmp_path), - mapping=Mapping())) + 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) From ab95fcb80a5959aa594093e787c044f5d4846d94 Mon Sep 17 00:00:00 2001 From: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com> Date: Wed, 15 Jul 2026 02:00:32 -0700 Subject: [PATCH 3/4] [None][feat] expose layer-wise loading as prototype config Signed-off-by: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com> --- docs/source/features/checkpoint-loading.md | 38 ++ ...layerwise_hf_safetensors_loading_design.md | 337 ------------------ .../checkpoints/hf/checkpoint_loader.py | 5 - .../models/checkpoints/hf/weight_loader.py | 7 - .../_torch/pyexecutor/model_loader.py | 10 +- tensorrt_llm/llmapi/llm_args.py | 9 + .../usage/llm_args_golden_manifest.json | 7 + .../executor/test_model_loader_layerwise.py | 25 +- .../checkpoints/hf/test_weight_loader.py | 17 - .../api_stability/references/llm.yaml | 4 + .../llmapi/test_layerwise_loading_args.py | 32 ++ 11 files changed, 117 insertions(+), 374 deletions(-) delete mode 100644 docs/source/torch/features/layerwise_hf_safetensors_loading_design.md create mode 100644 tests/unittest/llmapi/test_layerwise_loading_args.py 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/docs/source/torch/features/layerwise_hf_safetensors_loading_design.md b/docs/source/torch/features/layerwise_hf_safetensors_loading_design.md deleted file mode 100644 index 318a397d71d7..000000000000 --- a/docs/source/torch/features/layerwise_hf_safetensors_loading_design.md +++ /dev/null @@ -1,337 +0,0 @@ - - -# Layer-wise Hugging Face Safetensors Loading for DeepSeek V4 - -## Status - -This document describes the experimental implementation on branch -`fix/dsv4-mixed-precision`. The implementation is opt-in, supports the PyTorch -backend, and currently enables semantic-layer loading only for DeepSeek V4. - -The immediate goal is to avoid host OOM while loading the -DeepSeek-V4-Pro-NVFP4 checkpoint on a 952 GiB GB300 node. It is not yet intended -to be a generic streaming-loader API for every model or checkpoint format. - -## Problem statement - -The normal Hugging Face loader materializes the complete checkpoint as one -host-side weight dictionary before the model-specific loader remaps, splits, -concatenates, transforms, and copies weights to the device. For a large NVFP4 -checkpoint, all four tensor-parallel ranks can therefore retain both raw -checkpoint tensors and transformation intermediates at the same time. The -observed aggregate demand exceeded node memory. - -DeepSeek V4 transformations are not all tensor-local. Examples include fused -attention projections, compressor `wkv` plus `wgate`, fused MoE projections, -scale conversion, TP slicing, and MTP name remapping. Streaming one tensor at a -time would break these fusion contracts. The smallest practical unit is a -semantic model layer. - -## Goals - -- Bound live raw checkpoint storage to top-level tensors or one semantic layer. -- Preserve the existing DeepSeek V4 remap, fusion, quantization, TP/EP split, - and module-copy logic. -- Keep all tensors needed by a same-layer fusion in the same bucket, even when - the tensors reside in different safetensors shards. -- Release a bucket only after asynchronous device consumers have completed. -- Leave the default eager loading path unchanged unless explicitly enabled. -- Fail before silently producing a partially or incorrectly loaded model. - -## Non-goals - -- Streaming `.bin` or `.pth` checkpoints. -- Supporting a separate speculative draft checkpoint in the first version. -- Eliminating the Linux page cache populated by checkpoint reads. -- Eliminating every host-side transformation temporary inside one layer. -- Providing a public LLM configuration field. The current switch is an - environment variable to minimize API surface while the design is reviewed. -- Making arbitrary models layer-wise loadable without model-specific support. - -## User-facing activation and compatibility boundary - -Set: - -```text -TRTLLM_HF_LAYERWISE_SAFETENSORS=1 -``` - -Accepted true values are `1`, `true`, `yes`, and `on`, case-insensitively. The -feature remains disabled by default. - -When enabled, `ModelLoader` verifies that `model.load_weights` explicitly -accepts `initial_bucket_loading` before opening the first bucket. Consequently, -an unsupported model fails early instead of receiving partial weights. The -DeepSeek V4 implementation is currently the only intended consumer. - -## Architecture - -```text -HF safetensors index - | - v -HfWeightLoader: build metadata-only bucket plan - | top -> layer 0 -> ... -> layer N-1 -> MTP - v -HfCheckpointLoader: expose capability and bucket iterator - | - v -ModelLoader: initialize mapper once, then iterate buckets - | - v -DeepseekV4WeightLoader: remap and load only the active semantic layer - | - v -CUDA synchronize -> drop bucket -> close safetensors handles - | - v -Run the existing global post-load path once -``` - -The implementation changes four runtime layers: - -1. `HfWeightLoader` plans and materializes semantic buckets. -2. `HfCheckpointLoader` exposes the opt-in capability and iterator. -3. `ModelLoader` owns bucket lifetime and synchronization. -4. `DeepseekV4WeightLoader` scopes existing loading logic to the active bucket. - -## Bucket planning and lifetime - -### Metadata discovery - -`HfWeightLoader.iter_layer_weight_buckets` first discovers `.safetensors` -files. For an ordinary sharded checkpoint, it reads the first -`*.safetensors.index.json` as metadata and builds `(tensor name, shard path)` -entries without loading tensor storage. - -Before model mutation, it verifies that every shard referenced by the index is -present and that the index has a non-empty `weight_map`. This prevents an -incomplete checkpoint directory from being interpreted as a valid partial -model. - -For checkpoints without an index, the loader enumerates keys from each -safetensors file and rejects duplicate keys across files. Consolidated files -are intentionally handled independently of an ordinary sharded index when -`use_consolidated=True`. - -### Semantic classification - -Names are assigned to buckets as follows: - -| Checkpoint name | Bucket | -|---|---| -| `layers..*` | `("layer", n)` | -| `model.layers..*` | `("layer", n)` | -| `mtp..*` | `("mtp", n)` | -| Everything else | `("top", 0)` | - -Buckets are yielded in deterministic order: top-level weights, numeric base -layers, then numeric MTP layers. Classification is independent of shard -boundaries, so all same-layer fusion inputs remain atomic. - -### Storage lifetime - -Only the current bucket is materialized. An `ExitStack` keeps every -`safe_open` handle used by that bucket alive across the `yield`. Tensor storage -therefore remains valid while the model-specific loader consumes it. Handles -close when the caller advances the iterator or unwinds it after an exception. - -After `model.load_weights` returns, `ModelLoader` calls -`torch.cuda.synchronize()`. This is required because non-blocking copies or -backend transforms may still reference mmap-backed source storage. The bucket -reference is deleted only after synchronization. - -The Linux page cache may remain large after handles close. It is reclaimable -file-backed memory and is distinct from Python-owned anonymous tensor storage. - -## Model-loader orchestration - -The eager path remains the default. In layer-wise mode, `ModelLoader`: - -1. Resolves `model.llm_checkpoint_dir` when present, otherwise the ordinary - checkpoint directory. -2. Rejects a separate speculative draft checkpoint. -3. Checks the model loading capability before consuming the iterator. -4. Initializes the existing checkpoint weight mapper once. -5. Calls `model.load_weights(bucket, initial_bucket_loading=True)` for each - bucket. -6. Synchronizes CUDA and releases the bucket. -7. Continues through the existing post-load path after all buckets complete. - -`post_load_weights` is deliberately not run per bucket. Global alias setup, -derived state, and other finalization still occur once after the complete model -has been populated. - -An error after earlier buckets have loaded leaves the in-process model -partially mutated, but startup fails and the model is never published. The -implementation does not attempt transactional rollback. - -## DeepSeek V4 model behavior - -`DeepseekV4WeightLoader.load_weights` accepts -`initial_bucket_loading=False` by default. When true, it derives one -`active_layer` from checkpoint keys and validates that the bucket is either: - -- top-level weights only, or -- exactly one base/MTP semantic layer with no top-level weights mixed in. - -The existing raw-key remap still performs the same fusion and quantization -logic. Named-parameter default synthesis and named-module traversal are scoped -to the active layer: - -- A top bucket skips every `model.layers.*` module. -- A base-layer bucket visits only the matching runtime layer. -- An MTP bucket visits runtime MTP replicas and reuses the existing modulo - mapping back to checkpoint MTP layers. -- Synthesized indexer defaults are created only for the active base layer. - -No known DeepSeek V4 checkpoint fusion crosses a semantic layer boundary. The -attention, compressor, MoE, scale, and TP transformations used by this loader -consume tensors from one layer or from the top-level bucket. - -The same patch also fixes routed-expert scale-layout detection for an isolated -`mtp.0.*` bucket. Without recognizing the MTP prefix, MTP-only loading could -select the wrong MXFP4/NVFP4 interpretation. This is a correctness fix rather -than an intentional numerical change. - -## Numerical-equivalence argument - -Layer-wise loading changes scheduling and lifetime, not the mathematical -operations. It uses the same checkpoint tensors, mapper, remap functions, -fusion order, TP/EP slicing, scale transforms, destination dtypes, and module -copy routines as eager loading. CUDA synchronization adds an ordering barrier -but does not change values. - -The equivalence argument depends on these invariants: - -- Every fusion input is in the same bucket. -- The model-specific loader never reads a different base layer while loading - the active layer. -- Global post-load operations run once after all buckets. -- Missing or malformed buckets fail loudly rather than being silently skipped. -- MTP runtime replicas map to the same checkpoint MTP layer as eager loading. - -These invariants require unit tests and an end-to-end accuracy gate; successful -model startup alone is not sufficient proof of numerical equivalence. - -## Unit-test design - -### Tests already implemented - -`tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py` contains: - -- `test_layerwise_safetensors_keeps_cross_shard_layer_atomic`: verifies that - same-layer tensors split across shards are yielded together and that order is - top, base layers, MTP. -- `test_layerwise_safetensors_without_index_discovers_keys`: verifies fallback - key discovery and support for `model.layers.` names. -- `test_layerwise_safetensors_rejects_missing_index_shard`: verifies that all - index-referenced shards are validated before the first bucket is yielded. -- Parameterized environment tests cover the disabled default and all accepted - true spellings. -- Numeric-order, empty-index, duplicate-key, consolidated-selection, tensor - value, and handle-lifetime tests cover metadata and resource boundaries. - -`tests/unittest/_torch/executor/test_model_loader_layerwise.py` contains fake -model and checkpoint-loader orchestration tests. They verify mapper ordering, -one call and one CUDA barrier per successful bucket, keyword-only capability -detection, `llm_checkpoint_dir`, iterator cleanup after consumer failure, and -early rejection of unsupported models and separate draft checkpoints. - -`tests/unittest/_torch/modeling/test_modeling_deepseekv4.py` contains: - -- `test_deepseek_v4_weight_remap_for_mtp_mxfp4_routed_experts`: verifies MTP - routed-expert scale-layout detection and remapped dtypes. -- Multi-MTP remapping and malformed semantic-bucket rejection tests. - -These tests are CPU-only in intent. Full test-file collection still requires a -Linux TensorRT-LLM development environment because module imports depend on the -TensorRT-LLM runtime stack. - -### Remaining tests before upstreaming - -#### DeepSeek V4 bucket scoping - -- A top bucket never visits layer modules. -- A base-layer bucket visits only its matching runtime layer. -- Synthesized indexer defaults are limited to the active layer. -- An MTP bucket reaches every intended runtime MTP replica and preserves modulo - checkpoint mapping. -- Missing members of each fused group fail with the expected key/error rather - than silently loading a partial fused tensor. -- A small synthetic model loaded eagerly and bucket-by-bucket produces an - identical `state_dict`. Device-only transforms may be mocked for a CPU test, - followed by a real-GPU equivalence test. - -### Accuracy and integration gates - -The following tests complement, but do not replace, unit tests: - -1. Load the real checkpoint with TEP4 in CTX-only and GEN-only configurations. -2. Confirm all ranks traverse top, layers 0 through 60, and MTP without OOM. -3. Compare eager and layer-wise loading on a smaller compatible fixture using - exact state-dict equality after post-load. -4. With fixed prompt, seed, and greedy decoding, compare generated token IDs. -5. If logits are available, compare first-token logits or top-k logits under an - agreed tolerance and check for NaN/Inf. -6. Record per-task `MaxRSS`, job cgroup `memory.peak`, current anonymous memory, - and file cache separately. -7. Run CTX/GEN throughput smoke tests to catch runtime-only fusion, shape, or - communication failures. - -## Current validation evidence - -As of 2026-07-14: - -- Static checkpoint-index validation found 285,660 tensors in 63 semantic - buckets: one top bucket, 61 base layers, and one MTP bucket. -- `py_compile` passed for the modified source and test files. -- `git diff --check` passed. -- Real TEP4 CTX-only and disaggregated CTX/GEN workers loaded every bucket and - reached server/benchmark startup without host OOM. -- The observed node cgroup peak was approximately 874-878 GiB on a 952 GiB - node. After load, anonymous memory dropped substantially; most remaining - cgroup memory was reclaimable checkpoint page cache. -- In the branch-matched development image on a Slurm GPU node, the complete - DeepSeek V4 modeling file passed (41 tests). The complete HF weight-loader, - layer-wise ModelLoader, and existing MX/GMS ModelLoader files passed together - (62 tests). - -## Known limitations and review focus - -- All ranks still read the full checkpoint namespace; this is layer-bounded, - not rank-presharded I/O. -- Peak host memory is now below node capacity but remains high. Per-rank - historical `MaxRSS` was approximately 227-235 GiB in the real checkpoint - test. Review should separate anonymous transformation storage from mmap/file - cache accounting. -- Bucket classification is name-pattern based and currently understands top, - base layers, and MTP. New DeepSeek V4 checkpoint namespaces must be audited. -- Capability discovery uses Python signature inspection. Decorated loaders or - future keyword-only signatures should be considered during review. -- The iterator API currently lives on the HF concrete loader rather than a - generic base interface. -- There is no fused-group preflight manifest. Completeness currently relies on - index validation plus strict failures in existing model-specific transforms. -- The opt-in is process-global. It intentionally fails early for unsupported - models, but a future public configuration may need per-model scope. -- Page cache remains charged to the job cgroup until reclaimed by the kernel. - This is expected and must not be confused with live Python-owned weights. - -## Reviewer checklist - -- Verify eager behavior is byte-for-byte unchanged when the environment - variable is unset. -- Audit every DeepSeek V4 fusion for cross-layer dependencies. -- Audit bucket release ordering around asynchronous CUDA work. -- Verify top, base-layer, and MTP name classification against all supported - checkpoint naming variants. -- Check MTP replica routing for `num_nextn_predict_layers > 1`. -- Check exception behavior for missing shards, missing fused members, and - mid-stream consumer failure. -- Confirm no post-load transform is incorrectly executed once per bucket. -- Require eager-versus-layer-wise state/output equivalence before generalizing - the feature to additional models. diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py index c488c6e794de..04f6dce612b4 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py @@ -77,11 +77,6 @@ def config_loader(self) -> Optional[BaseConfigLoader]: def checkpoint_format(self) -> str: return self._checkpoint_format - def is_layerwise_loading_enabled(self) -> bool: - """Return whether opt-in HF safetensors layer-wise loading is enabled.""" - return (hasattr(self._weight_loader, "is_layerwise_safetensors_enabled") - and self._weight_loader.is_layerwise_safetensors_enabled()) - def iter_layer_weight_buckets( self, checkpoint_dir: str, **kwargs: Any) -> Iterator[ConsumableWeightsDict]: diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 97976cc27bcb..04055d7e2e81 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -39,7 +39,6 @@ from tensorrt_llm.mapping import Mapping _WEIGHT_CACHE_ENV = "TRTLLM_HF_WEIGHT_CACHE" -_LAYERWISE_SAFETENSORS_ENV = "TRTLLM_HF_LAYERWISE_SAFETENSORS" _WEIGHT_CACHE_MAX_ENTRIES_ENV = "TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES" # Default to a single cached checkpoint: each entry pins a full copy of the # raw weights in CPU RAM, so callers wanting cross-model caching must opt in @@ -58,12 +57,6 @@ class HfWeightLoader(BaseWeightLoader): Loads weights from SafeTensors/bin/pth files. """ - @staticmethod - def is_layerwise_safetensors_enabled() -> bool: - """Return whether semantic-layer safetensors loading is enabled.""" - return os.environ.get(_LAYERWISE_SAFETENSORS_ENV, - "0").lower() in ("1", "true", "yes", "on") - @staticmethod def _layer_bucket_name(weight_name: str) -> Tuple[str, int]: """Return a stable, layer-atomic bucket for a checkpoint tensor.""" diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 4d95bdde7f0d..596c189403b4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -560,10 +560,10 @@ def init_meta_tensor(t: torch.Tensor): model_checkpoint_dir = (model.llm_checkpoint_dir if hasattr( model, 'llm_checkpoint_dir') else checkpoint_dir) - layerwise_loading = ( - checkpoint_loader.checkpoint_format == "HF" and hasattr( - checkpoint_loader, 'is_layerwise_loading_enabled') - and checkpoint_loader.is_layerwise_loading_enabled()) + 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( @@ -576,7 +576,7 @@ def init_meta_tensor(t: torch.Tensor): model.load_weights).parameters if "initial_bucket_loading" not in load_parameters: raise RuntimeError( - "TRTLLM_HF_LAYERWISE_SAFETENSORS is enabled, but " + "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( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index cc6a66e53a08..d11436ad26ca 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4486,6 +4486,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 027cfd4e6a58..8d99efa2d80f 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -396,6 +396,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 index ff42881e52c7..1c0438078896 100644 --- a/tests/unittest/_torch/executor/test_model_loader_layerwise.py +++ b/tests/unittest/_torch/executor/test_model_loader_layerwise.py @@ -48,9 +48,12 @@ def _moe_context(_config, _mapping): yield None -def _make_loader(monkeypatch, model, *, spec_config=None): +def _make_loader(monkeypatch, model, *, spec_config=None, layerwise_loading=True): loader = ModelLoader( - llm_args=SimpleNamespace(load_format=LoadFormat.AUTO), + 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, @@ -79,7 +82,6 @@ def _make_loader(monkeypatch, model, *, spec_config=None): def _make_checkpoint_loader(events, buckets): checkpoint_loader = MagicMock(name="checkpoint_loader") checkpoint_loader.checkpoint_format = "HF" - checkpoint_loader.is_layerwise_loading_enabled.return_value = True mapper = MagicMock(name="weight_mapper") def initialize_mapper(_model, _config): @@ -134,6 +136,23 @@ def test_layerwise_model_loader_orders_mapper_buckets_sync_and_post_load(monkeyp 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)) 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 82e2a39add9f..30d5deed18b1 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py @@ -101,23 +101,6 @@ def test_layerwise_safetensors_rejects_missing_index_shard(tmp_path): next(loader.iter_layer_weight_buckets(str(tmp_path), mapping=Mapping())) -@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "Yes", "on", "ON"]) -def test_layerwise_safetensors_environment_true_values(monkeypatch, value): - monkeypatch.setenv("TRTLLM_HF_LAYERWISE_SAFETENSORS", value) - - assert HfWeightLoader.is_layerwise_safetensors_enabled() - - -@pytest.mark.parametrize("value", [None, "0", "false", "no", "off", "unexpected"]) -def test_layerwise_safetensors_environment_false_values(monkeypatch, value): - if value is None: - monkeypatch.delenv("TRTLLM_HF_LAYERWISE_SAFETENSORS", raising=False) - else: - monkeypatch.setenv("TRTLLM_HF_LAYERWISE_SAFETENSORS", value) - - assert not HfWeightLoader.is_layerwise_safetensors_enabled() - - def test_layerwise_safetensors_uses_natural_layer_order_and_preserves_values(tmp_path): save_file( { diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 0db864d9e540..addd12a49c7a 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 From e1898ae78cc96da4ccd6b41a7cfa5ba3b66fca92 Mon Sep 17 00:00:00 2001 From: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com> Date: Thu, 16 Jul 2026 02:36:58 -0700 Subject: [PATCH 4/4] [None][fix] exclude base layers from DeepSeek V4 MTP buckets Signed-off-by: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com> --- .../_torch/models/modeling_deepseekv4.py | 42 +++++++++++++------ .../modeling/test_modeling_deepseekv4.py | 42 +++++++++++++++++++ 2 files changed, 71 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 740e6e7ac567..e07be988d661 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -287,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: @@ -877,19 +900,12 @@ def load_flat_hc_weights(module, names: List[str]) -> bool: is_target_module = False if name.startswith("model.layers."): module_layer = int(name.split(".", 3)[2]) - if active_layer >= self.config.num_hidden_layers: - # Runtime may instantiate multiple MTP replicas. - # The existing name rewrite below maps each replica - # back to the checkpoint MTP layer modulo the number - # of checkpoint next-token layers. - checkpoint_mtp_idx = active_layer - self.config.num_hidden_layers - runtime_mtp_idx = module_layer - self.config.num_hidden_layers - is_target_module = ( - runtime_mtp_idx % self.config.num_nextn_predict_layers - == checkpoint_mtp_idx - ) - else: - is_target_module = module_layer == active_layer + 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(".") diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 14f893072466..f81c245611a0 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -35,6 +35,7 @@ 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, ) @@ -240,6 +241,47 @@ def test_deepseek_v4_layerwise_loading_rejects_non_atomic_bucket(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))