[WIP][None][feat] add layer-wise safetensors loading for DeepSeek V4#16437
Open
RockmanXiao wants to merge 6 commits into
Open
[WIP][None][feat] add layer-wise safetensors loading for DeepSeek V4#16437RockmanXiao wants to merge 6 commits into
RockmanXiao wants to merge 6 commits into
Conversation
Signed-off-by: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com>
Signed-off-by: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com>
Signed-off-by: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com>
RockmanXiao
changed the title
RockmanXiao
marked this pull request as draft
July 15, 2026 15:43
Contributor
📝 WalkthroughWalkthroughAdds prototype layer-wise Hugging Face safetensors loading for DeepSeek V4, including configuration, shard-aware bucket iteration, model integration, ModelLoader orchestration, documentation, and unit tests. ChangesLayer-wise checkpoint loading
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ModelLoader
participant HfCheckpointLoader
participant DeepseekV4ForCausalLM
participant CUDA
ModelLoader->>HfCheckpointLoader: Iterate checkpoint weight buckets
HfCheckpointLoader-->>ModelLoader: Return one bucket
ModelLoader->>DeepseekV4ForCausalLM: Load bucket with initial_bucket_loading
ModelLoader->>CUDA: Synchronize between buckets
ModelLoader->>HfCheckpointLoader: Request next bucket
HfCheckpointLoader-->>ModelLoader: Return remaining buckets
ModelLoader->>DeepseekV4ForCausalLM: Run post-load processing
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (1)
tests/unittest/_torch/modeling/test_modeling_deepseekv4.py (1)
218-240: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd successful atomic-bucket coverage.
Coverage is insufficient because only rejection paths call
load_weights(..., initial_bucket_loading=True). Add tests for one valid base layer, one MTP layer, and one top-level bucket, asserting traversal is restricted to that bucket.Suggested names:
test_deepseek_v4_layerwise_loading_accepts_single_base_layertest_deepseek_v4_layerwise_loading_accepts_single_mtp_layertest_deepseek_v4_layerwise_loading_accepts_top_level_bucketAs per path instructions,
tests/**feedback must state whether coverage is sufficient and suggest concrete test names.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modeling/test_modeling_deepseekv4.py` around lines 218 - 240, Add successful atomic-bucket coverage; current tests only cover rejection paths. In the DeepSeek V4 loader tests, add test_deepseek_v4_layerwise_loading_accepts_single_base_layer, test_deepseek_v4_layerwise_loading_accepts_single_mtp_layer, and test_deepseek_v4_layerwise_loading_accepts_top_level_bucket, each calling load_weights with initial_bucket_loading=True and asserting traversal is restricted to the supplied bucket.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py`:
- Around line 71-75: Annotate every changed callable with precise types: in
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py:71-75 type
iter_layer_weight_buckets kwargs and at :151-153 type bucket_sort_key; annotate
all new tests and nested recording-handle methods in
tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py:21-256, new
tests and parameterized weights in
tests/unittest/_torch/modeling/test_modeling_deepseekv4.py:192-240, model-loader
arguments with -> None in
tensorrt_llm/_torch/pyexecutor/model_loader.py:1331-1364, and all new helpers,
model methods, and tests in
tests/unittest/_torch/executor/test_model_loader_layerwise.py:18-208, preferring
precise built-in generic types.
- Around line 114-128: Update the index-selection logic around use_index in
weight_loader.py to filter safetensors indexes by checkpoint kind, reject
ambiguous matching indexes, and validate every weight_map key against the
available checkpoint entries before appending entries. Add regression tests in
test_weight_loader.py covering multiple matching indexes and a consolidated
index mismatched with the requested checkpoint kind.
- Around line 133-149: Reject empty unindexed safetensors checkpoints by adding
an empty-entries guard after the no-index loading loop in the checkpoint loader
before bucket construction; update
test_layerwise_safetensors_rejects_empty_unindexed_checkpoint in
tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py to verify this
failure, while preserving existing behavior for non-empty checkpoints.
In `@tensorrt_llm/_torch/models/modeling_deepseekv4.py`:
- Around line 872-894: Add an MTP-range guard in the active-layer branch of the
model-layer filtering logic so modulo matching runs only when module_layer is at
least self.config.num_hidden_layers. Keep base decoder layers excluded for MTP
buckets, while preserving the existing modulo mapping for valid MTP replica
layers and direct equality matching for regular layers.
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 1355-1363: Replace the assertion checks in the model-loading
option handling with explicit RuntimeError raises when allow_partial_loading or
initial_bucket_loading is enabled but unsupported by the model, preserving the
existing error messages. Add direct tests covering both unsupported-option paths
in test_model_loader_layerwise.py.
- Around line 575-592: Guard the self.weight_mapper initialization in the
enable_hf_layerwise_loading branch so DeepseekV4ForCausalLM does not call
checkpoint_loader.get_initialized_weight_mapper when its load_weights lacks
mapper support. Reuse the existing loader capability or registration check,
while preserving mapper initialization for models that accept and use it.
In `@tests/unittest/_torch/executor/test_model_loader_layerwise.py`:
- Around line 156-166: Coverage is insufficient because the current test only
exercises the outer layer-wise precheck, not _call_load_weights. Add
test_call_load_weights_rejects_unsupported_options to invoke _call_load_weights
directly against a method declaring neither option, and parameterize or
separately cover allow_partial_loading=True and initial_bucket_loading=True,
asserting both are rejected.
In `@tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py`:
- Around line 126-144: Coverage is missing for an empty safetensors checkpoint
when no index file exists. Add
test_layerwise_safetensors_rejects_empty_unindexed_checkpoint alongside the
existing layerwise safetensors rejection tests, creating an empty checkpoint
without model.safetensors.index.json and asserting iter_layer_weight_buckets
raises the expected empty-checkpoint RuntimeError; coverage is sufficient once
both indexed and unindexed empty cases are exercised.
- Around line 147-175: Add regression tests for layerwise safetensors index
selection: implement test_layerwise_safetensors_rejects_multiple_indexes to
verify multiple index files raise the expected error, and
test_layerwise_safetensors_ignores_consolidated_index_for_ordinary_weights to
verify ordinary loading ignores a consolidated index and selects the correct
ordinary weights. Existing coverage is insufficient for these index-file
scenarios; keep the tests focused on HfWeightLoader.iter_layer_weight_buckets.
---
Nitpick comments:
In `@tests/unittest/_torch/modeling/test_modeling_deepseekv4.py`:
- Around line 218-240: Add successful atomic-bucket coverage; current tests only
cover rejection paths. In the DeepSeek V4 loader tests, add
test_deepseek_v4_layerwise_loading_accepts_single_base_layer,
test_deepseek_v4_layerwise_loading_accepts_single_mtp_layer, and
test_deepseek_v4_layerwise_loading_accepts_top_level_bucket, each calling
load_weights with initial_bucket_loading=True and asserting traversal is
restricted to the supplied bucket.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c265d126-0ecb-43f5-bf4a-11bf6d8afec4
📒 Files selected for processing (12)
docs/source/features/checkpoint-loading.mdtensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.pytensorrt_llm/_torch/models/checkpoints/hf/weight_loader.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/unittest/_torch/executor/test_model_loader_layerwise.pytests/unittest/_torch/modeling/test_modeling_deepseekv4.pytests/unittest/_torch/models/checkpoints/hf/test_weight_loader.pytests/unittest/api_stability/references/llm.yamltests/unittest/llmapi/test_layerwise_loading_args.py
Comment on lines
+71
to
+75
| def iter_layer_weight_buckets(self, | ||
| checkpoint_dir: str, | ||
| mapping: Mapping, | ||
| use_consolidated: bool = False, | ||
| **kwargs) -> Iterator[ConsumableWeightsDict]: |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
Add the required annotations to changed callables.
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L71-L75: annotate**kwargs.tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L151-L153: annotatebucket_sort_key.tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py#L21-L256: annotate new tests and nested recording-handle methods.tests/unittest/_torch/modeling/test_modeling_deepseekv4.py#L192-L240: annotate new tests and the parameterizedweights.tensorrt_llm/_torch/pyexecutor/model_loader.py#L1331-L1364: precisely type arguments and-> None.tests/unittest/_torch/executor/test_model_loader_layerwise.py#L18-L208: annotate new helpers, model methods, and tests.
As per coding guidelines, “Annotate every function” and prefer precise built-in generic types.
📍 Affects 5 files
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L71-L75(this comment)tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L151-L153tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py#L21-L256tests/unittest/_torch/modeling/test_modeling_deepseekv4.py#L192-L240tensorrt_llm/_torch/pyexecutor/model_loader.py#L1331-L1364tests/unittest/_torch/executor/test_model_loader_layerwise.py#L18-L208
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py` around lines 71 -
75, Annotate every changed callable with precise types: in
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py:71-75 type
iter_layer_weight_buckets kwargs and at :151-153 type bucket_sort_key; annotate
all new tests and nested recording-handle methods in
tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py:21-256, new
tests and parameterized weights in
tests/unittest/_torch/modeling/test_modeling_deepseekv4.py:192-240, model-loader
arguments with -> None in
tensorrt_llm/_torch/pyexecutor/model_loader.py:1331-1364, and all new helpers,
model methods, and tests in
tests/unittest/_torch/executor/test_model_loader_layerwise.py:18-208, preferring
precise built-in generic types.
Source: Coding guidelines
Comment on lines
+114
to
+128
| 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])) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Select and validate the correct safetensors index before bucket loading.
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L114-L128: filter indexes by checkpoint kind, reject ambiguity, and validate mapped key membership before yielding.tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py#L147-L175: add multiple-index and mismatched consolidated-index regression tests.
🧰 Tools
🪛 ast-grep (0.44.1)
[warning] 119-119: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(index_files[0], encoding="utf-8")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
📍 Affects 2 files
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L114-L128(this comment)tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py#L147-L175
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py` around lines 114
- 128, Update the index-selection logic around use_index in weight_loader.py to
filter safetensors indexes by checkpoint kind, reject ambiguous matching
indexes, and validate every weight_map key against the available checkpoint
entries before appending entries. Add regression tests in test_weight_loader.py
covering multiple matching indexes and a consolidated index mismatched with the
requested checkpoint kind.
Source: Path instructions
Comment on lines
+133
to
+149
| 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)) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
sed -n '1,240p' tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 10435
🏁 Script executed:
sed -n '1,220p' tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py && python3 - <<'PY'
# Minimal behavioral probe based on the discovered control flow:
# if no tensors are present in any safetensors file and no index is used,
# entries stays empty and the function yields nothing.
entries = []
buckets = {}
for weight_name, path in entries:
bucket = ("top", 0)
buckets.setdefault(bucket, []).append((weight_name, path))
print("entries_empty:", not entries)
print("buckets_empty:", not buckets)
print("would_yield_anything:", bool(sorted(buckets)))
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 7706
Reject empty unindexed checkpoints. The no-index path can still return no buckets when every .safetensors file is empty, so an invalid checkpoint is accepted silently. Add the same empty-entries guard here and cover it with test_layerwise_safetensors_rejects_empty_unindexed_checkpoint.
📍 Affects 2 files
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L133-L149(this comment)tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py#L126-L144
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py` around lines 133
- 149, Reject empty unindexed safetensors checkpoints by adding an empty-entries
guard after the no-index loading loop in the checkpoint loader before bucket
construction; update
test_layerwise_safetensors_rejects_empty_unindexed_checkpoint in
tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py to verify this
failure, while preserving existing behavior for non-empty checkpoints.
Source: Path instructions
Comment on lines
+872
to
+894
| 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 |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Missing range check lets base decoder-layer modules incorrectly match an MTP bucket.
When active_layer >= self.config.num_hidden_layers (processing an MTP bucket), runtime_mtp_idx = module_layer - self.config.num_hidden_layers is computed and compared via % self.config.num_nextn_predict_layers for every model.layers.* module, without first checking that module_layer itself is in the MTP range. For module_layer < num_hidden_layers (a genuine base decoder layer), runtime_mtp_idx is negative, and Python's modulo (sign follows the divisor) can still equal checkpoint_mtp_idx.
Concrete counterexample: num_hidden_layers=4, num_nextn_predict_layers=2, active_layer=4 (checkpoint_mtp_idx=0). For module_layer=0: runtime_mtp_idx=-4, and -4 % 2 == 0, so is_target_module incorrectly becomes True for base layer 0 — even though the MTP-only bucket's weights dict has no model.layers.0.* keys. Downstream direct key lookups (e.g. weights[f"{module_name}.{weight_name}"] for kv_b_proj/kv_a_proj_with_mqa, or module_weights[n] in the generic per-parameter path) will then raise KeyError, crashing the load for essentially any DeepSeek V4 checkpoint that has MTP layers.
Add a guard so the modulo comparison only applies when module_layer is itself in the MTP range.
🐛 Proposed fix
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
- )
+ if module_layer >= self.config.num_hidden_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📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| 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 | |
| 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. | |
| if module_layer >= self.config.num_hidden_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 |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/models/modeling_deepseekv4.py` around lines 872 - 894,
Add an MTP-range guard in the active-layer branch of the model-layer filtering
logic so modulo matching runs only when module_layer is at least
self.config.num_hidden_layers. Keep base decoder layers excluded for MTP
buckets, while preserving the existing modulo mapping for valid MTP replica
layers and direct equality matching for regular layers.
Comment on lines
+575
to
+592
| load_parameters = inspect.signature( | ||
| model.load_weights).parameters | ||
| if "initial_bucket_loading" not in load_parameters: | ||
| raise RuntimeError( | ||
| "enable_hf_layerwise_loading is enabled, but " | ||
| f"{type(model).__name__}.load_weights does not support " | ||
| "initial_bucket_loading.") | ||
| self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( | ||
| model, config) | ||
| logger.info( | ||
| "Loading HF checkpoint with one semantic layer per host-memory bucket." | ||
| ) | ||
| for weight_bucket in checkpoint_loader.iter_layer_weight_buckets( | ||
| model_checkpoint_dir, **load_weights_kwargs): | ||
| self._call_load_weights(model.load_weights, | ||
| weight_bucket, | ||
| self.weight_mapper, | ||
| initial_bucket_loading=True) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -C4 \
'AutoCheckpointMapper\.(register|register_mapper)|DeepseekV4|DeepSeekV4|deepseek_v4' \
tensorrt_llm/_torch/models/checkpoints tensorrt_llm/_torch/models/modeling_deepseekv4.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 40591
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== model_loader outline ==\n'
ast-grep outline tensorrt_llm/_torch/pyexecutor/model_loader.py --view expanded | sed -n '1,260p'
printf '\n== relevant model_loader slice ==\n'
sed -n '540,610p' tensorrt_llm/_torch/pyexecutor/model_loader.py | cat -n
printf '\n== checkpoint loader search ==\n'
rg -n -C3 'def get_initialized_weight_mapper|def iter_layer_weight_buckets|weight_mapper|weight_map' tensorrt_llm/_torch/models/checkpoints tensorrt_llm/_torch/pyexecutor/model_loader.py
printf '\n== DeepSeek V4 mapper registration search ==\n'
rg -n -C3 'register_auto_model|register_mapper|AutoCheckpointMapper|weight_mapper' tensorrt_llm/_torch/models tensorrt_llm/_torch/models/checkpointsRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the exact implementation of the mapper initialization and the
# load-weights call path around the layer-wise branch.
sed -n '1,240p' tensorrt_llm/_torch/models/checkpoints/__init__.py | cat -n
printf '\n== checkpoint loader implementation ==\n'
rg -n -C4 'class .*CheckpointLoader|def get_initialized_weight_mapper|def iter_layer_weight_buckets|def _call_load_weights|weight_mapper' tensorrt_llm/_torch/models/checkpoints tensorrt_llm/_torch/pyexecutor/model_loader.py
printf '\n== DeepSeek V4 load_weights signature and uses ==\n'
sed -n '2636,2650p' tensorrt_llm/_torch/models/modeling_deepseekv4.py | cat -n
printf '\n== any DeepSeek V4 checkpoint-specific mapper file ==\n'
fd -a 'deepseek*' tensorrt_llm/_torch/models/checkpoints tensorrt_llm/_torch/models | sed -n '1,120p'Repository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== model_loader around the branch ==\n'
sed -n '560,600p' tensorrt_llm/_torch/pyexecutor/model_loader.py | cat -n
printf '\n== inspect _call_load_weights ==\n'
rg -n -C6 'def _call_load_weights|initial_bucket_loading|weight_mapper' tensorrt_llm/_torch/pyexecutor/model_loader.py
printf '\n== inspect mapper init ==\n'
rg -n -C6 'def get_initialized_weight_mapper|class .*Mapper|register_mapper|AutoCheckpointMapper' tensorrt_llm/_torch/models/checkpoints
printf '\n== repo-wide DeepSeek V4 checkpoint mapper references ==\n'
rg -n -C4 'DeepseekV4|DeepSeekV4|deepseek_v4|weight_mapper' tensorrt_llm -g '!**/__pycache__/**'Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '== DeepSeek V4 load_weights signature ==\n'
rg -n -C2 'def load_weights\(self, weights: Dict, initial_bucket_loading: bool = False\)' \
tensorrt_llm/_torch/models/modeling_deepseekv4.py
printf '\n== DeepSeek V4 mapper registration hits ==\n'
rg -n -C2 'register_mapper\("HF", ".*Deepseek|register_mapper\("HF", "Deepseek|DeepseekV4.*weight_mapper|deepseek_v4.*weight_mapper' \
tensorrt_llm/_torch/models
printf '\n== layer-wise loader call site ==\n'
sed -n '576,592p' tensorrt_llm/_torch/pyexecutor/model_loader.py | cat -nRepository: NVIDIA/TensorRT-LLM
Length of output: 503
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '== get_initialized_weight_mapper ==\n'
rg -n -C8 'def get_initialized_weight_mapper|AutoCheckpointMapper\.get\(' \
tensorrt_llm/_torch/pyexecutor/model_loader.py \
tensorrt_llm/_torch/models/checkpoints \
tensorrt_llm/_torch/models/modeling_utils.py
printf '\n== all DeepSeek V4 references in checkpoint mappers ==\n'
rg -n -C3 'DeepseekV4ForCausalLM|DeepseekV4|DeepSeekV4' \
tensorrt_llm/_torch/models/checkpoints \
tensorrt_llm/_torch/models/modeling_deepseekv4.py
printf '\n== nearby mapper registrations in hf checkpoint mappers ==\n'
rg -n -C2 '`@register_mapper`\("HF", ".*"\)' \
tensorrt_llm/_torch/models/checkpoints/hf/*.py | sed -n '1,220p'Repository: NVIDIA/TensorRT-LLM
Length of output: 43918
Guard the layer-wise mapper init for DeepSeek V4. DeepseekV4ForCausalLM.load_weights() doesn’t accept weight_mapper, and there’s no HF mapper registration for DeepseekV4ForCausalLM, so this branch can fail before the first bucket. Gate get_initialized_weight_mapper() behind loaders that actually use it, or add the missing mapper registration.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py` around lines 575 - 592, Guard
the self.weight_mapper initialization in the enable_hf_layerwise_loading branch
so DeepseekV4ForCausalLM does not call
checkpoint_loader.get_initialized_weight_mapper when its load_weights lacks
mapper support. Reuse the existing loader capability or registration check,
while preserving mapper initialization for models that accept and use it.
Comment on lines
1355
to
+1363
| if "allow_partial_loading" in args: | ||
| 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" |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Enforce unsupported loading options with explicit exceptions.
tensorrt_llm/_torch/pyexecutor/model_loader.py#L1355-L1363: replace optimization-sensitive assertions withRuntimeError.tests/unittest/_torch/executor/test_model_loader_layerwise.py#L156-L166: directly test both unsupported option paths.
📍 Affects 2 files
tensorrt_llm/_torch/pyexecutor/model_loader.py#L1355-L1363(this comment)tests/unittest/_torch/executor/test_model_loader_layerwise.py#L156-L166
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py` around lines 1355 - 1363,
Replace the assertion checks in the model-loading option handling with explicit
RuntimeError raises when allow_partial_loading or initial_bucket_loading is
enabled but unsupported by the model, preserving the existing error messages.
Add direct tests covering both unsupported-option paths in
test_model_loader_layerwise.py.
Source: Path instructions
Comment on lines
+156
to
+166
| 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 == [] |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Test _call_load_weights option rejection directly.
The current test exits through the outer layer-wise precheck. Add test_call_load_weights_rejects_unsupported_options, covering both allow_partial_loading=True and initial_bucket_loading=True against a method that declares neither option.
As per path instructions, tests/** feedback must state whether coverage is sufficient and suggest concrete test names.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/executor/test_model_loader_layerwise.py` around lines
156 - 166, Coverage is insufficient because the current test only exercises the
outer layer-wise precheck, not _call_load_weights. Add
test_call_load_weights_rejects_unsupported_options to invoke _call_load_weights
directly against a method declaring neither option, and parameterize or
separately cover allow_partial_loading=True and initial_bucket_loading=True,
asserting both are rejected.
Source: Path instructions
Comment on lines
+126
to
+144
| 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())) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Cover an empty checkpoint without an index.
The empty-index test does not exercise the no-index branch. Add test_layerwise_safetensors_rejects_empty_unindexed_checkpoint.
As per path instructions, tests/** feedback must state whether coverage is sufficient and suggest concrete test names.
🧰 Tools
🪛 ast-grep (0.44.1)
[info] 128-128: use jsonify instead of json.dumps for JSON output
Context: json.dumps({"weight_map": {}})
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py` around
lines 126 - 144, Coverage is missing for an empty safetensors checkpoint when no
index file exists. Add
test_layerwise_safetensors_rejects_empty_unindexed_checkpoint alongside the
existing layerwise safetensors rejection tests, creating an empty checkpoint
without model.safetensors.index.json and asserting iter_layer_weight_buckets
raises the expected empty-checkpoint RuntimeError; coverage is sufficient once
both indexed and unindexed empty cases are exercised.
Source: Path instructions
Comment on lines
+147
to
+175
| def test_layerwise_safetensors_strictly_selects_consolidated_files(tmp_path): | ||
| save_file({"layers.0.weight": torch.tensor([1.0])}, str(tmp_path / "model.safetensors")) | ||
| save_file( | ||
| {"layers.1.weight": torch.tensor([2.0])}, | ||
| str(tmp_path / "model-consolidated.safetensors"), | ||
| ) | ||
|
|
||
| ordinary = list(HfWeightLoader().iter_layer_weight_buckets(str(tmp_path), mapping=Mapping())) | ||
| consolidated = list( | ||
| HfWeightLoader().iter_layer_weight_buckets( | ||
| str(tmp_path), mapping=Mapping(), use_consolidated=True | ||
| ) | ||
| ) | ||
|
|
||
| assert [set(bucket) for bucket in ordinary] == [{"layers.0.weight"}] | ||
| assert [set(bucket) for bucket in consolidated] == [{"layers.1.weight"}] | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("use_consolidated", [False, True]) | ||
| def test_layerwise_safetensors_does_not_fallback_to_wrong_file_kind(tmp_path, use_consolidated): | ||
| filename = "model.safetensors" if use_consolidated else "model-consolidated.safetensors" | ||
| save_file({"layers.0.weight": torch.tensor([1.0])}, str(tmp_path / filename)) | ||
|
|
||
| with pytest.raises(RuntimeError, match="requires .*safetensors weights"): | ||
| next( | ||
| HfWeightLoader().iter_layer_weight_buckets( | ||
| str(tmp_path), mapping=Mapping(), use_consolidated=use_consolidated | ||
| ) | ||
| ) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add index-selection regression coverage.
Coverage is insufficient when multiple or mismatched index files exist. Add test_layerwise_safetensors_rejects_multiple_indexes and test_layerwise_safetensors_ignores_consolidated_index_for_ordinary_weights.
As per path instructions, tests/** feedback must state whether coverage is sufficient and suggest concrete test names.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 170-170: Pattern passed to match= contains metacharacters but is neither escaped nor raw
(RUF043)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py` around
lines 147 - 175, Add regression tests for layerwise safetensors index selection:
implement test_layerwise_safetensors_rejects_multiple_indexes to verify multiple
index files raise the expected error, and
test_layerwise_safetensors_ignores_consolidated_index_for_ordinary_weights to
verify ordinary loading ignores a consolidated index and selects the correct
ordinary weights. Existing coverage is insufficient for these index-file
scenarios; keep the tests focused on HfWeightLoader.iter_layer_weight_buckets.
Source: Path instructions
Signed-off-by: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com>
Signed-off-by: RockmanXiao <35864035+RockmanXiao@users.noreply.github.com>
RockmanXiao
marked this pull request as ready for review
July 16, 2026 10:01
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Reduce peak host memory during DeepSeek V4 startup by loading Hugging Face
safetensors checkpoints one semantic model layer at a time.
This change:
shard boundaries.
bucket.
transformations.
TorchLlmArgs.enable_hf_layerwise_loadingfield.docs/source/features/checkpoint-loading.md.Enable the feature with:
Current limitations:
cache.
The new field updates the non-committed API schema and usage telemetry
manifest. This PR should use the api-compatible label and requires the
expected telemetry/privacy CODEOWNER review.
Test Coverage
Tests cover:
Cluster results:
Summary by CodeRabbit
New Features
enable_hf_layerwise_loadingconfiguration option.Documentation
Tests