Unverified Commit 398a596e authored by weiguihua2's avatar weiguihua2 Committed by GitHub
Browse files

[MP executor] fix get device count for multi node of mp executor feature (#30042)


Signed-off-by: default avatarweiguihua2 <weiguihua2@huawei.com>
parent 67312cad
...@@ -27,6 +27,7 @@ from zmq import ( # type: ignore ...@@ -27,6 +27,7 @@ from zmq import ( # type: ignore
import vllm.envs as envs import vllm.envs as envs
from vllm.distributed.utils import StatelessProcessGroup, sched_yield from vllm.distributed.utils import StatelessProcessGroup, sched_yield
from vllm.logger import init_logger from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.network_utils import ( from vllm.utils.network_utils import (
get_ip, get_ip,
get_open_port, get_open_port,
...@@ -632,7 +633,7 @@ class MessageQueue: ...@@ -632,7 +633,7 @@ class MessageQueue:
The MessageQueue instance for the calling process, The MessageQueue instance for the calling process,
and a list of handles (only non-empty for the reader process). and a list of handles (only non-empty for the reader process).
""" """
local_size = torch.cuda.device_count() local_size = current_platform.device_count()
rank = dist.get_rank() rank = dist.get_rank()
same_node = rank // local_size == reader_rank // local_size same_node = rank // local_size == reader_rank // local_size
buffer_io = MessageQueue( buffer_io = MessageQueue(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment