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 15, 21:00 UTC
[Misc] Accurately capture the time of loading weights (#14063)
Signed-off-by: Jun Duan <jun.duan.phd@outlook.com>
Jun Duan committed
Mar 1, 2025 at 20:20 UTC
82fbeae92b86e404829a01441334a9505e8b190d
3 files changed
+13
-2
vllm/model_executor/model_loader/loader.py
+11
@@ -10,6 +10,7 @@ import inspect
10
import itertools
11
import math
12
import os
13
+import time
14
import warnings
15
from abc import ABC, abstractmethod
16
from contextlib import contextmanager
@@ -216,6 +217,9 @@ class DefaultModelLoader(BaseModelLoader):
217
allow_patterns_overrides: Optional[list[str]] = None
218
"""If defined, weights will load exclusively using these patterns."""
219
220
+ counter_before_loading_weights: float = 0.0
221
+ counter_after_loading_weights: float = 0.0
222
+
223
def __init__(self, load_config: LoadConfig):
224
super().__init__(load_config)
225
if load_config.model_loader_extra_config:
@@ -368,6 +372,8 @@ class DefaultModelLoader(BaseModelLoader):
372
373
weights_iterator = _xla_weights_iterator(weights_iterator)
374
375
+ if self.counter_before_loading_weights == 0.0:
376
+ self.counter_before_loading_weights = time.perf_counter()
377
# Apply the prefix.
378
return ((source.prefix + name, tensor)
379
for (name, tensor) in weights_iterator)
@@ -412,6 +418,11 @@ class DefaultModelLoader(BaseModelLoader):
418
weights_to_load = {name for name, _ in model.named_parameters()}
419
loaded_weights = model.load_weights(
420
self._get_all_weights(model_config, model))
421
+ self.counter_after_loading_weights = time.perf_counter()
422
+ logger.info(
423
+ "Loading weights took %.2f seconds",
424
+ self.counter_after_loading_weights -
425
+ self.counter_before_loading_weights)
426
# We only enable strict check for non-quantized models
427
# that have loaded weights tracking currently.
428
if model_config.quantization is None and loaded_weights is not None:
vllm/v1/worker/gpu_model_runner.py
+1
-1
@@ -1061,7 +1061,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1061
self.device)
1062
time_after_load = time.perf_counter()
1063
self.model_memory_usage = m.consumed_memory
1064
- logger.info("Loading model weights took %.4f GB and %.6f seconds",
1064
+ logger.info("Model loading took %.4f GB and %.6f seconds",
1065
self.model_memory_usage / float(2**30),
1066
time_after_load - time_before_load)
1067
vllm/worker/model_runner.py
+1
-1
@@ -1114,7 +1114,7 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1114
time_after_load = time.perf_counter()
1115
1116
self.model_memory_usage = m.consumed_memory
1117
- logger.info("Loading model weights took %.4f GB and %.6f seconds",
1117
+ logger.info("Model loading took %.4f GB and %.6f seconds",
1118
self.model_memory_usage / float(2**30),
1119
time_after_load - time_before_load)
1120