mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-07-29 16:34:40 +08:00
[MP executor] fix get device count for multi node of mp executor feature (#30042)
Signed-off-by: weiguihua2 <weiguihua2@huawei.com>
This commit is contained in:
parent
67312cad11
commit
398a596ed2
@ -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:
|
|||||||
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(
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user