From 9cf4759493919580011f03812abf16387eafe18c Mon Sep 17 00:00:00 2001 From: Mengqing Cao Date: Tue, 11 Feb 2025 21:20:53 +0800 Subject: [PATCH] [executor] init `local_rank` as device index (#13027) Signed-off-by: Mengqing Cao --- vllm/executor/uniproc_executor.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/executor/uniproc_executor.py b/vllm/executor/uniproc_executor.py index e5464cafaecbf..94db232240d55 100644 --- a/vllm/executor/uniproc_executor.py +++ b/vllm/executor/uniproc_executor.py @@ -28,6 +28,11 @@ class UniProcExecutor(ExecutorBase): distributed_init_method = get_distributed_init_method( get_ip(), get_open_port()) local_rank = 0 + # set local rank as the device index if specified + device_info = self.vllm_config.device_config.device.__str__().split( + ":") + if len(device_info) > 1: + local_rank = int(device_info[1]) rank = 0 kwargs = dict( vllm_config=self.vllm_config,