| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | # ===================== import region ===================== |
| 6 | import contextlib |
| 7 | import threading |
| 8 | |
| 9 | import torch |
| 10 | import torch.distributed as dist |
| 11 | from torch.distributed import ProcessGroup, ReduceOp |
| 12 | |
| 13 | import vllm.envs as envs |
| 14 | from vllm.distributed.device_communicators.pynccl_wrapper import ( |
| 15 | NCCL_UNIQUE_ID_BYTES, |
| 16 | NCCLLibrary, |
| 17 | buffer_type, |
| 18 | cudaStream_t, |
| 19 | ncclComm_t, |
| 20 | ncclDataTypeEnum, |
| 21 | ncclRedOpTypeEnum, |
| 22 | ncclUniqueId, |
| 23 | ) |
| 24 | from vllm.distributed.utils import StatelessProcessGroup |
| 25 | from vllm.logger import init_logger |
| 26 | from vllm.utils.torch_utils import current_stream |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | _warmup_tls = threading.local() |
| 31 | |
| 32 | |
| 33 | def _warmup_deferred() -> bool: |
| 34 | return getattr(_warmup_tls, "deferred", False) |
| 35 | |
| 36 | |
| 37 | @contextlib.contextmanager |
| 38 | def defer_comm_warmup_on_rocm(): |
| 39 | """Build communicators on this thread without their warm-up on ROCm, where |
| 40 | the warm-up can deadlock if built next to a live engine whose existing |
| 41 | communicators still have collectives in flight; instead the warm-up is |
| 42 | deferred to when the engine is quiet (during commit). |
| 43 | """ |
| 44 | from vllm.platforms import current_platform |
| 45 | |
| 46 | if not current_platform.is_rocm(): |
| 47 | yield |
| 48 | return |
| 49 | prev = _warmup_deferred() |
| 50 | _warmup_tls.deferred = True |
| 51 | try: |
| 52 | yield |
| 53 | finally: |
| 54 | _warmup_tls.deferred = prev |
| 55 | |
| 56 | |
| 57 | _NCCL_SYMM_OPS_REGISTERED = False |
| 58 | |
| 59 | _NCCL_SUSPEND_MEM = 0x01 |
| 60 | |
| 61 | |
| 62 | def register_nccl_symmetric_ops(pynccl_comm): |
| 63 | from vllm.distributed.device_communicators.pynccl_allocator import ( |
| 64 | nccl_symm_mem_context, |
| 65 | ) |
| 66 | from vllm.utils.torch_utils import direct_register_custom_op |
| 67 | |
| 68 | global _NCCL_SYMM_OPS_REGISTERED |
| 69 | if _NCCL_SYMM_OPS_REGISTERED: |
| 70 | return |
| 71 | _NCCL_SYMM_OPS_REGISTERED = True |
| 72 | |
| 73 | def all_reduce_symmetric_with_copy_impl(input_tensor: torch.Tensor) -> torch.Tensor: |
| 74 | with nccl_symm_mem_context(pynccl_comm): |
| 75 | symm_input = torch.empty_like(input_tensor) |
| 76 | symm_output = torch.empty_like(input_tensor) |
| 77 | symm_input.copy_(input_tensor) |
| 78 | symm_output = pynccl_comm.all_reduce(symm_input, symm_output) |
| 79 | return symm_output |
| 80 | |
| 81 | def all_reduce_symmetric_with_copy_fake(input_tensor: torch.Tensor) -> torch.Tensor: |
| 82 | return torch.empty_like(input_tensor) |
| 83 | |
| 84 | direct_register_custom_op( |
| 85 | op_name="all_reduce_symmetric_with_copy", |
| 86 | op_func=all_reduce_symmetric_with_copy_impl, |
| 87 | fake_impl=all_reduce_symmetric_with_copy_fake, |
| 88 | ) |
| 89 | |
| 90 | |
| 91 | class PyNcclCommunicator: |
| 92 | # None for communicators built via `from_unique_id_bytes` (no process group). |
| 93 | group: ProcessGroup | StatelessProcessGroup | None |
| 94 | |
| 95 | def __init__( |
| 96 | self, |
| 97 | group: ProcessGroup | StatelessProcessGroup, |
| 98 | device: int | str | torch.device, |
| 99 | library_path: str | None = None, |
| 100 | ): |
| 101 | """Args: |
| 102 | group: the process group to work on. If None, it will use the |
| 103 | default process group. |
| 104 | device: the device to bind the PyNcclCommunicator to. If None, |
| 105 | it will be bound to f"cuda:{local_rank}". |
| 106 | library_path: the path to the NCCL library. If None, it will |
| 107 | use the default library path. |
| 108 | It is the caller's responsibility to make sure each communicator |
| 109 | is bind to a unique device. |
| 110 | |
| 111 | """ |
| 112 | if not isinstance(group, StatelessProcessGroup): |
| 113 | assert dist.is_initialized() |
| 114 | assert dist.get_backend(group) != dist.Backend.NCCL, ( |
| 115 | "PyNcclCommunicator should be attached to a non-NCCL group." |
| 116 | ) |
| 117 | # note: this rank is the rank in the group |
| 118 | self.rank = dist.get_rank(group) |
| 119 | self.world_size = dist.get_world_size(group) |
| 120 | else: |
| 121 | self.rank = group.rank |
| 122 | self.world_size = group.world_size |
| 123 | |
| 124 | self.group = group |
| 125 | |
| 126 | # if world_size == 1, no need to create communicator |
| 127 | if self.world_size == 1 or envs.VLLM_DISABLE_PYNCCL: |
| 128 | self.available = False |
| 129 | self.disabled = True |
| 130 | return |
| 131 | try: |
| 132 | self.nccl = NCCLLibrary(library_path) |
| 133 | except Exception: |
| 134 | # disable because of missing NCCL library |
| 135 | # e.g. in a non-GPU environment |
| 136 | self.available = False |
| 137 | self.disabled = True |
| 138 | return |
| 139 | |
| 140 | self.available = True |
| 141 | self.disabled = False |
| 142 | self._suspended = False |
| 143 | |
| 144 | self.nccl_version = self.nccl.ncclGetRawVersion() |
| 145 | if self.rank == 0: |
| 146 | # get the unique id from NCCL |
| 147 | self.unique_id = self.nccl.ncclGetUniqueId() |
| 148 | logger.info_once("vLLM is using nccl==%s", self.nccl.ncclGetVersion()) |
| 149 | else: |
| 150 | # construct an empty unique id |
| 151 | self.unique_id = ncclUniqueId() |
| 152 | |
| 153 | if not isinstance(group, StatelessProcessGroup): |
| 154 | tensor = torch.ByteTensor(list(self.unique_id.internal)) |
| 155 | ranks = dist.get_process_group_ranks(group) |
| 156 | # arg `src` in `broadcast` is the global rank |
| 157 | dist.broadcast(tensor, src=ranks[0], group=group) |
| 158 | byte_list = tensor.tolist() |
| 159 | for i, byte in enumerate(byte_list): |
| 160 | self.unique_id.internal[i] = byte |
| 161 | else: |
| 162 | self.unique_id = group.broadcast_obj(self.unique_id, src=0) |
| 163 | self._init_comm(device) |
| 164 | |
| 165 | def _init_comm(self, device: int | str | torch.device) -> None: |
| 166 | """Create the communicator on `device` from the already-resolved |
| 167 | `self.unique_id` / `self.rank` / `self.world_size`, then run the |
| 168 | one-element warm-up all_reduce. Shared by `__init__` and |
| 169 | `from_unique_id_bytes` so the init handshake stays identical on both. |
| 170 | """ |
| 171 | if isinstance(device, int): |
| 172 | device = torch.device(f"cuda:{device}") |
| 173 | elif isinstance(device, str): |
| 174 | device = torch.device(device) |
| 175 | # now `device` is a `torch.device` object |
| 176 | assert isinstance(device, torch.device) |
| 177 | self.device = device |
| 178 | # nccl communicator and stream will use this device |
| 179 | with torch.accelerator.device_index(device.index): |
| 180 | self.comm: ncclComm_t = self.nccl.ncclCommInitRank( |
| 181 | self.world_size, self.unique_id, self.rank |
| 182 | ) |
| 183 | |
| 184 | if _warmup_deferred(): |
| 185 | return |
| 186 | |
| 187 | stream = current_stream() |
| 188 | # A small all_reduce for warmup. |
| 189 | data = torch.zeros(1, device=device) |
| 190 | self.all_reduce(data) |
| 191 | stream.synchronize() |
| 192 | del data |
| 193 | |
| 194 | @classmethod |
| 195 | def from_unique_id_bytes( |
| 196 | cls, |
| 197 | unique_id_bytes: bytes, |
| 198 | rank: int, |
| 199 | world_size: int, |
| 200 | device: int | str | torch.device, |
| 201 | library_path: str | None = None, |
| 202 | ) -> "PyNcclCommunicator": |
| 203 | """Build a communicator from pre-shared ``ncclUniqueId`` bytes. |
| 204 | |
| 205 | For peers that cannot join a ``StatelessProcessGroup`` / TCPStore (e.g. a |
| 206 | torch-free JAX trainer): every rank passes the same id, minted once via |
| 207 | ``ncclGetUniqueId`` and shared out of band. There is no barrier, so all |
| 208 | ranks must enter init concurrently or ``ncclCommInitRank`` hangs. |
| 209 | |
| 210 | Warm-up handshake: immediately after ``ncclCommInitRank`` this issues a |
| 211 | one-element ``all_reduce`` (mirroring ``__init__``). It is a collective, |
| 212 | so every peer -- including a foreign, non-vLLM rank -- must issue a |
| 213 | matching one-element ``all_reduce`` before any other collective, or all |
| 214 | ranks deadlock. |
| 215 | """ |
| 216 | if len(unique_id_bytes) != NCCL_UNIQUE_ID_BYTES: |
| 217 | raise ValueError( |
| 218 | f"expected a {NCCL_UNIQUE_ID_BYTES}-byte NCCL unique id, " |
| 219 | f"got {len(unique_id_bytes)} bytes" |
| 220 | ) |
| 221 | if not 0 <= rank < world_size: |
| 222 | raise ValueError(f"rank {rank} out of range for world_size {world_size}") |
| 223 | |
| 224 | self = cls.__new__(cls) |
| 225 | self.rank = rank |
| 226 | self.world_size = world_size |
| 227 | self.group = None |
| 228 | |
| 229 | if self.world_size == 1 or envs.VLLM_DISABLE_PYNCCL: |
| 230 | self.available = False |
| 231 | self.disabled = True |
| 232 | return self |
| 233 | try: |
| 234 | self.nccl = NCCLLibrary(library_path) |
| 235 | except Exception as e: |
| 236 | # Unlike the TCPStore path, silently disabling here leaves the peer |
| 237 | # blocked in ncclCommInitRank until timeout. The caller explicitly |
| 238 | # asked to join from a unique id, so fail loudly instead. |
| 239 | raise RuntimeError( |
| 240 | "failed to load the NCCL library for unique-id rendezvous" |
| 241 | ) from e |
| 242 | |
| 243 | self.available = True |
| 244 | self.disabled = False |
| 245 | self.nccl_version = self.nccl.ncclGetRawVersion() |
| 246 | self.unique_id = self.nccl.unique_id_from_bytes(unique_id_bytes) |
| 247 | self._init_comm(device) |
| 248 | return self |
| 249 | |
| 250 | def destroy(self): |
| 251 | if self.available and not self.disabled: |
| 252 | # ncclCommAbort can block until all CUDA graphs that |
| 253 | # captured NCCL ops on this comm are destroyed — and |
| 254 | # those graphs are released later in this same main- |
| 255 | # thread teardown, so a direct call here self-deadlocks. |
| 256 | # Run it in a daemon thread with a timeout: the main |
| 257 | # thread proceeds, the graphs drop, and the abort returns. |
| 258 | def _abort(): |
| 259 | with torch.accelerator.device_index(self.device.index): |
| 260 | self.nccl.ncclCommAbort(self.comm) |
| 261 | |
| 262 | abort_thread = threading.Thread(target=_abort, daemon=True) |
| 263 | abort_thread.start() |
| 264 | abort_thread.join(timeout=5.0) |
| 265 | self.available = False |
| 266 | self.disabled = True |
| 267 | |
| 268 | def all_reduce( |
| 269 | self, |
| 270 | in_tensor: torch.Tensor, |
| 271 | out_tensor: torch.Tensor = None, |
| 272 | op: ReduceOp = ReduceOp.SUM, |
| 273 | stream=None, |
| 274 | ) -> torch.Tensor: |
| 275 | if self.disabled: |
| 276 | return None |
| 277 | # nccl communicator created on a specific device |
| 278 | # will only work on tensors on the same device |
| 279 | # otherwise it will cause "illegal memory access" |
| 280 | assert in_tensor.device == self.device, ( |
| 281 | f"this nccl communicator is created to work on {self.device}, " |
| 282 | f"but the input tensor is on {in_tensor.device}" |
| 283 | ) |
| 284 | |
| 285 | if out_tensor is None: |
| 286 | out_tensor = torch.empty_like(in_tensor) |
| 287 | |
| 288 | if stream is None: |
| 289 | stream = current_stream() |
| 290 | self.nccl.ncclAllReduce( |
| 291 | buffer_type(in_tensor.data_ptr()), |
| 292 | buffer_type(out_tensor.data_ptr()), |
| 293 | in_tensor.numel(), |
| 294 | ncclDataTypeEnum.from_torch(in_tensor.dtype), |
| 295 | ncclRedOpTypeEnum.from_torch(op), |
| 296 | self.comm, |
| 297 | cudaStream_t(stream.cuda_stream), |
| 298 | ) |
| 299 | return out_tensor |
| 300 | |
| 301 | def all_gather( |
| 302 | self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, stream=None |
| 303 | ): |
| 304 | if self.disabled: |
| 305 | return |
| 306 | # nccl communicator created on a specific device |
| 307 | # will only work on tensors on the same device |
| 308 | # otherwise it will cause "illegal memory access" |
| 309 | assert input_tensor.device == self.device, ( |
| 310 | f"this nccl communicator is created to work on {self.device}, " |
| 311 | f"but the input tensor is on {input_tensor.device}" |
| 312 | ) |
| 313 | if stream is None: |
| 314 | stream = current_stream() |
| 315 | self.nccl.ncclAllGather( |
| 316 | buffer_type(input_tensor.data_ptr()), |
| 317 | buffer_type(output_tensor.data_ptr()), |
| 318 | input_tensor.numel(), |
| 319 | ncclDataTypeEnum.from_torch(input_tensor.dtype), |
| 320 | self.comm, |
| 321 | cudaStream_t(stream.cuda_stream), |
| 322 | ) |
| 323 | |
| 324 | def all_gatherv( |
| 325 | self, |
| 326 | output_tensor: torch.Tensor, |
| 327 | input_tensor: torch.Tensor, |
| 328 | sizes: list[int], |
| 329 | stream=None, |
| 330 | ): |
| 331 | if self.disabled: |
| 332 | return |
| 333 | # nccl communicator created on a specific device |
| 334 | # will only work on tensors on the same device |
| 335 | # otherwise it will cause "illegal memory access" |
| 336 | assert input_tensor.device == self.device, ( |
| 337 | f"this nccl communicator is created to work on {self.device}, " |
| 338 | f"but the input tensor is on {input_tensor.device}" |
| 339 | ) |
| 340 | if stream is None: |
| 341 | stream = current_stream() |
| 342 | assert output_tensor.shape[0] == sum(sizes) |
| 343 | split_offset = 0 |
| 344 | self.nccl.ncclGroupStart() |
| 345 | for root, split_size in enumerate(sizes): |
| 346 | dst_slice = output_tensor[split_offset : split_offset + split_size] |
| 347 | self.nccl.ncclBroadcast( |
| 348 | buffer_type(input_tensor.data_ptr()), |
| 349 | buffer_type(dst_slice.data_ptr()), |
| 350 | dst_slice.numel(), |
| 351 | ncclDataTypeEnum.from_torch(input_tensor.dtype), |
| 352 | root, |
| 353 | self.comm, |
| 354 | cudaStream_t(stream.cuda_stream), |
| 355 | ) |
| 356 | split_offset += split_size |
| 357 | self.nccl.ncclGroupEnd() |
| 358 | |
| 359 | def reduce_scatter( |
| 360 | self, |
| 361 | output_tensor: torch.Tensor, |
| 362 | input_tensor: torch.Tensor, |
| 363 | op: ReduceOp = ReduceOp.SUM, |
| 364 | stream=None, |
| 365 | ): |
| 366 | if self.disabled: |
| 367 | return |
| 368 | # nccl communicator created on a specific device |
| 369 | # will only work on tensors on the same device |
| 370 | # otherwise it will cause "illegal memory access" |
| 371 | assert input_tensor.device == self.device, ( |
| 372 | f"this nccl communicator is created to work on {self.device}, " |
| 373 | f"but the input tensor is on {input_tensor.device}" |
| 374 | ) |
| 375 | if stream is None: |
| 376 | stream = current_stream() |
| 377 | self.nccl.ncclReduceScatter( |
| 378 | buffer_type(input_tensor.data_ptr()), |
| 379 | buffer_type(output_tensor.data_ptr()), |
| 380 | output_tensor.numel(), |
| 381 | ncclDataTypeEnum.from_torch(input_tensor.dtype), |
| 382 | ncclRedOpTypeEnum.from_torch(op), |
| 383 | self.comm, |
| 384 | cudaStream_t(stream.cuda_stream), |
| 385 | ) |
| 386 | |
| 387 | def reduce_scatterv( |
| 388 | self, |
| 389 | output_tensor: torch.Tensor, |
| 390 | input_tensor: torch.Tensor, |
| 391 | sizes: list[int], |
| 392 | op: ReduceOp = ReduceOp.SUM, |
| 393 | stream=None, |
| 394 | ): |
| 395 | if self.disabled: |
| 396 | return |
| 397 | # nccl communicator created on a specific device |
| 398 | # will only work on tensors on the same device |
| 399 | # otherwise it will cause "illegal memory access" |
| 400 | assert input_tensor.device == self.device, ( |
| 401 | f"this nccl communicator is created to work on {self.device}, " |
| 402 | f"but the input tensor is on {input_tensor.device}" |
| 403 | ) |
| 404 | if stream is None: |
| 405 | stream = current_stream() |
| 406 | |
| 407 | split_offset = 0 |
| 408 | self.nccl.ncclGroupStart() |
| 409 | for root, split_size in enumerate(sizes): |
| 410 | chunk = input_tensor[split_offset : split_offset + split_size, ...] |
| 411 | self.nccl.ncclReduce( |
| 412 | buffer_type(chunk.data_ptr()), |
| 413 | buffer_type(output_tensor.data_ptr()), |
| 414 | chunk.numel(), |
| 415 | ncclDataTypeEnum.from_torch(input_tensor.dtype), |
| 416 | ncclRedOpTypeEnum.from_torch(op), |
| 417 | root, |
| 418 | self.comm, |
| 419 | cudaStream_t(stream.cuda_stream), |
| 420 | ) |
| 421 | split_offset += split_size |
| 422 | self.nccl.ncclGroupEnd() |
| 423 | |
| 424 | def reduce( |
| 425 | self, |
| 426 | output_tensor: torch.Tensor, |
| 427 | input_tensor: torch.Tensor, |
| 428 | root: int, |
| 429 | op: ReduceOp = ReduceOp.SUM, |
| 430 | stream=None, |
| 431 | ): |
| 432 | if self.disabled: |
| 433 | return |
| 434 | assert input_tensor.device == self.device, ( |
| 435 | f"this nccl communicator is created to work on {self.device}, " |
| 436 | f"but the input tensor is on {input_tensor.device}" |
| 437 | ) |
| 438 | if stream is None: |
| 439 | stream = current_stream() |
| 440 | self.nccl.ncclReduce( |
| 441 | buffer_type(input_tensor.data_ptr()), |
| 442 | buffer_type(output_tensor.data_ptr()), |
| 443 | input_tensor.numel(), |
| 444 | ncclDataTypeEnum.from_torch(input_tensor.dtype), |
| 445 | ncclRedOpTypeEnum.from_torch(op), |
| 446 | root, |
| 447 | self.comm, |
| 448 | cudaStream_t(stream.cuda_stream), |
| 449 | ) |
| 450 | |
| 451 | def scatter( |
| 452 | self, |
| 453 | output_tensor: torch.Tensor, |
| 454 | input_tensor: torch.Tensor, |
| 455 | sizes: list[int], |
| 456 | root: int = 0, |
| 457 | stream=None, |
| 458 | ): |
| 459 | if self.disabled: |
| 460 | return |
| 461 | assert output_tensor.device == self.device, ( |
| 462 | f"this nccl communicator is created to work on {self.device}, " |
| 463 | f"but the output tensor is on {output_tensor.device}" |
| 464 | ) |
| 465 | if stream is None: |
| 466 | stream = current_stream() |
| 467 | self.nccl.ncclGroupStart() |
| 468 | if self.rank == root: |
| 469 | split_offset = 0 |
| 470 | for dst, split_size in enumerate(sizes): |
| 471 | if split_size == 0: |
| 472 | continue |
| 473 | |
| 474 | chunk = input_tensor[split_offset : split_offset + split_size, ...] |
| 475 | if dst == root: |
| 476 | output_tensor.copy_(chunk) |
| 477 | else: |
| 478 | self.send(chunk, dst, stream) |
| 479 | split_offset += split_size |
| 480 | elif sizes[self.rank] > 0: |
| 481 | self.recv(output_tensor, root, stream) |
| 482 | self.nccl.ncclGroupEnd() |
| 483 | |
| 484 | def send(self, tensor: torch.Tensor, dst: int, stream=None): |
| 485 | if self.disabled: |
| 486 | return |
| 487 | assert tensor.device == self.device, ( |
| 488 | f"this nccl communicator is created to work on {self.device}, " |
| 489 | f"but the input tensor is on {tensor.device}" |
| 490 | ) |
| 491 | if stream is None: |
| 492 | stream = current_stream() |
| 493 | if tensor.dtype in [ |
| 494 | torch.float8_e5m2, |
| 495 | torch.float8_e4m3fn, |
| 496 | torch.float8_e4m3fnuz, |
| 497 | torch.float8_e5m2fnuz, |
| 498 | ]: |
| 499 | nccl_dtype = ncclDataTypeEnum.from_torch(torch.uint8) |
| 500 | else: |
| 501 | nccl_dtype = ncclDataTypeEnum.from_torch(tensor.dtype) |
| 502 | self.nccl.ncclSend( |
| 503 | buffer_type(tensor.data_ptr()), |
| 504 | tensor.numel(), |
| 505 | nccl_dtype, |
| 506 | dst, |
| 507 | self.comm, |
| 508 | cudaStream_t(stream.cuda_stream), |
| 509 | ) |
| 510 | |
| 511 | def recv(self, tensor: torch.Tensor, src: int, stream=None): |
| 512 | if self.disabled: |
| 513 | return |
| 514 | assert tensor.device == self.device, ( |
| 515 | f"this nccl communicator is created to work on {self.device}, " |
| 516 | f"but the input tensor is on {tensor.device}" |
| 517 | ) |
| 518 | if stream is None: |
| 519 | stream = current_stream() |
| 520 | if tensor.dtype in [ |
| 521 | torch.float8_e5m2, |
| 522 | torch.float8_e4m3fn, |
| 523 | torch.float8_e4m3fnuz, |
| 524 | torch.float8_e5m2fnuz, |
| 525 | ]: |
| 526 | nccl_dtype = ncclDataTypeEnum.from_torch(torch.uint8) |
| 527 | else: |
| 528 | nccl_dtype = ncclDataTypeEnum.from_torch(tensor.dtype) |
| 529 | self.nccl.ncclRecv( |
| 530 | buffer_type(tensor.data_ptr()), |
| 531 | tensor.numel(), |
| 532 | nccl_dtype, |
| 533 | src, |
| 534 | self.comm, |
| 535 | cudaStream_t(stream.cuda_stream), |
| 536 | ) |
| 537 | |
| 538 | def broadcast(self, tensor: torch.Tensor, src: int, stream=None): |
| 539 | if self.disabled: |
| 540 | return |
| 541 | assert tensor.device == self.device, ( |
| 542 | f"this nccl communicator is created to work on {self.device}, " |
| 543 | f"but the input tensor is on {tensor.device}" |
| 544 | ) |
| 545 | if stream is None: |
| 546 | stream = current_stream() |
| 547 | if src == self.rank: |
| 548 | sendbuff = buffer_type(tensor.data_ptr()) |
| 549 | # NCCL requires the sender also to have a receive buffer |
| 550 | recvbuff = buffer_type(tensor.data_ptr()) |
| 551 | else: |
| 552 | sendbuff = buffer_type() |
| 553 | recvbuff = buffer_type(tensor.data_ptr()) |
| 554 | self.nccl.ncclBroadcast( |
| 555 | sendbuff, |
| 556 | recvbuff, |
| 557 | tensor.numel(), |
| 558 | ncclDataTypeEnum.from_torch(tensor.dtype), |
| 559 | src, |
| 560 | self.comm, |
| 561 | cudaStream_t(stream.cuda_stream), |
| 562 | ) |
| 563 | |
| 564 | def group_start(self): |
| 565 | self.nccl.ncclGroupStart() |
| 566 | |
| 567 | def group_end(self): |
| 568 | self.nccl.ncclGroupEnd() |
| 569 | |
| 570 | def register_comm_window(self, tensor: torch.Tensor): |
| 571 | return self.nccl.ncclCommWindowRegister( |
| 572 | self.comm, |
| 573 | buffer_type(tensor.data_ptr()), |
| 574 | tensor.numel() * tensor.element_size(), |
| 575 | 1, |
| 576 | ) |
| 577 | |
| 578 | def register_comm_window_raw(self, ptr: int, size: int): |
| 579 | return self.nccl.ncclCommWindowRegister(self.comm, buffer_type(ptr), size, 1) |
| 580 | |
| 581 | def deregister_comm_window(self, window): |
| 582 | return self.nccl.ncclCommWindowDeregister(self.comm, window) |
| 583 | |
| 584 | def suspend(self): |
| 585 | """Release comm GPU memory (collective, idempotent); keeps topology.""" |
| 586 | if self.disabled or self._suspended: |
| 587 | return |
| 588 | if not self.nccl.has_symbol("ncclCommSuspend"): |
| 589 | logger.warning_once( |
| 590 | "ncclCommSuspend is not available in the loaded NCCL/RCCL " |
| 591 | "library (requires NCCL >= 2.29.7); skipping communicator " |
| 592 | "memory suspension." |
| 593 | ) |
| 594 | return |
| 595 | self.nccl.ncclCommSuspend(self.comm, _NCCL_SUSPEND_MEM) |
| 596 | self._suspended = True |
| 597 | |
| 598 | def resume(self): |
| 599 | """Restore a suspended comm (collective); no-op unless suspended.""" |
| 600 | if self.disabled or not self._suspended: |
| 601 | return |
| 602 | self.nccl.ncclCommResume(self.comm) |
| 603 | self._suspended = False |
| 604 | |
| 605 | def batch_isend_irecv(self, p2p_ops: list, stream=None): |
| 606 | if self.disabled: |
| 607 | return |
| 608 | if stream is None: |
| 609 | stream = current_stream() |
| 610 | self.group_start() |
| 611 | for op in p2p_ops: |
| 612 | if op.op is torch.distributed.isend: |
| 613 | self.send(op.tensor, op.group_peer, stream) |
| 614 | elif op.op is torch.distributed.irecv: |
| 615 | self.recv(op.tensor, op.group_peer, stream) |
| 616 | |
| 617 | self.group_end() |