iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
main
vLLM/vllm/distributed/parallel_state.py
T
Nick HillhjjqClaude Opus 5.5
 de7d70676d [Perf] Low-latency CPU all-reduce for DP metadata syncs (#61767) 
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: hjjq <50634613+hjjq@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-15 15:29:13 -07:00

2751 lines
100 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Copyright 2023 The vLLM team.
# Adapted from
# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
"""vLLM distributed state.
It takes over the control of the distributed environment from PyTorch.
The typical workflow is:
- call `init_distributed_environment` to initialize the distributed environment.
- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to
initialize the model parallel groups.
- any code dealing with the distributed stuff
- call `destroy_model_parallel` to destroy the model parallel groups.
- call `destroy_distributed_environment` to destroy the distributed environment.
If you only need to use the distributed environment without model/pipeline
parallelism, you can skip the model parallel initialization and destruction
steps.
"""
import contextlib
import gc
import pickle
import weakref
from collections import deque, namedtuple
from collections.abc import Callable
from contextlib import AbstractContextManager, contextmanager, nullcontext
from dataclasses import dataclass
from datetime import timedelta
from math import gcd
from multiprocessing import shared_memory
from typing import TYPE_CHECKING, Any, Protocol
from unittest.mock import patch
import torch
import torch.distributed
import torch.distributed._functional_collectives as funcol
import torch.distributed._symmetric_memory
from torch.distributed import Backend, ProcessGroup, Store
from torch.distributed.constants import default_pg_timeout
import vllm.envs as envs
from vllm.distributed.device_communicators.base_device_communicator import (
DeviceCommunicatorBase,
)
from vllm.distributed.utils import (
StatelessProcessGroup,
cpu_all_reduce_,
create_tcp_store,
get_cached_tcp_store_client,
)
from vllm.logger import init_logger
from vllm.utils.import_utils import resolve_obj_by_qualname
from vllm.utils.network_utils import get_distributed_init_method
from vllm.utils.system_utils import suppress_stdout
from vllm.utils.torch_utils import (
direct_register_custom_op,
)
if TYPE_CHECKING:
from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
@dataclass
class GraphCaptureContext:
stream: torch.cuda.Stream
TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
class Handle(Protocol):
"""Minimal async work handle used by P2P send/recv methods."""
def is_completed(self) -> bool: ...
def wait(self) -> None: ...
class _RetainedHandle:
"""Handle that retains source tensors until all posted work completes.
Used by ``isend_object`` to keep the serialized CPU tensors alive while
gloo copies them in the background; ``wait`` drops the refs once drained.
``wait`` must be idempotent: gloo ``Work.wait()`` is single-shot for
point-to-point operations — calling it a second time on an already
completed send blocks forever (verified empirically). Both the worker's
top-of-step drain of ``_pp_send_work`` and ``_reap_completed_isends``
may wait the same entry's metadata handle, so the second caller must
not re-wait the underlying works.
"""
def __init__(self, works: list[Any], retained: tuple[torch.Tensor, ...]) -> None:
self._works = works
self._retained = retained
self._waited = False
def is_completed(self) -> bool:
return all(w.is_completed() for w in self._works)
def wait(self) -> None:
if self._waited:
return
for w in self._works:
w.wait()
self._waited = True
self._retained = ()
def _split_tensor_dict(
tensor_dict: dict[str, torch.Tensor | Any],
) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:
"""Split the tensor dictionary into two parts:
1. A list of (key, value) pairs. If the value is a tensor, it is replaced
by its metadata.
2. A list of tensors.
"""
metadata_list: list[tuple[str, Any]] = []
tensor_list: list[torch.Tensor] = []
for key, value in tensor_dict.items():
if isinstance(value, torch.Tensor):
# Note: we cannot use `value.device` here,
# because it contains not only the device type but also the device
# index (e.g. "cuda:0"). We only need the device type.
# receiving side will set the device index.
device = value.device.type
metadata_list.append(
(key, TensorMetadata(device, value.dtype, value.size()))
)
tensor_list.append(value)
else:
metadata_list.append((key, value))
return metadata_list, tensor_list
_group_name_counter: dict[str, int] = {}
def _get_unique_name(name: str) -> str:
"""Get a unique name for the group.
Example:
_get_unique_name("tp") -> "tp:0"
_get_unique_name("tp") -> "tp:1"
"""
if name not in _group_name_counter:
_group_name_counter[name] = 0
newname = f"{name}:{_group_name_counter[name]}"
_group_name_counter[name] += 1
return newname
_groups: dict[str, Callable[[], "GroupCoordinator | None"]] = {}
def _register_group(group: "GroupCoordinator") -> None:
_groups[group.unique_name] = weakref.ref(group)
def _apply_to_device_comms(
action: Callable[[DeviceCommunicatorBase], None],
) -> None:
"""Apply ``action`` to every group's device communicator.
Walks the registered parallel groups and skips those without a device
communicator (absent at ``world_size == 1``).
"""
comms = []
for group_ref in _groups.values():
group = group_ref()
if group is None:
continue
dc = group.device_communicator
if dc is None:
continue
comms.append(dc)
for dc in comms:
action(dc)
def suspend_device_comms() -> None:
"""Release device communicator memory (collective; comms must be idle)."""
_apply_to_device_comms(lambda comm: comm.suspend())
def resume_device_comms() -> None:
"""Restore suspended device communicators (collective)."""
_apply_to_device_comms(lambda comm: comm.resume())
def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
assert group_name in _groups, f"Group {group_name} is not found."
group = _groups[group_name]()
if group is None:
raise ValueError(f"Group {group_name} is destroyed.")
return group._all_reduce_out_place(tensor)
def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
return torch.empty_like(tensor)
def reduce_scatter(
tensor: torch.Tensor, dim: int, world_size: int, group_name: str
) -> torch.Tensor:
assert group_name in _groups, f"Group {group_name} is not found."
group = _groups[group_name]()
if group is None:
raise ValueError(f"Group {group_name} is destroyed.")
return group._reduce_scatter_out_place(tensor, dim)
def reduce_scatter_fake(
tensor: torch.Tensor, dim: int, world_size: int, group_name: str
) -> torch.Tensor:
new_shape = list(tensor.shape)
new_shape[dim] = tensor.shape[dim] // world_size
return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device)
def all_gather(
tensor: torch.Tensor, dim: int, world_size: int, group_name: str
) -> torch.Tensor:
assert group_name in _groups, f"Group {group_name} is not found."
group = _groups[group_name]()
if group is None:
raise ValueError(f"Group {group_name} is destroyed.")
return group._all_gather_out_place(tensor, dim)
def all_gather_fake(
tensor: torch.Tensor, dim: int, world_size: int, group_name: str
) -> torch.Tensor:
new_shape = list(tensor.shape)
new_shape[dim] = tensor.shape[dim] * world_size
return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device)
def patched_fused_scaled_matmul_reduce_scatter_fake(
A: torch.Tensor,
B: torch.Tensor,
A_scale: torch.Tensor,
B_scale: torch.Tensor,
reduce_op: str,
orig_scatter_dim: int,
scatter_dim_after_maybe_reshape: int,
group_name: str,
output_shape: list[int],
bias: torch.Tensor | None = None,
result_scale: torch.Tensor | None = None,
out_dtype: torch.dtype | None = None,
use_fast_accum: bool = False,
) -> torch.Tensor:
# Copied from
# https://github.com/pytorch/pytorch/blob/50c338c2da905062449e4d9ac807832d1b5cd90e/torch/distributed/_symmetric_memory/__init__.py#L1189
if A_scale.numel() > 1:
if A_scale.shape[:-1] != A.shape[:-1]:
raise ValueError(
"For row-wise scaling, the leading dims of A_scale "
"must match the leading dims of A "
f"(A shape: {A.shape}, A_scale shape: {A_scale.shape})"
)
A_scale = A_scale.flatten(0, -2).contiguous()
elif A_scale.numel() != 1:
raise ValueError(
"Invalid A_scale shape "
f"(A shape: {A.shape}, A_scale shape: {A_scale.shape})"
)
C = torch._scaled_mm(
A.flatten(0, -2).contiguous(),
B,
A_scale,
B_scale,
bias,
result_scale,
out_dtype,
use_fast_accum,
)
C = C.view(*output_shape[:-1], B.shape[1])
res = funcol.reduce_scatter_tensor(
C,
reduce_op,
orig_scatter_dim, # need original scatter dim for 3D+ output tensor here
group_name,
)
res = funcol.wait_tensor(res)
return res
def _platform_device_type() -> str:
"""Return the device-type string (e.g. ``"cuda"``, ``"xpu"``, ``"cpu"``)
for the current platform, in the form expected by
``torch.distributed.init_process_group(backend=...)``.
"""
from vllm.platforms import current_platform
if current_platform.is_cuda_alike():
return "cuda"
elif current_platform.is_xpu():
return "xpu"
elif current_platform.is_out_of_tree():
return current_platform.device_name
else:
return "cpu"
def _device_backend_str(torch_distributed_backend: str | Backend) -> str:
"""Normalize ``torch_distributed_backend`` to the ``"<device>:<backend>"``
format required by ``split_group``'s ``backend`` argument.
Accepts either a bare backend name (e.g. ``"nccl"``) or an already-prefixed
string (e.g. ``"cuda:nccl"``).
"""
backend_str = str(torch_distributed_backend)
if ":" in backend_str:
return backend_str
return f"{_platform_device_type()}:{backend_str}"
def _create_subgroups_split_group(
group_ranks: list[list[int]],
group_name: str,
torch_distributed_backend: str | Backend,
) -> tuple[ProcessGroup, ProcessGroup]:
"""Create the device + CPU subgroups for ``GroupCoordinator`` via
``torch.distributed.split_group``.
``split_group`` is collective on the parent group, so every parent rank
must enter with the same ``split_ranks`` definition. Each rank receives
the subgroup it belongs to.
"""
from vllm.distributed.utils import (
get_cpu_distributed_timeout_or_none,
get_distributed_timeout_or_none,
)
device_backend_str = _device_backend_str(torch_distributed_backend)
self_device_group = torch.distributed.split_group(
split_ranks=group_ranks,
group_desc=f"{group_name}:device",
backend=device_backend_str,
timeout=get_distributed_timeout_or_none(),
)
# CPU subgroup: split_group requires the requested backend filter to
# include the parent's default device type (= the device the parent PG
# was bound to via ``device_id``), so a cpu-only filter is rejected.
# Include the device backend in the filter; only the gloo backend is
# actually used for CPU collectives on this group.
self_cpu_group = torch.distributed.split_group(
split_ranks=group_ranks,
group_desc=f"{group_name}:cpu",
backend=f"cpu:gloo,{device_backend_str}",
timeout=get_cpu_distributed_timeout_or_none(),
)
return self_device_group, self_cpu_group
def patched_fused_scaled_matmul_reduce_scatter(
A: torch.Tensor,
B: torch.Tensor,
A_scale: torch.Tensor,
B_scale: torch.Tensor,
reduce_op: str,
orig_scatter_dim: int,
scatter_dim_after_maybe_reshape: int,
group_name: str,
output_shape: list[int],
bias: torch.Tensor | None = None,
result_scale: torch.Tensor | None = None,
out_dtype: torch.dtype | None = None,
use_fast_accum: bool = False,
) -> torch.Tensor:
return torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter(
A,
B,
A_scale,
B_scale,
reduce_op,
orig_scatter_dim,
scatter_dim_after_maybe_reshape,
group_name,
output_shape,
bias,
result_scale,
out_dtype,
use_fast_accum,
)
direct_register_custom_op(
op_name="all_reduce",
op_func=all_reduce,
fake_impl=all_reduce_fake,
)
direct_register_custom_op(
op_name="reduce_scatter",
op_func=reduce_scatter,
fake_impl=reduce_scatter_fake,
)
direct_register_custom_op(
op_name="all_gather",
op_func=all_gather,
fake_impl=all_gather_fake,
)
# TODO: Remove this once the pytorch fix
# (https://github.com/pytorch/pytorch/pull/165086) gets released,
# in either 2.9.1 or 2.10
direct_register_custom_op(
op_name="patched_fused_scaled_matmul_reduce_scatter",
op_func=patched_fused_scaled_matmul_reduce_scatter,
fake_impl=patched_fused_scaled_matmul_reduce_scatter_fake,
)
class GroupCoordinator:
"""PyTorch ProcessGroup wrapper for a group of processes.
PyTorch ProcessGroup is bound to one specific communication backend,
e.g. NCCL, Gloo, MPI, etc.
GroupCoordinator takes charge of all the communication operations among
the processes in the group. It manages both CPU and device
communication.
"""
# available attributes:
rank: int # global rank
ranks: list[int] # global ranks in the group
world_size: int # size of the group
# difference between `local_rank` and `rank_in_group`:
# if we have a group of size 4 across two nodes:
# Process | Node | Rank | Local Rank | Rank in Group
# 0 | 0 | 0 | 0 | 0
# 1 | 0 | 1 | 1 | 1
# 2 | 1 | 2 | 0 | 2
# 3 | 1 | 3 | 1 | 3
local_rank: int # local rank used to assign devices
rank_in_group: int # rank inside the group
cpu_group: ProcessGroup # group for CPU communication
device_group: ProcessGroup # group for device communication
# device communicator (if use_device_communicator=True)
device_communicator: DeviceCommunicatorBase | None
mq_broadcaster: Any | None # shared memory broadcaster
def __init__(
self,
group_ranks: list[list[int]],
local_rank: int,
torch_distributed_backend: str | Backend,
use_device_communicator: bool, # whether to use device communicator
use_message_queue_broadcaster: bool = False,
group_name: str | None = None,
use_all2all: bool = False,
):
group_name = group_name or "anonymous"
self.unique_name = _get_unique_name(group_name)
_register_group(self)
self.rank = torch.distributed.get_rank()
self.local_rank = local_rank
self.device_index: int
assert local_rank >= 0, (
"local_rank must be provided when creating the world group"
)
self.device_index = local_rank
self_device_group = None
self_cpu_group = None
# VLLM_DISTRIBUTED_USE_SPLIT_GROUP gates the new ``split_group``
# codepath. Default (False) preserves the legacy ``new_group`` path.
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
self_device_group, self_cpu_group = _create_subgroups_split_group(
group_ranks, group_name, torch_distributed_backend
)
for ranks in group_ranks:
if self.rank in ranks:
self.ranks = ranks
self.world_size = len(ranks)
self.rank_in_group = ranks.index(self.rank)
break
else:
from vllm.distributed.utils import (
get_cpu_distributed_timeout_or_none,
get_distributed_timeout_or_none,
)
timeout = get_cpu_distributed_timeout_or_none()
device_timeout = get_distributed_timeout_or_none()
for ranks in group_ranks:
device_group = torch.distributed.new_group(
ranks,
backend=torch_distributed_backend,
timeout=device_timeout,
)
# a group with `gloo` backend, to allow direct coordination between
# processes through the CPU.
with suppress_stdout():
cpu_group = torch.distributed.new_group(
ranks, backend="gloo", timeout=timeout
)
if self.rank in ranks:
self.ranks = ranks
self.world_size = len(ranks)
self.rank_in_group = ranks.index(self.rank)
self_device_group = device_group
self_cpu_group = cpu_group
assert self_cpu_group is not None
assert self_device_group is not None
self.group_ranks = group_ranks
self.torch_distributed_backend = torch_distributed_backend
self.cpu_group = self_cpu_group
self.device_group = self_device_group
from vllm.platforms import current_platform
if current_platform.is_cuda_alike():
visible_device_index = (
current_platform.logical_device_id_to_visible_device_id(
self.device_index
)
)
self.device = torch.device(f"cuda:{visible_device_index}")
elif current_platform.is_xpu():
self.device = torch.device(f"xpu:{self.device_index}")
elif current_platform.is_out_of_tree():
self.device = torch.device(
f"{current_platform.device_name}:{self.device_index}"
)
else:
self.device = torch.device("cpu")
self.use_device_communicator = use_device_communicator
self.device_communicator = None
if use_device_communicator and self.world_size > 1:
device_comm_cls = resolve_obj_by_qualname(
current_platform.get_device_communicator_cls()
)
self.device_communicator = device_comm_cls(
cpu_group=self.cpu_group,
device=self.device,
device_group=self.device_group,
unique_name=self.unique_name,
use_all2all=use_all2all,
)
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
self.mq_broadcaster: MessageQueue | None = None
if use_message_queue_broadcaster and self.world_size > 1:
self.mq_broadcaster = MessageQueue.create_from_process_group(
self.cpu_group, 1 << 22, 6
)
# TODO(#35915): Remove is_tpu() check once tpu_inference
# overrides use_custom_op_collectives() to return True.
self.use_custom_op_call = (
current_platform.is_tpu() or current_platform.use_custom_op_collectives()
)
self.use_cpu_custom_send_recv = (
current_platform.is_cpu()
and self.device_communicator
and getattr(self.device_communicator, "supports_tensor_dict", False)
)
# FIFO of in-flight ``isend_tensor_dict`` calls: one entry per call,
# (handles, sent source tensors). Self-retention keeps the send
# buffers alive for fire-and-forget callers that drop the returned
# handles; entries are reaped lazily on the next ``isend_tensor_dict``
# (see ``_reap_completed_isends``). Lazily created instead of plain
# attribute init because some subclasses replicate ``__init__``
# without calling ``super().__init__``.
self._pending_isends: deque[tuple[list[Handle], list[torch.Tensor]]] = deque()
def make_sibling_device_group(self, group_desc: str | None = None) -> ProcessGroup:
"""Create a new device-side ProcessGroup with the same per-rank membership
as this coordinator's `device_group`, but backed by a distinct communicator.
This is a collective call: every world rank must invoke it. Used where we
want to issue ops that can run concurrently with ops on `device_group`.
"""
from vllm.distributed.utils import get_distributed_timeout_or_none
device_timeout = get_distributed_timeout_or_none()
sibling: ProcessGroup | None = None
for ranks in self.group_ranks:
pg = torch.distributed.new_group(
ranks,
backend=self.torch_distributed_backend,
group_desc=group_desc,
timeout=device_timeout,
)
if self.rank in ranks:
sibling = pg
assert sibling is not None
return sibling
def create_mq_broadcaster(
self, writer_rank=0, external_writer_handle=None, blocking=True
):
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
return MessageQueue.create_from_process_group(
self.cpu_group,
1 << 22,
6,
writer_rank=writer_rank,
external_writer_handle=external_writer_handle,
blocking=blocking,
)
def create_single_reader_mq_broadcasters(
self, reader_rank_in_group=0, blocking=False
):
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
return MessageQueue.create_from_process_group_single_reader(
self.cpu_group,
1 << 22,
6,
reader_rank=self.ranks[reader_rank_in_group],
blocking=blocking,
)
@property
def first_rank(self):
"""Return the global rank of the first process in the group."""
return self.ranks[0]
@property
def last_rank(self):
"""Return the global rank of the last process in the group."""
return self.ranks[-1]
@property
def is_first_rank(self):
"""Return whether the caller is the first process in the group."""
return self.rank == self.first_rank
@property
def is_last_rank(self):
"""Return whether the caller is the last process in the group."""
return self.rank == self.last_rank
@property
def next_rank(self):
"""Return the global rank of the process that follows the caller."""
rank_in_group = self.rank_in_group
world_size = self.world_size
return self.ranks[(rank_in_group + 1) % world_size]
@property
def prev_rank(self):
"""Return the global rank of the process that precedes the caller."""
rank_in_group = self.rank_in_group
world_size = self.world_size
return self.ranks[(rank_in_group - 1) % world_size]
@contextmanager
def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None):
if graph_capture_context is None:
stream = torch.cuda.Stream()
graph_capture_context = GraphCaptureContext(stream)
else:
stream = graph_capture_context.stream
# only cuda/rocm uses this function,
# so we don't abstract it into the base class
maybe_ca_context = nullcontext()
maybe_fi_pcie_ipc_context: AbstractContextManager[Any] = nullcontext()
maybe_aiter_ar_context = nullcontext()
from vllm.distributed.device_communicators.cuda_communicator import (
CudaCommunicator,
)
from vllm.distributed.device_communicators.xpu_communicator import (
XpuCommunicator,
)
if self.device_communicator is not None:
assert isinstance(
self.device_communicator,
(CudaCommunicator, XpuCommunicator),
)
ca_comm = self.device_communicator.ca_comm
if ca_comm is not None:
maybe_ca_context = ca_comm.capture() # type: ignore
if isinstance(self.device_communicator, CudaCommunicator):
fi_pcie_ipc_ar_comm = self.device_communicator.fi_pcie_ipc_ar_comm
if fi_pcie_ipc_ar_comm is not None:
maybe_fi_pcie_ipc_context = fi_pcie_ipc_ar_comm.capture()
# Capture each group's own comm. A global lookup would double-capture
aiter_ar_comm = getattr(self.device_communicator, "aiter_ar_comm", None)
if aiter_ar_comm is not None:
maybe_aiter_ar_context = aiter_ar_comm.capture() # type: ignore
# ensure all initialization operations complete before attempting to
# capture the graph on another stream
curr_stream = torch.cuda.current_stream()
if curr_stream != stream:
stream.wait_stream(curr_stream)
with (
torch.cuda.stream(stream),
maybe_ca_context,
maybe_fi_pcie_ipc_context,
maybe_aiter_ar_context,
):
yield graph_capture_context
def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
"""User-facing all-reduce function before we actually call the
all-reduce operation.
We need this because Dynamo does not support passing an arbitrary
object (`self` in this case) to a custom op. We need to pass the
group name as a string, and then look up the group coordinator from
the group name, dispatch the all-reduce operation to the group
coordinator.
In addition, PyTorch custom ops do not support mutation or returning
a new tensor in the same op. So we always make the all-reduce operation
out-of-place.
"""
# Bypass the function if we are using only 1 GPU.
if self.world_size == 1:
return input_
if self.use_custom_op_call:
return torch.ops.vllm.all_reduce(input_, group_name=self.unique_name)
else:
return self._all_reduce_out_place(input_)
def _all_reduce_out_place(self, input_: torch.Tensor) -> torch.Tensor:
if self.device_communicator is None:
raise ValueError("No device communicator found")
return self.device_communicator.all_reduce(input_)
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
world_size = self.world_size
# Bypass the function if we are using only 1 GPU.
if world_size == 1:
return input_
assert -input_.dim() <= dim < input_.dim(), (
f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
)
if self.use_custom_op_call:
return torch.ops.vllm.all_gather(
input_, dim, world_size, group_name=self.unique_name
)
else:
return self._all_gather_out_place(input_, dim)
def _all_gather_out_place(self, input_: torch.Tensor, dim: int) -> torch.Tensor:
if self.device_communicator is None:
raise ValueError("No device communicator found")
return self.device_communicator.all_gather(input_, dim)
def all_gatherv(
self,
input_: torch.Tensor | list[torch.Tensor],
dim: int = 0,
sizes: list[int] | None = None,
):
if self.device_communicator is None:
raise ValueError("No device communicator found")
return self.device_communicator.all_gatherv(input_, dim, sizes)
def reduce_scatter(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
world_size = self.world_size
# Bypass the function if we are using only 1 GPU.
if world_size == 1:
return input_
assert -input_.dim() <= dim < input_.dim(), (
f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
)
if self.use_custom_op_call:
return torch.ops.vllm.reduce_scatter(
input_, dim, world_size, group_name=self.unique_name
)
else:
return self._reduce_scatter_out_place(input_, dim)
def reduce_scatterv(
self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None
) -> torch.Tensor:
if self.device_communicator is None:
raise ValueError("No device communicator found")
return self.device_communicator.reduce_scatterv(input_, dim, sizes)
def reduce_scatterv_into_output(
self,
input_: torch.Tensor,
output: torch.Tensor,
dim: int = -1,
sizes: list[int] | None = None,
) -> torch.Tensor:
if self.device_communicator is None:
raise ValueError("No device communicator found")
return self.device_communicator.reduce_scatterv_into_output(
input_, output, dim, sizes
)
def _reduce_scatter_out_place(self, input_: torch.Tensor, dim: int) -> torch.Tensor:
if self.device_communicator is None:
raise ValueError("No device communicator found")
return self.device_communicator.reduce_scatter(input_, dim)
def gather(
self, input_: torch.Tensor, dst: int = 0, dim: int = -1
) -> torch.Tensor | None:
"""NOTE: We assume that the input tensor is on the same device across
all the ranks.
NOTE: `dst` is the local rank of the destination rank.
"""
world_size = self.world_size
# Bypass the function if we are using only 1 GPU.
if world_size == 1:
return input_
if self.device_communicator is None:
raise ValueError("No device communicator found")
return self.device_communicator.gather(input_, dst, dim)
def broadcast(self, input_: torch.Tensor, src: int = 0):
"""Broadcast the input tensor.
NOTE: `src` is the local rank of the source rank.
"""
assert src < self.world_size, f"Invalid src rank ({src})"
# Bypass the function if we are using only 1 GPU.
if self.world_size == 1:
return input_
# Broadcast.
torch.distributed.broadcast(
input_, src=self.ranks[src], group=self.device_group
)
return input_
def broadcast_object(self, obj: Any | None = None, src: int = 0):
"""Broadcast the input object.
NOTE: `src` is the local rank of the source rank.
"""
assert src < self.world_size, f"Invalid src rank ({src})"
# Bypass the function if we are using only 1 GPU.
if self.world_size == 1:
return obj
if self.mq_broadcaster is not None:
assert src == 0, "Message queue broadcaster only supports src=0"
return self.mq_broadcaster.broadcast_object(obj)
if self.rank_in_group == src:
torch.distributed.broadcast_object_list(
[obj], src=self.ranks[src], group=self.cpu_group
)
return obj
else:
recv = [None]
torch.distributed.broadcast_object_list(
recv, src=self.ranks[src], group=self.cpu_group
)
return recv[0]
def broadcast_object_list(
self, obj_list: list[Any], src: int = 0, group: ProcessGroup | None = None
):
"""Broadcast the input object list.
NOTE: `src` is the local rank of the source rank.
"""
assert src < self.world_size, f"Invalid src rank ({src})"
# Bypass the function if we are using only 1 GPU.
if self.world_size == 1:
return obj_list
# Broadcast.
torch.distributed.broadcast_object_list(
obj_list, src=self.ranks[src], group=self.device_group
)
return obj_list
def send_object(self, obj: Any, dst: int) -> None:
"""Send the input object list to the destination rank."""
"""NOTE: `dst` is the local rank of the destination rank."""
assert dst < self.world_size, f"Invalid dst rank ({dst})"
assert dst != self.rank_in_group, (
"Invalid destination rank. Destination rank is the same "
"as the current rank."
)
# Serialize object to tensor and get the size as well
object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8)
size_tensor = torch.tensor(
[object_tensor.numel()], dtype=torch.long, device="cpu"
)
# Send object size
torch.distributed.send(size_tensor, dst=self.ranks[dst], group=self.cpu_group)
# Send object
torch.distributed.send(object_tensor, dst=self.ranks[dst], group=self.cpu_group)
return None
def recv_object(self, src: int) -> Any:
"""Receive the input object list from the source rank."""
"""NOTE: `src` is the local rank of the source rank."""
assert src < self.world_size, f"Invalid src rank ({src})"
assert src != self.rank_in_group, (
"Invalid source rank. Source rank is the same as the current rank."
)
size_tensor = torch.empty(1, dtype=torch.long, device="cpu")
# Receive object size
rank_size = torch.distributed.recv(
size_tensor, src=self.ranks[src], group=self.cpu_group
)
# Tensor to receive serialized objects into.
object_tensor = torch.empty( # type: ignore[call-overload]
size_tensor.item(), # type: ignore[arg-type]
dtype=torch.uint8,
device="cpu",
)
rank_object = torch.distributed.recv(
object_tensor, src=self.ranks[src], group=self.cpu_group
)
assert rank_object == rank_size, (
"Received object sender rank does not match the size sender rank."
)
obj = pickle.loads(object_tensor.numpy().tobytes())
return obj
def isend_object(self, obj: Any, dst: int) -> Handle:
"""Non-blocking ``send_object``: isend size + pickled object on the
CPU group. Returns a handle that retains the serialized source
tensors until ``wait`` drains both sends.
"""
assert dst < self.world_size, f"Invalid dst rank ({dst})"
assert dst != self.rank_in_group, (
"Invalid destination rank. Destination rank is the same "
"as the current rank."
)
object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8)
size_tensor = torch.tensor(
[object_tensor.numel()], dtype=torch.long, device="cpu"
)
retained = (size_tensor, object_tensor)
works: list[Any] = []
for tensor in retained:
work = torch.distributed.isend(
tensor, dst=self.ranks[dst], group=self.cpu_group
)
if work is not None:
works.append(work)
return _RetainedHandle(works, retained)
def broadcast_tensor_dict(
self,
tensor_dict: dict[str, torch.Tensor | Any] | None = None,
src: int = 0,
group: ProcessGroup | None = None,
metadata_group: ProcessGroup | None = None,
) -> dict[str, torch.Tensor | Any] | None:
"""Broadcast the input tensor dictionary.
NOTE: `src` is the local rank of the source rank.
"""
# Bypass the function if we are using only 1 GPU.
if not torch.distributed.is_initialized() or self.world_size == 1:
return tensor_dict
group = self.device_group
metadata_group = self.cpu_group
assert src < self.world_size, f"Invalid src rank ({src})"
rank_in_group = self.rank_in_group
if rank_in_group == src:
metadata_list: list[tuple[Any, Any]] = []
assert isinstance(tensor_dict, dict), (
f"Expecting a dictionary, got {type(tensor_dict)}"
)
metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
# `metadata_list` lives in CPU memory.
# `broadcast_object_list` has serialization & deserialization,
# all happening on CPU. Therefore, we can use the CPU group.
self.broadcast_object(metadata_list, src=src)
async_handles = []
for tensor in tensor_list:
if tensor.numel() == 0:
# Skip broadcasting empty tensors.
continue
if tensor.is_cpu:
# use metadata_group for CPU tensors
handle = torch.distributed.broadcast(
tensor, src=self.ranks[src], group=metadata_group, async_op=True
)
else:
# use group for GPU tensors
handle = torch.distributed.broadcast(
tensor, src=self.ranks[src], group=group, async_op=True
)
async_handles.append(handle)
for async_handle in async_handles:
async_handle.wait()
else:
metadata_list = self.broadcast_object(None, src=src)
tensor_dict = {}
async_handles = []
for key, value in metadata_list:
if isinstance(value, TensorMetadata):
tensor = torch.empty(
value.size, dtype=value.dtype, device=value.device
)
if tensor.numel() == 0:
# Skip broadcasting empty tensors.
tensor_dict[key] = tensor
continue
if tensor.is_cpu:
# use metadata_group for CPU tensors
handle = torch.distributed.broadcast(
tensor,
src=self.ranks[src],
group=metadata_group,
async_op=True,
)
else:
# use group for GPU tensors
handle = torch.distributed.broadcast(
tensor, src=self.ranks[src], group=group, async_op=True
)
async_handles.append(handle)
tensor_dict[key] = tensor
else:
tensor_dict[key] = value
for async_handle in async_handles:
async_handle.wait()
return tensor_dict
def _should_use_all_gather(
self,
key: str,
numel: int,
all_gather_group: "GroupCoordinator | None",
all_gather_tensors: dict[str, bool] | None,
) -> bool:
if all_gather_group is None:
return False
use_all_gather = numel % all_gather_group.world_size == 0
if all_gather_tensors is not None:
use_all_gather = all_gather_tensors.get(key, use_all_gather)
return use_all_gather
def send_tensor_dict(
self,
tensor_dict: dict[str, torch.Tensor | Any],
dst: int | None = None,
all_gather_group: "GroupCoordinator | None" = None,
all_gather_tensors: dict[str, bool] | None = None,
) -> dict[str, torch.Tensor | Any] | None:
"""Send the input tensor dictionary.
NOTE: `dst` is the local rank of the source rank.
all_gather_group: The group for the all-gather operation. If provided,
an optimization is enabled where each rank in the group sends a
slice of a tensor and the receiver reconstructs it using an
all-gather, which can improve performance. This is typically the
tensor-parallel group.
all_gather_tensors: A dictionary to specify which tensors should use
the all-gather optimization, which is only effective when
`all_gather_group` is provided. By default, this optimization is
on for any tensor whose size is divisible by the
`all_gather_group`'s world size. However, it should be disabled
for tensors that are not fully replicated across the group (e.g.,
the residual tensor when sequence parallelism is enabled). This
dictionary allows overriding the default behavior on a per-tensor
basis.
"""
# Bypass the function if we are using only 1 GPU.
if not torch.distributed.is_initialized() or self.world_size == 1:
return tensor_dict
handles = self.isend_tensor_dict(
tensor_dict,
dst=dst,
all_gather_group=all_gather_group,
all_gather_tensors=all_gather_tensors,
)
for handle in handles:
handle.wait()
# The sync path waited every handle, so the self-retention entry
# appended by ``isend_tensor_dict`` is already drained; drop it
# instead of leaking the tensor refs until the next async send.
pending = getattr(self, "_pending_isends", None)
if pending and pending[-1][0] is handles:
pending.pop()
return None
def _reap_completed_isends(self) -> None:
"""Lazily drop self-retained ``isend_tensor_dict`` entries that have
completed, oldest first (FIFO).
gloo ``Work.is_completed()`` is unreliable (it can report completion
before the background copy of the source buffer finishes), so the
metadata handle (``handles[0]``, a gloo-backed ``_RetainedHandle``) is
never used as the completion gate. Instead we gate on the tensor-send
handles (``handles[1:]``), which run on the device backend where
``is_completed()`` is reliable. Once all tensor sends are done the
peer must already have posted the matching tensor recvs — and since
the receiver ingests metadata before tensors, the gloo metadata send
is then guaranteed complete, so ``wait()`` on it is ~instant and only
serves to drop the retained refs.
Metadata-only entries (``handles[1:]`` empty, e.g. an empty
``IntermediateTensors``) have no reliable completion signal; drop them
best-effort without blocking on the gloo metadata wait.
"""
pending = getattr(self, "_pending_isends", None)
if pending is None:
# Subclasses that replicate __init__ without calling
# super().__init__ may not have the attribute yet.
self._pending_isends = pending = deque()
while pending:
handles, _ = pending[0]
tensor_handles = handles[1:]
if not tensor_handles:
pending.popleft()
continue
if not all(h.is_completed() for h in tensor_handles):
# The oldest send is still in flight; FIFO ordering means the
# rest are no further along, so stop here.
break
handles[0].wait()
pending.popleft()
def isend_tensor_dict(
self,
tensor_dict: dict[str, torch.Tensor | Any],
dst: int | None = None,
all_gather_group: "GroupCoordinator | None" = None,
all_gather_tensors: dict[str, bool] | None = None,
) -> list[Handle]:
"""Send the input tensor dictionary asynchronously.
The returned handles are additionally self-retained in
``self._pending_isends`` (together with the sent source tensors) and
reaped lazily on subsequent calls, so fire-and-forget callers that
drop the handles cannot free or overwrite the send buffers while the
sends are still in flight. Callers that do keep the handles may
``wait()`` them as before; the retention entry is dropped on reap
either way.
"""
if self.world_size <= 1:
return []
if dst is None:
dst = (self.rank_in_group + 1) % self.world_size
assert dst < self.world_size, f"Invalid dst rank ({dst})"
if self.use_cpu_custom_send_recv:
if self.device_communicator is None:
raise ValueError("No device communicator found")
# custom device communicator path is synchronous
self.device_communicator.send_tensor_dict( # type: ignore
tensor_dict, dst
)
return []
all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size
all_gather_rank = (
0 if all_gather_group is None else all_gather_group.rank_in_group
)
group = self.device_group
metadata_group = self.cpu_group
metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
tensor_keys = [k for k, v in tensor_dict.items() if isinstance(v, torch.Tensor)]
assert len(tensor_keys) == len(tensor_list)
# Reap completed sends before posting new ones, bounding the
# self-retention FIFO.
self._reap_completed_isends()
handles: list[Handle] = [self.isend_object(metadata_list, dst=dst)]
sent_tensors: list[torch.Tensor] = []
for key, tensor in zip(tensor_keys, tensor_list):
if tensor.numel() == 0:
continue
if self._should_use_all_gather(
key, tensor.numel(), all_gather_group, all_gather_tensors
):
tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank]
comm_group = metadata_group if tensor.is_cpu else group
handle = torch.distributed.isend(
tensor, dst=self.ranks[dst], group=comm_group
)
if tensor.is_cuda:
tensor.record_stream(torch.cuda.current_stream(tensor.device))
handles.append(handle)
sent_tensors.append(tensor)
# Self-retain for fire-and-forget callers: keep the handles and
# the source tensors alive until the sends complete (reaped
# lazily by ``_reap_completed_isends``). Without this, dropping
# the returned handles would GC the ``_RetainedHandle``'s
# serialized metadata buffers mid-send, and non-CUDA devices
# have no ``record_stream`` equivalent protecting the tensor
# sources.
pending = getattr(self, "_pending_isends", None)
if pending is None:
self._pending_isends = pending = deque()
pending.append((handles, sent_tensors))
return handles
def recv_tensor_dict(
self,
src: int | None = None,
all_gather_group: "GroupCoordinator | None" = None,
all_gather_tensors: dict[str, bool] | None = None,
) -> dict[str, torch.Tensor | Any] | None:
"""Recv the input tensor dictionary.
NOTE: `src` is the local rank of the source rank.
all_gather_group: The group for the all-gather operation. If provided,
an optimization is enabled where each rank in the group sends a
slice of a tensor and the receiver reconstructs it using an
all-gather, which can improve performance. This is typically the
tensor-parallel group.
all_gather_tensors: A dictionary to specify which tensors should use
the all-gather optimization, which is only effective when
`all_gather_group` is provided. By default, this optimization is
on for any tensor whose size is divisible by the
`all_gather_group`'s world size. However, it should be disabled
for tensors that are not fully replicated across the group (e.g.,
the residual tensor when sequence parallelism is enabled). This
dictionary allows overriding the default behavior on a per-tensor
basis.
"""
# Bypass the function if we are using only 1 GPU.
if not torch.distributed.is_initialized() or self.world_size == 1:
return None
tensor_dict, handles, postprocess = self.irecv_tensor_dict(
src=src,
all_gather_group=all_gather_group,
all_gather_tensors=all_gather_tensors,
)
for handle in handles:
handle.wait()
for fn in postprocess:
fn()
return tensor_dict
def irecv_tensor_dict(
self,
src: int | None = None,
all_gather_group: "GroupCoordinator | None" = None,
all_gather_tensors: dict[str, bool] | None = None,
) -> tuple[
dict[str, torch.Tensor | Any] | None,
list[Handle],
list[Callable[[], None]],
]:
if not torch.distributed.is_initialized() or self.world_size == 1:
return None, [], []
if src is None:
src = (self.rank_in_group - 1) % self.world_size
assert src < self.world_size, f"Invalid src rank ({src})"
if self.use_cpu_custom_send_recv:
if self.device_communicator is None:
raise ValueError("No device communicator found")
# custom device communicator path is synchronous
sync_tensor_dict = self.device_communicator.recv_tensor_dict( # type: ignore
src
)
return sync_tensor_dict, [], []
all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size
all_gather_rank = (
0 if all_gather_group is None else all_gather_group.rank_in_group
)
group = self.device_group
metadata_group = self.cpu_group
recv_metadata_list = self.recv_object(src=src)
tensor_dict: dict[str, Any] = {}
handles: list[Handle] = []
postprocess: list[Callable[[], None]] = []
for key, value in recv_metadata_list:
if isinstance(value, TensorMetadata):
full_tensor = torch.empty(
value.size, dtype=value.dtype, device=value.device
)
if full_tensor.numel() == 0:
tensor_dict[key] = full_tensor
continue
if self._should_use_all_gather(
key, full_tensor.numel(), all_gather_group, all_gather_tensors
):
orig_shape = full_tensor.shape
slice_tensor = full_tensor.reshape(all_gather_size, -1)[
all_gather_rank
]
comm_group = metadata_group if slice_tensor.is_cpu else group
handle = torch.distributed.irecv(
slice_tensor, src=self.ranks[src], group=comm_group
)
handles.append(handle)
def _postprocess(
key: str = key,
slice_tensor: torch.Tensor = slice_tensor,
orig_shape: tuple[int, ...] = tuple(orig_shape),
all_gather_group=all_gather_group,
) -> None:
assert all_gather_group is not None
tensor_dict[key] = all_gather_group.all_gather(
slice_tensor, dim=0
).reshape(orig_shape)
postprocess.append(_postprocess)
tensor_dict[key] = slice_tensor
else:
comm_group = metadata_group if full_tensor.is_cpu else group
handle = torch.distributed.irecv(
full_tensor, src=self.ranks[src], group=comm_group
)
handles.append(handle)
tensor_dict[key] = full_tensor
else:
tensor_dict[key] = value
return tensor_dict, handles, postprocess
def barrier(self):
"""Barrier synchronization among the group.
NOTE: don't use `device_group` here! `barrier` in NCCL is
terrible because it is internally a broadcast operation with
secretly created GPU tensors. It is easy to mess up the current
device. Use the CPU group instead.
"""
torch.distributed.barrier(group=self.cpu_group)
def cpu_all_reduce_(self, tensor: torch.Tensor) -> torch.Tensor:
"""In-place SUM all-reduce of a small CPU tensor over `cpu_group`."""
return cpu_all_reduce_(tensor, self.cpu_group)
def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
"""Sends a tensor to the destination rank in a blocking way."""
"""NOTE: `dst` is the local rank of the destination rank."""
if self.device_communicator is None:
raise ValueError("No device communicator found")
self.device_communicator.send(tensor, dst)
def recv(
self, size: torch.Size, dtype: torch.dtype, src: int | None = None
) -> torch.Tensor:
"""Receives a tensor from the source rank."""
"""NOTE: `src` is the local rank of the source rank."""
if self.device_communicator is None:
raise ValueError("No device communicator found")
return self.device_communicator.recv(size, dtype, src)
def destroy(self):
# Device communicators can own collective workspaces whose teardown
# uses these process groups (for example FlashInfer PCIe IPC barriers).
# Release them before destroying the groups they depend on.
if self.device_communicator is not None:
self.device_communicator.destroy()
if hasattr(self, "device_group"):
torch.distributed.destroy_process_group(self.device_group)
del self.device_group
if hasattr(self, "cpu_group"):
torch.distributed.destroy_process_group(self.cpu_group)
del self.cpu_group
if self.mq_broadcaster is not None:
self.mq_broadcaster = None
def dispatch_router_logits(
self,
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
is_sequence_parallel: bool = False,
extra_tensors: list[torch.Tensor] | None = None,
) -> (
tuple[torch.Tensor, torch.Tensor]
| tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
):
if self.device_communicator is not None:
return self.device_communicator.dispatch_router_logits(
hidden_states,
router_logits,
is_sequence_parallel,
extra_tensors,
)
else:
return hidden_states, router_logits
def dispatch(
self,
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
is_sequence_parallel: bool = False,
extra_tensors: list[torch.Tensor] | None = None,
) -> (
tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
| tuple[torch.Tensor, torch.Tensor, torch.Tensor]
):
if self.device_communicator is not None:
return self.device_communicator.dispatch(
hidden_states,
topk_weights,
topk_ids,
is_sequence_parallel,
extra_tensors,
)
else:
return hidden_states, topk_weights, topk_ids
def combine(
self, hidden_states, is_sequence_parallel: bool = False
) -> torch.Tensor:
if self.device_communicator is not None:
return self.device_communicator.combine(hidden_states, is_sequence_parallel)
else:
return hidden_states
def allocate_combine_input(
self,
shape: tuple[int, ...],
dtype: torch.dtype,
device: torch.device,
is_sequence_parallel: bool = False,
) -> torch.Tensor | None:
if self.device_communicator is None:
return None
return self.device_communicator.allocate_combine_input(
shape, dtype, device, is_sequence_parallel
)
def combine_into_output(
self,
hidden_states: torch.Tensor,
output: torch.Tensor,
is_sequence_parallel: bool = False,
) -> torch.Tensor:
if self.device_communicator is None:
output.copy_(hidden_states)
return output
return self.device_communicator.combine_into_output(
hidden_states, output, is_sequence_parallel
)
_WORLD: GroupCoordinator | None = None
_INNER_DP_WORLD: GroupCoordinator | None = None
_NODE_COUNT: int | None = None
def get_world_group() -> GroupCoordinator:
assert _WORLD is not None, "world group is not initialized"
return _WORLD
def get_inner_dp_world_group() -> GroupCoordinator:
assert _INNER_DP_WORLD is not None, "inner dp world group is not initialized"
return _INNER_DP_WORLD
def init_world_group(
ranks: list[int], local_rank: int, backend: str
) -> GroupCoordinator:
return GroupCoordinator(
group_ranks=[ranks],
local_rank=local_rank,
torch_distributed_backend=backend,
use_device_communicator=False,
group_name="world",
)
def init_model_parallel_group(
group_ranks: list[list[int]],
local_rank: int,
backend: str,
use_message_queue_broadcaster: bool = False,
group_name: str | None = None,
use_device_communicator: bool = True,
use_all2all: bool = False,
) -> GroupCoordinator:
return GroupCoordinator(
group_ranks=group_ranks,
local_rank=local_rank,
torch_distributed_backend=backend,
use_device_communicator=use_device_communicator,
use_message_queue_broadcaster=use_message_queue_broadcaster,
group_name=group_name,
use_all2all=use_all2all,
)
def _init_stateless_group(
group_ranks: list[list[int]],
group_name: str,
host: str,
backend: str,
coord_store: Store,
use_device_communicator: bool = True,
use_all2all: bool = False,
) -> "StatelessGroupCoordinator":
"""Create a StatelessGroupCoordinator with the given parameters."""
from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
world = get_world_group()
return StatelessGroupCoordinator(
group_ranks=group_ranks,
local_rank=world.local_rank,
torch_distributed_backend=backend,
use_device_communicator=use_device_communicator,
group_name=group_name,
host=host,
coord_store=coord_store,
global_rank=world.rank,
global_world_size=world.world_size,
use_all2all=use_all2all,
)
def _replace_active_groups(
*,
world: GroupCoordinator | None,
dp: GroupCoordinator | None,
ep: GroupCoordinator | None,
eplb: GroupCoordinator | None,
node_count: int | None,
) -> tuple[GroupCoordinator | None, ...]:
"""Replace the active groups and return the groups they replaced.
The caller must destroy the returned DP, EP, WORLD, and EPLB groups
collectively and in that order. Pass all-``None`` to remove the active
groups without replacement.
"""
global _WORLD, _DP, _EP, _EPLB, _NODE_COUNT
old_groups = _DP, _EP, _WORLD, _EPLB
_WORLD = world
_DP = dp
_EP = ep
_EPLB = eplb
_NODE_COUNT = node_count
return old_groups
_TP: GroupCoordinator | None = None
def get_tp_group() -> GroupCoordinator:
assert _TP is not None, "tensor model parallel group is not initialized"
return _TP
_ETP: GroupCoordinator | None = None
def get_etp_group() -> GroupCoordinator:
"""Return the Engram Tensor Parallel (ETP) group that shards one embedding table."""
assert _ETP is not None, "Engram tensor-parallel group is not initialized"
return _ETP
_ENGRAM_DP: GroupCoordinator | None = None
def get_engram_dp_group() -> GroupCoordinator | None:
"""Return the DP replicas that share one engram embedding table.
None when every replica holds a full (TP-sharded) copy, which is the case
for models without engram layers and for replicas that span nodes.
"""
return _ENGRAM_DP
def get_engram_dp_size() -> int:
"""Number of DP replicas one engram embedding table is sharded over."""
return _ENGRAM_DP.world_size if _ENGRAM_DP is not None else 1
_DCP: GroupCoordinator | None = None
def get_dcp_group() -> GroupCoordinator:
assert _DCP is not None, "decode context model parallel group is not initialized"
return _DCP
def get_dcp_world_size_and_rank(enabled: bool = True) -> tuple[int, int]:
"""Return ``(world_size, rank)`` in the DCP group, or ``(1, 0)`` when disabled
(e.g. a replicated draft cache) or the group is uninitialized (unit tests)."""
if not enabled or _DCP is None:
return 1, 0
return _DCP.world_size, _DCP.rank_in_group
_PP: GroupCoordinator | None = None
def get_pp_group() -> GroupCoordinator:
assert _PP is not None, "pipeline model parallel group is not initialized"
return _PP
_DP: GroupCoordinator | None = None
def get_dp_group() -> GroupCoordinator:
assert _DP is not None, "data parallel group is not initialized"
return _DP
_EP: GroupCoordinator | None = None
def get_ep_group() -> GroupCoordinator:
assert _EP is not None, (
"expert parallel group is not initialized. "
"EP group is only created for MoE models with num_experts > 0. "
"This function should only be called for MoE models."
)
return _EP
_EPLB: GroupCoordinator | None = None
def get_eplb_group() -> GroupCoordinator:
assert _EPLB is not None, (
"EPLB group is not initialized. "
"EPLB group is only created for MoE models when EPLB is enabled. "
"Ensure parallel_config.enable_eplb is True."
)
return _EPLB
_PCP: GroupCoordinator | None = None
def get_pcp_group() -> GroupCoordinator:
assert _PCP is not None, "prefill context parallel group is not initialized"
return _PCP
@contextmanager
def graph_capture(device: torch.device):
"""`graph_capture` is a context manager which should surround the code that
is capturing the CUDA graph. Its main purpose is to ensure that some
operations will be run after the graph is captured, before the graph
is replayed. It returns a `GraphCaptureContext` object which contains the
necessary data for the graph capture. Currently, it only contains the
stream that the graph capture is running on. This stream is set to the
current CUDA stream when the context manager is entered and reset to the
default stream when the context manager is exited. This is to ensure that
the graph capture is running on a separate stream from the default stream,
in order to explicitly distinguish the kernels to capture
from other kernels possibly launched on background in the default stream.
"""
context = GraphCaptureContext(torch.cuda.Stream(device=device))
with (
get_tp_group().graph_capture(context),
get_pp_group().graph_capture(context),
get_dp_group().graph_capture(context),
):
yield context
logger = init_logger(__name__)
_ENABLE_CUSTOM_ALL_REDUCE = True
def set_custom_all_reduce(enable: bool):
global _ENABLE_CUSTOM_ALL_REDUCE
_ENABLE_CUSTOM_ALL_REDUCE = enable
def _init_process_group_for_split_group(
*,
backend: str,
distributed_init_method: str,
store: Store | None,
world_size: int,
rank: int,
local_rank: int,
timeout: timedelta | None,
) -> None:
"""Initialize the default PG with both CPU (gloo) and device (e.g. nccl)
backends and an eager ``device_id`` binding so that subgroups can be
created via ``split_group`` (which requires the parent communicator to
be eagerly initialized). Falls back to ``gloo`` on CPU-only systems.
"""
if torch.accelerator.is_available() and backend != "gloo":
init_backend = "cpu:gloo,cuda:nccl"
from vllm.platforms import current_platform
visible_device_index = current_platform.logical_device_id_to_visible_device_id(
local_rank
)
device_id: torch.device | None = torch.device(f"cuda:{visible_device_index}")
else:
init_backend = "gloo"
device_id = None
torch.distributed.init_process_group(
backend=init_backend,
init_method=distributed_init_method if store is None else None,
store=store,
world_size=world_size,
rank=rank,
timeout=timeout,
device_id=device_id,
)
def _validate_default_pg_for_split_group() -> None:
"""When an external launcher (e.g. ``torchrun``) initialized the default
PG, ``GroupCoordinator`` cannot patch in additional backends or change
the eager-init behavior — ``split_group`` only selects subsets of an
existing parent. Validate that the parent has both ``device_id`` and a
CPU (gloo) backend, and emit a descriptive error pointing at the exact
init call to update otherwise.
"""
default_pg = torch.distributed.distributed_c10d._get_default_group()
assert default_pg.bound_device_id is not None, (
"External launcher initialized the default process group "
"without device_id. vLLM requires the default PG to be device-"
"bound for split_group. Pass device_id=torch.device(f'cuda:"
"{local_rank}') to torch.distributed.init_process_group()."
)
try:
default_pg._get_backend(torch.device("cpu"))
except RuntimeError as e:
raise RuntimeError(
"External launcher initialized the default process group "
"without a CPU (gloo) backend. vLLM requires both CPU and "
"device backends. Pass backend='cpu:gloo,cuda:nccl' to "
"torch.distributed.init_process_group()."
) from e
def _init_elastic_ep_world(
config, local_rank: int, backend: str, rank: int, world_size: int
) -> None:
from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
global _WORLD, _NODE_COUNT
assert _WORLD is None, "world group already initialized"
parallel_config = config.parallel_config
global_rank = parallel_config.data_parallel_rank * world_size + rank
global_world_size = parallel_config.world_size_across_dp
all_ranks = list(range(global_world_size))
group_ranks = [all_ranks[i : i + 1] for i in range(global_world_size)]
if global_rank in all_ranks:
group_ranks = [all_ranks]
coord_store = get_cached_tcp_store_client(
parallel_config.data_parallel_master_ip, parallel_config._coord_store_port
)
world = StatelessGroupCoordinator(
group_ranks=group_ranks,
local_rank=local_rank,
torch_distributed_backend=backend,
use_device_communicator=False,
group_name="world",
host=parallel_config.data_parallel_master_ip,
coord_store=coord_store,
global_rank=global_rank,
global_world_size=global_world_size,
)
assert parallel_config.nnodes_within_dp == 1, (
"Elastic EP is not supported with multi-node TP/PP"
)
_NODE_COUNT = _node_count(world.tcp_store_group)
_WORLD = world
def init_distributed_environment(
world_size: int = -1,
rank: int = -1,
distributed_init_method: str = "env://",
local_rank: int = -1,
backend: str = "nccl",
timeout: timedelta | None = None,
):
logger.debug(
"world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s",
world_size,
rank,
local_rank,
distributed_init_method,
backend,
)
from vllm.config import get_current_vllm_config_or_none
config = get_current_vllm_config_or_none()
enable_elastic_ep = config is not None and config.parallel_config.enable_elastic_ep
world_pg_store: Store | None = None
if (
config is not None
and config.parallel_config.distributed_executor_backend != "external_launcher"
and (
config.parallel_config.nnodes > 1
or config.parallel_config.data_parallel_size > 1
)
and not enable_elastic_ep
):
parallel_config = config.parallel_config
# adjust to take into account data parallelism
# offset the rank by the data parallel rank
rank = parallel_config.data_parallel_rank * world_size + rank
# adjust the world size to take into account data parallelism
world_size = parallel_config.world_size_across_dp
# Use appropriate IP and port based on configuration
if parallel_config.nnodes > 1:
ip = parallel_config.master_addr
port = parallel_config.master_port
distributed_init_method = get_distributed_init_method(ip, port)
else:
ip = parallel_config.data_parallel_master_ip
if (
parallel_config._coord_store_port
and not torch.distributed.is_initialized()
):
# Group rank 0 lets the world-group TCPStore bind port 0
# and publishes the kernel-assigned port via the
# coordination store. A one-shot client of our own: the
# cached one is inherited across fork and shared with
# sibling workers.
coord_store = create_tcp_store(
ip,
parallel_config._coord_store_port,
is_master=False,
wait_for_workers=False,
)
if rank == 0:
# Workers can only connect once the port is published.
world_pg_store = create_tcp_store(
ip,
0,
world_size=world_size,
is_master=True,
wait_for_workers=False,
timeout=timeout or default_pg_timeout,
multi_tenant=True,
)
port = world_pg_store.port
coord_store.set("world_pg_port", str(port).encode())
else:
port = int(coord_store.get("world_pg_port").decode())
world_pg_store = create_tcp_store(
ip,
port,
world_size=world_size,
is_master=False,
timeout=timeout or default_pg_timeout,
multi_tenant=True,
)
else:
port = parallel_config.get_next_dp_init_port()
distributed_init_method = get_distributed_init_method(ip, port)
logger.debug(
"Adjusting world_size=%d rank=%d distributed_init_method=%s for DP",
world_size,
rank,
distributed_init_method,
)
if not torch.distributed.is_initialized():
logger.info(
"world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s",
world_size,
rank,
local_rank,
distributed_init_method,
backend,
)
assert distributed_init_method is not None, (
"distributed_init_method must be provided when initializing "
"distributed environment"
)
if not torch.distributed.is_backend_available(backend):
logger.warning(
"Distributed backend %s is not available; falling back to gloo.",
backend,
)
assert torch.distributed.is_gloo_available(), (
"Fallback Gloo backend is not available."
)
backend = "gloo"
store = world_pg_store
if store is None and distributed_init_method.startswith("file://"):
store = torch.distributed.FileStore(
distributed_init_method.removeprefix("file://"), world_size
)
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
# split_group needs local_rank early to compute device_id for
# the eager init. local_rank is not available in torch
# ProcessGroup, see https://github.com/pytorch/pytorch/issues/122816
if local_rank == -1:
local_rank = (
int(envs.LOCAL_RANK)
if distributed_init_method == "env://"
else rank
)
_init_process_group_for_split_group(
backend=backend,
distributed_init_method=distributed_init_method,
store=store,
world_size=world_size,
rank=rank,
local_rank=local_rank,
timeout=timeout,
)
else:
# this backend is used for WORLD
torch.distributed.init_process_group(
backend=backend,
init_method=distributed_init_method if store is None else None,
store=store,
world_size=world_size,
rank=rank,
timeout=timeout,
)
if enable_elastic_ep:
tp_pp_cpu_group = torch.distributed.new_group(
backend="gloo", timeout=timeout
)
if _node_count(tp_pp_cpu_group) > 1:
# NOTE(yongji): StatelessGroupCoordinator uses data_parallel_master_ip
# to initialize all DP/EP groups, hence all ranks within TP/PP group
# must reside on the same node
raise RuntimeError(
"Elastic EP is not yet supported with multi-node TP/PP"
)
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP and torch.accelerator.is_available():
_validate_default_pg_for_split_group()
# set the local rank
# local_rank is not available in torch ProcessGroup,
# see https://github.com/pytorch/pytorch/issues/122816
if local_rank == -1:
# local rank not set, this usually happens in single-node
# setting, where we can use rank as local rank
local_rank = envs.LOCAL_RANK if distributed_init_method == "env://" else rank
global _WORLD, _NODE_COUNT, _INNER_DP_WORLD
if enable_elastic_ep:
_init_elastic_ep_world(config, local_rank, backend, rank, world_size)
return
if _WORLD is None:
ranks = list(range(torch.distributed.get_world_size()))
_WORLD = init_world_group(ranks, local_rank, backend)
if config is not None and config.parallel_config.nnodes > 1:
_NODE_COUNT = config.parallel_config.nnodes
else:
_NODE_COUNT = _node_count(_WORLD.cpu_group)
logger.debug("Detected %d nodes in the distributed environment", _NODE_COUNT)
else:
assert _WORLD.world_size == torch.distributed.get_world_size(), (
"world group already initialized with a different world size"
)
if config is not None and config.parallel_config.nnodes_within_dp > 1:
if parallel_config.data_parallel_size > 1:
world_size_inner_dp = parallel_config.world_size
group_ranks = [
[dp_rank * world_size_inner_dp + i for i in range(world_size_inner_dp)]
for dp_rank in range(parallel_config.data_parallel_size)
]
_INNER_DP_WORLD = init_model_parallel_group(
group_ranks,
get_world_group().local_rank,
backend,
use_message_queue_broadcaster=True,
group_name="inner_dp_world",
use_device_communicator=False,
)
else:
_INNER_DP_WORLD = _WORLD
def _elastic_join_warmup_ctx() -> AbstractContextManager[Any]:
"""Defer communicator warm-up when joining a live elastic-EP cluster on ROCm.
The existing ranks defer their matching warm-up in create_standby_groups(),
so a joining worker must too or the warm-up collective has no peers. Both
sides warm up at commit.
"""
from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm
if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
return defer_comm_warmup_on_rocm()
return nullcontext()
def _engram_dp_shard_size(
world_size: int,
data_parallel_size: int,
replica_size: int,
) -> int:
"""Size of DP shards aligned with complete replicas and node boundaries."""
node_count = get_node_count()
replicas_per_node, remainder = divmod(world_size, node_count * replica_size)
if remainder:
return 1
# Shards must divide both the DP group and each node's replica count.
shard_size = gcd(data_parallel_size, replicas_per_node)
if shard_size == 1 or node_count == 1:
return shard_size
ranks_per_node = replicas_per_node * replica_size
for start in range(0, world_size, ranks_per_node):
same_node = in_the_same_node_as(get_world_group().cpu_group, start)
node_ranks = [rank for rank, local in enumerate(same_node) if local]
if node_ranks != list(range(start, start + ranks_per_node)):
return 1
return shard_size
def initialize_model_parallel(
tensor_model_parallel_size: int = 1,
pipeline_model_parallel_size: int = 1,
prefill_context_model_parallel_size: int = 1,
decode_context_model_parallel_size: int | None = 1,
backend: str | None = None,
) -> None:
"""Initialize model parallel groups.
Arguments:
tensor_model_parallel_size: number of GPUs used for tensor model
parallelism.
pipeline_model_parallel_size: number of GPUs used for pipeline model
parallelism.
prefill_context_model_parallel_size: number of GPUs used for context
parallelism during prefill.
decode_context_model_parallel_size: number of GPUs used for context
parallelism during decode.
backend: name of torch distributed communication backend.
Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
the model pipeline. The present function will
create 4 tensor model-parallel groups and 2 pipeline model-parallel groups:
4 tensor model-parallel groups:
[g0, g1], [g2, g3], [g4, g5], [g6, g7]
2 pipeline model-parallel groups:
[g0, g2, g4, g6], [g1, g3, g5, g7]
Note that for efficiency, the caller should make sure adjacent ranks
are on the same DGX box. For example if we are using 2 DGX-1 boxes
with a total of 16 GPUs, rank 0 to 7 belong to the first box and
ranks 8 to 15 belong to the second box.
"""
# Get world size and rank. Ensure some consistencies.
assert torch.distributed.is_initialized()
from vllm.config import get_current_vllm_config
config = get_current_vllm_config()
data_parallel_size = config.parallel_config.data_parallel_size
enable_elastic_ep = config.parallel_config.enable_elastic_ep
parallel_config = config.parallel_config
coord_store: Store | None = None
if enable_elastic_ep:
coord_store = get_cached_tcp_store_client(
parallel_config.data_parallel_master_ip,
parallel_config._coord_store_port,
)
# Use stateless world group for global information
world_size = get_world_group().world_size
rank = get_world_group().rank
backend = backend or "nccl"
tp_pp_pcp_size = (
tensor_model_parallel_size
* pipeline_model_parallel_size
* prefill_context_model_parallel_size
)
local_all_ranks = torch.arange(tp_pp_pcp_size).reshape(
pipeline_model_parallel_size,
prefill_context_model_parallel_size,
tensor_model_parallel_size,
)
else:
world_size = torch.distributed.get_world_size()
rank = torch.distributed.get_rank()
backend = backend or torch.distributed.get_backend(
get_world_group().device_group
)
# the layout order is: ExternalDP x DP x PP x PCP x TP
# ExternalDP is the data parallel group that is not part of the model,
# every dp rank can generate independently (in verl integration).
# DP is the data parallel group that is part of the model,
# all the ranks in the same DP group should generate simultaneously,
# i.e. the `generate` call in the same DP group should be called together,
# otherwise it will cause deadlock.
# to get group_ranks for each dimension, transpose that dimension to the
# last dimension, then reshape to 2D, then unbind the last dimension
all_ranks = torch.arange(world_size).reshape(
-1,
data_parallel_size,
pipeline_model_parallel_size,
prefill_context_model_parallel_size,
tensor_model_parallel_size,
) # noqa
# Build the tensor model-parallel groups.
global _TP
assert _TP is None, "tensor model parallel group is already initialized"
group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0)
group_ranks = [x.tolist() for x in group_ranks]
if enable_elastic_ep:
group_ranks = local_all_ranks.view(-1, tensor_model_parallel_size).unbind(0)
group_ranks = [x.tolist() for x in group_ranks]
# message queue broadcaster is only used in tensor model parallel group
_TP = init_model_parallel_group(
group_ranks,
get_world_group().local_rank,
backend,
use_message_queue_broadcaster=True,
group_name="tp",
)
global _ETP
assert _ETP is None, "Engram tensor-parallel group is already initialized"
engram_tensor_parallel_size = (
config.engram_config.get_parallel_size(parallel_config)
if config.engram_config is not None
else tensor_model_parallel_size
)
if engram_tensor_parallel_size == tensor_model_parallel_size:
_ETP = _TP
else:
group_ranks = (
all_ranks.permute(0, 2, 3, 1, 4)
.reshape(-1, engram_tensor_parallel_size)
.unbind(0)
)
_ETP = init_model_parallel_group(
[ranks.tolist() for ranks in group_ranks],
get_world_group().local_rank,
backend,
group_name="etp",
)
# Build the node-local DP group used to shard or share an Engram table.
global _ENGRAM_DP
assert _ENGRAM_DP is None, "engram data parallel group is already initialized"
engram_dp_size = 1
engram_config = config.engram_config
if (
engram_config is not None
and config.model_config is not None
and config.model_config.architecture == "DeepseekV41ForCausalLM"
and not enable_elastic_ep
):
engram_dp_size = _engram_dp_shard_size(
world_size,
data_parallel_size,
tensor_model_parallel_size
* pipeline_model_parallel_size
* prefill_context_model_parallel_size,
)
if engram_dp_size > 1:
group_ranks = (
all_ranks.permute(0, 2, 3, 4, 1).reshape(-1, engram_dp_size).unbind(0)
)
_ENGRAM_DP = init_model_parallel_group(
[x.tolist() for x in group_ranks],
get_world_group().local_rank,
backend,
group_name="edp",
)
logger.info_once(
"Engram node-local DP group size: %d (TP=%d)",
engram_dp_size,
tensor_model_parallel_size,
)
# Build the DCP model-parallel groups.
global _DCP
assert _DCP is None, "decode context model parallel group is already initialized"
dcp_size = decode_context_model_parallel_size or 1
dcp_ranks = local_all_ranks if enable_elastic_ep else all_ranks
if dcp_size > 1:
# DCP spans PCP first, then TP for full TP x PCP groups.
dcp_ranks = dcp_ranks.transpose(-1, -2)
group_ranks = dcp_ranks.reshape(-1, dcp_size).unbind(0)
group_ranks = [x.tolist() for x in group_ranks]
_DCP = init_model_parallel_group(
group_ranks,
get_world_group().local_rank,
backend,
use_message_queue_broadcaster=True,
group_name="dcp",
)
global _PCP
assert _PCP is None, "prefill context parallel group is already initialized"
group_ranks = (
all_ranks.transpose(3, 4)
.reshape(-1, prefill_context_model_parallel_size)
.unbind(0)
)
group_ranks = [x.tolist() for x in group_ranks]
if enable_elastic_ep:
group_ranks = (
local_all_ranks.transpose(1, 2)
.reshape(-1, prefill_context_model_parallel_size)
.unbind(0)
)
group_ranks = [x.tolist() for x in group_ranks]
_PCP = init_model_parallel_group(
group_ranks, get_world_group().local_rank, backend, group_name="pcp"
)
# Build the pipeline model-parallel groups.
global _PP
assert _PP is None, "pipeline model parallel group is already initialized"
group_ranks = (
all_ranks.transpose(2, 4).reshape(-1, pipeline_model_parallel_size).unbind(0)
)
group_ranks = [x.tolist() for x in group_ranks]
if enable_elastic_ep:
group_ranks = (
local_all_ranks.transpose(0, 2)
.reshape(-1, pipeline_model_parallel_size)
.unbind(0)
)
group_ranks = [x.tolist() for x in group_ranks]
_PP = init_model_parallel_group(
group_ranks, get_world_group().local_rank, backend, group_name="pp"
)
global _DP
assert _DP is None, "data parallel group is already initialized"
group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0)
group_ranks = [x.tolist() for x in group_ranks]
if enable_elastic_ep:
with _elastic_join_warmup_ctx():
_DP = _init_stateless_group(
group_ranks,
"dp",
parallel_config.data_parallel_master_ip,
backend,
coord_store=coord_store,
)
else:
_DP = init_model_parallel_group(
group_ranks, get_world_group().local_rank, backend, group_name="dp"
)
global _EP
assert _EP is None, "expert parallel group is already initialized"
# Don't create EP group for dense models.
if config.model_config is None or config.model_config.is_moe:
group_ranks = (
all_ranks.transpose(1, 2)
.reshape(
-1,
data_parallel_size
* prefill_context_model_parallel_size
* tensor_model_parallel_size,
)
.unbind(0)
)
group_ranks = [x.tolist() for x in group_ranks]
use_all2all = parallel_config.use_all2all
if enable_elastic_ep:
with _elastic_join_warmup_ctx():
_EP = _init_stateless_group(
group_ranks,
"ep",
parallel_config.data_parallel_master_ip,
backend,
coord_store=coord_store,
use_all2all=use_all2all,
)
else:
_EP = init_model_parallel_group(
group_ranks,
get_world_group().local_rank,
backend,
group_name="ep",
use_all2all=use_all2all,
)
# Create EPLB group with the same ranks as EP if EPLB is enabled.
# This is a separate process group to isolate EPLB communications
# from MoE forward pass collectives and prevent deadlocks when
# using torch.distributed in execution with torch.distributed in EPLB.
global _EPLB
assert _EPLB is None, "EPLB group is already initialized"
if config.parallel_config.enable_eplb:
if enable_elastic_ep:
with _elastic_join_warmup_ctx():
_EPLB = _init_stateless_group(
group_ranks,
"eplb",
parallel_config.data_parallel_master_ip,
backend,
coord_store=coord_store,
)
else:
_EPLB = init_model_parallel_group(
group_ranks,
get_world_group().local_rank,
backend,
group_name="eplb",
)
# If no EP group needed, _EP remains None
# If no EPLB group needed, _EPLB remains None
logger.info_once(
"rank %s in world size %s is assigned as "
"DP rank %s, PP rank %s, PCP rank %s, "
"TP rank %s, ETP rank %s, EP rank %s, EPLB rank %s",
rank,
world_size,
_DP.rank_in_group,
_PP.rank_in_group,
_PCP.rank_in_group,
_TP.rank_in_group,
_ETP.rank_in_group,
_EP.rank_in_group if _EP is not None else "N/A",
_EPLB.rank_in_group if _EPLB is not None else "N/A",
)
def ensure_model_parallel_initialized(
tensor_model_parallel_size: int,
pipeline_model_parallel_size: int,
prefill_context_model_parallel_size: int = 1,
decode_context_model_parallel_size: int | None = 1,
backend: str | None = None,
) -> None:
"""Helper to initialize model parallel groups if they are not initialized,
or ensure tensor-parallel and pipeline-parallel sizes are equal to expected
values if the model parallel groups are initialized.
"""
world_group = get_world_group()
if hasattr(world_group, "backend"):
backend = backend or world_group.backend
else:
backend = backend or torch.distributed.get_backend(world_group.device_group)
if not model_parallel_is_initialized():
initialize_model_parallel(
tensor_model_parallel_size,
pipeline_model_parallel_size,
prefill_context_model_parallel_size,
decode_context_model_parallel_size,
backend,
)
return
assert get_tensor_model_parallel_world_size() == tensor_model_parallel_size, (
"tensor parallel group already initialized, but of unexpected size. "
f"got: {get_tensor_model_parallel_world_size()=} vs. "
f"wanted: {tensor_model_parallel_size=}"
)
pp_world_size = get_pp_group().world_size
assert pp_world_size == pipeline_model_parallel_size, (
"pipeline parallel group already initialized, but of unexpected size. "
f"got: {pp_world_size=} vs. "
f"wanted: {pipeline_model_parallel_size=}"
)
pcp_world_size = get_pcp_group().world_size
assert pcp_world_size == prefill_context_model_parallel_size, (
"prefill context parallel group already initialized, but of unexpected size: "
f"{pcp_world_size=} vs. "
f"{prefill_context_model_parallel_size=}"
)
dcp_world_size = get_dcp_group().world_size
dcp_model_parallel_size = decode_context_model_parallel_size or 1
assert dcp_world_size == dcp_model_parallel_size, (
"decode context parallel group already initialized, but of unexpected size: "
f"{dcp_world_size=} vs. "
f"{dcp_model_parallel_size=}"
)
def checkpoint_prepare_distributed_state() -> None:
"""Prepare every device communicator for a process checkpoint."""
torch.accelerator.synchronize()
_apply_to_device_comms(lambda comm: comm.checkpoint_prepare())
torch.accelerator.synchronize()
def checkpoint_restore_distributed_state() -> None:
"""Restore every device communicator after a process checkpoint."""
torch.accelerator.synchronize()
_apply_to_device_comms(lambda comm: comm.checkpoint_restore())
torch.accelerator.synchronize()
def model_parallel_is_initialized():
"""Check if tensor and pipeline parallel groups are initialized."""
return _TP is not None and _PP is not None
def get_tensor_model_parallel_world_size() -> int:
"""Return world size for the tensor model parallel group."""
return get_tp_group().world_size
def get_tensor_model_parallel_rank() -> int:
"""Return my rank for the tensor model parallel group."""
return get_tp_group().rank_in_group
def get_pipeline_model_parallel_rank() -> int:
"""Return my rank for the pipeline model parallel group."""
return get_pp_group().rank_in_group
def get_node_count() -> int:
"""Return the total number of nodes in the distributed environment."""
assert _NODE_COUNT is not None, "distributed environment is not initialized"
return _NODE_COUNT
def destroy_model_parallel():
"""Set the groups to none and destroy them."""
global _TP, _ETP
if _ETP and _ETP is not _TP:
_ETP.destroy()
_ETP = None
if _TP:
_TP.destroy()
_TP = None
global _ENGRAM_DP
if _ENGRAM_DP:
_ENGRAM_DP.destroy()
_ENGRAM_DP = None
global _DCP
if _DCP:
_DCP.destroy()
_DCP = None
global _PCP
if _PCP:
_PCP.destroy()
_PCP = None
global _PP
if _PP:
_PP.destroy()
_PP = None
global _DP
if _DP:
_DP.destroy()
_DP = None
global _EP
if _EP:
_EP.destroy()
_EP = None
global _EPLB
if _EPLB:
_EPLB.destroy()
_EPLB = None
def destroy_distributed_environment():
global _WORLD, _NODE_COUNT
if _WORLD:
_WORLD.destroy()
_WORLD = None
_NODE_COUNT = None
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()
def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
logger.debug(
"[shutdown] Distributed: cleanup start shutdown_ray=%s",
shutdown_ray,
)
# Reset environment variable cache
envs.disable_envs_cache()
# Reset rocm_aiter_ops class variables to match current os.environ.
# These are class-level attributes that persist across tests and are
# NOT restored by monkeypatch (which only restores os.environ).
from vllm.platforms import current_platform
if current_platform.is_rocm():
from vllm._aiter_ops import rocm_aiter_ops
rocm_aiter_ops.refresh_env_variables()
# Ensure all objects are not frozen before cleanup
gc.unfreeze()
destroy_model_parallel()
destroy_distributed_environment()
if shutdown_ray:
import ray # Lazy import Ray
ray.shutdown()
gc.collect()
from vllm.platforms import current_platform
if not current_platform.is_cpu():
from vllm.triton_utils import HAS_TRITON
if HAS_TRITON:
from vllm.v1.sample.ops.topk_topp_triton import reset_buffer_cache
reset_buffer_cache()
torch.accelerator.empty_cache()
try:
torch.accelerator.empty_host_cache()
except AttributeError:
logger.warning(
"torch.accelerator.empty_host_cache() only available in Pytorch >=2.9"
)
logger.debug_once("[shutdown] Distributed: cleanup complete")
def in_the_same_node_as(
pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0
) -> list[bool]:
"""This is a collective operation that returns if each rank is in the same node
as the source rank. It tests if processes are attached to the same
memory system (shared access to shared memory).
"""
if isinstance(pg, ProcessGroup):
assert torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL, (
"in_the_same_node_as should be tested with a non-NCCL group."
)
# local rank inside the group
rank = torch.distributed.get_rank(group=pg)
world_size = torch.distributed.get_world_size(group=pg)
# global ranks of the processes in the group
ranks = torch.distributed.get_process_group_ranks(pg)
else:
rank = pg.rank
world_size = pg.world_size
ranks = list(range(world_size))
# local tensor in each process to store the result
is_in_the_same_node = torch.tensor(
[0] * world_size, dtype=torch.int32, device="cpu"
)
magic_message = b"magic_message"
shm = None
try:
with contextlib.suppress(OSError):
if rank == source_rank:
# create a shared memory segment
shm = shared_memory.SharedMemory(create=True, size=128)
assert shm.buf is not None, "Buffer was not created"
shm.buf[: len(magic_message)] = magic_message
if isinstance(pg, ProcessGroup):
torch.distributed.broadcast_object_list(
[shm.name], src=ranks[source_rank], group=pg
)
else:
pg.broadcast_obj(shm.name, src=source_rank)
is_in_the_same_node[rank] = 1
else:
# try to open the shared memory segment
if isinstance(pg, ProcessGroup):
recv = [None]
torch.distributed.broadcast_object_list(
recv, src=ranks[source_rank], group=pg
)
name = recv[0]
else:
name = pg.broadcast_obj(None, src=source_rank)
# fix to https://stackoverflow.com/q/62748654/9191338
# Python incorrectly tracks shared memory even if it is not
# created by the process. The following patch is a workaround.
with patch(
"multiprocessing.resource_tracker.register",
lambda *args, **kwargs: None,
):
shm = shared_memory.SharedMemory(name=name)
assert shm.buf is not None, "Buffer was not opened"
if shm.buf[: len(magic_message)] == magic_message:
is_in_the_same_node[rank] = 1
except Exception as e:
logger.error("Error ignored in is_in_the_same_node: %s", e)
finally:
if shm:
shm.close()
if isinstance(pg, ProcessGroup):
torch.distributed.barrier(group=pg)
else:
pg.barrier()
# clean up the shared memory segment
with contextlib.suppress(OSError):
if rank == source_rank and shm:
shm.unlink()
if isinstance(pg, ProcessGroup):
torch.distributed.all_reduce(is_in_the_same_node, group=pg)
aggregated_data = is_in_the_same_node
else:
aggregated_data = torch.zeros_like(is_in_the_same_node)
for i in range(world_size):
rank_data = pg.broadcast_obj(is_in_the_same_node, src=i)
aggregated_data += rank_data
return [x == 1 for x in aggregated_data.tolist()]
def is_global_first_rank() -> bool:
"""Check if the current process is the first rank globally across all
parallelism strategies (PP, TP, DP, EP, etc.).
Unlike group-specific checks like `get_tensor_model_parallel_rank() == 0`
or `get_pp_group().is_first_rank`, this function checks the global rank
across all parallelism dimensions.
Returns:
bool: True if this is the global first rank (rank 0), False otherwise.
Returns True if distributed is not initialized (single process).
"""
try:
# If world group is available, use it for the most accurate check
global _WORLD
if _WORLD is not None:
return _WORLD.is_first_rank
# If torch distributed is not initialized, assume single process
if not torch.distributed.is_initialized():
return True
# Fallback to torch's global rank
return torch.distributed.get_rank() == 0
except Exception:
# If anything goes wrong, assume this is the first rank
return True
def is_local_first_rank() -> bool:
"""Check if the current process is the first local rank (rank 0 on its node)."""
try:
# prefer the initialized world group if available
global _WORLD
if _WORLD is not None:
return _WORLD.local_rank == 0
if not torch.distributed.is_initialized():
return True
# fallback to environment-provided local rank if available
# note: envs.LOCAL_RANK is set when using env:// launchers (e.g., torchrun)
try:
return int(envs.LOCAL_RANK) == 0 # type: ignore[arg-type]
except Exception:
return torch.distributed.get_rank() == 0
except Exception:
return True
def _node_count(pg: ProcessGroup | StatelessProcessGroup) -> int:
"""Returns the total number of nodes in the process group.
Args:
pg: The process group to analyze
Returns:
int: The total number of nodes
"""
if isinstance(pg, ProcessGroup):
world_size = torch.distributed.get_world_size(group=pg)
else:
world_size = pg.world_size
if world_size == 1:
return 1
# Build node assignment map
node_assignment = [0] * world_size # rank -> node_id
next_node_id = 0
for current_rank in range(world_size):
if node_assignment[current_rank] != 0:
continue # Already assigned to a node
# Assign current rank to a new node
next_node_id += 1
node_assignment[current_rank] = next_node_id
# Find all ranks on the same node as current_rank
same_node_flags = in_the_same_node_as(pg, current_rank)
for other_rank, is_same_node in enumerate(same_node_flags):
if is_same_node and node_assignment[other_rank] == 0:
node_assignment[other_rank] = next_node_id
return next_node_id
Reference in New Issue View Git Blame Copy Permalink