| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """IPC-based weight transfer engine using CUDA IPC for communication.""" |
| 4 | |
| 5 | import pickle |
| 6 | from dataclasses import asdict, dataclass |
| 7 | from typing import TYPE_CHECKING, Any, ClassVar |
| 8 | |
| 9 | import pybase64 as base64 |
| 10 | import torch |
| 11 | from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor |
| 12 | from typing_extensions import Self |
| 13 | |
| 14 | from vllm import envs |
| 15 | from vllm.config.weight_transfer import WeightTransferConfig |
| 16 | from vllm.distributed.weight_transfer.base import ( |
| 17 | TrainerInitInfo, |
| 18 | TrainerWeightTransferEngine, |
| 19 | VLLMWeightSyncClient, |
| 20 | WeightSource, |
| 21 | WeightTransferEngine, |
| 22 | WeightTransferInitInfo, |
| 23 | WeightTransferUpdateInfo, |
| 24 | ) |
| 25 | |
| 26 | if TYPE_CHECKING: |
| 27 | from vllm.config import VllmConfig |
| 28 | from vllm.distributed.weight_transfer.packed_tensor import ( |
| 29 | DEFAULT_PACKED_BUFFER_SIZE_BYTES, |
| 30 | PackedBufferImporter, |
| 31 | packed_ipc_consumer, |
| 32 | packed_ipc_producer, |
| 33 | ) |
| 34 | |
| 35 | |
| 36 | @dataclass |
| 37 | class IPCWeightTransferInitInfo(WeightTransferInitInfo): |
| 38 | """Worker-side init info for IPC weight transfer. No rendezvous needed. |
| 39 | |
| 40 | `packed` is a must-agree wire param: the trainer ships it here at the init |
| 41 | handshake so the worker decodes with the same setting the trainer encoded |
| 42 | with. The consumer rebuilds from the IPC handle + `tensor_sizes`, so it does |
| 43 | not need the buffer size (producer-only).""" |
| 44 | |
| 45 | packed: bool = False |
| 46 | |
| 47 | |
| 48 | @dataclass |
| 49 | class IPCTrainerInitInfo(TrainerInitInfo): |
| 50 | """Trainer-side init info for IPC weight transfer. No rendezvous needed; |
| 51 | `rank` (from `TrainerInitInfo`) identifies this trainer process — rank 0 |
| 52 | ships the merged IPC handles. All ranks still join the handle all-gather. |
| 53 | |
| 54 | `packed` / `packed_buffer_size_bytes` are the transfer's wire params. The |
| 55 | trainer propagates `packed` to the worker at `trainer_init` so the two sides |
| 56 | cannot disagree. `backend` is the factory dispatch key.""" |
| 57 | |
| 58 | backend: ClassVar[str] = "ipc" |
| 59 | packed: bool = False |
| 60 | packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES |
| 61 | |
| 62 | |
| 63 | @dataclass |
| 64 | class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo): |
| 65 | """Per-round update info for the IPC weight transfer backend.""" |
| 66 | |
| 67 | names: list[str] |
| 68 | dtype_names: list[str] |
| 69 | shapes: list[list[int]] |
| 70 | ipc_handles: list[dict[str, tuple]] | dict[str, tuple] | None = None |
| 71 | """IPC handles mapping physical GPU UUID to rebuild_cuda_tensor args. |
| 72 | For non-packed mode: list of per-parameter handle dicts. |
| 73 | For packed mode: single handle dict for the packed buffer.""" |
| 74 | ipc_handles_pickled: str | None = None |
| 75 | """Base64-encoded pickled IPC handles, used for HTTP transport.""" |
| 76 | tensor_sizes: list[int] | None = None |
| 77 | """Per-parameter sizes in bytes within the packed buffer. |
| 78 | Required when packed=True, unused otherwise.""" |
| 79 | |
| 80 | def __post_init__(self): |
| 81 | if self.ipc_handles_pickled is not None: |
| 82 | if self.ipc_handles is not None: |
| 83 | raise ValueError( |
| 84 | "Cannot specify both `ipc_handles` and `ipc_handles_pickled`" |
| 85 | ) |
| 86 | |
| 87 | if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION: |
| 88 | raise ValueError( |
| 89 | "Refusing to deserialize `ipc_handles_pickled` without " |
| 90 | "VLLM_ALLOW_INSECURE_SERIALIZATION=1" |
| 91 | ) |
| 92 | |
| 93 | self.ipc_handles = pickle.loads(base64.b64decode(self.ipc_handles_pickled)) |
| 94 | self.ipc_handles_pickled = None |
| 95 | |
| 96 | if self.ipc_handles is None: |
| 97 | raise ValueError( |
| 98 | "Either `ipc_handles` or `ipc_handles_pickled` must be provided" |
| 99 | ) |
| 100 | num_params = len(self.names) |
| 101 | if len(self.dtype_names) != num_params: |
| 102 | raise ValueError( |
| 103 | f"`dtype_names` should be of the same size as `names`: " |
| 104 | f"got {len(self.dtype_names)} and {len(self.names)}" |
| 105 | ) |
| 106 | if len(self.shapes) != num_params: |
| 107 | raise ValueError( |
| 108 | f"`shapes` should be of the same size as `names`: " |
| 109 | f"got {len(self.shapes)} and {len(self.names)}" |
| 110 | ) |
| 111 | # Unpacked transfers carry a per-parameter handle dict (a list); packed |
| 112 | # transfers carry a single dict for the whole buffer, so the list check |
| 113 | # only applies to the list case. |
| 114 | if isinstance(self.ipc_handles, list) and len(self.ipc_handles) != num_params: |
| 115 | raise ValueError( |
| 116 | f"`ipc_handles` should be of the same size as `names`: " |
| 117 | f"got {len(self.ipc_handles)} and {len(self.names)}" |
| 118 | ) |
| 119 | |
| 120 | |
| 121 | class IPCWeightTransferEngine( |
| 122 | WeightTransferEngine[IPCWeightTransferInitInfo, IPCWeightTransferUpdateInfo] |
| 123 | ): |
| 124 | """Weight transfer engine using CUDA IPC between trainer and workers. |
| 125 | |
| 126 | This implementation uses CUDA IPC to transfer weights from the trainer (rank 0) |
| 127 | to all inference workers in a process group. IPC handles are used to share |
| 128 | memory between processes on the same node. |
| 129 | """ |
| 130 | |
| 131 | # Define backend-specific dataclass types |
| 132 | init_info_cls = IPCWeightTransferInitInfo |
| 133 | update_info_cls = IPCWeightTransferUpdateInfo |
| 134 | |
| 135 | def __init__( |
| 136 | self, |
| 137 | config: WeightTransferConfig, |
| 138 | vllm_config: "VllmConfig", |
| 139 | device: torch.device, |
| 140 | model: torch.nn.Module, |
| 141 | ) -> None: |
| 142 | super().__init__(config, vllm_config, device, model) |
| 143 | # Set from the trainer-supplied init info at the handshake; defaults are |
| 144 | # only for the (unreachable) receive-before-init case. |
| 145 | self.packed = False |
| 146 | # Shared across all chunks of every packed transfer this engine |
| 147 | # receives; see PackedBufferImporter for the refcount contract. |
| 148 | self._packed_importer = PackedBufferImporter() |
| 149 | |
| 150 | def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None: |
| 151 | """Initialize the weight transfer mechanism. No data-plane rendezvous is |
| 152 | needed for IPC; this just records the trainer-supplied wire params so the |
| 153 | worker decodes exactly as the trainer encoded. |
| 154 | |
| 155 | Args: |
| 156 | init_info: IPC initialization info (carries `packed`). |
| 157 | |
| 158 | """ |
| 159 | self.packed = init_info.packed |
| 160 | |
| 161 | def start_weight_update(self) -> None: |
| 162 | """Initialize layerwise reloading for the incoming checkpoint weights.""" |
| 163 | from vllm.model_executor.model_loader.reload import ( |
| 164 | initialize_layerwise_reload, |
| 165 | ) |
| 166 | |
| 167 | initialize_layerwise_reload(self.model) |
| 168 | |
| 169 | def finish_weight_update(self) -> None: |
| 170 | """Finalize layerwise reloading after all weights have been received.""" |
| 171 | from vllm.model_executor.model_loader.reload import ( |
| 172 | finalize_layerwise_reload, |
| 173 | ) |
| 174 | |
| 175 | finalize_layerwise_reload(self.model, self.model_config) |
| 176 | # Every reduce_tensor call is a fresh export with its own refcount |
| 177 | # slot, so releasing once per update always balances this update's |
| 178 | # export and lets the trainer reclaim its staging buffer. Callers |
| 179 | # that skip finish are still covered by the replace-on-next-export |
| 180 | # path inside the importer. |
| 181 | self._packed_importer.close() |
| 182 | |
| 183 | def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None: |
| 184 | """Receive weights from the trainer via CUDA IPC handles and load them. |
| 185 | |
| 186 | Whether the transfer is packed is read from `self.packed`, set at the |
| 187 | init handshake from the trainer's init info, so it is guaranteed to |
| 188 | match how the trainer encoded. |
| 189 | |
| 190 | Args: |
| 191 | update_info: IPC update info containing parameter names, dtypes, shapes, |
| 192 | and IPC handles. Each IPC handle is a mapping between physical |
| 193 | GPU UUID and the rebuild_cuda_tensor args tuple. |
| 194 | |
| 195 | """ |
| 196 | # Use the worker's assigned device rather than the ambient current |
| 197 | # device: the receive path is no longer wrapped in |
| 198 | # `with torch.device(self.device)` by the caller, so the current device |
| 199 | # is not guaranteed to match self.device. The IPC tensors must be |
| 200 | # rebuilt on the device the model lives on. |
| 201 | device_index = self.device.index |
| 202 | |
| 203 | if self.packed: |
| 204 | if update_info.tensor_sizes is None: |
| 205 | raise ValueError("`tensor_sizes` is required when packed=True") |
| 206 | assert isinstance(update_info.ipc_handles, dict) |
| 207 | weights = packed_ipc_consumer( |
| 208 | ipc_handle=update_info.ipc_handles, |
| 209 | names=update_info.names, |
| 210 | shapes=update_info.shapes, |
| 211 | dtype_names=update_info.dtype_names, |
| 212 | tensor_sizes=update_info.tensor_sizes, |
| 213 | device_index=device_index, |
| 214 | importer=self._packed_importer, |
| 215 | ) |
| 216 | else: |
| 217 | assert isinstance(update_info.ipc_handles, list) |
| 218 | weights = [] |
| 219 | for name, ipc_handle in zip( |
| 220 | update_info.names, |
| 221 | update_info.ipc_handles, |
| 222 | ): |
| 223 | props = torch.cuda.get_device_properties(device_index) |
| 224 | physical_gpu_id = str(props.uuid) |
| 225 | |
| 226 | if physical_gpu_id not in ipc_handle: |
| 227 | raise ValueError( |
| 228 | f"IPC handle not found for GPU UUID " |
| 229 | f"{physical_gpu_id}. " |
| 230 | f"Available UUIDs: {list(ipc_handle.keys())}" |
| 231 | ) |
| 232 | |
| 233 | args = ipc_handle[physical_gpu_id] |
| 234 | list_args = list(args) |
| 235 | # Index 6 of the args from reduce_tensor is the device_index. |
| 236 | # We need to overwrite it with the receiver's device index. |
| 237 | list_args[6] = device_index |
| 238 | weight = rebuild_cuda_tensor(*list_args) |
| 239 | weights.append((name, weight)) |
| 240 | |
| 241 | from vllm.model_executor.model_loader.mtp_validation import ( |
| 242 | disable_mtp_completeness_check, |
| 243 | ) |
| 244 | |
| 245 | with disable_mtp_completeness_check(): |
| 246 | self.model.load_weights(weights) |
| 247 | |
| 248 | def shutdown(self) -> None: |
| 249 | self._packed_importer.close() |
| 250 | |
| 251 | |
| 252 | class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitInfo]): |
| 253 | """Trainer-side CUDA IPC weight transfer engine. |
| 254 | |
| 255 | Called on every trainer rank. For multi-rank (e.g. FSDP) trainers all ranks |
| 256 | iterate the source (materializing each tensor) and contribute to the |
| 257 | IPC-handle all-gather; only the sender (rank 0) ships the merged handles to |
| 258 | the inference side. IPC transfer is straight-line (no concurrent broadcast |
| 259 | like NCCL): `update_weights` *is* the transfer, and it rides the client, so |
| 260 | it no-ops on non-senders. |
| 261 | |
| 262 | `packed` / `packed_buffer_size_bytes` come from `IPCTrainerInitInfo`; the |
| 263 | sender propagates `packed` to the worker at `trainer_init`. |
| 264 | """ |
| 265 | |
| 266 | init_info_cls = IPCTrainerInitInfo |
| 267 | |
| 268 | def __init__( |
| 269 | self, |
| 270 | *, |
| 271 | client: VLLMWeightSyncClient, |
| 272 | source: WeightSource, |
| 273 | is_sender: bool = True, |
| 274 | packed: bool = False, |
| 275 | packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, |
| 276 | ) -> None: |
| 277 | super().__init__(client=client, source=source, is_sender=is_sender) |
| 278 | self.packed = packed |
| 279 | self.packed_buffer_size_bytes = packed_buffer_size_bytes |
| 280 | self.device_index = torch.accelerator.current_device_index() |
| 281 | self.gpu_uuid = str(torch.cuda.get_device_properties(self.device_index).uuid) |
| 282 | |
| 283 | @classmethod |
| 284 | def trainer_init( |
| 285 | cls, |
| 286 | init_info: IPCTrainerInitInfo, |
| 287 | *, |
| 288 | client: VLLMWeightSyncClient, |
| 289 | source: WeightSource | None = None, |
| 290 | ) -> Self: |
| 291 | if source is None: |
| 292 | raise ValueError("IPC trainer weight transfer requires a WeightSource.") |
| 293 | engine = cls( |
| 294 | client=client, |
| 295 | source=source, |
| 296 | is_sender=init_info.is_sender, |
| 297 | packed=init_info.packed, |
| 298 | packed_buffer_size_bytes=init_info.packed_buffer_size_bytes, |
| 299 | ) |
| 300 | # IPC needs no data-plane rendezvous. The sender ships the must-agree |
| 301 | # `packed` flag so the worker decodes exactly as this trainer encodes. |
| 302 | if engine.is_sender: |
| 303 | engine.client.init_weight_transfer_engine({"packed": init_info.packed}) |
| 304 | return engine |
| 305 | |
| 306 | def send_weights(self) -> None: |
| 307 | assert self.source is not None |
| 308 | source = self.source |
| 309 | if self.is_sender: |
| 310 | self.client.start_weight_update() |
| 311 | # Unpacked returns strong refs to its IPC-shared copies; they must stay |
| 312 | # alive across `finish_weight_update` too. |
| 313 | weight_refs = self._send(source) |
| 314 | if self.is_sender: |
| 315 | self.client.finish_weight_update() |
| 316 | self._post_send_sync() |
| 317 | del weight_refs |
| 318 | |
| 319 | # ---- data plane (runs on all ranks; only the sender ships) ---- |
| 320 | |
| 321 | def _send(self, source: WeightSource) -> list[torch.Tensor] | None: |
| 322 | if self.packed: |
| 323 | self._send_packed(source) |
| 324 | return None |
| 325 | return self._send_unpacked(source) |
| 326 | |
| 327 | def _all_gather_and_merge_handles( |
| 328 | self, |
| 329 | handles: list[dict[str, tuple]], |
| 330 | ) -> list[dict[str, tuple]]: |
| 331 | """All-gather and merge IPC handle dicts across ranks in one call. |
| 332 | |
| 333 | Each rank contributes a list of {gpu_uuid: ipc_args} dicts (one |
| 334 | per parameter or one per chunk). A single all_gather_object |
| 335 | collects every rank's full list, then the sender merges per-index so |
| 336 | each dict maps every GPU UUID to its args. |
| 337 | |
| 338 | The all-gather runs over the *default* process group; this assumes the |
| 339 | default group is exactly the set of colocated trainer ranks and that the |
| 340 | sender is a member. No-op (returns handles unchanged) when no distributed |
| 341 | group exists. |
| 342 | """ |
| 343 | if ( |
| 344 | not torch.distributed.is_initialized() |
| 345 | or torch.distributed.get_world_size() == 1 |
| 346 | ): |
| 347 | return handles |
| 348 | |
| 349 | world_size = torch.distributed.get_world_size() |
| 350 | gathered: list[list[dict[str, tuple]] | None] = [None] * world_size |
| 351 | torch.distributed.all_gather_object(gathered, handles) |
| 352 | torch.distributed.barrier() |
| 353 | torch.cuda.synchronize() |
| 354 | |
| 355 | if self.is_sender: |
| 356 | merged: list[dict[str, tuple]] = [] |
| 357 | for param_idx in range(len(handles)): |
| 358 | m: dict[str, tuple] = {} |
| 359 | for rank_handles in gathered: |
| 360 | if rank_handles is not None: |
| 361 | m.update(rank_handles[param_idx]) |
| 362 | merged.append(m) |
| 363 | return merged |
| 364 | return [{} for _ in handles] |
| 365 | |
| 366 | @staticmethod |
| 367 | def _post_send_sync() -> None: |
| 368 | """Barrier + ipc_collect after a send; no-op if single-GPU.""" |
| 369 | if ( |
| 370 | torch.distributed.is_initialized() |
| 371 | and torch.distributed.get_world_size() > 1 |
| 372 | ): |
| 373 | torch.distributed.barrier() |
| 374 | torch.cuda.ipc_collect() |
| 375 | |
| 376 | def _send_unpacked(self, source: WeightSource) -> list[torch.Tensor]: |
| 377 | """Iterate the source, build one IPC handle per param, all-gather the |
| 378 | handles across ranks, and (sender) ship them in one update call. |
| 379 | |
| 380 | Returns the strong refs to every contiguous copy. reduce_tensor's args |
| 381 | do NOT keep storage alive, and non-contiguous inputs allocate fresh |
| 382 | storage in .contiguous(); the caller must keep these alive until the |
| 383 | post-send barrier (past `finish`) so the consumer's IPC views stay valid. |
| 384 | """ |
| 385 | names: list[str] = [] |
| 386 | dtype_names: list[str] = [] |
| 387 | shapes: list[list[int]] = [] |
| 388 | ipc_handles: list[dict[str, tuple]] = [] |
| 389 | weight_refs: list[torch.Tensor] = [] |
| 390 | |
| 391 | for name, tensor in source: |
| 392 | # FSDP shards were gathered by the source; ensure contiguity here. |
| 393 | names.append(name) |
| 394 | dtype_names.append(str(tensor.dtype).split(".")[-1]) |
| 395 | shapes.append(list(tensor.shape)) |
| 396 | |
| 397 | weight = tensor.detach().contiguous() |
| 398 | weight_refs.append(weight) |
| 399 | _, ipc_args = reduce_tensor(weight) |
| 400 | ipc_handles.append({self.gpu_uuid: ipc_args}) |
| 401 | |
| 402 | ipc_handles = self._all_gather_and_merge_handles(ipc_handles) |
| 403 | self._do_send( |
| 404 | names=names, |
| 405 | dtype_names=dtype_names, |
| 406 | shapes=shapes, |
| 407 | ipc_handles=ipc_handles, |
| 408 | ) |
| 409 | return weight_refs |
| 410 | |
| 411 | def _send_packed(self, source: WeightSource) -> None: |
| 412 | """Send weights in bounded-memory chunks (packed mode).""" |
| 413 | for chunk in packed_ipc_producer( |
| 414 | iterator=iter(source), |
| 415 | gpu_uuid=self.gpu_uuid, |
| 416 | post_iter_func=lambda item: item[1], |
| 417 | buffer_size_bytes=self.packed_buffer_size_bytes, |
| 418 | ): |
| 419 | ipc_handle = self._all_gather_and_merge_handles([chunk.ipc_handle])[0] |
| 420 | self._do_send( |
| 421 | names=chunk.names, |
| 422 | dtype_names=chunk.dtype_names, |
| 423 | shapes=chunk.shapes, |
| 424 | ipc_handles=ipc_handle, |
| 425 | tensor_sizes=chunk.tensor_sizes, |
| 426 | ) |
| 427 | # Per-chunk barrier: the producer reuses a single IPC buffer across |
| 428 | # chunks, but only the sender waits for the consumers (via _do_send). |
| 429 | # Without syncing every rank here, non-sender ranks race ahead and |
| 430 | # overwrite their buffer while their colocated worker is still |
| 431 | # reading the current chunk, silently corrupting the transfer. |
| 432 | self._post_send_sync() |
| 433 | |
| 434 | def _do_send( |
| 435 | self, |
| 436 | names: list[str], |
| 437 | dtype_names: list[str], |
| 438 | shapes: list[list[int]], |
| 439 | ipc_handles: list[dict[str, tuple]] | dict[str, tuple], |
| 440 | tensor_sizes: list[int] | None = None, |
| 441 | ) -> None: |
| 442 | """Build one update payload and ship it via the client. Only the sender |
| 443 | ships (non-sender ranks already contributed to the handle all-gather). |
| 444 | |
| 445 | Emits raw `ipc_handles`; transports that cannot carry them natively |
| 446 | (HTTP/JSON) pickle them in their client (see `HTTPVLLMWeightSyncClient`). |
| 447 | """ |
| 448 | if not self.is_sender: |
| 449 | return |
| 450 | update_fields: dict[str, Any] = { |
| 451 | "names": names, |
| 452 | "dtype_names": dtype_names, |
| 453 | "shapes": shapes, |
| 454 | "ipc_handles": ipc_handles, |
| 455 | } |
| 456 | if tensor_sizes is not None: |
| 457 | update_fields["tensor_sizes"] = tensor_sizes |
| 458 | |
| 459 | update_info = IPCWeightTransferUpdateInfo(**update_fields) |
| 460 | self.client.update_weights(asdict(update_info)) |