Move driver worker to device

vwxyzjn · Aug 18, 2024 · 7341071 · 7341071
1 parent e680349
commit 7341071
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm/executor/gpu_executor.py b/vllm/executor/gpu_executor.py
@@ -31,7 +31,7 @@ def _init_executor(self) -> None:
         assert self.parallel_config.world_size == 1, (
             "GPUExecutor only supports single GPU.")
 
-        self.driver_worker = self._create_worker()
+        self.driver_worker = self._create_worker(local_rank=self.device_config.device.index)
         self.driver_worker.init_device()
         self.driver_worker.load_model()