| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright 2023 The vLLM team. |
| 5 | # Adapted from |
| 6 | # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py |
| 7 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. |
| 8 | """vLLM distributed state. |
| 9 | It takes over the control of the distributed environment from PyTorch. |
| 10 | The typical workflow is: |
| 11 | |
| 12 | - call `init_distributed_environment` to initialize the distributed environment. |
| 13 | - call `initialize_model_parallel` or `ensure_model_parallel_initialized` to |
| 14 | initialize the model parallel groups. |
| 15 | |
| 16 | - any code dealing with the distributed stuff |
| 17 | |
| 18 | - call `destroy_model_parallel` to destroy the model parallel groups. |
| 19 | - call `destroy_distributed_environment` to destroy the distributed environment. |
| 20 | |
| 21 | If you only need to use the distributed environment without model/pipeline |
| 22 | parallelism, you can skip the model parallel initialization and destruction |
| 23 | steps. |
| 24 | """ |
| 25 | |
| 26 | import contextlib |
| 27 | import gc |
| 28 | import pickle |
| 29 | import weakref |
| 30 | from collections import deque, namedtuple |
| 31 | from collections.abc import Callable |
| 32 | from contextlib import AbstractContextManager, contextmanager, nullcontext |
| 33 | from dataclasses import dataclass |
| 34 | from datetime import timedelta |
| 35 | from math import gcd |
| 36 | from multiprocessing import shared_memory |
| 37 | from typing import TYPE_CHECKING, Any, Protocol |
| 38 | from unittest.mock import patch |
| 39 | |
| 40 | import torch |
| 41 | import torch.distributed |
| 42 | import torch.distributed._functional_collectives as funcol |
| 43 | import torch.distributed._symmetric_memory |
| 44 | from torch.distributed import Backend, ProcessGroup, Store |
| 45 | |
| 46 | import vllm.envs as envs |
| 47 | from vllm.distributed.device_communicators.base_device_communicator import ( |
| 48 | DeviceCommunicatorBase, |
| 49 | ) |
| 50 | from vllm.distributed.utils import ( |
| 51 | StatelessProcessGroup, |
| 52 | get_cached_tcp_store_client, |
| 53 | ) |
| 54 | from vllm.logger import init_logger |
| 55 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 56 | from vllm.utils.network_utils import get_distributed_init_method |
| 57 | from vllm.utils.system_utils import suppress_stdout |
| 58 | from vllm.utils.torch_utils import ( |
| 59 | direct_register_custom_op, |
| 60 | ) |
| 61 | |
| 62 | if TYPE_CHECKING: |
| 63 | from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator |
| 64 | |
| 65 | |
| 66 | @dataclass |
| 67 | class GraphCaptureContext: |
| 68 | stream: torch.cuda.Stream |
| 69 | |
| 70 | |
| 71 | TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) |
| 72 | |
| 73 | |
| 74 | class Handle(Protocol): |
| 75 | """Minimal async work handle used by P2P send/recv methods.""" |
| 76 | |
| 77 | def is_completed(self) -> bool: ... |
| 78 | |
| 79 | def wait(self) -> None: ... |
| 80 | |
| 81 | |
| 82 | class _RetainedHandle: |
| 83 | """Handle that retains source tensors until all posted work completes. |
| 84 | |
| 85 | Used by ``isend_object`` to keep the serialized CPU tensors alive while |
| 86 | gloo copies them in the background; ``wait`` drops the refs once drained. |
| 87 | |
| 88 | ``wait`` must be idempotent: gloo ``Work.wait()`` is single-shot for |
| 89 | point-to-point operations — calling it a second time on an already |
| 90 | completed send blocks forever (verified empirically). Both the worker's |
| 91 | top-of-step drain of ``_pp_send_work`` and ``_reap_completed_isends`` |
| 92 | may wait the same entry's metadata handle, so the second caller must |
| 93 | not re-wait the underlying works. |
| 94 | """ |
| 95 | |
| 96 | def __init__(self, works: list[Any], retained: tuple[torch.Tensor, ...]) -> None: |
| 97 | self._works = works |
| 98 | self._retained = retained |
| 99 | self._waited = False |
| 100 | |
| 101 | def is_completed(self) -> bool: |
| 102 | return all(w.is_completed() for w in self._works) |
| 103 | |
| 104 | def wait(self) -> None: |
| 105 | if self._waited: |
| 106 | return |
| 107 | for w in self._works: |
| 108 | w.wait() |
| 109 | self._waited = True |
| 110 | self._retained = () |
| 111 | |
| 112 | |
| 113 | def _split_tensor_dict( |
| 114 | tensor_dict: dict[str, torch.Tensor | Any], |
| 115 | ) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]: |
| 116 | """Split the tensor dictionary into two parts: |
| 117 | 1. A list of (key, value) pairs. If the value is a tensor, it is replaced |
| 118 | by its metadata. |
| 119 | 2. A list of tensors. |
| 120 | """ |
| 121 | metadata_list: list[tuple[str, Any]] = [] |
| 122 | tensor_list: list[torch.Tensor] = [] |
| 123 | for key, value in tensor_dict.items(): |
| 124 | if isinstance(value, torch.Tensor): |
| 125 | # Note: we cannot use `value.device` here, |
| 126 | # because it contains not only the device type but also the device |
| 127 | # index (e.g. "cuda:0"). We only need the device type. |
| 128 | # receiving side will set the device index. |
| 129 | device = value.device.type |
| 130 | metadata_list.append( |
| 131 | (key, TensorMetadata(device, value.dtype, value.size())) |
| 132 | ) |
| 133 | tensor_list.append(value) |
| 134 | else: |
| 135 | metadata_list.append((key, value)) |
| 136 | return metadata_list, tensor_list |
| 137 | |
| 138 | |
| 139 | _group_name_counter: dict[str, int] = {} |
| 140 | |
| 141 | |
| 142 | def _get_unique_name(name: str) -> str: |
| 143 | """Get a unique name for the group. |
| 144 | |
| 145 | Example: |
| 146 | _get_unique_name("tp") -> "tp:0" |
| 147 | _get_unique_name("tp") -> "tp:1" |
| 148 | |
| 149 | """ |
| 150 | if name not in _group_name_counter: |
| 151 | _group_name_counter[name] = 0 |
| 152 | newname = f"{name}:{_group_name_counter[name]}" |
| 153 | _group_name_counter[name] += 1 |
| 154 | return newname |
| 155 | |
| 156 | |
| 157 | _groups: dict[str, Callable[[], "GroupCoordinator | None"]] = {} |
| 158 | |
| 159 | |
| 160 | def _register_group(group: "GroupCoordinator") -> None: |
| 161 | _groups[group.unique_name] = weakref.ref(group) |
| 162 | |
| 163 | |
| 164 | def _apply_to_device_comms( |
| 165 | action: Callable[[DeviceCommunicatorBase], None], |
| 166 | ) -> None: |
| 167 | """Apply ``action`` to every group's device communicator. |
| 168 | |
| 169 | Walks the registered parallel groups and skips those without a device |
| 170 | communicator (absent at ``world_size == 1``). |
| 171 | """ |
| 172 | comms = [] |
| 173 | for group_ref in _groups.values(): |
| 174 | group = group_ref() |
| 175 | if group is None: |
| 176 | continue |
| 177 | dc = group.device_communicator |
| 178 | if dc is None: |
| 179 | continue |
| 180 | comms.append(dc) |
| 181 | |
| 182 | for dc in comms: |
| 183 | action(dc) |
| 184 | |
| 185 | |
| 186 | def suspend_device_comms() -> None: |
| 187 | """Release device communicator memory (collective; comms must be idle).""" |
| 188 | _apply_to_device_comms(lambda comm: comm.suspend()) |
| 189 | |
| 190 | |
| 191 | def resume_device_comms() -> None: |
| 192 | """Restore suspended device communicators (collective).""" |
| 193 | _apply_to_device_comms(lambda comm: comm.resume()) |
| 194 | |
| 195 | |
| 196 | def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor: |
| 197 | assert group_name in _groups, f"Group {group_name} is not found." |
| 198 | group = _groups[group_name]() |
| 199 | if group is None: |
| 200 | raise ValueError(f"Group {group_name} is destroyed.") |
| 201 | return group._all_reduce_out_place(tensor) |
| 202 | |
| 203 | |
| 204 | def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor: |
| 205 | return torch.empty_like(tensor) |
| 206 | |
| 207 | |
| 208 | def reduce_scatter( |
| 209 | tensor: torch.Tensor, dim: int, world_size: int, group_name: str |
| 210 | ) -> torch.Tensor: |
| 211 | assert group_name in _groups, f"Group {group_name} is not found." |
| 212 | group = _groups[group_name]() |
| 213 | if group is None: |
| 214 | raise ValueError(f"Group {group_name} is destroyed.") |
| 215 | return group._reduce_scatter_out_place(tensor, dim) |
| 216 | |
| 217 | |
| 218 | def reduce_scatter_fake( |
| 219 | tensor: torch.Tensor, dim: int, world_size: int, group_name: str |
| 220 | ) -> torch.Tensor: |
| 221 | new_shape = list(tensor.shape) |
| 222 | new_shape[dim] = tensor.shape[dim] // world_size |
| 223 | return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device) |
| 224 | |
| 225 | |
| 226 | def all_gather( |
| 227 | tensor: torch.Tensor, dim: int, world_size: int, group_name: str |
| 228 | ) -> torch.Tensor: |
| 229 | assert group_name in _groups, f"Group {group_name} is not found." |
| 230 | group = _groups[group_name]() |
| 231 | if group is None: |
| 232 | raise ValueError(f"Group {group_name} is destroyed.") |
| 233 | return group._all_gather_out_place(tensor, dim) |
| 234 | |
| 235 | |
| 236 | def all_gather_fake( |
| 237 | tensor: torch.Tensor, dim: int, world_size: int, group_name: str |
| 238 | ) -> torch.Tensor: |
| 239 | new_shape = list(tensor.shape) |
| 240 | new_shape[dim] = tensor.shape[dim] * world_size |
| 241 | return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device) |
| 242 | |
| 243 | |
| 244 | def patched_fused_scaled_matmul_reduce_scatter_fake( |
| 245 | A: torch.Tensor, |
| 246 | B: torch.Tensor, |
| 247 | A_scale: torch.Tensor, |
| 248 | B_scale: torch.Tensor, |
| 249 | reduce_op: str, |
| 250 | orig_scatter_dim: int, |
| 251 | scatter_dim_after_maybe_reshape: int, |
| 252 | group_name: str, |
| 253 | output_shape: list[int], |
| 254 | bias: torch.Tensor | None = None, |
| 255 | result_scale: torch.Tensor | None = None, |
| 256 | out_dtype: torch.dtype | None = None, |
| 257 | use_fast_accum: bool = False, |
| 258 | ) -> torch.Tensor: |
| 259 | # Copied from |
| 260 | # https://github.com/pytorch/pytorch/blob/50c338c2da905062449e4d9ac807832d1b5cd90e/torch/distributed/_symmetric_memory/__init__.py#L1189 |
| 261 | if A_scale.numel() > 1: |
| 262 | if A_scale.shape[:-1] != A.shape[:-1]: |
| 263 | raise ValueError( |
| 264 | "For row-wise scaling, the leading dims of A_scale " |
| 265 | "must match the leading dims of A " |
| 266 | f"(A shape: {A.shape}, A_scale shape: {A_scale.shape})" |
| 267 | ) |
| 268 | A_scale = A_scale.flatten(0, -2).contiguous() |
| 269 | elif A_scale.numel() != 1: |
| 270 | raise ValueError( |
| 271 | "Invalid A_scale shape " |
| 272 | f"(A shape: {A.shape}, A_scale shape: {A_scale.shape})" |
| 273 | ) |
| 274 | |
| 275 | C = torch._scaled_mm( |
| 276 | A.flatten(0, -2).contiguous(), |
| 277 | B, |
| 278 | A_scale, |
| 279 | B_scale, |
| 280 | bias, |
| 281 | result_scale, |
| 282 | out_dtype, |
| 283 | use_fast_accum, |
| 284 | ) |
| 285 | C = C.view(*output_shape[:-1], B.shape[1]) |
| 286 | res = funcol.reduce_scatter_tensor( |
| 287 | C, |
| 288 | reduce_op, |
| 289 | orig_scatter_dim, # need original scatter dim for 3D+ output tensor here |
| 290 | group_name, |
| 291 | ) |
| 292 | res = funcol.wait_tensor(res) |
| 293 | return res |
| 294 | |
| 295 | |
| 296 | def _platform_device_type() -> str: |
| 297 | """Return the device-type string (e.g. ``"cuda"``, ``"xpu"``, ``"cpu"``) |
| 298 | for the current platform, in the form expected by |
| 299 | ``torch.distributed.init_process_group(backend=...)``. |
| 300 | """ |
| 301 | from vllm.platforms import current_platform |
| 302 | |
| 303 | if current_platform.is_cuda_alike(): |
| 304 | return "cuda" |
| 305 | elif current_platform.is_xpu(): |
| 306 | return "xpu" |
| 307 | elif current_platform.is_out_of_tree(): |
| 308 | return current_platform.device_name |
| 309 | else: |
| 310 | return "cpu" |
| 311 | |
| 312 | |
| 313 | def _device_backend_str(torch_distributed_backend: str | Backend) -> str: |
| 314 | """Normalize ``torch_distributed_backend`` to the ``"<device>:<backend>"`` |
| 315 | format required by ``split_group``'s ``backend`` argument. |
| 316 | |
| 317 | Accepts either a bare backend name (e.g. ``"nccl"``) or an already-prefixed |
| 318 | string (e.g. ``"cuda:nccl"``). |
| 319 | """ |
| 320 | backend_str = str(torch_distributed_backend) |
| 321 | if ":" in backend_str: |
| 322 | return backend_str |
| 323 | return f"{_platform_device_type()}:{backend_str}" |
| 324 | |
| 325 | |
| 326 | def _create_subgroups_split_group( |
| 327 | group_ranks: list[list[int]], |
| 328 | group_name: str, |
| 329 | torch_distributed_backend: str | Backend, |
| 330 | ) -> tuple[ProcessGroup, ProcessGroup]: |
| 331 | """Create the device + CPU subgroups for ``GroupCoordinator`` via |
| 332 | ``torch.distributed.split_group``. |
| 333 | |
| 334 | ``split_group`` is collective on the parent group, so every parent rank |
| 335 | must enter with the same ``split_ranks`` definition. Each rank receives |
| 336 | the subgroup it belongs to. |
| 337 | """ |
| 338 | from vllm.distributed.utils import ( |
| 339 | get_cpu_distributed_timeout_or_none, |
| 340 | get_distributed_timeout_or_none, |
| 341 | ) |
| 342 | |
| 343 | device_backend_str = _device_backend_str(torch_distributed_backend) |
| 344 | self_device_group = torch.distributed.split_group( |
| 345 | split_ranks=group_ranks, |
| 346 | group_desc=f"{group_name}:device", |
| 347 | backend=device_backend_str, |
| 348 | timeout=get_distributed_timeout_or_none(), |
| 349 | ) |
| 350 | # CPU subgroup: split_group requires the requested backend filter to |
| 351 | # include the parent's default device type (= the device the parent PG |
| 352 | # was bound to via ``device_id``), so a cpu-only filter is rejected. |
| 353 | # Include the device backend in the filter; only the gloo backend is |
| 354 | # actually used for CPU collectives on this group. |
| 355 | self_cpu_group = torch.distributed.split_group( |
| 356 | split_ranks=group_ranks, |
| 357 | group_desc=f"{group_name}:cpu", |
| 358 | backend=f"cpu:gloo,{device_backend_str}", |
| 359 | timeout=get_cpu_distributed_timeout_or_none(), |
| 360 | ) |
| 361 | return self_device_group, self_cpu_group |
| 362 | |
| 363 | |
| 364 | def patched_fused_scaled_matmul_reduce_scatter( |
| 365 | A: torch.Tensor, |
| 366 | B: torch.Tensor, |
| 367 | A_scale: torch.Tensor, |
| 368 | B_scale: torch.Tensor, |
| 369 | reduce_op: str, |
| 370 | orig_scatter_dim: int, |
| 371 | scatter_dim_after_maybe_reshape: int, |
| 372 | group_name: str, |
| 373 | output_shape: list[int], |
| 374 | bias: torch.Tensor | None = None, |
| 375 | result_scale: torch.Tensor | None = None, |
| 376 | out_dtype: torch.dtype | None = None, |
| 377 | use_fast_accum: bool = False, |
| 378 | ) -> torch.Tensor: |
| 379 | return torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter( |
| 380 | A, |
| 381 | B, |
| 382 | A_scale, |
| 383 | B_scale, |
| 384 | reduce_op, |
| 385 | orig_scatter_dim, |
| 386 | scatter_dim_after_maybe_reshape, |
| 387 | group_name, |
| 388 | output_shape, |
| 389 | bias, |
| 390 | result_scale, |
| 391 | out_dtype, |
| 392 | use_fast_accum, |
| 393 | ) |
| 394 | |
| 395 | |
| 396 | direct_register_custom_op( |
| 397 | op_name="all_reduce", |
| 398 | op_func=all_reduce, |
| 399 | fake_impl=all_reduce_fake, |
| 400 | ) |
| 401 | |
| 402 | direct_register_custom_op( |
| 403 | op_name="reduce_scatter", |
| 404 | op_func=reduce_scatter, |
| 405 | fake_impl=reduce_scatter_fake, |
| 406 | ) |
| 407 | |
| 408 | direct_register_custom_op( |
| 409 | op_name="all_gather", |
| 410 | op_func=all_gather, |
| 411 | fake_impl=all_gather_fake, |
| 412 | ) |
| 413 | |
| 414 | # TODO: Remove this once the pytorch fix |
| 415 | # (https://github.com/pytorch/pytorch/pull/165086) gets released, |
| 416 | # in either 2.9.1 or 2.10 |
| 417 | direct_register_custom_op( |
| 418 | op_name="patched_fused_scaled_matmul_reduce_scatter", |
| 419 | op_func=patched_fused_scaled_matmul_reduce_scatter, |
| 420 | fake_impl=patched_fused_scaled_matmul_reduce_scatter_fake, |
| 421 | ) |
| 422 | |
| 423 | |
| 424 | class GroupCoordinator: |
| 425 | """PyTorch ProcessGroup wrapper for a group of processes. |
| 426 | PyTorch ProcessGroup is bound to one specific communication backend, |
| 427 | e.g. NCCL, Gloo, MPI, etc. |
| 428 | GroupCoordinator takes charge of all the communication operations among |
| 429 | the processes in the group. It manages both CPU and device |
| 430 | communication. |
| 431 | """ |
| 432 | |
| 433 | # available attributes: |
| 434 | rank: int # global rank |
| 435 | ranks: list[int] # global ranks in the group |
| 436 | world_size: int # size of the group |
| 437 | # difference between `local_rank` and `rank_in_group`: |
| 438 | # if we have a group of size 4 across two nodes: |
| 439 | # Process | Node | Rank | Local Rank | Rank in Group |
| 440 | # 0 | 0 | 0 | 0 | 0 |
| 441 | # 1 | 0 | 1 | 1 | 1 |
| 442 | # 2 | 1 | 2 | 0 | 2 |
| 443 | # 3 | 1 | 3 | 1 | 3 |
| 444 | local_rank: int # local rank used to assign devices |
| 445 | rank_in_group: int # rank inside the group |
| 446 | cpu_group: ProcessGroup # group for CPU communication |
| 447 | device_group: ProcessGroup # group for device communication |
| 448 | # device communicator (if use_device_communicator=True) |
| 449 | device_communicator: DeviceCommunicatorBase | None |
| 450 | mq_broadcaster: Any | None # shared memory broadcaster |
| 451 | |
| 452 | def __init__( |
| 453 | self, |
| 454 | group_ranks: list[list[int]], |
| 455 | local_rank: int, |
| 456 | torch_distributed_backend: str | Backend, |
| 457 | use_device_communicator: bool, # whether to use device communicator |
| 458 | use_message_queue_broadcaster: bool = False, |
| 459 | group_name: str | None = None, |
| 460 | use_all2all: bool = False, |
| 461 | ): |
| 462 | group_name = group_name or "anonymous" |
| 463 | self.unique_name = _get_unique_name(group_name) |
| 464 | _register_group(self) |
| 465 | |
| 466 | self.rank = torch.distributed.get_rank() |
| 467 | self.local_rank = local_rank |
| 468 | self.device_index: int |
| 469 | assert local_rank >= 0, ( |
| 470 | "local_rank must be provided when creating the world group" |
| 471 | ) |
| 472 | self.device_index = local_rank |
| 473 | |
| 474 | self_device_group = None |
| 475 | self_cpu_group = None |
| 476 | |
| 477 | # VLLM_DISTRIBUTED_USE_SPLIT_GROUP gates the new ``split_group`` |
| 478 | # codepath. Default (False) preserves the legacy ``new_group`` path. |
| 479 | if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP: |
| 480 | self_device_group, self_cpu_group = _create_subgroups_split_group( |
| 481 | group_ranks, group_name, torch_distributed_backend |
| 482 | ) |
| 483 | for ranks in group_ranks: |
| 484 | if self.rank in ranks: |
| 485 | self.ranks = ranks |
| 486 | self.world_size = len(ranks) |
| 487 | self.rank_in_group = ranks.index(self.rank) |
| 488 | break |
| 489 | else: |
| 490 | from vllm.distributed.utils import ( |
| 491 | get_cpu_distributed_timeout_or_none, |
| 492 | get_distributed_timeout_or_none, |
| 493 | ) |
| 494 | |
| 495 | timeout = get_cpu_distributed_timeout_or_none() |
| 496 | device_timeout = get_distributed_timeout_or_none() |
| 497 | |
| 498 | for ranks in group_ranks: |
| 499 | device_group = torch.distributed.new_group( |
| 500 | ranks, |
| 501 | backend=torch_distributed_backend, |
| 502 | timeout=device_timeout, |
| 503 | ) |
| 504 | # a group with `gloo` backend, to allow direct coordination between |
| 505 | # processes through the CPU. |
| 506 | with suppress_stdout(): |
| 507 | cpu_group = torch.distributed.new_group( |
| 508 | ranks, backend="gloo", timeout=timeout |
| 509 | ) |
| 510 | if self.rank in ranks: |
| 511 | self.ranks = ranks |
| 512 | self.world_size = len(ranks) |
| 513 | self.rank_in_group = ranks.index(self.rank) |
| 514 | self_device_group = device_group |
| 515 | self_cpu_group = cpu_group |
| 516 | |
| 517 | assert self_cpu_group is not None |
| 518 | assert self_device_group is not None |
| 519 | |
| 520 | self.group_ranks = group_ranks |
| 521 | self.torch_distributed_backend = torch_distributed_backend |
| 522 | |
| 523 | self.cpu_group = self_cpu_group |
| 524 | self.device_group = self_device_group |
| 525 | |
| 526 | from vllm.platforms import current_platform |
| 527 | |
| 528 | if current_platform.is_cuda_alike(): |
| 529 | visible_device_index = ( |
| 530 | current_platform.logical_device_id_to_visible_device_id( |
| 531 | self.device_index |
| 532 | ) |
| 533 | ) |
| 534 | self.device = torch.device(f"cuda:{visible_device_index}") |
| 535 | elif current_platform.is_xpu(): |
| 536 | self.device = torch.device(f"xpu:{self.device_index}") |
| 537 | elif current_platform.is_out_of_tree(): |
| 538 | self.device = torch.device( |
| 539 | f"{current_platform.device_name}:{self.device_index}" |
| 540 | ) |
| 541 | else: |
| 542 | self.device = torch.device("cpu") |
| 543 | |
| 544 | self.use_device_communicator = use_device_communicator |
| 545 | self.device_communicator = None |
| 546 | if use_device_communicator and self.world_size > 1: |
| 547 | device_comm_cls = resolve_obj_by_qualname( |
| 548 | current_platform.get_device_communicator_cls() |
| 549 | ) |
| 550 | self.device_communicator = device_comm_cls( |
| 551 | cpu_group=self.cpu_group, |
| 552 | device=self.device, |
| 553 | device_group=self.device_group, |
| 554 | unique_name=self.unique_name, |
| 555 | use_all2all=use_all2all, |
| 556 | ) |
| 557 | |
| 558 | from vllm.distributed.device_communicators.shm_broadcast import MessageQueue |
| 559 | |
| 560 | self.mq_broadcaster: MessageQueue | None = None |
| 561 | if use_message_queue_broadcaster and self.world_size > 1: |
| 562 | self.mq_broadcaster = MessageQueue.create_from_process_group( |
| 563 | self.cpu_group, 1 << 22, 6 |
| 564 | ) |
| 565 | |
| 566 | # TODO(#35915): Remove is_tpu() check once tpu_inference |
| 567 | # overrides use_custom_op_collectives() to return True. |
| 568 | self.use_custom_op_call = ( |
| 569 | current_platform.is_tpu() or current_platform.use_custom_op_collectives() |
| 570 | ) |
| 571 | |
| 572 | self.use_cpu_custom_send_recv = ( |
| 573 | current_platform.is_cpu() |
| 574 | and self.device_communicator |
| 575 | and getattr(self.device_communicator, "supports_tensor_dict", False) |
| 576 | ) |
| 577 | |
| 578 | # FIFO of in-flight ``isend_tensor_dict`` calls: one entry per call, |
| 579 | # (handles, sent source tensors). Self-retention keeps the send |
| 580 | # buffers alive for fire-and-forget callers that drop the returned |
| 581 | # handles; entries are reaped lazily on the next ``isend_tensor_dict`` |
| 582 | # (see ``_reap_completed_isends``). Lazily created instead of plain |
| 583 | # attribute init because some subclasses replicate ``__init__`` |
| 584 | # without calling ``super().__init__``. |
| 585 | self._pending_isends: deque[tuple[list[Handle], list[torch.Tensor]]] = deque() |
| 586 | |
| 587 | def make_sibling_device_group(self, group_desc: str | None = None) -> ProcessGroup: |
| 588 | """Create a new device-side ProcessGroup with the same per-rank membership |
| 589 | as this coordinator's `device_group`, but backed by a distinct communicator. |
| 590 | This is a collective call: every world rank must invoke it. Used where we |
| 591 | want to issue ops that can run concurrently with ops on `device_group`. |
| 592 | """ |
| 593 | from vllm.distributed.utils import get_distributed_timeout_or_none |
| 594 | |
| 595 | device_timeout = get_distributed_timeout_or_none() |
| 596 | sibling: ProcessGroup | None = None |
| 597 | for ranks in self.group_ranks: |
| 598 | pg = torch.distributed.new_group( |
| 599 | ranks, |
| 600 | backend=self.torch_distributed_backend, |
| 601 | group_desc=group_desc, |
| 602 | timeout=device_timeout, |
| 603 | ) |
| 604 | if self.rank in ranks: |
| 605 | sibling = pg |
| 606 | assert sibling is not None |
| 607 | return sibling |
| 608 | |
| 609 | def create_mq_broadcaster( |
| 610 | self, writer_rank=0, external_writer_handle=None, blocking=True |
| 611 | ): |
| 612 | from vllm.distributed.device_communicators.shm_broadcast import MessageQueue |
| 613 | |
| 614 | return MessageQueue.create_from_process_group( |
| 615 | self.cpu_group, |
| 616 | 1 << 22, |
| 617 | 6, |
| 618 | writer_rank=writer_rank, |
| 619 | external_writer_handle=external_writer_handle, |
| 620 | blocking=blocking, |
| 621 | ) |
| 622 | |
| 623 | def create_single_reader_mq_broadcasters( |
| 624 | self, reader_rank_in_group=0, blocking=False |
| 625 | ): |
| 626 | from vllm.distributed.device_communicators.shm_broadcast import MessageQueue |
| 627 | |
| 628 | return MessageQueue.create_from_process_group_single_reader( |
| 629 | self.cpu_group, |
| 630 | 1 << 22, |
| 631 | 6, |
| 632 | reader_rank=self.ranks[reader_rank_in_group], |
| 633 | blocking=blocking, |
| 634 | ) |
| 635 | |
| 636 | @property |
| 637 | def first_rank(self): |
| 638 | """Return the global rank of the first process in the group.""" |
| 639 | return self.ranks[0] |
| 640 | |
| 641 | @property |
| 642 | def last_rank(self): |
| 643 | """Return the global rank of the last process in the group.""" |
| 644 | return self.ranks[-1] |
| 645 | |
| 646 | @property |
| 647 | def is_first_rank(self): |
| 648 | """Return whether the caller is the first process in the group.""" |
| 649 | return self.rank == self.first_rank |
| 650 | |
| 651 | @property |
| 652 | def is_last_rank(self): |
| 653 | """Return whether the caller is the last process in the group.""" |
| 654 | return self.rank == self.last_rank |
| 655 | |
| 656 | @property |
| 657 | def next_rank(self): |
| 658 | """Return the global rank of the process that follows the caller.""" |
| 659 | rank_in_group = self.rank_in_group |
| 660 | world_size = self.world_size |
| 661 | return self.ranks[(rank_in_group + 1) % world_size] |
| 662 | |
| 663 | @property |
| 664 | def prev_rank(self): |
| 665 | """Return the global rank of the process that precedes the caller.""" |
| 666 | rank_in_group = self.rank_in_group |
| 667 | world_size = self.world_size |
| 668 | return self.ranks[(rank_in_group - 1) % world_size] |
| 669 | |
| 670 | @contextmanager |
| 671 | def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None): |
| 672 | if graph_capture_context is None: |
| 673 | stream = torch.cuda.Stream() |
| 674 | graph_capture_context = GraphCaptureContext(stream) |
| 675 | else: |
| 676 | stream = graph_capture_context.stream |
| 677 | |
| 678 | # only cuda/rocm uses this function, |
| 679 | # so we don't abstract it into the base class |
| 680 | maybe_ca_context = nullcontext() |
| 681 | maybe_fi_pcie_ipc_context: AbstractContextManager[Any] = nullcontext() |
| 682 | maybe_aiter_ar_context = nullcontext() |
| 683 | from vllm.distributed.device_communicators.cuda_communicator import ( |
| 684 | CudaCommunicator, |
| 685 | ) |
| 686 | from vllm.distributed.device_communicators.xpu_communicator import ( |
| 687 | XpuCommunicator, |
| 688 | ) |
| 689 | |
| 690 | if self.device_communicator is not None: |
| 691 | assert isinstance( |
| 692 | self.device_communicator, |
| 693 | (CudaCommunicator, XpuCommunicator), |
| 694 | ) |
| 695 | ca_comm = self.device_communicator.ca_comm |
| 696 | if ca_comm is not None: |
| 697 | maybe_ca_context = ca_comm.capture() # type: ignore |
| 698 | if isinstance(self.device_communicator, CudaCommunicator): |
| 699 | fi_pcie_ipc_ar_comm = self.device_communicator.fi_pcie_ipc_ar_comm |
| 700 | if fi_pcie_ipc_ar_comm is not None: |
| 701 | maybe_fi_pcie_ipc_context = fi_pcie_ipc_ar_comm.capture() |
| 702 | |
| 703 | # Capture each group's own comm. A global lookup would double-capture |
| 704 | aiter_ar_comm = getattr(self.device_communicator, "aiter_ar_comm", None) |
| 705 | if aiter_ar_comm is not None: |
| 706 | maybe_aiter_ar_context = aiter_ar_comm.capture() # type: ignore |
| 707 | |
| 708 | # ensure all initialization operations complete before attempting to |
| 709 | # capture the graph on another stream |
| 710 | curr_stream = torch.cuda.current_stream() |
| 711 | if curr_stream != stream: |
| 712 | stream.wait_stream(curr_stream) |
| 713 | |
| 714 | with ( |
| 715 | torch.cuda.stream(stream), |
| 716 | maybe_ca_context, |
| 717 | maybe_fi_pcie_ipc_context, |
| 718 | maybe_aiter_ar_context, |
| 719 | ): |
| 720 | yield graph_capture_context |
| 721 | |
| 722 | def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: |
| 723 | """User-facing all-reduce function before we actually call the |
| 724 | all-reduce operation. |
| 725 | |
| 726 | We need this because Dynamo does not support passing an arbitrary |
| 727 | object (`self` in this case) to a custom op. We need to pass the |
| 728 | group name as a string, and then look up the group coordinator from |
| 729 | the group name, dispatch the all-reduce operation to the group |
| 730 | coordinator. |
| 731 | |
| 732 | In addition, PyTorch custom ops do not support mutation or returning |
| 733 | a new tensor in the same op. So we always make the all-reduce operation |
| 734 | out-of-place. |
| 735 | """ |
| 736 | # Bypass the function if we are using only 1 GPU. |
| 737 | if self.world_size == 1: |
| 738 | return input_ |
| 739 | |
| 740 | if self.use_custom_op_call: |
| 741 | return torch.ops.vllm.all_reduce(input_, group_name=self.unique_name) |
| 742 | else: |
| 743 | return self._all_reduce_out_place(input_) |
| 744 | |
| 745 | def _all_reduce_out_place(self, input_: torch.Tensor) -> torch.Tensor: |
| 746 | if self.device_communicator is None: |
| 747 | raise ValueError("No device communicator found") |
| 748 | return self.device_communicator.all_reduce(input_) |
| 749 | |
| 750 | def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: |
| 751 | world_size = self.world_size |
| 752 | # Bypass the function if we are using only 1 GPU. |
| 753 | if world_size == 1: |
| 754 | return input_ |
| 755 | assert -input_.dim() <= dim < input_.dim(), ( |
| 756 | f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" |
| 757 | ) |
| 758 | |
| 759 | if self.use_custom_op_call: |
| 760 | return torch.ops.vllm.all_gather( |
| 761 | input_, dim, world_size, group_name=self.unique_name |
| 762 | ) |
| 763 | else: |
| 764 | return self._all_gather_out_place(input_, dim) |
| 765 | |
| 766 | def _all_gather_out_place(self, input_: torch.Tensor, dim: int) -> torch.Tensor: |
| 767 | if self.device_communicator is None: |
| 768 | raise ValueError("No device communicator found") |
| 769 | return self.device_communicator.all_gather(input_, dim) |
| 770 | |
| 771 | def all_gatherv( |
| 772 | self, |
| 773 | input_: torch.Tensor | list[torch.Tensor], |
| 774 | dim: int = 0, |
| 775 | sizes: list[int] | None = None, |
| 776 | ): |
| 777 | if self.device_communicator is None: |
| 778 | raise ValueError("No device communicator found") |
| 779 | return self.device_communicator.all_gatherv(input_, dim, sizes) |
| 780 | |
| 781 | def reduce_scatter(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: |
| 782 | world_size = self.world_size |
| 783 | # Bypass the function if we are using only 1 GPU. |
| 784 | if world_size == 1: |
| 785 | return input_ |
| 786 | assert -input_.dim() <= dim < input_.dim(), ( |
| 787 | f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" |
| 788 | ) |
| 789 | |
| 790 | if self.use_custom_op_call: |
| 791 | return torch.ops.vllm.reduce_scatter( |
| 792 | input_, dim, world_size, group_name=self.unique_name |
| 793 | ) |
| 794 | else: |
| 795 | return self._reduce_scatter_out_place(input_, dim) |
| 796 | |
| 797 | def reduce_scatterv( |
| 798 | self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None |
| 799 | ) -> torch.Tensor: |
| 800 | if self.device_communicator is None: |
| 801 | raise ValueError("No device communicator found") |
| 802 | return self.device_communicator.reduce_scatterv(input_, dim, sizes) |
| 803 | |
| 804 | def _reduce_scatter_out_place(self, input_: torch.Tensor, dim: int) -> torch.Tensor: |
| 805 | if self.device_communicator is None: |
| 806 | raise ValueError("No device communicator found") |
| 807 | return self.device_communicator.reduce_scatter(input_, dim) |
| 808 | |
| 809 | def gather( |
| 810 | self, input_: torch.Tensor, dst: int = 0, dim: int = -1 |
| 811 | ) -> torch.Tensor | None: |
| 812 | """NOTE: We assume that the input tensor is on the same device across |
| 813 | all the ranks. |
| 814 | NOTE: `dst` is the local rank of the destination rank. |
| 815 | """ |
| 816 | world_size = self.world_size |
| 817 | # Bypass the function if we are using only 1 GPU. |
| 818 | if world_size == 1: |
| 819 | return input_ |
| 820 | if self.device_communicator is None: |
| 821 | raise ValueError("No device communicator found") |
| 822 | return self.device_communicator.gather(input_, dst, dim) |
| 823 | |
| 824 | def broadcast(self, input_: torch.Tensor, src: int = 0): |
| 825 | """Broadcast the input tensor. |
| 826 | NOTE: `src` is the local rank of the source rank. |
| 827 | """ |
| 828 | assert src < self.world_size, f"Invalid src rank ({src})" |
| 829 | |
| 830 | # Bypass the function if we are using only 1 GPU. |
| 831 | if self.world_size == 1: |
| 832 | return input_ |
| 833 | # Broadcast. |
| 834 | torch.distributed.broadcast( |
| 835 | input_, src=self.ranks[src], group=self.device_group |
| 836 | ) |
| 837 | return input_ |
| 838 | |
| 839 | def broadcast_object(self, obj: Any | None = None, src: int = 0): |
| 840 | """Broadcast the input object. |
| 841 | NOTE: `src` is the local rank of the source rank. |
| 842 | """ |
| 843 | assert src < self.world_size, f"Invalid src rank ({src})" |
| 844 | |
| 845 | # Bypass the function if we are using only 1 GPU. |
| 846 | if self.world_size == 1: |
| 847 | return obj |
| 848 | if self.mq_broadcaster is not None: |
| 849 | assert src == 0, "Message queue broadcaster only supports src=0" |
| 850 | return self.mq_broadcaster.broadcast_object(obj) |
| 851 | if self.rank_in_group == src: |
| 852 | torch.distributed.broadcast_object_list( |
| 853 | [obj], src=self.ranks[src], group=self.cpu_group |
| 854 | ) |
| 855 | return obj |
| 856 | else: |
| 857 | recv = [None] |
| 858 | torch.distributed.broadcast_object_list( |
| 859 | recv, src=self.ranks[src], group=self.cpu_group |
| 860 | ) |
| 861 | return recv[0] |
| 862 | |
| 863 | def broadcast_object_list( |
| 864 | self, obj_list: list[Any], src: int = 0, group: ProcessGroup | None = None |
| 865 | ): |
| 866 | """Broadcast the input object list. |
| 867 | NOTE: `src` is the local rank of the source rank. |
| 868 | """ |
| 869 | assert src < self.world_size, f"Invalid src rank ({src})" |
| 870 | |
| 871 | # Bypass the function if we are using only 1 GPU. |
| 872 | if self.world_size == 1: |
| 873 | return obj_list |
| 874 | # Broadcast. |
| 875 | torch.distributed.broadcast_object_list( |
| 876 | obj_list, src=self.ranks[src], group=self.device_group |
| 877 | ) |
| 878 | return obj_list |
| 879 | |
| 880 | def send_object(self, obj: Any, dst: int) -> None: |
| 881 | """Send the input object list to the destination rank.""" |
| 882 | """NOTE: `dst` is the local rank of the destination rank.""" |
| 883 | |
| 884 | assert dst < self.world_size, f"Invalid dst rank ({dst})" |
| 885 | |
| 886 | assert dst != self.rank_in_group, ( |
| 887 | "Invalid destination rank. Destination rank is the same " |
| 888 | "as the current rank." |
| 889 | ) |
| 890 | |
| 891 | # Serialize object to tensor and get the size as well |
| 892 | object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8) |
| 893 | |
| 894 | size_tensor = torch.tensor( |
| 895 | [object_tensor.numel()], dtype=torch.long, device="cpu" |
| 896 | ) |
| 897 | |
| 898 | # Send object size |
| 899 | |
| 900 | torch.distributed.send(size_tensor, dst=self.ranks[dst], group=self.cpu_group) |
| 901 | |
| 902 | # Send object |
| 903 | torch.distributed.send(object_tensor, dst=self.ranks[dst], group=self.cpu_group) |
| 904 | |
| 905 | return None |
| 906 | |
| 907 | def recv_object(self, src: int) -> Any: |
| 908 | """Receive the input object list from the source rank.""" |
| 909 | """NOTE: `src` is the local rank of the source rank.""" |
| 910 | |
| 911 | assert src < self.world_size, f"Invalid src rank ({src})" |
| 912 | |
| 913 | assert src != self.rank_in_group, ( |
| 914 | "Invalid source rank. Source rank is the same as the current rank." |
| 915 | ) |
| 916 | |
| 917 | size_tensor = torch.empty(1, dtype=torch.long, device="cpu") |
| 918 | |
| 919 | # Receive object size |
| 920 | rank_size = torch.distributed.recv( |
| 921 | size_tensor, src=self.ranks[src], group=self.cpu_group |
| 922 | ) |
| 923 | |
| 924 | # Tensor to receive serialized objects into. |
| 925 | object_tensor = torch.empty( # type: ignore[call-overload] |
| 926 | size_tensor.item(), # type: ignore[arg-type] |
| 927 | dtype=torch.uint8, |
| 928 | device="cpu", |
| 929 | ) |
| 930 | |
| 931 | rank_object = torch.distributed.recv( |
| 932 | object_tensor, src=self.ranks[src], group=self.cpu_group |
| 933 | ) |
| 934 | |
| 935 | assert rank_object == rank_size, ( |
| 936 | "Received object sender rank does not match the size sender rank." |
| 937 | ) |
| 938 | |
| 939 | obj = pickle.loads(object_tensor.numpy().tobytes()) |
| 940 | |
| 941 | return obj |
| 942 | |
| 943 | def isend_object(self, obj: Any, dst: int) -> Handle: |
| 944 | """Non-blocking ``send_object``: isend size + pickled object on the |
| 945 | CPU group. Returns a handle that retains the serialized source |
| 946 | tensors until ``wait`` drains both sends. |
| 947 | """ |
| 948 | assert dst < self.world_size, f"Invalid dst rank ({dst})" |
| 949 | assert dst != self.rank_in_group, ( |
| 950 | "Invalid destination rank. Destination rank is the same " |
| 951 | "as the current rank." |
| 952 | ) |
| 953 | |
| 954 | object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8) |
| 955 | size_tensor = torch.tensor( |
| 956 | [object_tensor.numel()], dtype=torch.long, device="cpu" |
| 957 | ) |
| 958 | |
| 959 | retained = (size_tensor, object_tensor) |
| 960 | works: list[Any] = [] |
| 961 | for tensor in retained: |
| 962 | work = torch.distributed.isend( |
| 963 | tensor, dst=self.ranks[dst], group=self.cpu_group |
| 964 | ) |
| 965 | if work is not None: |
| 966 | works.append(work) |
| 967 | |
| 968 | return _RetainedHandle(works, retained) |
| 969 | |
| 970 | def broadcast_tensor_dict( |
| 971 | self, |
| 972 | tensor_dict: dict[str, torch.Tensor | Any] | None = None, |
| 973 | src: int = 0, |
| 974 | group: ProcessGroup | None = None, |
| 975 | metadata_group: ProcessGroup | None = None, |
| 976 | ) -> dict[str, torch.Tensor | Any] | None: |
| 977 | """Broadcast the input tensor dictionary. |
| 978 | NOTE: `src` is the local rank of the source rank. |
| 979 | """ |
| 980 | # Bypass the function if we are using only 1 GPU. |
| 981 | if not torch.distributed.is_initialized() or self.world_size == 1: |
| 982 | return tensor_dict |
| 983 | |
| 984 | group = self.device_group |
| 985 | metadata_group = self.cpu_group |
| 986 | assert src < self.world_size, f"Invalid src rank ({src})" |
| 987 | |
| 988 | rank_in_group = self.rank_in_group |
| 989 | if rank_in_group == src: |
| 990 | metadata_list: list[tuple[Any, Any]] = [] |
| 991 | assert isinstance(tensor_dict, dict), ( |
| 992 | f"Expecting a dictionary, got {type(tensor_dict)}" |
| 993 | ) |
| 994 | metadata_list, tensor_list = _split_tensor_dict(tensor_dict) |
| 995 | # `metadata_list` lives in CPU memory. |
| 996 | # `broadcast_object_list` has serialization & deserialization, |
| 997 | # all happening on CPU. Therefore, we can use the CPU group. |
| 998 | self.broadcast_object(metadata_list, src=src) |
| 999 | async_handles = [] |
| 1000 | for tensor in tensor_list: |
| 1001 | if tensor.numel() == 0: |
| 1002 | # Skip broadcasting empty tensors. |
| 1003 | continue |
| 1004 | if tensor.is_cpu: |
| 1005 | # use metadata_group for CPU tensors |
| 1006 | handle = torch.distributed.broadcast( |
| 1007 | tensor, src=self.ranks[src], group=metadata_group, async_op=True |
| 1008 | ) |
| 1009 | else: |
| 1010 | # use group for GPU tensors |
| 1011 | handle = torch.distributed.broadcast( |
| 1012 | tensor, src=self.ranks[src], group=group, async_op=True |
| 1013 | ) |
| 1014 | async_handles.append(handle) |
| 1015 | for async_handle in async_handles: |
| 1016 | async_handle.wait() |
| 1017 | |
| 1018 | else: |
| 1019 | metadata_list = self.broadcast_object(None, src=src) |
| 1020 | tensor_dict = {} |
| 1021 | async_handles = [] |
| 1022 | for key, value in metadata_list: |
| 1023 | if isinstance(value, TensorMetadata): |
| 1024 | tensor = torch.empty( |
| 1025 | value.size, dtype=value.dtype, device=value.device |
| 1026 | ) |
| 1027 | if tensor.numel() == 0: |
| 1028 | # Skip broadcasting empty tensors. |
| 1029 | tensor_dict[key] = tensor |
| 1030 | continue |
| 1031 | if tensor.is_cpu: |
| 1032 | # use metadata_group for CPU tensors |
| 1033 | handle = torch.distributed.broadcast( |
| 1034 | tensor, |
| 1035 | src=self.ranks[src], |
| 1036 | group=metadata_group, |
| 1037 | async_op=True, |
| 1038 | ) |
| 1039 | else: |
| 1040 | # use group for GPU tensors |
| 1041 | handle = torch.distributed.broadcast( |
| 1042 | tensor, src=self.ranks[src], group=group, async_op=True |
| 1043 | ) |
| 1044 | async_handles.append(handle) |
| 1045 | tensor_dict[key] = tensor |
| 1046 | else: |
| 1047 | tensor_dict[key] = value |
| 1048 | for async_handle in async_handles: |
| 1049 | async_handle.wait() |
| 1050 | return tensor_dict |
| 1051 | |
| 1052 | def _should_use_all_gather( |
| 1053 | self, |
| 1054 | key: str, |
| 1055 | numel: int, |
| 1056 | all_gather_group: "GroupCoordinator | None", |
| 1057 | all_gather_tensors: dict[str, bool] | None, |
| 1058 | ) -> bool: |
| 1059 | if all_gather_group is None: |
| 1060 | return False |
| 1061 | use_all_gather = numel % all_gather_group.world_size == 0 |
| 1062 | if all_gather_tensors is not None: |
| 1063 | use_all_gather = all_gather_tensors.get(key, use_all_gather) |
| 1064 | return use_all_gather |
| 1065 | |
| 1066 | def send_tensor_dict( |
| 1067 | self, |
| 1068 | tensor_dict: dict[str, torch.Tensor | Any], |
| 1069 | dst: int | None = None, |
| 1070 | all_gather_group: "GroupCoordinator | None" = None, |
| 1071 | all_gather_tensors: dict[str, bool] | None = None, |
| 1072 | ) -> dict[str, torch.Tensor | Any] | None: |
| 1073 | """Send the input tensor dictionary. |
| 1074 | NOTE: `dst` is the local rank of the source rank. |
| 1075 | |
| 1076 | all_gather_group: The group for the all-gather operation. If provided, |
| 1077 | an optimization is enabled where each rank in the group sends a |
| 1078 | slice of a tensor and the receiver reconstructs it using an |
| 1079 | all-gather, which can improve performance. This is typically the |
| 1080 | tensor-parallel group. |
| 1081 | all_gather_tensors: A dictionary to specify which tensors should use |
| 1082 | the all-gather optimization, which is only effective when |
| 1083 | `all_gather_group` is provided. By default, this optimization is |
| 1084 | on for any tensor whose size is divisible by the |
| 1085 | `all_gather_group`'s world size. However, it should be disabled |
| 1086 | for tensors that are not fully replicated across the group (e.g., |
| 1087 | the residual tensor when sequence parallelism is enabled). This |
| 1088 | dictionary allows overriding the default behavior on a per-tensor |
| 1089 | basis. |
| 1090 | """ |
| 1091 | # Bypass the function if we are using only 1 GPU. |
| 1092 | if not torch.distributed.is_initialized() or self.world_size == 1: |
| 1093 | return tensor_dict |
| 1094 | handles = self.isend_tensor_dict( |
| 1095 | tensor_dict, |
| 1096 | dst=dst, |
| 1097 | all_gather_group=all_gather_group, |
| 1098 | all_gather_tensors=all_gather_tensors, |
| 1099 | ) |
| 1100 | for handle in handles: |
| 1101 | handle.wait() |
| 1102 | # The sync path waited every handle, so the self-retention entry |
| 1103 | # appended by ``isend_tensor_dict`` is already drained; drop it |
| 1104 | # instead of leaking the tensor refs until the next async send. |
| 1105 | pending = getattr(self, "_pending_isends", None) |
| 1106 | if pending and pending[-1][0] is handles: |
| 1107 | pending.pop() |
| 1108 | return None |
| 1109 | |
| 1110 | def _reap_completed_isends(self) -> None: |
| 1111 | """Lazily drop self-retained ``isend_tensor_dict`` entries that have |
| 1112 | completed, oldest first (FIFO). |
| 1113 | |
| 1114 | gloo ``Work.is_completed()`` is unreliable (it can report completion |
| 1115 | before the background copy of the source buffer finishes), so the |
| 1116 | metadata handle (``handles[0]``, a gloo-backed ``_RetainedHandle``) is |
| 1117 | never used as the completion gate. Instead we gate on the tensor-send |
| 1118 | handles (``handles[1:]``), which run on the device backend where |
| 1119 | ``is_completed()`` is reliable. Once all tensor sends are done the |
| 1120 | peer must already have posted the matching tensor recvs — and since |
| 1121 | the receiver ingests metadata before tensors, the gloo metadata send |
| 1122 | is then guaranteed complete, so ``wait()`` on it is ~instant and only |
| 1123 | serves to drop the retained refs. |
| 1124 | |
| 1125 | Metadata-only entries (``handles[1:]`` empty, e.g. an empty |
| 1126 | ``IntermediateTensors``) have no reliable completion signal; drop them |
| 1127 | best-effort without blocking on the gloo metadata wait. |
| 1128 | """ |
| 1129 | pending = getattr(self, "_pending_isends", None) |
| 1130 | if pending is None: |
| 1131 | # Subclasses that replicate __init__ without calling |
| 1132 | # super().__init__ may not have the attribute yet. |
| 1133 | self._pending_isends = pending = deque() |
| 1134 | while pending: |
| 1135 | handles, _ = pending[0] |
| 1136 | tensor_handles = handles[1:] |
| 1137 | if not tensor_handles: |
| 1138 | pending.popleft() |
| 1139 | continue |
| 1140 | if not all(h.is_completed() for h in tensor_handles): |
| 1141 | # The oldest send is still in flight; FIFO ordering means the |
| 1142 | # rest are no further along, so stop here. |
| 1143 | break |
| 1144 | handles[0].wait() |
| 1145 | pending.popleft() |
| 1146 | |
| 1147 | def isend_tensor_dict( |
| 1148 | self, |
| 1149 | tensor_dict: dict[str, torch.Tensor | Any], |
| 1150 | dst: int | None = None, |
| 1151 | all_gather_group: "GroupCoordinator | None" = None, |
| 1152 | all_gather_tensors: dict[str, bool] | None = None, |
| 1153 | ) -> list[Handle]: |
| 1154 | """Send the input tensor dictionary asynchronously. |
| 1155 | |
| 1156 | The returned handles are additionally self-retained in |
| 1157 | ``self._pending_isends`` (together with the sent source tensors) and |
| 1158 | reaped lazily on subsequent calls, so fire-and-forget callers that |
| 1159 | drop the handles cannot free or overwrite the send buffers while the |
| 1160 | sends are still in flight. Callers that do keep the handles may |
| 1161 | ``wait()`` them as before; the retention entry is dropped on reap |
| 1162 | either way. |
| 1163 | """ |
| 1164 | if self.world_size <= 1: |
| 1165 | return [] |
| 1166 | |
| 1167 | if dst is None: |
| 1168 | dst = (self.rank_in_group + 1) % self.world_size |
| 1169 | assert dst < self.world_size, f"Invalid dst rank ({dst})" |
| 1170 | |
| 1171 | if self.use_cpu_custom_send_recv: |
| 1172 | if self.device_communicator is None: |
| 1173 | raise ValueError("No device communicator found") |
| 1174 | # custom device communicator path is synchronous |
| 1175 | self.device_communicator.send_tensor_dict( # type: ignore |
| 1176 | tensor_dict, dst |
| 1177 | ) |
| 1178 | return [] |
| 1179 | |
| 1180 | all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size |
| 1181 | all_gather_rank = ( |
| 1182 | 0 if all_gather_group is None else all_gather_group.rank_in_group |
| 1183 | ) |
| 1184 | |
| 1185 | group = self.device_group |
| 1186 | metadata_group = self.cpu_group |
| 1187 | |
| 1188 | metadata_list, tensor_list = _split_tensor_dict(tensor_dict) |
| 1189 | |
| 1190 | tensor_keys = [k for k, v in tensor_dict.items() if isinstance(v, torch.Tensor)] |
| 1191 | assert len(tensor_keys) == len(tensor_list) |
| 1192 | |
| 1193 | # Reap completed sends before posting new ones, bounding the |
| 1194 | # self-retention FIFO. |
| 1195 | self._reap_completed_isends() |
| 1196 | handles: list[Handle] = [self.isend_object(metadata_list, dst=dst)] |
| 1197 | |
| 1198 | sent_tensors: list[torch.Tensor] = [] |
| 1199 | for key, tensor in zip(tensor_keys, tensor_list): |
| 1200 | if tensor.numel() == 0: |
| 1201 | continue |
| 1202 | |
| 1203 | if self._should_use_all_gather( |
| 1204 | key, tensor.numel(), all_gather_group, all_gather_tensors |
| 1205 | ): |
| 1206 | tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] |
| 1207 | |
| 1208 | comm_group = metadata_group if tensor.is_cpu else group |
| 1209 | handle = torch.distributed.isend( |
| 1210 | tensor, dst=self.ranks[dst], group=comm_group |
| 1211 | ) |
| 1212 | if tensor.is_cuda: |
| 1213 | tensor.record_stream(torch.cuda.current_stream(tensor.device)) |
| 1214 | handles.append(handle) |
| 1215 | sent_tensors.append(tensor) |
| 1216 | |
| 1217 | # Self-retain for fire-and-forget callers: keep the handles and |
| 1218 | # the source tensors alive until the sends complete (reaped |
| 1219 | # lazily by ``_reap_completed_isends``). Without this, dropping |
| 1220 | # the returned handles would GC the ``_RetainedHandle``'s |
| 1221 | # serialized metadata buffers mid-send, and non-CUDA devices |
| 1222 | # have no ``record_stream`` equivalent protecting the tensor |
| 1223 | # sources. |
| 1224 | pending = getattr(self, "_pending_isends", None) |
| 1225 | if pending is None: |
| 1226 | self._pending_isends = pending = deque() |
| 1227 | pending.append((handles, sent_tensors)) |
| 1228 | |
| 1229 | return handles |
| 1230 | |
| 1231 | def recv_tensor_dict( |
| 1232 | self, |
| 1233 | src: int | None = None, |
| 1234 | all_gather_group: "GroupCoordinator | None" = None, |
| 1235 | all_gather_tensors: dict[str, bool] | None = None, |
| 1236 | ) -> dict[str, torch.Tensor | Any] | None: |
| 1237 | """Recv the input tensor dictionary. |
| 1238 | NOTE: `src` is the local rank of the source rank. |
| 1239 | |
| 1240 | all_gather_group: The group for the all-gather operation. If provided, |
| 1241 | an optimization is enabled where each rank in the group sends a |
| 1242 | slice of a tensor and the receiver reconstructs it using an |
| 1243 | all-gather, which can improve performance. This is typically the |
| 1244 | tensor-parallel group. |
| 1245 | all_gather_tensors: A dictionary to specify which tensors should use |
| 1246 | the all-gather optimization, which is only effective when |
| 1247 | `all_gather_group` is provided. By default, this optimization is |
| 1248 | on for any tensor whose size is divisible by the |
| 1249 | `all_gather_group`'s world size. However, it should be disabled |
| 1250 | for tensors that are not fully replicated across the group (e.g., |
| 1251 | the residual tensor when sequence parallelism is enabled). This |
| 1252 | dictionary allows overriding the default behavior on a per-tensor |
| 1253 | basis. |
| 1254 | """ |
| 1255 | # Bypass the function if we are using only 1 GPU. |
| 1256 | if not torch.distributed.is_initialized() or self.world_size == 1: |
| 1257 | return None |
| 1258 | tensor_dict, handles, postprocess = self.irecv_tensor_dict( |
| 1259 | src=src, |
| 1260 | all_gather_group=all_gather_group, |
| 1261 | all_gather_tensors=all_gather_tensors, |
| 1262 | ) |
| 1263 | for handle in handles: |
| 1264 | handle.wait() |
| 1265 | for fn in postprocess: |
| 1266 | fn() |
| 1267 | return tensor_dict |
| 1268 | |
| 1269 | def irecv_tensor_dict( |
| 1270 | self, |
| 1271 | src: int | None = None, |
| 1272 | all_gather_group: "GroupCoordinator | None" = None, |
| 1273 | all_gather_tensors: dict[str, bool] | None = None, |
| 1274 | ) -> tuple[ |
| 1275 | dict[str, torch.Tensor | Any] | None, |
| 1276 | list[Handle], |
| 1277 | list[Callable[[], None]], |
| 1278 | ]: |
| 1279 | if not torch.distributed.is_initialized() or self.world_size == 1: |
| 1280 | return None, [], [] |
| 1281 | |
| 1282 | if src is None: |
| 1283 | src = (self.rank_in_group - 1) % self.world_size |
| 1284 | assert src < self.world_size, f"Invalid src rank ({src})" |
| 1285 | |
| 1286 | if self.use_cpu_custom_send_recv: |
| 1287 | if self.device_communicator is None: |
| 1288 | raise ValueError("No device communicator found") |
| 1289 | # custom device communicator path is synchronous |
| 1290 | sync_tensor_dict = self.device_communicator.recv_tensor_dict( # type: ignore |
| 1291 | src |
| 1292 | ) |
| 1293 | return sync_tensor_dict, [], [] |
| 1294 | |
| 1295 | all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size |
| 1296 | all_gather_rank = ( |
| 1297 | 0 if all_gather_group is None else all_gather_group.rank_in_group |
| 1298 | ) |
| 1299 | |
| 1300 | group = self.device_group |
| 1301 | metadata_group = self.cpu_group |
| 1302 | |
| 1303 | recv_metadata_list = self.recv_object(src=src) |
| 1304 | tensor_dict: dict[str, Any] = {} |
| 1305 | handles: list[Handle] = [] |
| 1306 | postprocess: list[Callable[[], None]] = [] |
| 1307 | |
| 1308 | for key, value in recv_metadata_list: |
| 1309 | if isinstance(value, TensorMetadata): |
| 1310 | full_tensor = torch.empty( |
| 1311 | value.size, dtype=value.dtype, device=value.device |
| 1312 | ) |
| 1313 | if full_tensor.numel() == 0: |
| 1314 | tensor_dict[key] = full_tensor |
| 1315 | continue |
| 1316 | |
| 1317 | if self._should_use_all_gather( |
| 1318 | key, full_tensor.numel(), all_gather_group, all_gather_tensors |
| 1319 | ): |
| 1320 | orig_shape = full_tensor.shape |
| 1321 | slice_tensor = full_tensor.reshape(all_gather_size, -1)[ |
| 1322 | all_gather_rank |
| 1323 | ] |
| 1324 | comm_group = metadata_group if slice_tensor.is_cpu else group |
| 1325 | handle = torch.distributed.irecv( |
| 1326 | slice_tensor, src=self.ranks[src], group=comm_group |
| 1327 | ) |
| 1328 | handles.append(handle) |
| 1329 | |
| 1330 | def _postprocess( |
| 1331 | key: str = key, |
| 1332 | slice_tensor: torch.Tensor = slice_tensor, |
| 1333 | orig_shape: tuple[int, ...] = tuple(orig_shape), |
| 1334 | all_gather_group=all_gather_group, |
| 1335 | ) -> None: |
| 1336 | assert all_gather_group is not None |
| 1337 | tensor_dict[key] = all_gather_group.all_gather( |
| 1338 | slice_tensor, dim=0 |
| 1339 | ).reshape(orig_shape) |
| 1340 | |
| 1341 | postprocess.append(_postprocess) |
| 1342 | tensor_dict[key] = slice_tensor |
| 1343 | else: |
| 1344 | comm_group = metadata_group if full_tensor.is_cpu else group |
| 1345 | handle = torch.distributed.irecv( |
| 1346 | full_tensor, src=self.ranks[src], group=comm_group |
| 1347 | ) |
| 1348 | handles.append(handle) |
| 1349 | tensor_dict[key] = full_tensor |
| 1350 | else: |
| 1351 | tensor_dict[key] = value |
| 1352 | |
| 1353 | return tensor_dict, handles, postprocess |
| 1354 | |
| 1355 | def barrier(self): |
| 1356 | """Barrier synchronization among the group. |
| 1357 | NOTE: don't use `device_group` here! `barrier` in NCCL is |
| 1358 | terrible because it is internally a broadcast operation with |
| 1359 | secretly created GPU tensors. It is easy to mess up the current |
| 1360 | device. Use the CPU group instead. |
| 1361 | """ |
| 1362 | torch.distributed.barrier(group=self.cpu_group) |
| 1363 | |
| 1364 | def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: |
| 1365 | """Sends a tensor to the destination rank in a blocking way.""" |
| 1366 | """NOTE: `dst` is the local rank of the destination rank.""" |
| 1367 | if self.device_communicator is None: |
| 1368 | raise ValueError("No device communicator found") |
| 1369 | self.device_communicator.send(tensor, dst) |
| 1370 | |
| 1371 | def recv( |
| 1372 | self, size: torch.Size, dtype: torch.dtype, src: int | None = None |
| 1373 | ) -> torch.Tensor: |
| 1374 | """Receives a tensor from the source rank.""" |
| 1375 | """NOTE: `src` is the local rank of the source rank.""" |
| 1376 | if self.device_communicator is None: |
| 1377 | raise ValueError("No device communicator found") |
| 1378 | return self.device_communicator.recv(size, dtype, src) |
| 1379 | |
| 1380 | def destroy(self): |
| 1381 | # Device communicators can own collective workspaces whose teardown |
| 1382 | # uses these process groups (for example FlashInfer PCIe IPC barriers). |
| 1383 | # Release them before destroying the groups they depend on. |
| 1384 | if self.device_communicator is not None: |
| 1385 | self.device_communicator.destroy() |
| 1386 | if hasattr(self, "device_group"): |
| 1387 | torch.distributed.destroy_process_group(self.device_group) |
| 1388 | del self.device_group |
| 1389 | if hasattr(self, "cpu_group"): |
| 1390 | torch.distributed.destroy_process_group(self.cpu_group) |
| 1391 | del self.cpu_group |
| 1392 | if self.mq_broadcaster is not None: |
| 1393 | self.mq_broadcaster = None |
| 1394 | |
| 1395 | def dispatch_router_logits( |
| 1396 | self, |
| 1397 | hidden_states: torch.Tensor, |
| 1398 | router_logits: torch.Tensor, |
| 1399 | is_sequence_parallel: bool = False, |
| 1400 | extra_tensors: list[torch.Tensor] | None = None, |
| 1401 | ) -> ( |
| 1402 | tuple[torch.Tensor, torch.Tensor] |
| 1403 | | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 1404 | ): |
| 1405 | if self.device_communicator is not None: |
| 1406 | return self.device_communicator.dispatch_router_logits( |
| 1407 | hidden_states, |
| 1408 | router_logits, |
| 1409 | is_sequence_parallel, |
| 1410 | extra_tensors, |
| 1411 | ) |
| 1412 | else: |
| 1413 | return hidden_states, router_logits |
| 1414 | |
| 1415 | def dispatch( |
| 1416 | self, |
| 1417 | hidden_states: torch.Tensor, |
| 1418 | topk_weights: torch.Tensor, |
| 1419 | topk_ids: torch.Tensor, |
| 1420 | is_sequence_parallel: bool = False, |
| 1421 | extra_tensors: list[torch.Tensor] | None = None, |
| 1422 | ) -> ( |
| 1423 | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 1424 | | tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 1425 | ): |
| 1426 | if self.device_communicator is not None: |
| 1427 | return self.device_communicator.dispatch( |
| 1428 | hidden_states, |
| 1429 | topk_weights, |
| 1430 | topk_ids, |
| 1431 | is_sequence_parallel, |
| 1432 | extra_tensors, |
| 1433 | ) |
| 1434 | else: |
| 1435 | return hidden_states, topk_weights, topk_ids |
| 1436 | |
| 1437 | def combine( |
| 1438 | self, hidden_states, is_sequence_parallel: bool = False |
| 1439 | ) -> torch.Tensor: |
| 1440 | if self.device_communicator is not None: |
| 1441 | return self.device_communicator.combine(hidden_states, is_sequence_parallel) |
| 1442 | else: |
| 1443 | return hidden_states |
| 1444 | |
| 1445 | |
| 1446 | _WORLD: GroupCoordinator | None = None |
| 1447 | _INNER_DP_WORLD: GroupCoordinator | None = None |
| 1448 | _NODE_COUNT: int | None = None |
| 1449 | |
| 1450 | |
| 1451 | def get_world_group() -> GroupCoordinator: |
| 1452 | assert _WORLD is not None, "world group is not initialized" |
| 1453 | return _WORLD |
| 1454 | |
| 1455 | |
| 1456 | def get_inner_dp_world_group() -> GroupCoordinator: |
| 1457 | assert _INNER_DP_WORLD is not None, "inner dp world group is not initialized" |
| 1458 | return _INNER_DP_WORLD |
| 1459 | |
| 1460 | |
| 1461 | def init_world_group( |
| 1462 | ranks: list[int], local_rank: int, backend: str |
| 1463 | ) -> GroupCoordinator: |
| 1464 | return GroupCoordinator( |
| 1465 | group_ranks=[ranks], |
| 1466 | local_rank=local_rank, |
| 1467 | torch_distributed_backend=backend, |
| 1468 | use_device_communicator=False, |
| 1469 | group_name="world", |
| 1470 | ) |
| 1471 | |
| 1472 | |
| 1473 | def init_model_parallel_group( |
| 1474 | group_ranks: list[list[int]], |
| 1475 | local_rank: int, |
| 1476 | backend: str, |
| 1477 | use_message_queue_broadcaster: bool = False, |
| 1478 | group_name: str | None = None, |
| 1479 | use_device_communicator: bool = True, |
| 1480 | use_all2all: bool = False, |
| 1481 | ) -> GroupCoordinator: |
| 1482 | return GroupCoordinator( |
| 1483 | group_ranks=group_ranks, |
| 1484 | local_rank=local_rank, |
| 1485 | torch_distributed_backend=backend, |
| 1486 | use_device_communicator=use_device_communicator, |
| 1487 | use_message_queue_broadcaster=use_message_queue_broadcaster, |
| 1488 | group_name=group_name, |
| 1489 | use_all2all=use_all2all, |
| 1490 | ) |
| 1491 | |
| 1492 | |
| 1493 | def _init_stateless_group( |
| 1494 | group_ranks: list[list[int]], |
| 1495 | group_name: str, |
| 1496 | host: str, |
| 1497 | backend: str, |
| 1498 | coord_store: Store, |
| 1499 | use_device_communicator: bool = True, |
| 1500 | use_all2all: bool = False, |
| 1501 | ) -> "StatelessGroupCoordinator": |
| 1502 | """Create a StatelessGroupCoordinator with the given parameters.""" |
| 1503 | from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator |
| 1504 | |
| 1505 | world = get_world_group() |
| 1506 | return StatelessGroupCoordinator( |
| 1507 | group_ranks=group_ranks, |
| 1508 | local_rank=world.local_rank, |
| 1509 | torch_distributed_backend=backend, |
| 1510 | use_device_communicator=use_device_communicator, |
| 1511 | group_name=group_name, |
| 1512 | host=host, |
| 1513 | coord_store=coord_store, |
| 1514 | global_rank=world.rank, |
| 1515 | global_world_size=world.world_size, |
| 1516 | use_all2all=use_all2all, |
| 1517 | ) |
| 1518 | |
| 1519 | |
| 1520 | def _replace_active_groups( |
| 1521 | *, |
| 1522 | world: GroupCoordinator | None, |
| 1523 | dp: GroupCoordinator | None, |
| 1524 | ep: GroupCoordinator | None, |
| 1525 | eplb: GroupCoordinator | None, |
| 1526 | node_count: int | None, |
| 1527 | ) -> tuple[GroupCoordinator | None, ...]: |
| 1528 | """Replace the active groups and return the groups they replaced. |
| 1529 | |
| 1530 | The caller must destroy the returned DP, EP, WORLD, and EPLB groups |
| 1531 | collectively and in that order. Pass all-``None`` to remove the active |
| 1532 | groups without replacement. |
| 1533 | """ |
| 1534 | global _WORLD, _DP, _EP, _EPLB, _NODE_COUNT |
| 1535 | old_groups = _DP, _EP, _WORLD, _EPLB |
| 1536 | _WORLD = world |
| 1537 | _DP = dp |
| 1538 | _EP = ep |
| 1539 | _EPLB = eplb |
| 1540 | _NODE_COUNT = node_count |
| 1541 | return old_groups |
| 1542 | |
| 1543 | |
| 1544 | _TP: GroupCoordinator | None = None |
| 1545 | |
| 1546 | |
| 1547 | def get_tp_group() -> GroupCoordinator: |
| 1548 | assert _TP is not None, "tensor model parallel group is not initialized" |
| 1549 | return _TP |
| 1550 | |
| 1551 | |
| 1552 | _ETP: GroupCoordinator | None = None |
| 1553 | |
| 1554 | |
| 1555 | def get_etp_group() -> GroupCoordinator: |
| 1556 | """Return the Engram Tensor Parallel (ETP) group that shards one embedding table.""" |
| 1557 | assert _ETP is not None, "Engram tensor-parallel group is not initialized" |
| 1558 | return _ETP |
| 1559 | |
| 1560 | |
| 1561 | _ENGRAM_DP: GroupCoordinator | None = None |
| 1562 | |
| 1563 | |
| 1564 | def get_engram_dp_group() -> GroupCoordinator | None: |
| 1565 | """Return the DP replicas that share one engram embedding table. |
| 1566 | |
| 1567 | None when every replica holds a full (TP-sharded) copy, which is the case |
| 1568 | for models without engram layers and for replicas that span nodes. |
| 1569 | """ |
| 1570 | return _ENGRAM_DP |
| 1571 | |
| 1572 | |
| 1573 | def get_engram_dp_size() -> int: |
| 1574 | """Number of DP replicas one engram embedding table is sharded over.""" |
| 1575 | return _ENGRAM_DP.world_size if _ENGRAM_DP is not None else 1 |
| 1576 | |
| 1577 | |
| 1578 | _DCP: GroupCoordinator | None = None |
| 1579 | |
| 1580 | |
| 1581 | def get_dcp_group() -> GroupCoordinator: |
| 1582 | assert _DCP is not None, "decode context model parallel group is not initialized" |
| 1583 | return _DCP |
| 1584 | |
| 1585 | |
| 1586 | def get_dcp_world_size_and_rank(enabled: bool = True) -> tuple[int, int]: |
| 1587 | """Return ``(world_size, rank)`` in the DCP group, or ``(1, 0)`` when disabled |
| 1588 | (e.g. a replicated draft cache) or the group is uninitialized (unit tests).""" |
| 1589 | if not enabled or _DCP is None: |
| 1590 | return 1, 0 |
| 1591 | return _DCP.world_size, _DCP.rank_in_group |
| 1592 | |
| 1593 | |
| 1594 | _PP: GroupCoordinator | None = None |
| 1595 | |
| 1596 | |
| 1597 | def get_pp_group() -> GroupCoordinator: |
| 1598 | assert _PP is not None, "pipeline model parallel group is not initialized" |
| 1599 | return _PP |
| 1600 | |
| 1601 | |
| 1602 | _DP: GroupCoordinator | None = None |
| 1603 | |
| 1604 | |
| 1605 | def get_dp_group() -> GroupCoordinator: |
| 1606 | assert _DP is not None, "data parallel group is not initialized" |
| 1607 | return _DP |
| 1608 | |
| 1609 | |
| 1610 | _EP: GroupCoordinator | None = None |
| 1611 | |
| 1612 | |
| 1613 | def get_ep_group() -> GroupCoordinator: |
| 1614 | assert _EP is not None, ( |
| 1615 | "expert parallel group is not initialized. " |
| 1616 | "EP group is only created for MoE models with num_experts > 0. " |
| 1617 | "This function should only be called for MoE models." |
| 1618 | ) |
| 1619 | return _EP |
| 1620 | |
| 1621 | |
| 1622 | _EPLB: GroupCoordinator | None = None |
| 1623 | |
| 1624 | |
| 1625 | def get_eplb_group() -> GroupCoordinator: |
| 1626 | assert _EPLB is not None, ( |
| 1627 | "EPLB group is not initialized. " |
| 1628 | "EPLB group is only created for MoE models when EPLB is enabled. " |
| 1629 | "Ensure parallel_config.enable_eplb is True." |
| 1630 | ) |
| 1631 | return _EPLB |
| 1632 | |
| 1633 | |
| 1634 | _PCP: GroupCoordinator | None = None |
| 1635 | |
| 1636 | |
| 1637 | def get_pcp_group() -> GroupCoordinator: |
| 1638 | assert _PCP is not None, "prefill context parallel group is not initialized" |
| 1639 | return _PCP |
| 1640 | |
| 1641 | |
| 1642 | @contextmanager |
| 1643 | def graph_capture(device: torch.device): |
| 1644 | """`graph_capture` is a context manager which should surround the code that |
| 1645 | is capturing the CUDA graph. Its main purpose is to ensure that some |
| 1646 | operations will be run after the graph is captured, before the graph |
| 1647 | is replayed. It returns a `GraphCaptureContext` object which contains the |
| 1648 | necessary data for the graph capture. Currently, it only contains the |
| 1649 | stream that the graph capture is running on. This stream is set to the |
| 1650 | current CUDA stream when the context manager is entered and reset to the |
| 1651 | default stream when the context manager is exited. This is to ensure that |
| 1652 | the graph capture is running on a separate stream from the default stream, |
| 1653 | in order to explicitly distinguish the kernels to capture |
| 1654 | from other kernels possibly launched on background in the default stream. |
| 1655 | """ |
| 1656 | context = GraphCaptureContext(torch.cuda.Stream(device=device)) |
| 1657 | with ( |
| 1658 | get_tp_group().graph_capture(context), |
| 1659 | get_pp_group().graph_capture(context), |
| 1660 | get_dp_group().graph_capture(context), |
| 1661 | ): |
| 1662 | yield context |
| 1663 | |
| 1664 | |
| 1665 | logger = init_logger(__name__) |
| 1666 | |
| 1667 | _ENABLE_CUSTOM_ALL_REDUCE = True |
| 1668 | |
| 1669 | |
| 1670 | def set_custom_all_reduce(enable: bool): |
| 1671 | global _ENABLE_CUSTOM_ALL_REDUCE |
| 1672 | _ENABLE_CUSTOM_ALL_REDUCE = enable |
| 1673 | |
| 1674 | |
| 1675 | def _init_process_group_for_split_group( |
| 1676 | *, |
| 1677 | backend: str, |
| 1678 | distributed_init_method: str, |
| 1679 | store: Store | None, |
| 1680 | world_size: int, |
| 1681 | rank: int, |
| 1682 | local_rank: int, |
| 1683 | timeout: timedelta | None, |
| 1684 | ) -> None: |
| 1685 | """Initialize the default PG with both CPU (gloo) and device (e.g. nccl) |
| 1686 | backends and an eager ``device_id`` binding so that subgroups can be |
| 1687 | created via ``split_group`` (which requires the parent communicator to |
| 1688 | be eagerly initialized). Falls back to ``gloo`` on CPU-only systems. |
| 1689 | """ |
| 1690 | if torch.accelerator.is_available() and backend != "gloo": |
| 1691 | init_backend = "cpu:gloo,cuda:nccl" |
| 1692 | from vllm.platforms import current_platform |
| 1693 | |
| 1694 | visible_device_index = current_platform.logical_device_id_to_visible_device_id( |
| 1695 | local_rank |
| 1696 | ) |
| 1697 | device_id: torch.device | None = torch.device(f"cuda:{visible_device_index}") |
| 1698 | else: |
| 1699 | init_backend = "gloo" |
| 1700 | device_id = None |
| 1701 | torch.distributed.init_process_group( |
| 1702 | backend=init_backend, |
| 1703 | init_method=distributed_init_method if store is None else None, |
| 1704 | store=store, |
| 1705 | world_size=world_size, |
| 1706 | rank=rank, |
| 1707 | timeout=timeout, |
| 1708 | device_id=device_id, |
| 1709 | ) |
| 1710 | |
| 1711 | |
| 1712 | def _validate_default_pg_for_split_group() -> None: |
| 1713 | """When an external launcher (e.g. ``torchrun``) initialized the default |
| 1714 | PG, ``GroupCoordinator`` cannot patch in additional backends or change |
| 1715 | the eager-init behavior — ``split_group`` only selects subsets of an |
| 1716 | existing parent. Validate that the parent has both ``device_id`` and a |
| 1717 | CPU (gloo) backend, and emit a descriptive error pointing at the exact |
| 1718 | init call to update otherwise. |
| 1719 | """ |
| 1720 | default_pg = torch.distributed.distributed_c10d._get_default_group() |
| 1721 | assert default_pg.bound_device_id is not None, ( |
| 1722 | "External launcher initialized the default process group " |
| 1723 | "without device_id. vLLM requires the default PG to be device-" |
| 1724 | "bound for split_group. Pass device_id=torch.device(f'cuda:" |
| 1725 | "{local_rank}') to torch.distributed.init_process_group()." |
| 1726 | ) |
| 1727 | try: |
| 1728 | default_pg._get_backend(torch.device("cpu")) |
| 1729 | except RuntimeError as e: |
| 1730 | raise RuntimeError( |
| 1731 | "External launcher initialized the default process group " |
| 1732 | "without a CPU (gloo) backend. vLLM requires both CPU and " |
| 1733 | "device backends. Pass backend='cpu:gloo,cuda:nccl' to " |
| 1734 | "torch.distributed.init_process_group()." |
| 1735 | ) from e |
| 1736 | |
| 1737 | |
| 1738 | def _init_elastic_ep_world( |
| 1739 | config, local_rank: int, backend: str, rank: int, world_size: int |
| 1740 | ) -> None: |
| 1741 | from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator |
| 1742 | |
| 1743 | global _WORLD, _NODE_COUNT |
| 1744 | assert _WORLD is None, "world group already initialized" |
| 1745 | parallel_config = config.parallel_config |
| 1746 | global_rank = parallel_config.data_parallel_rank * world_size + rank |
| 1747 | global_world_size = parallel_config.world_size_across_dp |
| 1748 | all_ranks = list(range(global_world_size)) |
| 1749 | group_ranks = [all_ranks[i : i + 1] for i in range(global_world_size)] |
| 1750 | if global_rank in all_ranks: |
| 1751 | group_ranks = [all_ranks] |
| 1752 | coord_store = get_cached_tcp_store_client( |
| 1753 | parallel_config.data_parallel_master_ip, parallel_config._coord_store_port |
| 1754 | ) |
| 1755 | world = StatelessGroupCoordinator( |
| 1756 | group_ranks=group_ranks, |
| 1757 | local_rank=local_rank, |
| 1758 | torch_distributed_backend=backend, |
| 1759 | use_device_communicator=False, |
| 1760 | group_name="world", |
| 1761 | host=parallel_config.data_parallel_master_ip, |
| 1762 | coord_store=coord_store, |
| 1763 | global_rank=global_rank, |
| 1764 | global_world_size=global_world_size, |
| 1765 | ) |
| 1766 | assert parallel_config.nnodes_within_dp == 1, ( |
| 1767 | "Elastic EP is not supported with multi-node TP/PP" |
| 1768 | ) |
| 1769 | _NODE_COUNT = _node_count(world.tcp_store_group) |
| 1770 | _WORLD = world |
| 1771 | |
| 1772 | |
| 1773 | def init_distributed_environment( |
| 1774 | world_size: int = -1, |
| 1775 | rank: int = -1, |
| 1776 | distributed_init_method: str = "env://", |
| 1777 | local_rank: int = -1, |
| 1778 | backend: str = "nccl", |
| 1779 | timeout: timedelta | None = None, |
| 1780 | ): |
| 1781 | logger.debug( |
| 1782 | "world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s", |
| 1783 | world_size, |
| 1784 | rank, |
| 1785 | local_rank, |
| 1786 | distributed_init_method, |
| 1787 | backend, |
| 1788 | ) |
| 1789 | from vllm.config import get_current_vllm_config_or_none |
| 1790 | |
| 1791 | config = get_current_vllm_config_or_none() |
| 1792 | enable_elastic_ep = config is not None and config.parallel_config.enable_elastic_ep |
| 1793 | if ( |
| 1794 | config is not None |
| 1795 | and config.parallel_config.distributed_executor_backend != "external_launcher" |
| 1796 | and ( |
| 1797 | config.parallel_config.nnodes > 1 |
| 1798 | or config.parallel_config.data_parallel_size > 1 |
| 1799 | ) |
| 1800 | and not enable_elastic_ep |
| 1801 | ): |
| 1802 | parallel_config = config.parallel_config |
| 1803 | # adjust to take into account data parallelism |
| 1804 | # offset the rank by the data parallel rank |
| 1805 | rank = parallel_config.data_parallel_rank * world_size + rank |
| 1806 | # adjust the world size to take into account data parallelism |
| 1807 | world_size = parallel_config.world_size_across_dp |
| 1808 | |
| 1809 | # Use appropriate IP and port based on configuration |
| 1810 | if parallel_config.nnodes > 1: |
| 1811 | ip = parallel_config.master_addr |
| 1812 | port = parallel_config.master_port |
| 1813 | distributed_init_method = get_distributed_init_method(ip, port) |
| 1814 | else: |
| 1815 | ip = parallel_config.data_parallel_master_ip |
| 1816 | port = parallel_config.get_next_dp_init_port() |
| 1817 | distributed_init_method = get_distributed_init_method(ip, port) |
| 1818 | logger.debug( |
| 1819 | "Adjusting world_size=%d rank=%d distributed_init_method=%s for DP", |
| 1820 | world_size, |
| 1821 | rank, |
| 1822 | distributed_init_method, |
| 1823 | ) |
| 1824 | if not torch.distributed.is_initialized(): |
| 1825 | logger.info( |
| 1826 | "world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s", |
| 1827 | world_size, |
| 1828 | rank, |
| 1829 | local_rank, |
| 1830 | distributed_init_method, |
| 1831 | backend, |
| 1832 | ) |
| 1833 | assert distributed_init_method is not None, ( |
| 1834 | "distributed_init_method must be provided when initializing " |
| 1835 | "distributed environment" |
| 1836 | ) |
| 1837 | if not torch.distributed.is_backend_available(backend): |
| 1838 | logger.warning( |
| 1839 | "Distributed backend %s is not available; falling back to gloo.", |
| 1840 | backend, |
| 1841 | ) |
| 1842 | assert torch.distributed.is_gloo_available(), ( |
| 1843 | "Fallback Gloo backend is not available." |
| 1844 | ) |
| 1845 | backend = "gloo" |
| 1846 | store = None |
| 1847 | if distributed_init_method.startswith("file://"): |
| 1848 | store = torch.distributed.FileStore( |
| 1849 | distributed_init_method.removeprefix("file://"), world_size |
| 1850 | ) |
| 1851 | if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP: |
| 1852 | # split_group needs local_rank early to compute device_id for |
| 1853 | # the eager init. local_rank is not available in torch |
| 1854 | # ProcessGroup, see https://github.com/pytorch/pytorch/issues/122816 |
| 1855 | if local_rank == -1: |
| 1856 | local_rank = ( |
| 1857 | int(envs.LOCAL_RANK) |
| 1858 | if distributed_init_method == "env://" |
| 1859 | else rank |
| 1860 | ) |
| 1861 | _init_process_group_for_split_group( |
| 1862 | backend=backend, |
| 1863 | distributed_init_method=distributed_init_method, |
| 1864 | store=store, |
| 1865 | world_size=world_size, |
| 1866 | rank=rank, |
| 1867 | local_rank=local_rank, |
| 1868 | timeout=timeout, |
| 1869 | ) |
| 1870 | else: |
| 1871 | # this backend is used for WORLD |
| 1872 | torch.distributed.init_process_group( |
| 1873 | backend=backend, |
| 1874 | init_method=distributed_init_method if store is None else None, |
| 1875 | store=store, |
| 1876 | world_size=world_size, |
| 1877 | rank=rank, |
| 1878 | timeout=timeout, |
| 1879 | ) |
| 1880 | if enable_elastic_ep: |
| 1881 | tp_pp_cpu_group = torch.distributed.new_group( |
| 1882 | backend="gloo", timeout=timeout |
| 1883 | ) |
| 1884 | if _node_count(tp_pp_cpu_group) > 1: |
| 1885 | # NOTE(yongji): StatelessGroupCoordinator uses data_parallel_master_ip |
| 1886 | # to initialize all DP/EP groups, hence all ranks within TP/PP group |
| 1887 | # must reside on the same node |
| 1888 | raise RuntimeError( |
| 1889 | "Elastic EP is not yet supported with multi-node TP/PP" |
| 1890 | ) |
| 1891 | |
| 1892 | if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP and torch.accelerator.is_available(): |
| 1893 | _validate_default_pg_for_split_group() |
| 1894 | |
| 1895 | # set the local rank |
| 1896 | # local_rank is not available in torch ProcessGroup, |
| 1897 | # see https://github.com/pytorch/pytorch/issues/122816 |
| 1898 | if local_rank == -1: |
| 1899 | # local rank not set, this usually happens in single-node |
| 1900 | # setting, where we can use rank as local rank |
| 1901 | local_rank = envs.LOCAL_RANK if distributed_init_method == "env://" else rank |
| 1902 | |
| 1903 | global _WORLD, _NODE_COUNT, _INNER_DP_WORLD |
| 1904 | if enable_elastic_ep: |
| 1905 | _init_elastic_ep_world(config, local_rank, backend, rank, world_size) |
| 1906 | return |
| 1907 | if _WORLD is None: |
| 1908 | ranks = list(range(torch.distributed.get_world_size())) |
| 1909 | _WORLD = init_world_group(ranks, local_rank, backend) |
| 1910 | if config is not None and config.parallel_config.nnodes > 1: |
| 1911 | _NODE_COUNT = config.parallel_config.nnodes |
| 1912 | else: |
| 1913 | _NODE_COUNT = _node_count(_WORLD.cpu_group) |
| 1914 | logger.debug("Detected %d nodes in the distributed environment", _NODE_COUNT) |
| 1915 | else: |
| 1916 | assert _WORLD.world_size == torch.distributed.get_world_size(), ( |
| 1917 | "world group already initialized with a different world size" |
| 1918 | ) |
| 1919 | if config is not None and config.parallel_config.nnodes_within_dp > 1: |
| 1920 | if parallel_config.data_parallel_size > 1: |
| 1921 | world_size_inner_dp = parallel_config.world_size |
| 1922 | group_ranks = [ |
| 1923 | [dp_rank * world_size_inner_dp + i for i in range(world_size_inner_dp)] |
| 1924 | for dp_rank in range(parallel_config.data_parallel_size) |
| 1925 | ] |
| 1926 | _INNER_DP_WORLD = init_model_parallel_group( |
| 1927 | group_ranks, |
| 1928 | get_world_group().local_rank, |
| 1929 | backend, |
| 1930 | use_message_queue_broadcaster=True, |
| 1931 | group_name="inner_dp_world", |
| 1932 | use_device_communicator=False, |
| 1933 | ) |
| 1934 | else: |
| 1935 | _INNER_DP_WORLD = _WORLD |
| 1936 | |
| 1937 | |
| 1938 | def _elastic_join_warmup_ctx() -> AbstractContextManager[Any]: |
| 1939 | """Defer communicator warm-up when joining a live elastic-EP cluster on ROCm. |
| 1940 | |
| 1941 | The existing ranks defer their matching warm-up in create_standby_groups(), |
| 1942 | so a joining worker must too or the warm-up collective has no peers. Both |
| 1943 | sides warm up at commit. |
| 1944 | """ |
| 1945 | from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm |
| 1946 | |
| 1947 | if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: |
| 1948 | return defer_comm_warmup_on_rocm() |
| 1949 | return nullcontext() |
| 1950 | |
| 1951 | |
| 1952 | def _engram_dp_shard_size( |
| 1953 | world_size: int, |
| 1954 | data_parallel_size: int, |
| 1955 | replica_size: int, |
| 1956 | ) -> int: |
| 1957 | """Size of DP shards aligned with complete replicas and node boundaries.""" |
| 1958 | node_count = get_node_count() |
| 1959 | replicas_per_node, remainder = divmod(world_size, node_count * replica_size) |
| 1960 | if remainder: |
| 1961 | return 1 |
| 1962 | # Shards must divide both the DP group and each node's replica count. |
| 1963 | shard_size = gcd(data_parallel_size, replicas_per_node) |
| 1964 | if shard_size == 1 or node_count == 1: |
| 1965 | return shard_size |
| 1966 | ranks_per_node = replicas_per_node * replica_size |
| 1967 | for start in range(0, world_size, ranks_per_node): |
| 1968 | same_node = in_the_same_node_as(get_world_group().cpu_group, start) |
| 1969 | node_ranks = [rank for rank, local in enumerate(same_node) if local] |
| 1970 | if node_ranks != list(range(start, start + ranks_per_node)): |
| 1971 | return 1 |
| 1972 | return shard_size |
| 1973 | |
| 1974 | |
| 1975 | def initialize_model_parallel( |
| 1976 | tensor_model_parallel_size: int = 1, |
| 1977 | pipeline_model_parallel_size: int = 1, |
| 1978 | prefill_context_model_parallel_size: int = 1, |
| 1979 | decode_context_model_parallel_size: int | None = 1, |
| 1980 | backend: str | None = None, |
| 1981 | ) -> None: |
| 1982 | """Initialize model parallel groups. |
| 1983 | |
| 1984 | Arguments: |
| 1985 | tensor_model_parallel_size: number of GPUs used for tensor model |
| 1986 | parallelism. |
| 1987 | pipeline_model_parallel_size: number of GPUs used for pipeline model |
| 1988 | parallelism. |
| 1989 | prefill_context_model_parallel_size: number of GPUs used for context |
| 1990 | parallelism during prefill. |
| 1991 | decode_context_model_parallel_size: number of GPUs used for context |
| 1992 | parallelism during decode. |
| 1993 | backend: name of torch distributed communication backend. |
| 1994 | |
| 1995 | Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we |
| 1996 | use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize |
| 1997 | the model pipeline. The present function will |
| 1998 | create 4 tensor model-parallel groups and 2 pipeline model-parallel groups: |
| 1999 | 4 tensor model-parallel groups: |
| 2000 | [g0, g1], [g2, g3], [g4, g5], [g6, g7] |
| 2001 | 2 pipeline model-parallel groups: |
| 2002 | [g0, g2, g4, g6], [g1, g3, g5, g7] |
| 2003 | Note that for efficiency, the caller should make sure adjacent ranks |
| 2004 | are on the same DGX box. For example if we are using 2 DGX-1 boxes |
| 2005 | with a total of 16 GPUs, rank 0 to 7 belong to the first box and |
| 2006 | ranks 8 to 15 belong to the second box. |
| 2007 | |
| 2008 | """ |
| 2009 | # Get world size and rank. Ensure some consistencies. |
| 2010 | assert torch.distributed.is_initialized() |
| 2011 | |
| 2012 | from vllm.config import get_current_vllm_config |
| 2013 | |
| 2014 | config = get_current_vllm_config() |
| 2015 | data_parallel_size = config.parallel_config.data_parallel_size |
| 2016 | enable_elastic_ep = config.parallel_config.enable_elastic_ep |
| 2017 | parallel_config = config.parallel_config |
| 2018 | coord_store: Store | None = None |
| 2019 | if enable_elastic_ep: |
| 2020 | coord_store = get_cached_tcp_store_client( |
| 2021 | parallel_config.data_parallel_master_ip, |
| 2022 | parallel_config._coord_store_port, |
| 2023 | ) |
| 2024 | # Use stateless world group for global information |
| 2025 | world_size = get_world_group().world_size |
| 2026 | rank = get_world_group().rank |
| 2027 | backend = backend or "nccl" |
| 2028 | tp_pp_pcp_size = ( |
| 2029 | tensor_model_parallel_size |
| 2030 | * pipeline_model_parallel_size |
| 2031 | * prefill_context_model_parallel_size |
| 2032 | ) |
| 2033 | local_all_ranks = torch.arange(tp_pp_pcp_size).reshape( |
| 2034 | pipeline_model_parallel_size, |
| 2035 | prefill_context_model_parallel_size, |
| 2036 | tensor_model_parallel_size, |
| 2037 | ) |
| 2038 | else: |
| 2039 | world_size = torch.distributed.get_world_size() |
| 2040 | rank = torch.distributed.get_rank() |
| 2041 | backend = backend or torch.distributed.get_backend( |
| 2042 | get_world_group().device_group |
| 2043 | ) |
| 2044 | |
| 2045 | # the layout order is: ExternalDP x DP x PP x PCP x TP |
| 2046 | # ExternalDP is the data parallel group that is not part of the model, |
| 2047 | # every dp rank can generate independently (in verl integration). |
| 2048 | # DP is the data parallel group that is part of the model, |
| 2049 | # all the ranks in the same DP group should generate simultaneously, |
| 2050 | # i.e. the `generate` call in the same DP group should be called together, |
| 2051 | # otherwise it will cause deadlock. |
| 2052 | # to get group_ranks for each dimension, transpose that dimension to the |
| 2053 | # last dimension, then reshape to 2D, then unbind the last dimension |
| 2054 | all_ranks = torch.arange(world_size).reshape( |
| 2055 | -1, |
| 2056 | data_parallel_size, |
| 2057 | pipeline_model_parallel_size, |
| 2058 | prefill_context_model_parallel_size, |
| 2059 | tensor_model_parallel_size, |
| 2060 | ) # noqa |
| 2061 | |
| 2062 | # Build the tensor model-parallel groups. |
| 2063 | global _TP |
| 2064 | assert _TP is None, "tensor model parallel group is already initialized" |
| 2065 | group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0) |
| 2066 | group_ranks = [x.tolist() for x in group_ranks] |
| 2067 | if enable_elastic_ep: |
| 2068 | group_ranks = local_all_ranks.view(-1, tensor_model_parallel_size).unbind(0) |
| 2069 | group_ranks = [x.tolist() for x in group_ranks] |
| 2070 | # message queue broadcaster is only used in tensor model parallel group |
| 2071 | _TP = init_model_parallel_group( |
| 2072 | group_ranks, |
| 2073 | get_world_group().local_rank, |
| 2074 | backend, |
| 2075 | use_message_queue_broadcaster=True, |
| 2076 | group_name="tp", |
| 2077 | ) |
| 2078 | |
| 2079 | global _ETP |
| 2080 | assert _ETP is None, "Engram tensor-parallel group is already initialized" |
| 2081 | engram_tensor_parallel_size = ( |
| 2082 | config.engram_config.get_parallel_size(parallel_config) |
| 2083 | if config.engram_config is not None |
| 2084 | else tensor_model_parallel_size |
| 2085 | ) |
| 2086 | if engram_tensor_parallel_size == tensor_model_parallel_size: |
| 2087 | _ETP = _TP |
| 2088 | else: |
| 2089 | group_ranks = ( |
| 2090 | all_ranks.permute(0, 2, 3, 1, 4) |
| 2091 | .reshape(-1, engram_tensor_parallel_size) |
| 2092 | .unbind(0) |
| 2093 | ) |
| 2094 | _ETP = init_model_parallel_group( |
| 2095 | [ranks.tolist() for ranks in group_ranks], |
| 2096 | get_world_group().local_rank, |
| 2097 | backend, |
| 2098 | group_name="etp", |
| 2099 | ) |
| 2100 | |
| 2101 | # Build the node-local DP group used to shard or share an Engram table. |
| 2102 | global _ENGRAM_DP |
| 2103 | assert _ENGRAM_DP is None, "engram data parallel group is already initialized" |
| 2104 | engram_dp_size = 1 |
| 2105 | engram_config = config.engram_config |
| 2106 | if ( |
| 2107 | engram_config is not None |
| 2108 | and config.model_config is not None |
| 2109 | and config.model_config.architecture == "DeepseekV41ForCausalLM" |
| 2110 | and not enable_elastic_ep |
| 2111 | ): |
| 2112 | engram_dp_size = _engram_dp_shard_size( |
| 2113 | world_size, |
| 2114 | data_parallel_size, |
| 2115 | tensor_model_parallel_size |
| 2116 | * pipeline_model_parallel_size |
| 2117 | * prefill_context_model_parallel_size, |
| 2118 | ) |
| 2119 | if engram_dp_size > 1: |
| 2120 | group_ranks = ( |
| 2121 | all_ranks.permute(0, 2, 3, 4, 1).reshape(-1, engram_dp_size).unbind(0) |
| 2122 | ) |
| 2123 | _ENGRAM_DP = init_model_parallel_group( |
| 2124 | [x.tolist() for x in group_ranks], |
| 2125 | get_world_group().local_rank, |
| 2126 | backend, |
| 2127 | group_name="edp", |
| 2128 | ) |
| 2129 | logger.info_once( |
| 2130 | "Engram node-local DP group size: %d (TP=%d)", |
| 2131 | engram_dp_size, |
| 2132 | tensor_model_parallel_size, |
| 2133 | ) |
| 2134 | |
| 2135 | # Build the DCP model-parallel groups. |
| 2136 | global _DCP |
| 2137 | assert _DCP is None, "decode context model parallel group is already initialized" |
| 2138 | dcp_size = decode_context_model_parallel_size or 1 |
| 2139 | dcp_ranks = local_all_ranks if enable_elastic_ep else all_ranks |
| 2140 | if dcp_size > 1: |
| 2141 | # DCP spans PCP first, then TP for full TP x PCP groups. |
| 2142 | dcp_ranks = dcp_ranks.transpose(-1, -2) |
| 2143 | group_ranks = dcp_ranks.reshape(-1, dcp_size).unbind(0) |
| 2144 | group_ranks = [x.tolist() for x in group_ranks] |
| 2145 | _DCP = init_model_parallel_group( |
| 2146 | group_ranks, |
| 2147 | get_world_group().local_rank, |
| 2148 | backend, |
| 2149 | use_message_queue_broadcaster=True, |
| 2150 | group_name="dcp", |
| 2151 | ) |
| 2152 | |
| 2153 | global _PCP |
| 2154 | assert _PCP is None, "prefill context parallel group is already initialized" |
| 2155 | group_ranks = ( |
| 2156 | all_ranks.transpose(3, 4) |
| 2157 | .reshape(-1, prefill_context_model_parallel_size) |
| 2158 | .unbind(0) |
| 2159 | ) |
| 2160 | group_ranks = [x.tolist() for x in group_ranks] |
| 2161 | if enable_elastic_ep: |
| 2162 | group_ranks = ( |
| 2163 | local_all_ranks.transpose(1, 2) |
| 2164 | .reshape(-1, prefill_context_model_parallel_size) |
| 2165 | .unbind(0) |
| 2166 | ) |
| 2167 | group_ranks = [x.tolist() for x in group_ranks] |
| 2168 | _PCP = init_model_parallel_group( |
| 2169 | group_ranks, get_world_group().local_rank, backend, group_name="pcp" |
| 2170 | ) |
| 2171 | |
| 2172 | # Build the pipeline model-parallel groups. |
| 2173 | global _PP |
| 2174 | assert _PP is None, "pipeline model parallel group is already initialized" |
| 2175 | group_ranks = ( |
| 2176 | all_ranks.transpose(2, 4).reshape(-1, pipeline_model_parallel_size).unbind(0) |
| 2177 | ) |
| 2178 | group_ranks = [x.tolist() for x in group_ranks] |
| 2179 | if enable_elastic_ep: |
| 2180 | group_ranks = ( |
| 2181 | local_all_ranks.transpose(0, 2) |
| 2182 | .reshape(-1, pipeline_model_parallel_size) |
| 2183 | .unbind(0) |
| 2184 | ) |
| 2185 | group_ranks = [x.tolist() for x in group_ranks] |
| 2186 | _PP = init_model_parallel_group( |
| 2187 | group_ranks, get_world_group().local_rank, backend, group_name="pp" |
| 2188 | ) |
| 2189 | |
| 2190 | global _DP |
| 2191 | assert _DP is None, "data parallel group is already initialized" |
| 2192 | group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0) |
| 2193 | group_ranks = [x.tolist() for x in group_ranks] |
| 2194 | if enable_elastic_ep: |
| 2195 | with _elastic_join_warmup_ctx(): |
| 2196 | _DP = _init_stateless_group( |
| 2197 | group_ranks, |
| 2198 | "dp", |
| 2199 | parallel_config.data_parallel_master_ip, |
| 2200 | backend, |
| 2201 | coord_store=coord_store, |
| 2202 | ) |
| 2203 | else: |
| 2204 | _DP = init_model_parallel_group( |
| 2205 | group_ranks, get_world_group().local_rank, backend, group_name="dp" |
| 2206 | ) |
| 2207 | |
| 2208 | global _EP |
| 2209 | assert _EP is None, "expert parallel group is already initialized" |
| 2210 | # Don't create EP group for dense models. |
| 2211 | if config.model_config is None or config.model_config.is_moe: |
| 2212 | group_ranks = ( |
| 2213 | all_ranks.transpose(1, 2) |
| 2214 | .reshape( |
| 2215 | -1, |
| 2216 | data_parallel_size |
| 2217 | * prefill_context_model_parallel_size |
| 2218 | * tensor_model_parallel_size, |
| 2219 | ) |
| 2220 | .unbind(0) |
| 2221 | ) |
| 2222 | group_ranks = [x.tolist() for x in group_ranks] |
| 2223 | use_all2all = parallel_config.use_all2all |
| 2224 | if enable_elastic_ep: |
| 2225 | with _elastic_join_warmup_ctx(): |
| 2226 | _EP = _init_stateless_group( |
| 2227 | group_ranks, |
| 2228 | "ep", |
| 2229 | parallel_config.data_parallel_master_ip, |
| 2230 | backend, |
| 2231 | coord_store=coord_store, |
| 2232 | use_all2all=use_all2all, |
| 2233 | ) |
| 2234 | else: |
| 2235 | _EP = init_model_parallel_group( |
| 2236 | group_ranks, |
| 2237 | get_world_group().local_rank, |
| 2238 | backend, |
| 2239 | group_name="ep", |
| 2240 | use_all2all=use_all2all, |
| 2241 | ) |
| 2242 | |
| 2243 | # Create EPLB group with the same ranks as EP if EPLB is enabled. |
| 2244 | # This is a separate process group to isolate EPLB communications |
| 2245 | # from MoE forward pass collectives and prevent deadlocks when |
| 2246 | # using torch.distributed in execution with torch.distributed in EPLB. |
| 2247 | global _EPLB |
| 2248 | assert _EPLB is None, "EPLB group is already initialized" |
| 2249 | if config.parallel_config.enable_eplb: |
| 2250 | if enable_elastic_ep: |
| 2251 | with _elastic_join_warmup_ctx(): |
| 2252 | _EPLB = _init_stateless_group( |
| 2253 | group_ranks, |
| 2254 | "eplb", |
| 2255 | parallel_config.data_parallel_master_ip, |
| 2256 | backend, |
| 2257 | coord_store=coord_store, |
| 2258 | ) |
| 2259 | else: |
| 2260 | _EPLB = init_model_parallel_group( |
| 2261 | group_ranks, |
| 2262 | get_world_group().local_rank, |
| 2263 | backend, |
| 2264 | group_name="eplb", |
| 2265 | ) |
| 2266 | # If no EP group needed, _EP remains None |
| 2267 | # If no EPLB group needed, _EPLB remains None |
| 2268 | |
| 2269 | logger.info_once( |
| 2270 | "rank %s in world size %s is assigned as " |
| 2271 | "DP rank %s, PP rank %s, PCP rank %s, " |
| 2272 | "TP rank %s, ETP rank %s, EP rank %s, EPLB rank %s", |
| 2273 | rank, |
| 2274 | world_size, |
| 2275 | _DP.rank_in_group, |
| 2276 | _PP.rank_in_group, |
| 2277 | _PCP.rank_in_group, |
| 2278 | _TP.rank_in_group, |
| 2279 | _ETP.rank_in_group, |
| 2280 | _EP.rank_in_group if _EP is not None else "N/A", |
| 2281 | _EPLB.rank_in_group if _EPLB is not None else "N/A", |
| 2282 | ) |
| 2283 | |
| 2284 | |
| 2285 | def ensure_model_parallel_initialized( |
| 2286 | tensor_model_parallel_size: int, |
| 2287 | pipeline_model_parallel_size: int, |
| 2288 | prefill_context_model_parallel_size: int = 1, |
| 2289 | decode_context_model_parallel_size: int | None = 1, |
| 2290 | backend: str | None = None, |
| 2291 | ) -> None: |
| 2292 | """Helper to initialize model parallel groups if they are not initialized, |
| 2293 | or ensure tensor-parallel and pipeline-parallel sizes are equal to expected |
| 2294 | values if the model parallel groups are initialized. |
| 2295 | """ |
| 2296 | world_group = get_world_group() |
| 2297 | if hasattr(world_group, "backend"): |
| 2298 | backend = backend or world_group.backend |
| 2299 | else: |
| 2300 | backend = backend or torch.distributed.get_backend(world_group.device_group) |
| 2301 | if not model_parallel_is_initialized(): |
| 2302 | initialize_model_parallel( |
| 2303 | tensor_model_parallel_size, |
| 2304 | pipeline_model_parallel_size, |
| 2305 | prefill_context_model_parallel_size, |
| 2306 | decode_context_model_parallel_size, |
| 2307 | backend, |
| 2308 | ) |
| 2309 | return |
| 2310 | |
| 2311 | assert get_tensor_model_parallel_world_size() == tensor_model_parallel_size, ( |
| 2312 | "tensor parallel group already initialized, but of unexpected size. " |
| 2313 | f"got: {get_tensor_model_parallel_world_size()=} vs. " |
| 2314 | f"wanted: {tensor_model_parallel_size=}" |
| 2315 | ) |
| 2316 | pp_world_size = get_pp_group().world_size |
| 2317 | assert pp_world_size == pipeline_model_parallel_size, ( |
| 2318 | "pipeline parallel group already initialized, but of unexpected size. " |
| 2319 | f"got: {pp_world_size=} vs. " |
| 2320 | f"wanted: {pipeline_model_parallel_size=}" |
| 2321 | ) |
| 2322 | pcp_world_size = get_pcp_group().world_size |
| 2323 | assert pcp_world_size == prefill_context_model_parallel_size, ( |
| 2324 | "prefill context parallel group already initialized, but of unexpected size: " |
| 2325 | f"{pcp_world_size=} vs. " |
| 2326 | f"{prefill_context_model_parallel_size=}" |
| 2327 | ) |
| 2328 | dcp_world_size = get_dcp_group().world_size |
| 2329 | dcp_model_parallel_size = decode_context_model_parallel_size or 1 |
| 2330 | assert dcp_world_size == dcp_model_parallel_size, ( |
| 2331 | "decode context parallel group already initialized, but of unexpected size: " |
| 2332 | f"{dcp_world_size=} vs. " |
| 2333 | f"{dcp_model_parallel_size=}" |
| 2334 | ) |
| 2335 | |
| 2336 | |
| 2337 | def checkpoint_prepare_distributed_state() -> None: |
| 2338 | """Prepare every device communicator for a process checkpoint.""" |
| 2339 | torch.accelerator.synchronize() |
| 2340 | _apply_to_device_comms(lambda comm: comm.checkpoint_prepare()) |
| 2341 | torch.accelerator.synchronize() |
| 2342 | |
| 2343 | |
| 2344 | def checkpoint_restore_distributed_state() -> None: |
| 2345 | """Restore every device communicator after a process checkpoint.""" |
| 2346 | torch.accelerator.synchronize() |
| 2347 | _apply_to_device_comms(lambda comm: comm.checkpoint_restore()) |
| 2348 | torch.accelerator.synchronize() |
| 2349 | |
| 2350 | |
| 2351 | def model_parallel_is_initialized(): |
| 2352 | """Check if tensor and pipeline parallel groups are initialized.""" |
| 2353 | return _TP is not None and _PP is not None |
| 2354 | |
| 2355 | |
| 2356 | def get_tensor_model_parallel_world_size() -> int: |
| 2357 | """Return world size for the tensor model parallel group.""" |
| 2358 | return get_tp_group().world_size |
| 2359 | |
| 2360 | |
| 2361 | def get_tensor_model_parallel_rank() -> int: |
| 2362 | """Return my rank for the tensor model parallel group.""" |
| 2363 | return get_tp_group().rank_in_group |
| 2364 | |
| 2365 | |
| 2366 | def get_node_count() -> int: |
| 2367 | """Return the total number of nodes in the distributed environment.""" |
| 2368 | assert _NODE_COUNT is not None, "distributed environment is not initialized" |
| 2369 | return _NODE_COUNT |
| 2370 | |
| 2371 | |
| 2372 | def destroy_model_parallel(): |
| 2373 | """Set the groups to none and destroy them.""" |
| 2374 | global _TP, _ETP |
| 2375 | |
| 2376 | if _ETP and _ETP is not _TP: |
| 2377 | _ETP.destroy() |
| 2378 | _ETP = None |
| 2379 | |
| 2380 | if _TP: |
| 2381 | _TP.destroy() |
| 2382 | _TP = None |
| 2383 | |
| 2384 | global _ENGRAM_DP |
| 2385 | if _ENGRAM_DP: |
| 2386 | _ENGRAM_DP.destroy() |
| 2387 | _ENGRAM_DP = None |
| 2388 | |
| 2389 | global _DCP |
| 2390 | if _DCP: |
| 2391 | _DCP.destroy() |
| 2392 | _DCP = None |
| 2393 | |
| 2394 | global _PCP |
| 2395 | if _PCP: |
| 2396 | _PCP.destroy() |
| 2397 | _PCP = None |
| 2398 | |
| 2399 | global _PP |
| 2400 | if _PP: |
| 2401 | _PP.destroy() |
| 2402 | _PP = None |
| 2403 | |
| 2404 | global _DP |
| 2405 | if _DP: |
| 2406 | _DP.destroy() |
| 2407 | _DP = None |
| 2408 | |
| 2409 | global _EP |
| 2410 | if _EP: |
| 2411 | _EP.destroy() |
| 2412 | _EP = None |
| 2413 | |
| 2414 | global _EPLB |
| 2415 | if _EPLB: |
| 2416 | _EPLB.destroy() |
| 2417 | _EPLB = None |
| 2418 | |
| 2419 | |
| 2420 | def destroy_distributed_environment(): |
| 2421 | global _WORLD, _NODE_COUNT |
| 2422 | if _WORLD: |
| 2423 | _WORLD.destroy() |
| 2424 | _WORLD = None |
| 2425 | _NODE_COUNT = None |
| 2426 | if torch.distributed.is_initialized(): |
| 2427 | torch.distributed.destroy_process_group() |
| 2428 | |
| 2429 | |
| 2430 | def cleanup_dist_env_and_memory(shutdown_ray: bool = False): |
| 2431 | logger.debug( |
| 2432 | "[shutdown] Distributed: cleanup start shutdown_ray=%s", |
| 2433 | shutdown_ray, |
| 2434 | ) |
| 2435 | # Reset environment variable cache |
| 2436 | envs.disable_envs_cache() |
| 2437 | |
| 2438 | # Reset rocm_aiter_ops class variables to match current os.environ. |
| 2439 | # These are class-level attributes that persist across tests and are |
| 2440 | # NOT restored by monkeypatch (which only restores os.environ). |
| 2441 | from vllm.platforms import current_platform |
| 2442 | |
| 2443 | if current_platform.is_rocm(): |
| 2444 | from vllm._aiter_ops import rocm_aiter_ops |
| 2445 | |
| 2446 | rocm_aiter_ops.refresh_env_variables() |
| 2447 | |
| 2448 | # Ensure all objects are not frozen before cleanup |
| 2449 | gc.unfreeze() |
| 2450 | |
| 2451 | destroy_model_parallel() |
| 2452 | destroy_distributed_environment() |
| 2453 | if shutdown_ray: |
| 2454 | import ray # Lazy import Ray |
| 2455 | |
| 2456 | ray.shutdown() |
| 2457 | gc.collect() |
| 2458 | from vllm.platforms import current_platform |
| 2459 | |
| 2460 | if not current_platform.is_cpu(): |
| 2461 | from vllm.triton_utils import HAS_TRITON |
| 2462 | |
| 2463 | if HAS_TRITON: |
| 2464 | from vllm.v1.sample.ops.topk_topp_triton import reset_buffer_cache |
| 2465 | |
| 2466 | reset_buffer_cache() |
| 2467 | torch.accelerator.empty_cache() |
| 2468 | try: |
| 2469 | torch.accelerator.empty_host_cache() |
| 2470 | except AttributeError: |
| 2471 | logger.warning( |
| 2472 | "torch.accelerator.empty_host_cache() only available in Pytorch >=2.9" |
| 2473 | ) |
| 2474 | |
| 2475 | logger.debug_once("[shutdown] Distributed: cleanup complete") |
| 2476 | |
| 2477 | |
| 2478 | def in_the_same_node_as( |
| 2479 | pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0 |
| 2480 | ) -> list[bool]: |
| 2481 | """This is a collective operation that returns if each rank is in the same node |
| 2482 | as the source rank. It tests if processes are attached to the same |
| 2483 | memory system (shared access to shared memory). |
| 2484 | """ |
| 2485 | if isinstance(pg, ProcessGroup): |
| 2486 | assert torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL, ( |
| 2487 | "in_the_same_node_as should be tested with a non-NCCL group." |
| 2488 | ) |
| 2489 | # local rank inside the group |
| 2490 | rank = torch.distributed.get_rank(group=pg) |
| 2491 | world_size = torch.distributed.get_world_size(group=pg) |
| 2492 | |
| 2493 | # global ranks of the processes in the group |
| 2494 | ranks = torch.distributed.get_process_group_ranks(pg) |
| 2495 | else: |
| 2496 | rank = pg.rank |
| 2497 | world_size = pg.world_size |
| 2498 | ranks = list(range(world_size)) |
| 2499 | |
| 2500 | # local tensor in each process to store the result |
| 2501 | is_in_the_same_node = torch.tensor( |
| 2502 | [0] * world_size, dtype=torch.int32, device="cpu" |
| 2503 | ) |
| 2504 | |
| 2505 | magic_message = b"magic_message" |
| 2506 | shm = None |
| 2507 | |
| 2508 | try: |
| 2509 | with contextlib.suppress(OSError): |
| 2510 | if rank == source_rank: |
| 2511 | # create a shared memory segment |
| 2512 | shm = shared_memory.SharedMemory(create=True, size=128) |
| 2513 | assert shm.buf is not None, "Buffer was not created" |
| 2514 | shm.buf[: len(magic_message)] = magic_message |
| 2515 | if isinstance(pg, ProcessGroup): |
| 2516 | torch.distributed.broadcast_object_list( |
| 2517 | [shm.name], src=ranks[source_rank], group=pg |
| 2518 | ) |
| 2519 | else: |
| 2520 | pg.broadcast_obj(shm.name, src=source_rank) |
| 2521 | is_in_the_same_node[rank] = 1 |
| 2522 | else: |
| 2523 | # try to open the shared memory segment |
| 2524 | if isinstance(pg, ProcessGroup): |
| 2525 | recv = [None] |
| 2526 | torch.distributed.broadcast_object_list( |
| 2527 | recv, src=ranks[source_rank], group=pg |
| 2528 | ) |
| 2529 | name = recv[0] |
| 2530 | else: |
| 2531 | name = pg.broadcast_obj(None, src=source_rank) |
| 2532 | # fix to https://stackoverflow.com/q/62748654/9191338 |
| 2533 | # Python incorrectly tracks shared memory even if it is not |
| 2534 | # created by the process. The following patch is a workaround. |
| 2535 | with patch( |
| 2536 | "multiprocessing.resource_tracker.register", |
| 2537 | lambda *args, **kwargs: None, |
| 2538 | ): |
| 2539 | shm = shared_memory.SharedMemory(name=name) |
| 2540 | assert shm.buf is not None, "Buffer was not opened" |
| 2541 | if shm.buf[: len(magic_message)] == magic_message: |
| 2542 | is_in_the_same_node[rank] = 1 |
| 2543 | except Exception as e: |
| 2544 | logger.error("Error ignored in is_in_the_same_node: %s", e) |
| 2545 | finally: |
| 2546 | if shm: |
| 2547 | shm.close() |
| 2548 | |
| 2549 | if isinstance(pg, ProcessGroup): |
| 2550 | torch.distributed.barrier(group=pg) |
| 2551 | else: |
| 2552 | pg.barrier() |
| 2553 | |
| 2554 | # clean up the shared memory segment |
| 2555 | with contextlib.suppress(OSError): |
| 2556 | if rank == source_rank and shm: |
| 2557 | shm.unlink() |
| 2558 | |
| 2559 | if isinstance(pg, ProcessGroup): |
| 2560 | torch.distributed.all_reduce(is_in_the_same_node, group=pg) |
| 2561 | aggregated_data = is_in_the_same_node |
| 2562 | else: |
| 2563 | aggregated_data = torch.zeros_like(is_in_the_same_node) |
| 2564 | for i in range(world_size): |
| 2565 | rank_data = pg.broadcast_obj(is_in_the_same_node, src=i) |
| 2566 | aggregated_data += rank_data |
| 2567 | |
| 2568 | return [x == 1 for x in aggregated_data.tolist()] |
| 2569 | |
| 2570 | |
| 2571 | def is_global_first_rank() -> bool: |
| 2572 | """Check if the current process is the first rank globally across all |
| 2573 | parallelism strategies (PP, TP, DP, EP, etc.). |
| 2574 | |
| 2575 | Unlike group-specific checks like `get_tensor_model_parallel_rank() == 0` |
| 2576 | or `get_pp_group().is_first_rank`, this function checks the global rank |
| 2577 | across all parallelism dimensions. |
| 2578 | |
| 2579 | Returns: |
| 2580 | bool: True if this is the global first rank (rank 0), False otherwise. |
| 2581 | Returns True if distributed is not initialized (single process). |
| 2582 | |
| 2583 | """ |
| 2584 | try: |
| 2585 | # If world group is available, use it for the most accurate check |
| 2586 | global _WORLD |
| 2587 | if _WORLD is not None: |
| 2588 | return _WORLD.is_first_rank |
| 2589 | |
| 2590 | # If torch distributed is not initialized, assume single process |
| 2591 | if not torch.distributed.is_initialized(): |
| 2592 | return True |
| 2593 | |
| 2594 | # Fallback to torch's global rank |
| 2595 | return torch.distributed.get_rank() == 0 |
| 2596 | |
| 2597 | except Exception: |
| 2598 | # If anything goes wrong, assume this is the first rank |
| 2599 | return True |
| 2600 | |
| 2601 | |
| 2602 | def is_local_first_rank() -> bool: |
| 2603 | """Check if the current process is the first local rank (rank 0 on its node).""" |
| 2604 | try: |
| 2605 | # prefer the initialized world group if available |
| 2606 | global _WORLD |
| 2607 | if _WORLD is not None: |
| 2608 | return _WORLD.local_rank == 0 |
| 2609 | |
| 2610 | if not torch.distributed.is_initialized(): |
| 2611 | return True |
| 2612 | |
| 2613 | # fallback to environment-provided local rank if available |
| 2614 | # note: envs.LOCAL_RANK is set when using env:// launchers (e.g., torchrun) |
| 2615 | try: |
| 2616 | return int(envs.LOCAL_RANK) == 0 # type: ignore[arg-type] |
| 2617 | except Exception: |
| 2618 | return torch.distributed.get_rank() == 0 |
| 2619 | except Exception: |
| 2620 | return True |
| 2621 | |
| 2622 | |
| 2623 | def _node_count(pg: ProcessGroup | StatelessProcessGroup) -> int: |
| 2624 | """Returns the total number of nodes in the process group. |
| 2625 | |
| 2626 | Args: |
| 2627 | pg: The process group to analyze |
| 2628 | |
| 2629 | Returns: |
| 2630 | int: The total number of nodes |
| 2631 | |
| 2632 | """ |
| 2633 | if isinstance(pg, ProcessGroup): |
| 2634 | world_size = torch.distributed.get_world_size(group=pg) |
| 2635 | else: |
| 2636 | world_size = pg.world_size |
| 2637 | |
| 2638 | if world_size == 1: |
| 2639 | return 1 |
| 2640 | |
| 2641 | # Build node assignment map |
| 2642 | node_assignment = [0] * world_size # rank -> node_id |
| 2643 | next_node_id = 0 |
| 2644 | |
| 2645 | for current_rank in range(world_size): |
| 2646 | if node_assignment[current_rank] != 0: |
| 2647 | continue # Already assigned to a node |
| 2648 | |
| 2649 | # Assign current rank to a new node |
| 2650 | next_node_id += 1 |
| 2651 | node_assignment[current_rank] = next_node_id |
| 2652 | |
| 2653 | # Find all ranks on the same node as current_rank |
| 2654 | same_node_flags = in_the_same_node_as(pg, current_rank) |
| 2655 | for other_rank, is_same_node in enumerate(same_node_flags): |
| 2656 | if is_same_node and node_assignment[other_rank] == 0: |
| 2657 | node_assignment[other_rank] = next_node_id |
| 2658 | |
| 2659 | return next_node_id |