A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 18:30 UTC
[log] add weights loading time log to sharded_state loader (#28628)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Ning Xie committed
Nov 22, 2025 at 05:06 UTC
53a1ba6ec584ea93531a3195b3b9f8049786055b
1 file changed
+8
vllm/model_executor/model_loader/sharded_state_loader.py
+8
@@ -4,6 +4,7 @@
4
import collections
5
import glob
6
import os
7
+import time
8
from collections.abc import Generator
9
from typing import Any
10
@@ -132,6 +133,7 @@ class ShardedStateLoader(BaseModelLoader):
133
f"pre-sharded checkpoints are currently supported!"
134
)
135
state_dict = self._filter_subtensors(model.state_dict())
136
+ counter_before_loading_weights = time.perf_counter()
137
for key, tensor in self.iterate_over_files(filepaths):
138
# If loading with LoRA enabled, additional padding may
139
# be added to certain parameters. We only load into a
@@ -150,6 +152,12 @@ class ShardedStateLoader(BaseModelLoader):
152
)
153
param_data.copy_(tensor)
154
state_dict.pop(key)
155
+ counter_after_loading_weights = time.perf_counter()
156
+ logger.info_once(
157
+ "Loading weights took %.2f seconds",
158
+ counter_after_loading_weights - counter_before_loading_weights,
159
+ scope="local",
160
+ )
161
if state_dict:
162
raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!")
163