class PyNcclCommunicator:
# None for communicators built via `from_unique_id_bytes` (no process group).
group: ProcessGroup | StatelessProcessGroup | None
def __init__(
self,
group: ProcessGroup | StatelessProcessGroup,
device: int | str | torch.device,
library_path: str | None = None,
):
"""
Args:
group: the process group to work on. If None, it will use the
default process group.
device: the device to bind the PyNcclCommunicator to. If None,
it will be bound to f"cuda:{local_rank}".
library_path: the path to the NCCL library. If None, it will
use the default library path.
It is the caller's responsibility to make sure each communicator
is bind to a unique device.
"""
if not isinstance(group, StatelessProcessGroup):
assert dist.is_initialized()
assert dist.get_backend(group) != dist.Backend.NCCL, (
"PyNcclCommunicator should be attached to a non-NCCL group."
)
# note: this rank is the rank in the group
self.rank = dist.get_rank(group)
self.world_size = dist.get_world_size(group)
else:
self.rank = group.rank
self.world_size = group.world_size
self.group = group
# if world_size == 1, no need to create communicator
if self.world_size == 1 or envs.VLLM_DISABLE_PYNCCL:
self.available = False
self.disabled = True
return
try:
self.nccl = NCCLLibrary(library_path)
except Exception:
# disable because of missing NCCL library
# e.g. in a non-GPU environment
self.available = False
self.disabled = True
return
self.available = True
self.disabled = False
self._suspended = False
self.nccl_version = self.nccl.ncclGetRawVersion()
if self.rank == 0:
# get the unique id from NCCL
self.unique_id = self.nccl.ncclGetUniqueId()
logger.info_once("vLLM is using nccl==%s", self.nccl.ncclGetVersion())
else:
# construct an empty unique id
self.unique_id = ncclUniqueId()
if not isinstance(group, StatelessProcessGroup):
tensor = torch.ByteTensor(list(self.unique_id.internal))
ranks = dist.get_process_group_ranks(group)
# arg `src` in `broadcast` is the global rank
dist.broadcast(tensor, src=ranks[0], group=group)
byte_list = tensor.tolist()
for i, byte in enumerate(byte_list):
self.unique_id.internal[i] = byte
else:
self.unique_id = group.broadcast_obj(self.unique_id, src=0)
self._init_comm(device)
def _init_comm(self, device: int | str | torch.device) -> None:
"""Create the communicator on `device` from the already-resolved
`self.unique_id` / `self.rank` / `self.world_size`, then run the
one-element warm-up all_reduce. Shared by `__init__` and
`from_unique_id_bytes` so the init handshake stays identical on both.
"""
if isinstance(device, int):
device = torch.device(f"cuda:{device}")
elif isinstance(device, str):
device = torch.device(device)
# now `device` is a `torch.device` object
assert isinstance(device, torch.device)
self.device = device
# nccl communicator and stream will use this device
with torch.accelerator.device_index(device.index):
self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
self.world_size, self.unique_id, self.rank
)
stream = current_stream()
# A small all_reduce for warmup.
data = torch.zeros(1, device=device)
self.all_reduce(data)
stream.synchronize()
del data
@classmethod
def from_unique_id_bytes(
cls,
unique_id_bytes: bytes,
rank: int,
world_size: int,
device: int | str | torch.device,
library_path: str | None = None,
) -> "PyNcclCommunicator":
"""Build a communicator from pre-shared ``ncclUniqueId`` bytes.
For peers that cannot join a ``StatelessProcessGroup`` / TCPStore (e.g. a
torch-free JAX trainer): every rank passes the same id, minted once via
``ncclGetUniqueId`` and shared out of band. There is no barrier, so all
ranks must enter init concurrently or ``ncclCommInitRank`` hangs.
Warm-up handshake: immediately after ``ncclCommInitRank`` this issues a
one-element ``all_reduce`` (mirroring ``__init__``). It is a collective,
so every peer -- including a foreign, non-vLLM rank -- must issue a
matching one-element ``all_reduce`` before any other collective, or all
ranks deadlock.
"""
if len(unique_id_bytes) != NCCL_UNIQUE_ID_BYTES:
raise ValueError(
f"expected a {NCCL_UNIQUE_ID_BYTES}-byte NCCL unique id, "
f"got {len(unique_id_bytes)} bytes"
)
if not 0 <= rank < world_size:
raise ValueError(f"rank {rank} out of range for world_size {world_size}")
self = cls.__new__(cls)
self.rank = rank
self.world_size = world_size
self.group = None
if self.world_size == 1 or envs.VLLM_DISABLE_PYNCCL:
self.available = False
self.disabled = True
return self
try:
self.nccl = NCCLLibrary(library_path)
except Exception as e:
# Unlike the TCPStore path, silently disabling here leaves the peer
# blocked in ncclCommInitRank until timeout. The caller explicitly
# asked to join from a unique id, so fail loudly instead.
raise RuntimeError(
"failed to load the NCCL library for unique-id rendezvous"
) from e
self.available = True
self.disabled = False
self.nccl_version = self.nccl.ncclGetRawVersion()
self.unique_id = self.nccl.unique_id_from_bytes(unique_id_bytes)
self._init_comm(device)
return self
def destroy(self):
if self.available and not self.disabled:
# ncclCommAbort can block until all CUDA graphs that
# captured NCCL ops on this comm are destroyed — and
# those graphs are released later in this same main-
# thread teardown, so a direct call here self-deadlocks.
# Run it in a daemon thread with a timeout: the main
# thread proceeds, the graphs drop, and the abort returns.
def _abort():
with torch.accelerator.device_index(self.device.index):
self.nccl.ncclCommAbort(self.comm)
abort_thread = threading.Thread(target=_abort, daemon=True)
abort_thread.start()
abort_thread.join(timeout=5.0)
self.available = False
self.disabled = True
def all_reduce(
self,
in_tensor: torch.Tensor,
out_tensor: torch.Tensor = None,
op: ReduceOp = ReduceOp.SUM,
stream=None,
) -> torch.Tensor:
if self.disabled:
return None
# nccl communicator created on a specific device
# will only work on tensors on the same device
# otherwise it will cause "illegal memory access"
assert in_tensor.device == self.device, (
f"this nccl communicator is created to work on {self.device}, "
f"but the input tensor is on {in_tensor.device}"
)
if out_tensor is None:
out_tensor = torch.empty_like(in_tensor)
if stream is None:
stream = current_stream()
self.nccl.ncclAllReduce(
buffer_type(in_tensor.data_ptr()),
buffer_type(out_tensor.data_ptr()),
in_tensor.numel(),
ncclDataTypeEnum.from_torch(in_tensor.dtype),
ncclRedOpTypeEnum.from_torch(op),
self.comm,
cudaStream_t(stream.cuda_stream),
)
return out_tensor
def all_gather(
self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, stream=None
):
if self.disabled:
return
# nccl communicator created on a specific device
# will only work on tensors on the same device
# otherwise it will cause "illegal memory access"
assert input_tensor.device == self.device, (
f"this nccl communicator is created to work on {self.device}, "
f"but the input tensor is on {input_tensor.device}"
)
if stream is None:
stream = current_stream()
self.nccl.ncclAllGather(
buffer_type(input_tensor.data_ptr()),
buffer_type(output_tensor.data_ptr()),
input_tensor.numel(),
ncclDataTypeEnum.from_torch(input_tensor.dtype),
self.comm,
cudaStream_t(stream.cuda_stream),
)
def all_gatherv(
self,
output_tensor: torch.Tensor,
input_tensor: torch.Tensor,
sizes: list[int],
stream=None,
):
if self.disabled:
return
# nccl communicator created on a specific device
# will only work on tensors on the same device
# otherwise it will cause "illegal memory access"
assert input_tensor.device == self.device, (
f"this nccl communicator is created to work on {self.device}, "
f"but the input tensor is on {input_tensor.device}"
)
if stream is None:
stream = current_stream()
assert output_tensor.shape[0] == sum(sizes)
split_offset = 0
self.nccl.ncclGroupStart()
for root, split_size in enumerate(sizes):
dst_slice = output_tensor[split_offset : split_offset + split_size]
self.nccl.ncclBroadcast(
buffer_type(input_tensor.data_ptr()),
buffer_type(dst_slice.data_ptr()),
dst_slice.numel(),
ncclDataTypeEnum.from_torch(input_tensor.dtype),
root,
self.comm,
cudaStream_t(stream.cuda_stream),
)
split_offset += split_size
self.nccl.ncclGroupEnd()
def reduce_scatter(
self,
output_tensor: torch.Tensor,
input_tensor: torch.Tensor,
op: ReduceOp = ReduceOp.SUM,
stream=None,
):
if self.disabled:
return
# nccl communicator created on a specific device
# will only work on tensors on the same device
# otherwise it will cause "illegal memory access"
assert input_tensor.device == self.device, (
f"this nccl communicator is created to work on {self.device}, "
f"but the input tensor is on {input_tensor.device}"
)
if stream is None:
stream = current_stream()
self.nccl.ncclReduceScatter(
buffer_type(input_tensor.data_ptr()),
buffer_type(output_tensor.data_ptr()),
output_tensor.numel(),
ncclDataTypeEnum.from_torch(input_tensor.dtype),
ncclRedOpTypeEnum.from_torch(op),
self.comm,
cudaStream_t(stream.cuda_stream),
)
def reduce_scatterv(
self,
output_tensor: torch.Tensor,
input_tensor: torch.Tensor,
sizes: list[int],
op: ReduceOp = ReduceOp.SUM,
stream=None,
):
if self.disabled:
return
# nccl communicator created on a specific device
# will only work on tensors on the same device
# otherwise it will cause "illegal memory access"
assert input_tensor.device == self.device, (
f"this nccl communicator is created to work on {self.device}, "
f"but the input tensor is on {input_tensor.device}"
)
if stream is None:
stream = current_stream()
split_offset = 0
self.nccl.ncclGroupStart()
for root, split_size in enumerate(sizes):
chunk = input_tensor[split_offset : split_offset + split_size, ...]
self.nccl.ncclReduce(
buffer_type(chunk.data_ptr()),
buffer_type(output_tensor.data_ptr()),
chunk.numel(),
ncclDataTypeEnum.from_torch(input_tensor.dtype),
ncclRedOpTypeEnum.from_torch(op),
root,
self.comm,
cudaStream_t(stream.cuda_stream),
)
split_offset += split_size
self.nccl.ncclGroupEnd()
def reduce(
self,
output_tensor: torch.Tensor,
input_tensor: torch.Tensor,
root: int,
op: ReduceOp = ReduceOp.SUM,
stream=None,
):
if self.disabled:
return
assert input_tensor.device == self.device, (
f"this nccl communicator is created to work on {self.device}, "
f"but the input tensor is on {input_tensor.device}"
)
if stream is None:
stream = current_stream()
self.nccl.ncclReduce(
buffer_type(input_tensor.data_ptr()),
buffer_type(output_tensor.data_ptr()),
input_tensor.numel(),
ncclDataTypeEnum.from_torch(input_tensor.dtype),
ncclRedOpTypeEnum.from_torch(op),
root,
self.comm,
cudaStream_t(stream.cuda_stream),
)
def scatter(
self,
output_tensor: torch.Tensor,
input_tensor: torch.Tensor,
sizes: list[int],
root: int = 0,
stream=None,
):
if self.disabled:
return
assert output_tensor.device == self.device, (
f"this nccl communicator is created to work on {self.device}, "
f"but the output tensor is on {output_tensor.device}"
)
if stream is None:
stream = current_stream()
self.nccl.ncclGroupStart()
if self.rank == root:
split_offset = 0
for dst, split_size in enumerate(sizes):
if split_size == 0:
continue
chunk = input_tensor[split_offset : split_offset + split_size, ...]
if dst == root:
output_tensor.copy_(chunk)
else:
self.send(chunk, dst, stream)
split_offset += split_size
elif sizes[self.rank] > 0:
self.recv(output_tensor, root, stream)
self.nccl.ncclGroupEnd()
def send(self, tensor: torch.Tensor, dst: int, stream=None):
if self.disabled:
return
assert tensor.device == self.device, (
f"this nccl communicator is created to work on {self.device}, "
f"but the input tensor is on {tensor.device}"
)
if stream is None:
stream = current_stream()
if tensor.dtype in [
torch.float8_e5m2,
torch.float8_e4m3fn,
torch.float8_e4m3fnuz,
torch.float8_e5m2fnuz,
]:
nccl_dtype = ncclDataTypeEnum.from_torch(torch.uint8)
else:
nccl_dtype = ncclDataTypeEnum.from_torch(tensor.dtype)
self.nccl.ncclSend(
buffer_type(tensor.data_ptr()),
tensor.numel(),
nccl_dtype,
dst,
self.comm,
cudaStream_t(stream.cuda_stream),
)
def recv(self, tensor: torch.Tensor, src: int, stream=None):
if self.disabled:
return
assert tensor.device == self.device, (
f"this nccl communicator is created to work on {self.device}, "
f"but the input tensor is on {tensor.device}"
)
if stream is None:
stream = current_stream()
if tensor.dtype in [
torch.float8_e5m2,
torch.float8_e4m3fn,
torch.float8_e4m3fnuz,
torch.float8_e5m2fnuz,
]:
nccl_dtype = ncclDataTypeEnum.from_torch(torch.uint8)
else:
nccl_dtype = ncclDataTypeEnum.from_torch(tensor.dtype)
self.nccl.ncclRecv(
buffer_type(tensor.data_ptr()),
tensor.numel(),
nccl_dtype,
src,
self.comm,
cudaStream_t(stream.cuda_stream),
)
def broadcast(self, tensor: torch.Tensor, src: int, stream=None):
if self.disabled:
return
assert tensor.device == self.device, (
f"this nccl communicator is created to work on {self.device}, "
f"but the input tensor is on {tensor.device}"
)
if stream is None:
stream = current_stream()
if src == self.rank:
sendbuff = buffer_type(tensor.data_ptr())
# NCCL requires the sender also to have a receive buffer
recvbuff = buffer_type(tensor.data_ptr())
else:
sendbuff = buffer_type()
recvbuff = buffer_type(tensor.data_ptr())
self.nccl.ncclBroadcast(
sendbuff,
recvbuff,
tensor.numel(),
ncclDataTypeEnum.from_torch(tensor.dtype),
src,
self.comm,
cudaStream_t(stream.cuda_stream),
)
def group_start(self):
self.nccl.ncclGroupStart()
def group_end(self):
self.nccl.ncclGroupEnd()
def register_comm_window(self, tensor: torch.Tensor):
return self.nccl.ncclCommWindowRegister(
self.comm,
buffer_type(tensor.data_ptr()),
tensor.numel() * tensor.element_size(),
1,
)
def register_comm_window_raw(self, ptr: int, size: int):
return self.nccl.ncclCommWindowRegister(self.comm, buffer_type(ptr), size, 1)
def deregister_comm_window(self, window):
return self.nccl.ncclCommWindowDeregister(self.comm, window)
def suspend(self):
"""Release comm GPU memory (collective, idempotent); keeps topology."""
if self.disabled or self._suspended:
return
if not self.nccl.has_symbol("ncclCommSuspend"):
logger.warning_once(
"ncclCommSuspend is not available in the loaded NCCL/RCCL "
"library (requires NCCL >= 2.29.7); skipping communicator "
"memory suspension."
)
return
self.nccl.ncclCommSuspend(self.comm, _NCCL_SUSPEND_MEM)
self._suspended = True
def resume(self):
"""Restore a suspended comm (collective); no-op unless suspended."""
if self.disabled or not self._suspended:
return
self.nccl.ncclCommResume(self.comm)
self._suspended = False
def batch_isend_irecv(self, p2p_ops: list, stream=None):
if self.disabled:
return
if stream is None:
stream = current_stream()
self.group_start()
for op in p2p_ops:
if op.op is torch.distributed.isend:
self.send(op.tensor, op.group_peer, stream)
elif op.op is torch.distributed.irecv:
self.recv(op.tensor, op.group_peer, stream)
self.group_end()