| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import socket |
| 4 | import struct |
| 5 | from typing import Any, Optional |
| 6 | |
| 7 | import torch |
| 8 | from torch.distributed import Backend, ProcessGroup, Store |
| 9 | |
| 10 | from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator |
| 11 | from vllm.distributed.parallel_state import ( |
| 12 | GroupCoordinator, |
| 13 | TensorMetadata, |
| 14 | _get_unique_name, |
| 15 | _register_group, |
| 16 | _split_tensor_dict, |
| 17 | ) |
| 18 | from vllm.distributed.utils import ( |
| 19 | StatelessProcessGroup, |
| 20 | get_cpu_distributed_timeout_or_none, |
| 21 | get_distributed_timeout_or_none, |
| 22 | stateless_destroy_torch_distributed_process_group, |
| 23 | stateless_init_torch_distributed_process_group, |
| 24 | ) |
| 25 | from vllm.logger import init_logger |
| 26 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | _PORTS_FMT = "!3I" |
| 31 | |
| 32 | |
| 33 | def _allocate_group_ports( |
| 34 | key: str, |
| 35 | host: str, |
| 36 | coord_store: Store, |
| 37 | ) -> tuple[list[int], list[socket.socket]]: |
| 38 | """Bind 3 sockets and publish the ports to *coord_store*. |
| 39 | |
| 40 | Called by rank 0 only. Returns ``(ports, sockets)`` with the |
| 41 | sockets still open. |
| 42 | """ |
| 43 | socks: list[socket.socket] = [] |
| 44 | ports: list[int] = [] |
| 45 | for _ in range(3): |
| 46 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| 47 | s.bind((host, 0)) |
| 48 | s.listen() |
| 49 | socks.append(s) |
| 50 | ports.append(s.getsockname()[1]) |
| 51 | coord_store.set(key, struct.pack(_PORTS_FMT, *ports)) |
| 52 | return ports, socks |
| 53 | |
| 54 | |
| 55 | def _fetch_group_ports(key: str, coord_store: Store) -> list[int]: |
| 56 | """Read 3 ports published by rank 0 from *coord_store*. |
| 57 | |
| 58 | Blocks until the key is available. |
| 59 | """ |
| 60 | return list(struct.unpack(_PORTS_FMT, coord_store.get(key))) |
| 61 | |
| 62 | |
| 63 | class StatelessGroupCoordinator(GroupCoordinator): |
| 64 | """A stateless version of the GroupCoordinator class in parallel_state, |
| 65 | It will create CPU, device and TCPStore based communication groups |
| 66 | that are independent of PyTorch's WORLD group. Hence, |
| 67 | communication groups with a different set of participants GPUs |
| 68 | can be created without destroying the existing ones. |
| 69 | """ |
| 70 | |
| 71 | def __init__( |
| 72 | self, |
| 73 | group_ranks: list[list[int]], |
| 74 | local_rank: int, |
| 75 | torch_distributed_backend: str | Backend, |
| 76 | use_device_communicator: bool, |
| 77 | coord_store: Store, |
| 78 | use_message_queue_broadcaster: bool = False, |
| 79 | group_name: str | None = None, |
| 80 | host: str = "127.0.0.1", |
| 81 | global_rank: int = 0, |
| 82 | global_world_size: int = 1, |
| 83 | use_all2all: bool = False, |
| 84 | ): |
| 85 | group_name = group_name or "anonymous" |
| 86 | self.unique_name = _get_unique_name(group_name) |
| 87 | _register_group(self) |
| 88 | |
| 89 | self.rank = global_rank |
| 90 | self.local_rank = local_rank |
| 91 | from vllm.distributed.parallel_state import _WORLD |
| 92 | |
| 93 | if _WORLD is not None: |
| 94 | self.device_index = _WORLD.device_index |
| 95 | else: |
| 96 | assert local_rank >= 0, ( |
| 97 | "local_rank must be provided when creating the world group" |
| 98 | ) |
| 99 | self.device_index = local_rank |
| 100 | |
| 101 | self_device_group = None |
| 102 | self_cpu_group = None |
| 103 | self_tcp_store_group = None |
| 104 | |
| 105 | from vllm.platforms import current_platform |
| 106 | |
| 107 | backend = str(torch_distributed_backend) |
| 108 | self.backend = backend |
| 109 | cpu_timeout = get_cpu_distributed_timeout_or_none() |
| 110 | device_timeout = ( |
| 111 | cpu_timeout if backend == "gloo" else get_distributed_timeout_or_none() |
| 112 | ) |
| 113 | for idx, ranks in enumerate(group_ranks): |
| 114 | if self.rank in ranks: |
| 115 | self.ranks = ranks |
| 116 | self.world_size = len(ranks) |
| 117 | self.rank_in_group = ranks.index(self.rank) |
| 118 | |
| 119 | key = f"{group_name}_{idx}" |
| 120 | if self.rank_in_group == 0: |
| 121 | ports, socks = _allocate_group_ports( |
| 122 | key, |
| 123 | host, |
| 124 | coord_store, |
| 125 | ) |
| 126 | else: |
| 127 | ports = _fetch_group_ports(key, coord_store) |
| 128 | socks = [] |
| 129 | device_port, cpu_port, tcp_store_port = ports |
| 130 | |
| 131 | device_group = stateless_init_torch_distributed_process_group( |
| 132 | host=host, |
| 133 | port=device_port, |
| 134 | rank=self.rank_in_group, |
| 135 | world_size=self.world_size, |
| 136 | backend=backend, |
| 137 | group_name=f"{self.unique_name}_device", |
| 138 | listen_socket=socks[0] if socks else None, |
| 139 | timeout=device_timeout, |
| 140 | ) |
| 141 | cpu_group = stateless_init_torch_distributed_process_group( |
| 142 | host=host, |
| 143 | port=cpu_port, |
| 144 | rank=self.rank_in_group, |
| 145 | world_size=self.world_size, |
| 146 | backend="gloo", |
| 147 | group_name=f"{self.unique_name}_cpu", |
| 148 | listen_socket=socks[1] if socks else None, |
| 149 | timeout=cpu_timeout, |
| 150 | ) |
| 151 | tcp_store_group = StatelessProcessGroup.create( |
| 152 | host=host, |
| 153 | port=tcp_store_port, |
| 154 | rank=self.rank_in_group, |
| 155 | world_size=self.world_size, |
| 156 | listen_socket=socks[2] if socks else None, |
| 157 | ) |
| 158 | |
| 159 | self_device_group = device_group |
| 160 | self_cpu_group = cpu_group |
| 161 | self_tcp_store_group = tcp_store_group |
| 162 | |
| 163 | assert self_cpu_group is not None |
| 164 | assert self_device_group is not None |
| 165 | assert self_tcp_store_group is not None |
| 166 | |
| 167 | self.cpu_group = self_cpu_group |
| 168 | self.device_group = self_device_group |
| 169 | self.tcp_store_group = self_tcp_store_group |
| 170 | |
| 171 | if current_platform.is_cuda_alike(): |
| 172 | visible_device_index = ( |
| 173 | current_platform.logical_device_id_to_visible_device_id( |
| 174 | self.device_index |
| 175 | ) |
| 176 | ) |
| 177 | self.device = torch.device(f"cuda:{visible_device_index}") |
| 178 | elif current_platform.is_xpu(): |
| 179 | self.device = torch.device(f"xpu:{self.device_index}") |
| 180 | elif current_platform.is_out_of_tree(): |
| 181 | self.device = torch.device( |
| 182 | f"{current_platform.device_name}:{self.device_index}" |
| 183 | ) |
| 184 | else: |
| 185 | self.device = torch.device("cpu") |
| 186 | |
| 187 | self.use_device_communicator = use_device_communicator |
| 188 | self.device_communicator = None |
| 189 | if use_device_communicator and self.world_size > 1: |
| 190 | device_comm_cls = resolve_obj_by_qualname( |
| 191 | current_platform.get_device_communicator_cls() |
| 192 | ) |
| 193 | assert device_comm_cls == CudaCommunicator |
| 194 | self.device_communicator = CudaCommunicator( |
| 195 | cpu_group=self.cpu_group, |
| 196 | device=self.device, |
| 197 | device_group=self.device_group, |
| 198 | unique_name=self.unique_name, |
| 199 | global_ranks=self.ranks, |
| 200 | global_world_size=global_world_size, |
| 201 | tcp_store_group=self.tcp_store_group, |
| 202 | use_all2all=use_all2all, |
| 203 | ) |
| 204 | |
| 205 | self.mq_broadcaster = None |
| 206 | |
| 207 | self.use_custom_op_call = ( |
| 208 | current_platform.is_cuda_alike() or current_platform.is_tpu() |
| 209 | ) |
| 210 | self.use_cpu_custom_send_recv = False |
| 211 | |
| 212 | def destroy(self): |
| 213 | if self.device_communicator: |
| 214 | self.device_communicator.destroy() |
| 215 | if self.device_group: |
| 216 | stateless_destroy_torch_distributed_process_group(self.device_group) |
| 217 | if self.cpu_group: |
| 218 | stateless_destroy_torch_distributed_process_group(self.cpu_group) |
| 219 | |
| 220 | def size(self) -> int: |
| 221 | """Return the world size of this group.""" |
| 222 | return self.world_size |
| 223 | |
| 224 | def broadcast(self, input_: torch.Tensor, src: int = 0): |
| 225 | if self.world_size == 1: |
| 226 | return input_ |
| 227 | |
| 228 | if self.device_communicator and input_.is_cuda: |
| 229 | return self.device_communicator.broadcast(input_, src) |
| 230 | else: |
| 231 | return self.tcp_store_group.broadcast(input_, src) |
| 232 | |
| 233 | def broadcast_object(self, obj=None, src: int = 0): |
| 234 | if self.world_size == 1: |
| 235 | return obj |
| 236 | return self.tcp_store_group.broadcast_obj(obj, src) |
| 237 | |
| 238 | def broadcast_object_list( |
| 239 | self, obj_list: list[Any], src: int = 0, group: ProcessGroup | None = None |
| 240 | ): |
| 241 | assert src < self.world_size |
| 242 | |
| 243 | if self.world_size == 1: |
| 244 | return obj_list |
| 245 | |
| 246 | if self.rank_in_group == src: |
| 247 | for obj in obj_list: |
| 248 | self.tcp_store_group.broadcast_obj(obj, src) |
| 249 | else: |
| 250 | for i in range(len(obj_list)): |
| 251 | obj_list[i] = self.tcp_store_group.broadcast_obj(None, src) |
| 252 | |
| 253 | return obj_list |
| 254 | |
| 255 | def broadcast_tensor_dict( |
| 256 | self, |
| 257 | tensor_dict: dict[str, torch.Tensor | Any] | None = None, |
| 258 | src: int = 0, |
| 259 | group: ProcessGroup | None = None, |
| 260 | metadata_group: ProcessGroup | None = None, |
| 261 | ) -> dict[str, torch.Tensor | Any] | None: |
| 262 | if self.world_size == 1: |
| 263 | return tensor_dict |
| 264 | |
| 265 | if self.rank_in_group == src: |
| 266 | assert isinstance(tensor_dict, dict), ( |
| 267 | f"Expecting a dictionary, got {type(tensor_dict)}" |
| 268 | ) |
| 269 | metadata_list, tensor_list = _split_tensor_dict(tensor_dict) |
| 270 | else: |
| 271 | metadata_list = None |
| 272 | tensor_list = [] |
| 273 | |
| 274 | recv_metadata_list: list[tuple[str, Any]] = self.tcp_store_group.broadcast_obj( |
| 275 | metadata_list, src |
| 276 | ) |
| 277 | |
| 278 | if self.rank_in_group != src: |
| 279 | tensor_dict = {} |
| 280 | for key, value in recv_metadata_list: |
| 281 | if isinstance(value, TensorMetadata): |
| 282 | tensor = torch.empty( |
| 283 | value.size, dtype=value.dtype, device=value.device |
| 284 | ) |
| 285 | tensor_list.append(tensor) |
| 286 | tensor_dict[key] = tensor |
| 287 | else: |
| 288 | tensor_dict[key] = value |
| 289 | |
| 290 | for tensor in tensor_list: |
| 291 | if tensor.numel() == 0: |
| 292 | continue |
| 293 | if self.device_communicator and tensor.is_cuda: |
| 294 | tensor.copy_(self.device_communicator.broadcast(tensor, src)) |
| 295 | else: |
| 296 | tensor.copy_(self.tcp_store_group.broadcast(tensor, src)) |
| 297 | |
| 298 | return tensor_dict |
| 299 | |
| 300 | def send_object(self, obj, dst: int) -> None: |
| 301 | assert dst < self.world_size |
| 302 | assert dst != self.rank_in_group |
| 303 | self.tcp_store_group.send_obj(obj, dst) |
| 304 | |
| 305 | def recv_object(self, src: int): |
| 306 | assert src < self.world_size |
| 307 | assert src != self.rank_in_group |
| 308 | return self.tcp_store_group.recv_obj(src) |
| 309 | |
| 310 | def send_tensor_dict( |
| 311 | self, |
| 312 | tensor_dict: dict[str, torch.Tensor | Any], |
| 313 | dst: int | None = None, |
| 314 | all_gather_group: Optional["GroupCoordinator"] = None, |
| 315 | all_gather_tensors: dict[str, bool] | None = None, |
| 316 | ) -> dict[str, torch.Tensor | Any] | None: |
| 317 | if self.world_size == 1: |
| 318 | return tensor_dict |
| 319 | |
| 320 | if dst is None: |
| 321 | dst = (self.rank_in_group + 1) % self.world_size |
| 322 | assert dst < self.world_size |
| 323 | |
| 324 | metadata_list, tensor_list = _split_tensor_dict(tensor_dict) |
| 325 | self.tcp_store_group.send_obj(metadata_list, dst) |
| 326 | |
| 327 | for tensor in tensor_list: |
| 328 | if tensor.numel() == 0: |
| 329 | continue |
| 330 | if self.device_communicator and tensor.is_cuda: |
| 331 | self.device_communicator.send(tensor, dst) |
| 332 | else: |
| 333 | self.tcp_store_group.send(tensor, dst) |
| 334 | |
| 335 | return None |
| 336 | |
| 337 | def recv_tensor_dict( |
| 338 | self, |
| 339 | src: int | None = None, |
| 340 | all_gather_group: Optional["GroupCoordinator"] = None, |
| 341 | all_gather_tensors: dict[str, bool] | None = None, |
| 342 | ) -> dict[str, torch.Tensor | Any] | None: |
| 343 | if self.world_size == 1: |
| 344 | return None |
| 345 | |
| 346 | if src is None: |
| 347 | src = (self.rank_in_group - 1) % self.world_size |
| 348 | assert src < self.world_size |
| 349 | |
| 350 | recv_metadata_list = self.tcp_store_group.recv_obj(src) |
| 351 | tensor_dict = {} |
| 352 | for key, value in recv_metadata_list: |
| 353 | if isinstance(value, TensorMetadata): |
| 354 | tensor = torch.empty(value.size, dtype=value.dtype, device=value.device) |
| 355 | if tensor.numel() > 0: |
| 356 | if self.device_communicator and tensor.is_cuda: |
| 357 | tensor = self.device_communicator.recv( |
| 358 | tensor.size(), tensor.dtype, src |
| 359 | ) |
| 360 | else: |
| 361 | tensor = self.tcp_store_group.recv(tensor, src) |
| 362 | tensor_dict[key] = tensor |
| 363 | else: |
| 364 | tensor_dict[key] = value |
| 365 | return tensor_dict |
| 366 | |
| 367 | def barrier(self): |
| 368 | self.tcp_store_group.barrier() |
| 369 | |
| 370 | def gather( |
| 371 | self, input_: torch.Tensor, dst: int = 0, dim: int = -1 |
| 372 | ) -> torch.Tensor | None: |
| 373 | if self.world_size == 1: |
| 374 | return input_ |
| 375 | |
| 376 | if self.device_communicator is None: |
| 377 | raise ValueError("No device communicator found") |
| 378 | |
| 379 | if self.rank_in_group == dst: |
| 380 | gathered_list = [torch.empty_like(input_) for _ in range(self.world_size)] |
| 381 | gathered_list[self.rank_in_group] = input_ |
| 382 | for src_rank in range(self.world_size): |
| 383 | if src_rank != self.rank_in_group: |
| 384 | gathered_list[src_rank] = self.device_communicator.recv( |
| 385 | input_.size(), input_.dtype, src_rank |
| 386 | ) |
| 387 | return torch.cat(gathered_list, dim=dim) |
| 388 | else: |
| 389 | self.device_communicator.send(input_, dst) |
| 390 | return None |