[WIP][None][feat] add layer-wise safetensors loading for DeepSeek V4

docs/source/features/checkpoint-loading.md

@@ -1,3 +1,8 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+-->
+
 # 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 <model_path> --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

tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py

@@ -1,11 +1,14 @@
-from typing import Optional
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Iterator, Optional
 
 from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \
     BaseCheckpointLoader
 from tensorrt_llm._torch.models.checkpoints.base_config_loader import \
     BaseConfigLoader
-from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \
-    BaseWeightLoader
+from tensorrt_llm._torch.models.checkpoints.base_weight_loader import (
+    BaseWeightLoader, ConsumableWeightsDict)
 from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \
     BaseWeightMapper
 from tensorrt_llm._torch.models.checkpoints.hf.config_loader import \
@@ -73,3 +76,19 @@ def config_loader(self) -> Optional[BaseConfigLoader]:
     @property
     def checkpoint_format(self) -> str:
         return self._checkpoint_format
+
+    def iter_layer_weight_buckets(
+            self, checkpoint_dir: str,
+            **kwargs: Any) -> Iterator[ConsumableWeightsDict]:
+        """Yield semantic-layer weight buckets from the underlying HF loader.
+
+        Args:
+            checkpoint_dir: Directory containing Hugging Face checkpoint files.
+            **kwargs: Arguments forwarded to the concrete weight loader.
+
+        Yields:
+            A consumable dictionary containing top-level tensors or tensors
+            for exactly one base/MTP checkpoint layer.
+        """
+        return self._weight_loader.iter_layer_weight_buckets(
+            checkpoint_dir, **kwargs)

tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py

@@ -13,12 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import glob
+import json
 import multiprocessing
 import os
+import re
 import threading
 from collections import OrderedDict
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, List
+from contextlib import ExitStack
+from typing import Any, Iterator, List, Tuple
 
 import psutil
 import safetensors
@@ -54,6 +57,116 @@ class HfWeightLoader(BaseWeightLoader):
     Loads weights from SafeTensors/bin/pth files.
     """
 
+    @staticmethod
+    def _layer_bucket_name(weight_name: str) -> Tuple[str, int]:
+        """Return a stable, layer-atomic bucket for a checkpoint tensor."""
+        match = re.match(r"^(?:model\.)?layers\.(\d+)\.", weight_name)
+        if match is not None:
+            return ("layer", int(match.group(1)))
+        match = re.match(r"^mtp\.(\d+)\.", weight_name)
+        if match is not None:
+            return ("mtp", int(match.group(1)))
+        return ("top", 0)
+
+    def iter_layer_weight_buckets(self,
+                                  checkpoint_dir: str,
+                                  mapping: Mapping,
+                                  use_consolidated: bool = False,
+                                  **kwargs) -> Iterator[ConsumableWeightsDict]:
+        """Yield real CPU tensors one semantic layer at a time.
+
+        The safetensors index is used only as metadata. Tensor storage remains
+        mmap-backed and the file handles for a bucket stay alive until the
+        caller advances the iterator. This keeps cross-shard, same-layer
+        fusions atomic without materializing the complete checkpoint.
+
+        Args:
+            checkpoint_dir: Directory containing safetensors files and an
+                optional Hugging Face safetensors index.
+            mapping: Distributed mapping accepted for parity with eager
+                checkpoint loading. Bucketing itself is rank-independent.
+            use_consolidated: Select consolidated files instead of ordinary
+                Hugging Face shards.
+            **kwargs: Extra eager-loader arguments that do not affect bucket
+                planning.
+
+        Yields:
+            A consumable dictionary for all top-level tensors or one numeric
+            base/MTP layer. File handles remain open until the next iteration.
+
+        Raises:
+            RuntimeError: If the selected safetensors files are missing or
+                their index/key metadata is incomplete or ambiguous.
+        """
+        del mapping, kwargs
+        weight_files = glob.glob(f"{checkpoint_dir}/*.safetensors")
+        weight_files = [
+            path for path in weight_files
+            if ("consolidated" in os.path.basename(path)) == use_consolidated
+        ]
+        if not weight_files:
+            checkpoint_kind = "consolidated " if use_consolidated else ""
+            raise RuntimeError(
+                f"Layer-wise loading requires {checkpoint_kind}safetensors "
+                f"weights in {checkpoint_dir}.")
+
+        files_by_name = {os.path.basename(path): path for path in weight_files}
+        index_files = sorted(
+            glob.glob(f"{checkpoint_dir}/*.safetensors.index.json"))
+        entries: List[Tuple[str, str]] = []
+        use_index = bool(index_files) and not use_consolidated and all(
+            "consolidated" not in name for name in files_by_name)
+        if use_index:
+            with open(index_files[0], encoding="utf-8") as index_file:
+                weight_map = json.load(index_file).get("weight_map", {})
+            missing_files = set(weight_map.values()) - set(files_by_name)
+            if missing_files:
+                raise RuntimeError(
+                    f"Safetensors index {index_files[0]} references missing "
+                    f"checkpoint files: {sorted(missing_files)}")
+            for weight_name, file_name in weight_map.items():
+                entries.append((weight_name, files_by_name[file_name]))
+            if not entries:
+                raise RuntimeError(
+                    f"Safetensors index {index_files[0]} has an empty weight_map."
+                )
+        else:
+            seen_keys = set()
+            for path in sorted(weight_files):
+                with safetensors.safe_open(path, framework="pt",
+                                           device="cpu") as handle:
+                    for key in handle.keys():
+                        if key in seen_keys:
+                            raise RuntimeError(
+                                "Duplicate tensor key found across safetensors "
+                                f"files without an index: {key}")
+                        seen_keys.add(key)
+                        entries.append((key, path))
+
+        buckets = {}
+        for weight_name, path in entries:
+            bucket = self._layer_bucket_name(weight_name)
+            buckets.setdefault(bucket, []).append((weight_name, path))
+
+        def bucket_sort_key(bucket):
+            kind, index = bucket
+            return ({"top": 0, "layer": 1, "mtp": 2}[kind], index)
+
+        for bucket in sorted(buckets, key=bucket_sort_key):
+            tensors = {}
+            with ExitStack() as stack:
+                handles = {}
+                for weight_name, path in buckets[bucket]:
+                    if path not in handles:
+                        handles[path] = stack.enter_context(
+                            safetensors.safe_open(path,
+                                                  framework="pt",
+                                                  device="cpu"))
+                    tensors[weight_name] = handles[path].get_tensor(weight_name)
+                logger.info("Loading HF weight bucket %s with %d tensors",
+                            bucket, len(tensors))
+                yield ConsumableWeightsDict(tensors)
+
     @staticmethod
     def _is_weight_cache_enabled() -> bool:
         return os.environ.get(_WEIGHT_CACHE_ENV,