From 53a1ba6ec584ea93531a3195b3b9f8049786055b Mon Sep 17 00:00:00 2001 From: Ning Xie Date: Sat, 22 Nov 2025 05:06:09 +0800 Subject: [PATCH] [log] add weights loading time log to sharded_state loader (#28628) Signed-off-by: Andy Xie --- vllm/model_executor/model_loader/sharded_state_loader.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py index d94dbd9f06e0..1538f0c2af65 100644 --- a/vllm/model_executor/model_loader/sharded_state_loader.py +++ b/vllm/model_executor/model_loader/sharded_state_loader.py @@ -4,6 +4,7 @@ import collections import glob import os +import time from collections.abc import Generator from typing import Any @@ -132,6 +133,7 @@ class ShardedStateLoader(BaseModelLoader): f"pre-sharded checkpoints are currently supported!" ) state_dict = self._filter_subtensors(model.state_dict()) + counter_before_loading_weights = time.perf_counter() for key, tensor in self.iterate_over_files(filepaths): # If loading with LoRA enabled, additional padding may # be added to certain parameters. We only load into a @@ -150,6 +152,12 @@ class ShardedStateLoader(BaseModelLoader): ) param_data.copy_(tensor) state_dict.pop(key) + counter_after_loading_weights = time.perf_counter() + logger.info_once( + "Loading weights took %.2f seconds", + counter_after_loading_weights - counter_before_loading_weights, + scope="local", + ) if state_dict: raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!")