| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """EPLB communicator implementations and factory.""" |
| 4 | |
| 5 | import contextlib |
| 6 | import time |
| 7 | import uuid |
| 8 | from abc import ABC, abstractmethod |
| 9 | from collections.abc import Sequence |
| 10 | from datetime import timedelta |
| 11 | |
| 12 | import numpy as np |
| 13 | import torch |
| 14 | from torch.distributed import ( |
| 15 | P2POp, |
| 16 | ProcessGroup, |
| 17 | batch_isend_irecv, |
| 18 | ) |
| 19 | |
| 20 | import vllm.distributed.nixl_utils as nixl_utils |
| 21 | from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator |
| 22 | from vllm.distributed.device_communicators.pynccl_wrapper import ( |
| 23 | ncclDataTypeEnum, |
| 24 | ) |
| 25 | from vllm.distributed.eplb.eplb_utils import device_stream |
| 26 | from vllm.distributed.parallel_state import ( |
| 27 | GroupCoordinator, |
| 28 | get_pp_group, |
| 29 | is_local_first_rank, |
| 30 | ) |
| 31 | from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator |
| 32 | from vllm.distributed.utils import is_weak_contiguous |
| 33 | from vllm.logger import init_logger |
| 34 | from vllm.platforms import current_platform |
| 35 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 36 | from vllm.utils.torch_utils import PIN_MEMORY |
| 37 | |
| 38 | logger = init_logger(__name__) |
| 39 | |
| 40 | |
| 41 | def has_nixl() -> bool: |
| 42 | """Whether the optional NIXL package is available.""" |
| 43 | return nixl_utils.NixlWrapper is not None |
| 44 | |
| 45 | |
| 46 | class EplbCommunicator(ABC): |
| 47 | """Abstract EPLB communicator for expert weight transfers.""" |
| 48 | |
| 49 | @abstractmethod |
| 50 | def add_send( |
| 51 | self, |
| 52 | tensors: list[torch.Tensor], |
| 53 | dst_rank: int, |
| 54 | expert_id: int, |
| 55 | ) -> None: |
| 56 | pass |
| 57 | |
| 58 | @abstractmethod |
| 59 | def add_recv( |
| 60 | self, |
| 61 | tensors: list[torch.Tensor], |
| 62 | src_rank: int, |
| 63 | expert_id: int, |
| 64 | ) -> None: |
| 65 | pass |
| 66 | |
| 67 | @abstractmethod |
| 68 | def execute(self) -> None: |
| 69 | """Complete all enqueued transfers. |
| 70 | |
| 71 | Some backends perform communication here; others (e.g. NIXL) |
| 72 | issue transfers eagerly in add_recv and only wait here. |
| 73 | On return, all data is available in the destination buffers. |
| 74 | """ |
| 75 | |
| 76 | def set_transfer_context( # noqa: B027 |
| 77 | self, old_indices: np.ndarray, layer_idx: int |
| 78 | ) -> None: |
| 79 | """Pre-set layer context before add_recv calls. |
| 80 | |
| 81 | Default is a no-op; overridden by backends (e.g. NIXL) that need |
| 82 | layer-level context to issue transfers inside add_recv. |
| 83 | """ |
| 84 | |
| 85 | @property |
| 86 | def needs_profile_buffer_reservation(self) -> bool: |
| 87 | """Whether the profile path must run a dummy collective operation to reserve |
| 88 | communication buffers.""" |
| 89 | return True |
| 90 | |
| 91 | def set_stream(self, stream: torch.Stream | None) -> None: |
| 92 | self._stream = stream |
| 93 | |
| 94 | def _log_initialized(self) -> None: |
| 95 | if is_local_first_rank(): |
| 96 | logger.info("Initialized EPLB communicator: %s.", self.__class__.__name__) |
| 97 | |
| 98 | |
| 99 | class TorchDistNcclEplbCommunicator(EplbCommunicator): |
| 100 | """EPLB communicator backed by torch.distributed isend/irecv.""" |
| 101 | |
| 102 | def __init__( |
| 103 | self, |
| 104 | ep_group: ProcessGroup, |
| 105 | stream: torch.Stream | None = None, |
| 106 | ) -> None: |
| 107 | self._ep_group = ep_group |
| 108 | self._stream = stream |
| 109 | self._p2p_ops: list[P2POp] = [] |
| 110 | self._log_initialized() |
| 111 | |
| 112 | def add_send( |
| 113 | self, |
| 114 | tensors: list[torch.Tensor], |
| 115 | dst_rank: int, |
| 116 | expert_id: int, # unused by this backend |
| 117 | ) -> None: |
| 118 | for tensor in tensors: |
| 119 | self._p2p_ops.append( |
| 120 | P2POp( |
| 121 | torch.distributed.isend, |
| 122 | tensor, |
| 123 | dst_rank, |
| 124 | self._ep_group, |
| 125 | ) |
| 126 | ) |
| 127 | |
| 128 | def add_recv( |
| 129 | self, |
| 130 | tensors: list[torch.Tensor], |
| 131 | src_rank: int, |
| 132 | expert_id: int, # unused by this backend |
| 133 | ) -> None: |
| 134 | for tensor in tensors: |
| 135 | self._p2p_ops.append( |
| 136 | P2POp( |
| 137 | torch.distributed.irecv, |
| 138 | tensor, |
| 139 | src_rank, |
| 140 | self._ep_group, |
| 141 | ) |
| 142 | ) |
| 143 | |
| 144 | def execute(self) -> None: |
| 145 | if not self._p2p_ops: |
| 146 | return |
| 147 | try: |
| 148 | with device_stream(self._stream): |
| 149 | reqs = batch_isend_irecv(self._p2p_ops) |
| 150 | for req in reqs: |
| 151 | req.wait() |
| 152 | finally: |
| 153 | self._p2p_ops.clear() |
| 154 | |
| 155 | |
| 156 | class TorchDistGlooStagedEplbCommunicator(EplbCommunicator): |
| 157 | """EPLB communicator using gloo P2P with CPU staging.""" |
| 158 | |
| 159 | def __init__( |
| 160 | self, |
| 161 | cpu_group: ProcessGroup, |
| 162 | stream: torch.Stream | None = None, |
| 163 | ) -> None: |
| 164 | self._cpu_group = cpu_group |
| 165 | self._stream = stream |
| 166 | self._ops: list[tuple[str, torch.Tensor, int]] = [] |
| 167 | self._log_initialized() |
| 168 | |
| 169 | def add_send( |
| 170 | self, |
| 171 | tensors: list[torch.Tensor], |
| 172 | dst_rank: int, |
| 173 | expert_id: int, # unused by this backend |
| 174 | ) -> None: |
| 175 | for tensor in tensors: |
| 176 | self._ops.append(("send", tensor, dst_rank)) |
| 177 | |
| 178 | def add_recv( |
| 179 | self, |
| 180 | tensors: list[torch.Tensor], |
| 181 | src_rank: int, |
| 182 | expert_id: int, # unused by this backend |
| 183 | ) -> None: |
| 184 | for tensor in tensors: |
| 185 | self._ops.append(("recv", tensor, src_rank)) |
| 186 | |
| 187 | def execute(self) -> None: |
| 188 | if not self._ops: |
| 189 | return |
| 190 | |
| 191 | p2p_ops: list[P2POp] = [] |
| 192 | recv_staging: list[tuple[torch.Tensor, torch.Tensor]] = [] |
| 193 | |
| 194 | def build_ops() -> None: |
| 195 | for op, tensor, peer_rank in self._ops: |
| 196 | if op == "send": |
| 197 | cpu_tensor = tensor.to(device="cpu", non_blocking=True) |
| 198 | p2p_ops.append( |
| 199 | P2POp( |
| 200 | torch.distributed.isend, |
| 201 | cpu_tensor, |
| 202 | peer_rank, |
| 203 | self._cpu_group, |
| 204 | ) |
| 205 | ) |
| 206 | continue |
| 207 | cpu_tensor = torch.empty_like( |
| 208 | tensor, device="cpu", pin_memory=PIN_MEMORY |
| 209 | ) |
| 210 | p2p_ops.append( |
| 211 | P2POp( |
| 212 | torch.distributed.irecv, |
| 213 | cpu_tensor, |
| 214 | peer_rank, |
| 215 | self._cpu_group, |
| 216 | ) |
| 217 | ) |
| 218 | recv_staging.append((tensor, cpu_tensor)) |
| 219 | |
| 220 | try: |
| 221 | with device_stream(self._stream): |
| 222 | build_ops() |
| 223 | finally: |
| 224 | self._ops.clear() |
| 225 | |
| 226 | # Wait for all D2H copies to finish |
| 227 | # before issuing gloo batch_isend_irecv operations. |
| 228 | with gpu_sync_allowed(): |
| 229 | if self._stream is not None: |
| 230 | self._stream.synchronize() |
| 231 | else: |
| 232 | torch.accelerator.current_stream().synchronize() |
| 233 | |
| 234 | reqs = batch_isend_irecv(p2p_ops) |
| 235 | for req in reqs: |
| 236 | req.wait() |
| 237 | |
| 238 | if not recv_staging: |
| 239 | return |
| 240 | with device_stream(self._stream): |
| 241 | for dst_tensor, cpu_tensor in recv_staging: |
| 242 | dst_tensor.copy_(cpu_tensor, non_blocking=True) |
| 243 | |
| 244 | |
| 245 | class TorchDistXCCLStagedEplbCommunicator(EplbCommunicator): |
| 246 | """EPLB communicator using XCCL device-to-device P2P on XPU.""" |
| 247 | |
| 248 | def __init__( |
| 249 | self, |
| 250 | ep_group: ProcessGroup, |
| 251 | stream: torch.Stream | None = None, |
| 252 | ) -> None: |
| 253 | self._ep_group = ep_group |
| 254 | self._stream = stream |
| 255 | self._ops: list[tuple[str, torch.Tensor, int]] = [] |
| 256 | self._log_initialized() |
| 257 | |
| 258 | def add_send( |
| 259 | self, |
| 260 | tensors: list[torch.Tensor], |
| 261 | dst_rank: int, |
| 262 | expert_id: int, # unused by this backend |
| 263 | ) -> None: |
| 264 | for tensor in tensors: |
| 265 | self._ops.append(("send", tensor, dst_rank)) |
| 266 | |
| 267 | def add_recv( |
| 268 | self, |
| 269 | tensors: list[torch.Tensor], |
| 270 | src_rank: int, |
| 271 | expert_id: int, # unused by this backend |
| 272 | ) -> None: |
| 273 | for tensor in tensors: |
| 274 | self._ops.append(("recv", tensor, src_rank)) |
| 275 | |
| 276 | def execute(self) -> None: |
| 277 | if not self._ops: |
| 278 | return |
| 279 | |
| 280 | p2p_ops: list[P2POp] = [] |
| 281 | try: |
| 282 | for op, tensor, peer_rank in self._ops: |
| 283 | send_or_recv = ( |
| 284 | torch.distributed.isend if op == "send" else torch.distributed.irecv |
| 285 | ) |
| 286 | p2p_ops.append( |
| 287 | P2POp( |
| 288 | send_or_recv, |
| 289 | tensor, |
| 290 | peer_rank, |
| 291 | self._ep_group, |
| 292 | ) |
| 293 | ) |
| 294 | finally: |
| 295 | self._ops.clear() |
| 296 | |
| 297 | with device_stream(self._stream): |
| 298 | reqs = batch_isend_irecv(p2p_ops) |
| 299 | for req in reqs: |
| 300 | req.wait() |
| 301 | |
| 302 | |
| 303 | class NixlEplbCommunicator(EplbCommunicator): |
| 304 | """EPLB communicator backed by NIXL READ transfers.""" |
| 305 | |
| 306 | def __init__( |
| 307 | self, |
| 308 | cpu_group: ProcessGroup, |
| 309 | all_expert_weights: Sequence[Sequence[torch.Tensor]], |
| 310 | expert_buffer: Sequence[torch.Tensor], |
| 311 | ) -> None: |
| 312 | """Create a NIXL-backed EPLB communicator. |
| 313 | |
| 314 | Args: |
| 315 | cpu_group: CPU process group for metadata exchange. |
| 316 | all_expert_weights: Expert weight tensors for all MoE layers. |
| 317 | expert_buffer: Pre-allocated receive buffer tensors. |
| 318 | |
| 319 | """ |
| 320 | assert all_expert_weights, ( |
| 321 | "NixlEplbCommunicator requires non-empty all_expert_weights." |
| 322 | ) |
| 323 | assert expert_buffer, "NixlEplbCommunicator requires non-empty expert_buffer." |
| 324 | nixl_wrapper_cls = nixl_utils.NixlWrapper |
| 325 | if nixl_wrapper_cls is None: |
| 326 | raise RuntimeError("NIXL is unavailable.") |
| 327 | |
| 328 | self._cpu_group = cpu_group |
| 329 | self._world_size = cpu_group.size() |
| 330 | self._rank = cpu_group.rank() |
| 331 | |
| 332 | self._all_expert_weights = all_expert_weights |
| 333 | self._expert_buffer = expert_buffer |
| 334 | self._num_local_experts: int = all_expert_weights[0][0].shape[0] |
| 335 | self._device = all_expert_weights[0][0].device |
| 336 | |
| 337 | for layer_tensors in all_expert_weights: |
| 338 | for tensor in layer_tensors: |
| 339 | assert is_weak_contiguous(tensor), ( |
| 340 | "Expert weight tensors must be contiguous in memory" |
| 341 | ) |
| 342 | assert tensor.device == self._device, ( |
| 343 | "All local EPLB tensors are expected to be on the same " |
| 344 | f"device: expected={self._device}, got={tensor.device}" |
| 345 | ) |
| 346 | for tensor in expert_buffer: |
| 347 | assert is_weak_contiguous(tensor), ( |
| 348 | "expert_buffer tensors must be contiguous in memory" |
| 349 | ) |
| 350 | |
| 351 | # (local_dlist, remote_dlist, xfer_handle) for in-flight READs; |
| 352 | # accumulated by add_recv, drained by execute. |
| 353 | self._xfer_entries: list[tuple[int, int, int]] = [] |
| 354 | # Per-rank expert_id -> physical row; set by set_transfer_context. |
| 355 | self._expert_to_src_row: list[dict[int, int]] | None = None |
| 356 | self._layer_idx: int | None = None |
| 357 | |
| 358 | nixl_agent_config = nixl_utils.nixl_agent_config |
| 359 | config = ( |
| 360 | nixl_agent_config(capture_telemetry=False) |
| 361 | if nixl_agent_config is not None |
| 362 | else None |
| 363 | ) |
| 364 | self._nixl_wrapper = nixl_wrapper_cls(self._make_agent_name(), config) |
| 365 | self._nixl_memory_type = "VRAM" |
| 366 | # NIXL registration handles; deregistered in __del__. |
| 367 | self._registered_descs: list[object] = [] |
| 368 | self._remote_agents: dict[int, str] = {} |
| 369 | # peer -> (layer, tensor) -> (base_ptr, bytes_per_expert, dev_id). |
| 370 | self._remote_send_meta: dict[ |
| 371 | int, dict[tuple[int, int], tuple[int, int, int]] |
| 372 | ] = {} |
| 373 | |
| 374 | self._cuda_device_id = int(self._device.index or 0) |
| 375 | self._init_step("buffers", self._init_registered_buffers) |
| 376 | self._init_remote_state() |
| 377 | self._log_initialized() |
| 378 | |
| 379 | def _init_remote_state(self) -> None: |
| 380 | """Exchange NIXL agent metadata and RDMA pointer info with all peers. |
| 381 | |
| 382 | This is a collective operation (uses ``all_gather_object`` twice). |
| 383 | """ |
| 384 | self._init_step("agents", self._init_remote_agents) |
| 385 | self._init_step("send meta", self._exchange_remote_send_meta) |
| 386 | |
| 387 | @property |
| 388 | def needs_profile_buffer_reservation(self) -> bool: |
| 389 | return False |
| 390 | |
| 391 | @staticmethod |
| 392 | def _init_step(name: str, fn: object, *args: object, **kwargs: object) -> None: |
| 393 | try: |
| 394 | fn(*args, **kwargs) # type: ignore[operator] |
| 395 | except Exception as exc: |
| 396 | raise RuntimeError(f"NIXL EPLB init failed: {name}") from exc |
| 397 | |
| 398 | def _make_agent_name(self) -> str: |
| 399 | """Build a deployment-unique nixl agent name.""" |
| 400 | pp_size = get_pp_group().world_size |
| 401 | pp_suffix = f"-pp{get_pp_group().rank_in_group}" if pp_size > 1 else "" |
| 402 | uid = uuid.uuid4().hex[:8] |
| 403 | return f"eplb-{self._rank}{pp_suffix}-{uid}" |
| 404 | |
| 405 | def set_stream(self, stream: torch.Stream | None) -> None: |
| 406 | pass |
| 407 | |
| 408 | def add_send( |
| 409 | self, |
| 410 | tensors: list[torch.Tensor], |
| 411 | dst_rank: int, |
| 412 | expert_id: int, |
| 413 | ) -> None: |
| 414 | # No-op: NIXL READ is receiver-initiated. The sender's expert |
| 415 | # weights are pre-registered and always readable in-place. |
| 416 | pass |
| 417 | |
| 418 | def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None: |
| 419 | assert not self._xfer_entries, ( |
| 420 | f"set_transfer_context() called with {len(self._xfer_entries)} " |
| 421 | f"pending transfers from layer {self._layer_idx}; " |
| 422 | f"execute() was not called after previous add_recv() calls" |
| 423 | ) |
| 424 | self._layer_idx = layer_idx |
| 425 | n = self._num_local_experts |
| 426 | rank_experts = old_indices[: self._world_size * n].reshape(self._world_size, n) |
| 427 | self._expert_to_src_row = [ |
| 428 | {int(eid): i for i, eid in enumerate(row) if eid != -1} |
| 429 | for row in rank_experts |
| 430 | ] |
| 431 | |
| 432 | def add_recv( |
| 433 | self, |
| 434 | tensors: list[torch.Tensor], |
| 435 | src_rank: int, |
| 436 | expert_id: int, |
| 437 | ) -> None: |
| 438 | # Build NIXL descriptors and issue the RDMA READ immediately, |
| 439 | # overlapping the transfer with the remaining Python loop in |
| 440 | # move_to_buffer. |
| 441 | assert self._expert_to_src_row is not None and self._layer_idx is not None, ( |
| 442 | "set_transfer_context() must be called before add_recv()" |
| 443 | ) |
| 444 | src_row = self._expert_to_src_row[src_rank][expert_id] |
| 445 | layer_idx = self._layer_idx |
| 446 | |
| 447 | local_descs: list[tuple[int, int, int]] = [] |
| 448 | remote_descs: list[tuple[int, int, int]] = [] |
| 449 | for t_idx, t in enumerate(tensors): |
| 450 | send_base, send_stride, remote_dev = self._remote_send_meta[src_rank][ |
| 451 | (layer_idx, t_idx) |
| 452 | ] |
| 453 | assert t.nbytes == send_stride, ( |
| 454 | f"tensor {t_idx} size {t.nbytes} != remote stride {send_stride}" |
| 455 | ) |
| 456 | local_descs.append( |
| 457 | ( |
| 458 | t.data_ptr(), |
| 459 | t.nbytes, |
| 460 | self._cuda_device_id, |
| 461 | ) |
| 462 | ) |
| 463 | remote_descs.append( |
| 464 | ( |
| 465 | send_base + src_row * send_stride, |
| 466 | send_stride, |
| 467 | remote_dev, |
| 468 | ) |
| 469 | ) |
| 470 | |
| 471 | local_h, remote_h, xfer_h = self._create_peer_xfer( |
| 472 | src_rank, local_descs, remote_descs |
| 473 | ) |
| 474 | self._nixl_wrapper.transfer(xfer_h) |
| 475 | self._xfer_entries.append((local_h, remote_h, xfer_h)) |
| 476 | |
| 477 | def _init_remote_agents(self) -> None: |
| 478 | local_metadata = self._nixl_wrapper.get_agent_metadata() |
| 479 | gathered_metadata: list[bytes | None] = [None] * self._world_size |
| 480 | torch.distributed.all_gather_object( |
| 481 | gathered_metadata, local_metadata, group=self._cpu_group |
| 482 | ) |
| 483 | for peer in range(self._world_size): |
| 484 | if peer == self._rank: |
| 485 | continue |
| 486 | peer_metadata = gathered_metadata[peer] |
| 487 | assert peer_metadata is not None |
| 488 | self._remote_agents[peer] = self._nixl_wrapper.add_remote_agent( |
| 489 | peer_metadata |
| 490 | ) |
| 491 | |
| 492 | def _init_registered_buffers(self) -> None: |
| 493 | all_tensors: list[torch.Tensor] = [] |
| 494 | for layer_tensors in self._all_expert_weights: |
| 495 | all_tensors.extend(layer_tensors) |
| 496 | all_tensors.extend(self._expert_buffer) |
| 497 | |
| 498 | descs = self._nixl_wrapper.get_reg_descs(all_tensors) |
| 499 | self._nixl_wrapper.register_memory(descs) |
| 500 | self._registered_descs.append(descs) |
| 501 | |
| 502 | def _exchange_remote_send_meta(self) -> None: |
| 503 | """Exchange per-layer per-tensor metadata so receivers can compute |
| 504 | remote RDMA addresses at transfer time.""" |
| 505 | local_meta: dict[tuple[int, int], tuple[int, int, int]] = {} |
| 506 | for layer_idx, layer_tensors in enumerate(self._all_expert_weights): |
| 507 | for t_idx, t in enumerate(layer_tensors): |
| 508 | nbytes_per_expert = t.nbytes // self._num_local_experts |
| 509 | local_meta[(layer_idx, t_idx)] = ( |
| 510 | t.data_ptr(), |
| 511 | nbytes_per_expert, |
| 512 | self._cuda_device_id, |
| 513 | ) |
| 514 | |
| 515 | # Per-rank map: (layer_idx, tensor_idx) -> (base_ptr, bytes_per_expert, dev_id). |
| 516 | # add_recv uses base_ptr + src_row * bytes_per_expert to compute |
| 517 | # the remote RDMA address for each expert. |
| 518 | gathered_meta: list[dict[tuple[int, int], tuple[int, int, int]] | None] = [ |
| 519 | None |
| 520 | ] * self._world_size |
| 521 | torch.distributed.all_gather_object( |
| 522 | gathered_meta, local_meta, group=self._cpu_group |
| 523 | ) |
| 524 | |
| 525 | local_keys = set(local_meta.keys()) |
| 526 | for peer in self._remote_agents: |
| 527 | peer_meta = gathered_meta[peer] |
| 528 | assert peer_meta is not None |
| 529 | peer_keys = set(peer_meta.keys()) |
| 530 | if peer_keys != local_keys: |
| 531 | raise RuntimeError( |
| 532 | f"NIXL EPLB metadata key mismatch with rank {peer}: " |
| 533 | f"local={sorted(local_keys)}, peer={sorted(peer_keys)}" |
| 534 | ) |
| 535 | for key in local_keys: |
| 536 | _, local_stride, _ = local_meta[key] |
| 537 | _, peer_stride, _ = peer_meta[key] |
| 538 | if local_stride != peer_stride: |
| 539 | raise RuntimeError( |
| 540 | f"NIXL EPLB nbytes_per_expert mismatch for {key} " |
| 541 | f"with rank {peer}: " |
| 542 | f"local={local_stride}, peer={peer_stride}" |
| 543 | ) |
| 544 | self._remote_send_meta[peer] = peer_meta |
| 545 | |
| 546 | def _wait_for_all_transfers(self, handles: list[int]) -> None: |
| 547 | pending = set(handles) |
| 548 | while pending: |
| 549 | completed: list[int] = [] |
| 550 | for handle in pending: |
| 551 | state = self._nixl_wrapper.check_xfer_state(handle) |
| 552 | if state == "DONE": |
| 553 | completed.append(handle) |
| 554 | continue |
| 555 | if state != "PROC": |
| 556 | raise RuntimeError(f"NIXL transfer failed with state={state}") |
| 557 | for handle in completed: |
| 558 | pending.remove(handle) |
| 559 | if pending: |
| 560 | time.sleep(0.0005) |
| 561 | |
| 562 | def _create_peer_xfer( |
| 563 | self, |
| 564 | src: int, |
| 565 | local_descs: list[tuple[int, int, int]], |
| 566 | remote_descs: list[tuple[int, int, int]], |
| 567 | ) -> tuple[int, int, int]: |
| 568 | """Create a batched xfer for multiple descriptors from one peer. |
| 569 | |
| 570 | Each element in *local_descs* / *remote_descs* is an |
| 571 | ``(address, size, device_id)`` tuple. |
| 572 | |
| 573 | Returns ``(local_dlist, remote_dlist, xfer_handle)``. |
| 574 | """ |
| 575 | local_desc = self._nixl_wrapper.get_xfer_descs( |
| 576 | local_descs, self._nixl_memory_type |
| 577 | ) |
| 578 | local_handle = self._nixl_wrapper.prep_xfer_dlist( |
| 579 | "NIXL_INIT_AGENT", |
| 580 | local_desc, |
| 581 | ) |
| 582 | |
| 583 | remote_desc = self._nixl_wrapper.get_xfer_descs( |
| 584 | remote_descs, self._nixl_memory_type |
| 585 | ) |
| 586 | remote_handle = self._nixl_wrapper.prep_xfer_dlist( |
| 587 | self._remote_agents[src], |
| 588 | remote_desc, |
| 589 | ) |
| 590 | |
| 591 | indices = np.arange(len(local_descs), dtype=np.int32) |
| 592 | xfer_handle = self._nixl_wrapper.make_prepped_xfer( |
| 593 | "READ", |
| 594 | local_handle, |
| 595 | indices, |
| 596 | remote_handle, |
| 597 | indices, |
| 598 | ) |
| 599 | return (local_handle, remote_handle, xfer_handle) |
| 600 | |
| 601 | def _post_read_barrier(self) -> None: |
| 602 | """Correctness fence: prevents overwrite-while-remote-read race. |
| 603 | |
| 604 | We avoid ``torch.distributed.monitored_barrier`` because it |
| 605 | calls ``get_backend(group)`` which fails for stateless groups |
| 606 | (elastic EP). An async ``all_reduce`` + ``wait(timeout)`` |
| 607 | works with both regular and stateless groups and provides |
| 608 | equivalent timeout detection. |
| 609 | """ |
| 610 | _dummy = torch.zeros(1, dtype=torch.int32) |
| 611 | work = torch.distributed.all_reduce( |
| 612 | _dummy, group=self._cpu_group, async_op=True |
| 613 | ) |
| 614 | work.wait(timeout=timedelta(minutes=5)) |
| 615 | |
| 616 | def execute(self) -> None: |
| 617 | assert self._layer_idx is not None or not self._xfer_entries, ( |
| 618 | "set_transfer_context() must be called before execute() " |
| 619 | "if any add_recv() calls were made" |
| 620 | ) |
| 621 | try: |
| 622 | self._wait_for_all_transfers([x[2] for x in self._xfer_entries]) |
| 623 | |
| 624 | self._post_read_barrier() |
| 625 | finally: |
| 626 | for local_h, remote_h, xfer_h in self._xfer_entries: |
| 627 | with contextlib.suppress(Exception): |
| 628 | self._nixl_wrapper.release_xfer_handle(xfer_h) |
| 629 | with contextlib.suppress(Exception): |
| 630 | self._nixl_wrapper.release_dlist_handle(local_h) |
| 631 | with contextlib.suppress(Exception): |
| 632 | self._nixl_wrapper.release_dlist_handle(remote_h) |
| 633 | self._xfer_entries.clear() |
| 634 | self._expert_to_src_row = None |
| 635 | self._layer_idx = None |
| 636 | |
| 637 | def __del__(self) -> None: |
| 638 | with contextlib.suppress(Exception): |
| 639 | for local_h, remote_h, xfer_h in self._xfer_entries: |
| 640 | with contextlib.suppress(Exception): |
| 641 | self._nixl_wrapper.release_xfer_handle(xfer_h) |
| 642 | with contextlib.suppress(Exception): |
| 643 | self._nixl_wrapper.release_dlist_handle(local_h) |
| 644 | with contextlib.suppress(Exception): |
| 645 | self._nixl_wrapper.release_dlist_handle(remote_h) |
| 646 | with contextlib.suppress(Exception): |
| 647 | for descs in self._registered_descs: |
| 648 | with contextlib.suppress(Exception): |
| 649 | self._nixl_wrapper.deregister_memory(descs) |
| 650 | self._registered_descs.clear() |
| 651 | with contextlib.suppress(Exception): |
| 652 | for agent_name in self._remote_agents.values(): |
| 653 | with contextlib.suppress(Exception): |
| 654 | self._nixl_wrapper.remove_remote_agent(agent_name) |
| 655 | self._remote_agents.clear() |
| 656 | |
| 657 | |
| 658 | class PyNcclEplbCommunicator(EplbCommunicator): |
| 659 | """EPLB communicator backed by PyNcclCommunicator using ncclSend/ncclRecv.""" |
| 660 | |
| 661 | def __init__( |
| 662 | self, |
| 663 | pynccl_comm: PyNcclCommunicator, |
| 664 | stream: torch.Stream | None = None, |
| 665 | ) -> None: |
| 666 | self._pynccl_comm = pynccl_comm |
| 667 | self._stream = stream |
| 668 | self._group_started = False |
| 669 | self._log_initialized() |
| 670 | |
| 671 | def _ensure_group_started(self) -> None: |
| 672 | if not self._group_started: |
| 673 | self._pynccl_comm.group_start() |
| 674 | self._group_started = True |
| 675 | |
| 676 | def add_send( |
| 677 | self, |
| 678 | tensors: list[torch.Tensor], |
| 679 | dst_rank: int, |
| 680 | expert_id: int, # unused by this backend |
| 681 | ) -> None: |
| 682 | self._ensure_group_started() |
| 683 | for tensor in tensors: |
| 684 | self._pynccl_comm.send(tensor, dst_rank, stream=self._stream) |
| 685 | |
| 686 | def add_recv( |
| 687 | self, |
| 688 | tensors: list[torch.Tensor], |
| 689 | src_rank: int, |
| 690 | expert_id: int, # unused by this backend |
| 691 | ) -> None: |
| 692 | self._ensure_group_started() |
| 693 | for tensor in tensors: |
| 694 | self._pynccl_comm.recv(tensor, src_rank, stream=self._stream) |
| 695 | |
| 696 | def execute(self) -> None: |
| 697 | if self._group_started: |
| 698 | self._pynccl_comm.group_end() |
| 699 | self._group_started = False |
| 700 | |
| 701 | |
| 702 | def create_eplb_communicator( |
| 703 | group_coordinator: GroupCoordinator, |
| 704 | backend: str, |
| 705 | expert_weights: Sequence[Sequence[torch.Tensor]], |
| 706 | expert_buffer: Sequence[torch.Tensor], |
| 707 | ) -> EplbCommunicator: |
| 708 | """Create an EPLB communicator for the given backend. |
| 709 | |
| 710 | Args: |
| 711 | group_coordinator: Process-group coordinator that provides the |
| 712 | device and CPU communication groups. |
| 713 | backend: Communicator backend name (``"torch_nccl"``, |
| 714 | ``"torch_gloo"``, ``"pynccl"``, or ``"nixl"``). |
| 715 | Falls back to ``"torch_nccl"`` when *None*. |
| 716 | Stateless (elastic EP) groups support ``"torch_nccl"``, |
| 717 | ``"pynccl"``, and ``"nixl"``; ``"torch_nccl"`` is silently |
| 718 | promoted to ``"pynccl"``. When tensors reside on CPU, |
| 719 | ``"torch_gloo"`` or |
| 720 | ``"torch_nccl"`` are used via the CPU process group. |
| 721 | expert_weights: Expert weight tensors for *all* MoE layers. |
| 722 | Shape ``(num_layers)(num_tensors_per_layer)``. |
| 723 | NixlEplbCommunicator registers all layers with NIXL for |
| 724 | zero-copy RDMA reads. |
| 725 | expert_buffer: Pre-allocated receive buffer tensors (one per |
| 726 | weight tensor in a single layer). |
| 727 | |
| 728 | """ |
| 729 | first_layer = expert_weights[0] if expert_weights else [] |
| 730 | tensor_device_type = first_layer[0].device.type if first_layer else "cpu" |
| 731 | torch_group = ( |
| 732 | group_coordinator.cpu_group |
| 733 | if tensor_device_type == "cpu" |
| 734 | else group_coordinator.device_group |
| 735 | ) |
| 736 | |
| 737 | def _create_pynccl() -> EplbCommunicator: |
| 738 | if tensor_device_type == "cpu": |
| 739 | raise RuntimeError( |
| 740 | "EPLB communicator 'pynccl' supports only cuda-like devices " |
| 741 | f"(got {tensor_device_type})." |
| 742 | ) |
| 743 | unsupported_dtypes = sorted( |
| 744 | { |
| 745 | tensor.dtype |
| 746 | for tensor in first_layer |
| 747 | if not ncclDataTypeEnum.supports_torch_dtype(tensor.dtype) |
| 748 | }, |
| 749 | key=str, |
| 750 | ) |
| 751 | if unsupported_dtypes: |
| 752 | raise RuntimeError( |
| 753 | "EPLB communicator 'pynccl' requested but expert weights contain " |
| 754 | "unsupported dtypes: " |
| 755 | f"({', '.join(str(dtype) for dtype in unsupported_dtypes)})." |
| 756 | ) |
| 757 | |
| 758 | device_comm = group_coordinator.device_communicator |
| 759 | pynccl_comm = ( |
| 760 | getattr(device_comm, "pynccl_comm", None) |
| 761 | if device_comm is not None |
| 762 | else None |
| 763 | ) |
| 764 | if pynccl_comm is None or pynccl_comm.disabled or not pynccl_comm.available: |
| 765 | raise RuntimeError("EPLB communicator 'pynccl' requested but unavailable.") |
| 766 | try: |
| 767 | return PyNcclEplbCommunicator(pynccl_comm=pynccl_comm) |
| 768 | except Exception as exc: |
| 769 | raise RuntimeError( |
| 770 | f"Failed to initialize PyNcclEplbCommunicator ({exc})." |
| 771 | ) from exc |
| 772 | |
| 773 | is_stateless = isinstance(group_coordinator, StatelessGroupCoordinator) |
| 774 | if is_stateless and backend != "nixl": |
| 775 | if backend not in ("torch_nccl", "pynccl"): |
| 776 | raise ValueError( |
| 777 | f"Elastic EP requires 'torch_nccl', 'pynccl', or 'nixl' " |
| 778 | f"EPLB communicator (got '{backend}')." |
| 779 | ) |
| 780 | if backend == "torch_nccl": |
| 781 | logger.warning( |
| 782 | "Stateless elastic EP requires PyNCCL backend. " |
| 783 | "Forcing EPLB communicator to 'pynccl'." |
| 784 | ) |
| 785 | backend = "pynccl" |
| 786 | return _create_pynccl() |
| 787 | |
| 788 | if backend == "nixl": |
| 789 | if not has_nixl(): |
| 790 | raise RuntimeError( |
| 791 | "EPLB communicator 'nixl' requested but NIXL is unavailable." |
| 792 | ) |
| 793 | if not (current_platform.is_cuda_alike() and tensor_device_type != "cpu"): |
| 794 | raise RuntimeError( |
| 795 | "EPLB communicator 'nixl' supports only cuda-like devices " |
| 796 | f"(got {tensor_device_type})." |
| 797 | ) |
| 798 | try: |
| 799 | return NixlEplbCommunicator( |
| 800 | cpu_group=group_coordinator.cpu_group, |
| 801 | all_expert_weights=expert_weights, |
| 802 | expert_buffer=expert_buffer, |
| 803 | ) |
| 804 | except Exception as exc: |
| 805 | raise RuntimeError( |
| 806 | f"Failed to initialize NixlEplbCommunicator ({exc})." |
| 807 | ) from exc |
| 808 | elif backend == "torch_gloo": |
| 809 | return TorchDistGlooStagedEplbCommunicator( |
| 810 | cpu_group=group_coordinator.cpu_group, |
| 811 | ) |
| 812 | elif backend == "torch_xccl": |
| 813 | return TorchDistXCCLStagedEplbCommunicator(ep_group=torch_group) |
| 814 | elif backend == "torch_nccl": |
| 815 | return TorchDistNcclEplbCommunicator(ep_group=torch_group) |
| 816 | elif backend == "pynccl": |
| 817 | return _create_pynccl() |
| 818 | raise ValueError(f"Unknown EPLB communicator backend: {backend}") |