# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ===================== import region ===================== import threading import torch import torch.distributed as dist from torch.distributed import ProcessGroup, ReduceOp import vllm.envs as envs from vllm.distributed.device_communicators.pynccl_wrapper import ( NCCL_UNIQUE_ID_BYTES, NCCLLibrary, buffer_type, cudaStream_t, ncclComm_t, ncclDataTypeEnum, ncclRedOpTypeEnum, ncclUniqueId, ) from vllm.distributed.utils import StatelessProcessGroup from vllm.logger import init_logger from vllm.utils.torch_utils import current_stream logger = init_logger(__name__) _NCCL_SYMM_OPS_REGISTERED = False _NCCL_SUSPEND_MEM = 0x01 def register_nccl_symmetric_ops(pynccl_comm): from vllm.distributed.device_communicators.pynccl_allocator import ( nccl_symm_mem_context, ) from vllm.utils.torch_utils import direct_register_custom_op global _NCCL_SYMM_OPS_REGISTERED if _NCCL_SYMM_OPS_REGISTERED: return _NCCL_SYMM_OPS_REGISTERED = True def all_reduce_symmetric_with_copy_impl(input_tensor: torch.Tensor) -> torch.Tensor: with nccl_symm_mem_context(pynccl_comm): symm_input = torch.empty_like(input_tensor) symm_output = torch.empty_like(input_tensor) symm_input.copy_(input_tensor) symm_output = pynccl_comm.all_reduce(symm_input, symm_output) return symm_output def all_reduce_symmetric_with_copy_fake(input_tensor: torch.Tensor) -> torch.Tensor: return torch.empty_like(input_tensor) direct_register_custom_op( op_name="all_reduce_symmetric_with_copy", op_func=all_reduce_symmetric_with_copy_impl, fake_impl=all_reduce_symmetric_with_copy_fake, ) 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()