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, 22:15 UTC
use device param in load_model method (#13037)
Zhe Zhang committed
Feb 19, 2025 at 16:05 UTC
fdc5df6f54854ba437429fa0ff721939259364ad
1 file changed
+1
-1
vllm/worker/model_runner.py
+1
-1
@@ -1107,7 +1107,7 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1107
1108
def load_model(self) -> None:
1109
logger.info("Starting to load model %s...", self.model_config.model)
1110
- with DeviceMemoryProfiler() as m:
1110
+ with DeviceMemoryProfiler(self.device) as m:
1111
self.model = get_model(vllm_config=self.vllm_config)
1112
1113
self.model_memory_usage = m.consumed_memory