| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Packed tensor utilities for efficient weight transfer.""" |
| 4 | |
| 5 | import math |
| 6 | from collections.abc import Callable, Iterator |
| 7 | from dataclasses import dataclass |
| 8 | from functools import cache |
| 9 | from typing import Any |
| 10 | |
| 11 | import torch |
| 12 | from torch.multiprocessing.reductions import reduce_tensor |
| 13 | |
| 14 | # Default values for packed tensor transfer. |
| 15 | DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 # 1GB |
| 16 | DEFAULT_PACKED_NUM_BUFFERS = 2 |
| 17 | |
| 18 | |
| 19 | # Streams are cached across calls: the caching allocator keys cached blocks by the |
| 20 | # stream they were allocated on, so fresh streams at every call would strand the |
| 21 | # packed buffers' reserved memory per call. See gh-52950. |
| 22 | @cache |
| 23 | def _get_streams(device_idx: int, num_buffers: int) -> tuple[torch.cuda.Stream, ...]: |
| 24 | return tuple(torch.cuda.Stream(device=device_idx) for _ in range(num_buffers)) |
| 25 | |
| 26 | |
| 27 | def unpack_tensor( |
| 28 | packed_tensor: torch.Tensor, |
| 29 | names: list[str], |
| 30 | shapes: list[list[int]], |
| 31 | dtypes: list[torch.dtype], |
| 32 | tensor_sizes: list[int], |
| 33 | ) -> list[tuple[str, torch.Tensor]]: |
| 34 | """Unpack a packed uint8 tensor into a list of named tensors. |
| 35 | |
| 36 | The returned tensors are **views** of ``packed_tensor`` (the |
| 37 | ``.contiguous()`` call is a no-op on already-contiguous row-slices). |
| 38 | If ``packed_tensor`` lives in storage that may be reused — e.g. a |
| 39 | reused CUDA IPC buffer — callers must clone the results before the |
| 40 | underlying storage is overwritten. |
| 41 | |
| 42 | Args: |
| 43 | packed_tensor: The packed torch.uint8 tensor to unpack |
| 44 | names: List of tensor names |
| 45 | shapes: List of tensor shapes |
| 46 | dtypes: List of tensor dtypes |
| 47 | tensor_sizes: List of tensor sizes in bytes |
| 48 | |
| 49 | """ |
| 50 | unpacked_tensors = packed_tensor.split(tensor_sizes) |
| 51 | |
| 52 | return [ |
| 53 | (name, tensor.contiguous().view(dtype).view(*shape)) |
| 54 | for name, shape, dtype, tensor in zip(names, shapes, dtypes, unpacked_tensors) |
| 55 | ] |
| 56 | |
| 57 | |
| 58 | @dataclass |
| 59 | class PackedChunk: |
| 60 | """Result of packing tensors into a single contiguous uint8 buffer.""" |
| 61 | |
| 62 | packed_tensor: torch.Tensor |
| 63 | names: list[str] |
| 64 | shapes: list[list[int]] |
| 65 | dtypes: list[torch.dtype] |
| 66 | tensor_sizes: list[int] |
| 67 | |
| 68 | |
| 69 | def pack_tensors( |
| 70 | iterator: Iterator[tuple[str, torch.Tensor]], |
| 71 | post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor], |
| 72 | buffer_size_bytes: int, |
| 73 | tensor_list: list[torch.Tensor] | None = None, |
| 74 | current_size: int = 0, |
| 75 | ) -> PackedChunk | None: |
| 76 | """Pack tensors from an iterator into a single contiguous uint8 buffer. |
| 77 | |
| 78 | Consumes from the iterator until the accumulated size exceeds |
| 79 | buffer_size_bytes or the iterator is exhausted, then returns a |
| 80 | PackedChunk. Returns None if no tensors were consumed. |
| 81 | |
| 82 | Args: |
| 83 | iterator: Iterator of (name, tensor) pairs |
| 84 | post_iter_func: Applied to each item before linearizing to uint8 |
| 85 | buffer_size_bytes: Max bytes before flushing |
| 86 | tensor_list: Pre-existing tensor list to append to (for NCCL |
| 87 | multi-buffer reuse). If None, a fresh list is created. |
| 88 | current_size: Byte count already accumulated in tensor_list |
| 89 | |
| 90 | """ |
| 91 | if tensor_list is None: |
| 92 | tensor_list = [] |
| 93 | |
| 94 | names: list[str] = [] |
| 95 | shapes: list[list[int]] = [] |
| 96 | dtypes: list[torch.dtype] = [] |
| 97 | tensor_sizes: list[int] = [] |
| 98 | total_bytes = current_size |
| 99 | |
| 100 | while True: |
| 101 | try: |
| 102 | item = next(iterator) |
| 103 | except StopIteration: |
| 104 | break |
| 105 | |
| 106 | name, orig_tensor = item |
| 107 | # Apply post processing and convert to linearized uint8 tensor |
| 108 | tensor = post_iter_func(item).contiguous().view(torch.uint8).view(-1) |
| 109 | |
| 110 | if tensor.numel() > buffer_size_bytes: |
| 111 | import warnings |
| 112 | |
| 113 | warnings.warn( |
| 114 | f"Tensor '{name}' has size {tensor.numel()} bytes, which " |
| 115 | f"exceeds buffer_size_bytes={buffer_size_bytes}.", |
| 116 | stacklevel=2, |
| 117 | ) |
| 118 | |
| 119 | tensor_list.append(tensor) |
| 120 | names.append(name) |
| 121 | shapes.append(list(orig_tensor.shape)) |
| 122 | dtypes.append(orig_tensor.dtype) |
| 123 | tensor_sizes.append(tensor.numel()) |
| 124 | total_bytes += tensor.numel() |
| 125 | |
| 126 | if total_bytes > buffer_size_bytes: |
| 127 | break |
| 128 | |
| 129 | if not tensor_list: |
| 130 | return None |
| 131 | |
| 132 | packed = torch.cat(tensor_list, dim=0) |
| 133 | del tensor_list |
| 134 | return PackedChunk( |
| 135 | packed_tensor=packed, |
| 136 | names=names, |
| 137 | shapes=shapes, |
| 138 | dtypes=dtypes, |
| 139 | tensor_sizes=tensor_sizes, |
| 140 | ) |
| 141 | |
| 142 | |
| 143 | # ── NCCL packed broadcast ────────────────────────────────────────────── |
| 144 | |
| 145 | |
| 146 | def packed_nccl_broadcast_producer( |
| 147 | iterator: Iterator[tuple[str, torch.Tensor]], |
| 148 | group: Any, |
| 149 | src: int, |
| 150 | post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor], |
| 151 | buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, |
| 152 | num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS, |
| 153 | ) -> None: |
| 154 | """Broadcast tensors in a packed manner from trainer to workers. |
| 155 | |
| 156 | Args: |
| 157 | iterator: Iterator of model parameters. Returns a tuple of (name, tensor) |
| 158 | group: Process group (PyNcclCommunicator) |
| 159 | src: Source rank (0 in current implementation) |
| 160 | post_iter_func: Function to apply to each (name, tensor) pair before |
| 161 | packing, should return a tensor |
| 162 | buffer_size_bytes: Size in bytes for each packed tensor buffer. |
| 163 | Both producer and consumer must use the same value. |
| 164 | num_buffers: Number of buffers for double/triple buffering. |
| 165 | Both producer and consumer must use the same value. |
| 166 | |
| 167 | """ |
| 168 | streams = _get_streams(torch.accelerator.current_device_index(), num_buffers) |
| 169 | # Keep references to in-flight chunks so their packed_tensors |
| 170 | # aren't freed while an async broadcast is still reading them. |
| 171 | in_flight: list[PackedChunk | None] = [None] * num_buffers |
| 172 | buffer_idx = 0 |
| 173 | |
| 174 | while True: |
| 175 | # Synchronize the current stream |
| 176 | streams[buffer_idx].synchronize() |
| 177 | # Previous chunk on this buffer slot is now safe to free |
| 178 | in_flight[buffer_idx] = None |
| 179 | # Start tasks for the new buffer in a new stream |
| 180 | with torch.cuda.stream(streams[buffer_idx]): |
| 181 | chunk = pack_tensors(iterator, post_iter_func, buffer_size_bytes) |
| 182 | if chunk is None: |
| 183 | break |
| 184 | # Pack the tensors and call broadcast collective |
| 185 | group.broadcast(chunk.packed_tensor, src=src) |
| 186 | # Hold reference until this stream is synchronized |
| 187 | in_flight[buffer_idx] = chunk |
| 188 | # Move to the next buffer |
| 189 | buffer_idx = (buffer_idx + 1) % num_buffers |
| 190 | |
| 191 | # Drain every slot before returning. |
| 192 | for stream in streams: |
| 193 | stream.synchronize() |
| 194 | |
| 195 | |
| 196 | def packed_nccl_broadcast_consumer( |
| 197 | iterator: Iterator[tuple[str, tuple[list[int], torch.dtype]]], |
| 198 | group: Any, |
| 199 | src: int, |
| 200 | post_unpack_func: Callable[[list[tuple[str, torch.Tensor]]], None], |
| 201 | buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, |
| 202 | num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS, |
| 203 | device: torch.device | str = "cuda", |
| 204 | ) -> None: |
| 205 | """Consume packed tensors and unpack them into a list of tensors. |
| 206 | |
| 207 | Args: |
| 208 | iterator: Iterator of parameter metadata. Returns (name, (shape, dtype)) |
| 209 | group: Process group (PyNcclCommunicator) |
| 210 | src: Source rank (0 in current implementation) |
| 211 | post_unpack_func: Function to apply to each list of (name, tensor) after |
| 212 | unpacking |
| 213 | buffer_size_bytes: Size in bytes for each packed tensor buffer. |
| 214 | Both producer and consumer must use the same value. |
| 215 | num_buffers: Number of buffers for double/triple buffering. |
| 216 | Both producer and consumer must use the same value. |
| 217 | device: Device for the receive buffers. Must match the device the NCCL |
| 218 | communicator was created on (the worker's assigned device). |
| 219 | |
| 220 | """ |
| 221 | target_packed_tensor_size = buffer_size_bytes |
| 222 | |
| 223 | streams = _get_streams(torch.accelerator.current_device_index(), num_buffers) |
| 224 | buffer_idx = 0 |
| 225 | |
| 226 | packing_tensor_meta_data: list[list[tuple[str, list[int], torch.dtype, int]]] = [ |
| 227 | [] for _ in range(num_buffers) |
| 228 | ] |
| 229 | packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)] |
| 230 | packed_tensors: list[torch.Tensor] = [ |
| 231 | torch.empty(0, dtype=torch.uint8, device=device) for _ in range(num_buffers) |
| 232 | ] |
| 233 | |
| 234 | while True: |
| 235 | # Synchronize the current stream |
| 236 | streams[buffer_idx].synchronize() |
| 237 | with torch.cuda.stream(streams[buffer_idx]): |
| 238 | # Initialize the packing tensor meta data |
| 239 | packing_tensor_meta_data[buffer_idx] = [] |
| 240 | packing_tensor_sizes[buffer_idx] = 0 |
| 241 | try: |
| 242 | # Form a packed tensor |
| 243 | while True: |
| 244 | name, (shape, dtype) = next(iterator) |
| 245 | tensor_size = math.prod(shape) * dtype.itemsize |
| 246 | packing_tensor_meta_data[buffer_idx].append( |
| 247 | (name, shape, dtype, tensor_size) |
| 248 | ) |
| 249 | packing_tensor_sizes[buffer_idx] += tensor_size |
| 250 | if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size: |
| 251 | break |
| 252 | # Create a packed tensor and broadcast it |
| 253 | packed_tensors[buffer_idx] = torch.empty( |
| 254 | packing_tensor_sizes[buffer_idx], dtype=torch.uint8, device=device |
| 255 | ) |
| 256 | group.broadcast(packed_tensors[buffer_idx], src=src) |
| 257 | # Load the packed tensor into the model |
| 258 | names, shapes, dtypes, tensor_sizes = zip( |
| 259 | *packing_tensor_meta_data[buffer_idx] |
| 260 | ) |
| 261 | post_unpack_func( |
| 262 | unpack_tensor( |
| 263 | packed_tensors[buffer_idx], |
| 264 | list(names), |
| 265 | list(shapes), |
| 266 | list(dtypes), |
| 267 | list(tensor_sizes), |
| 268 | ) |
| 269 | ) |
| 270 | # Move to the next buffer |
| 271 | buffer_idx = (buffer_idx + 1) % num_buffers |
| 272 | except StopIteration: |
| 273 | # Do the last broadcast if there are remaining tensors |
| 274 | if len(packing_tensor_meta_data[buffer_idx]) > 0: |
| 275 | # Create a packed tensor and broadcast it |
| 276 | packed_tensors[buffer_idx] = torch.empty( |
| 277 | packing_tensor_sizes[buffer_idx], |
| 278 | dtype=torch.uint8, |
| 279 | device=device, |
| 280 | ) |
| 281 | group.broadcast(packed_tensors[buffer_idx], src=src) |
| 282 | # Load the packed tensor into the model |
| 283 | names, shapes, dtypes, tensor_sizes = zip( |
| 284 | *packing_tensor_meta_data[buffer_idx] |
| 285 | ) |
| 286 | post_unpack_func( |
| 287 | unpack_tensor( |
| 288 | packed_tensors[buffer_idx], |
| 289 | list(names), |
| 290 | list(shapes), |
| 291 | list(dtypes), |
| 292 | list(tensor_sizes), |
| 293 | ) |
| 294 | ) |
| 295 | break |
| 296 | |
| 297 | # Drain every slot before returning. |
| 298 | for stream in streams: |
| 299 | stream.synchronize() |
| 300 | |
| 301 | |
| 302 | # ── IPC packed transfer ──────────────────────────────────────────────── |
| 303 | |
| 304 | |
| 305 | @dataclass |
| 306 | class PackedIpcChunk: |
| 307 | """Metadata and IPC handle for a single packed chunk.""" |
| 308 | |
| 309 | names: list[str] |
| 310 | shapes: list[list[int]] |
| 311 | dtype_names: list[str] |
| 312 | tensor_sizes: list[int] |
| 313 | ipc_handle: dict[str, tuple] |
| 314 | |
| 315 | |
| 316 | def packed_ipc_producer( |
| 317 | iterator: Iterator[tuple[str, torch.Tensor]], |
| 318 | gpu_uuid: str, |
| 319 | post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor], |
| 320 | buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, |
| 321 | ) -> Iterator[PackedIpcChunk]: |
| 322 | """Pack tensors into a reusable IPC buffer and yield handles. |
| 323 | |
| 324 | Allocates a single GPU buffer of ``buffer_size_bytes`` and registers |
| 325 | it for IPC once via ``reduce_tensor``. Each chunk's packed data is |
| 326 | copied into this buffer before yielding, so only one IPC-shared |
| 327 | allocation exists for the lifetime of the transfer. |
| 328 | |
| 329 | Callers **must** ensure the consumer has finished reading the buffer |
| 330 | (e.g. ``ray.get`` returned) before resuming the generator for the |
| 331 | next chunk. |
| 332 | |
| 333 | Args: |
| 334 | iterator: Iterator of (name, tensor) pairs. |
| 335 | gpu_uuid: Physical GPU UUID string for this rank. |
| 336 | post_iter_func: Applied to each (name, tensor) before packing. |
| 337 | buffer_size_bytes: Exact capacity of the reusable IPC buffer. |
| 338 | Every chunk is guaranteed to fit within this size. A |
| 339 | ``ValueError`` is raised if any single tensor exceeds it. |
| 340 | |
| 341 | """ |
| 342 | ipc_buffer = torch.empty(buffer_size_bytes, dtype=torch.uint8, device="cuda") |
| 343 | _, ipc_args = reduce_tensor(ipc_buffer) |
| 344 | |
| 345 | names: list[str] = [] |
| 346 | shapes: list[list[int]] = [] |
| 347 | dtypes: list[torch.dtype] = [] |
| 348 | tensor_sizes: list[int] = [] |
| 349 | total_bytes = 0 |
| 350 | |
| 351 | for name, orig_tensor in iterator: |
| 352 | flat = ( |
| 353 | post_iter_func((name, orig_tensor)).contiguous().view(torch.uint8).view(-1) |
| 354 | ) |
| 355 | |
| 356 | if flat.numel() > buffer_size_bytes: |
| 357 | raise ValueError( |
| 358 | f"Tensor '{name}' has size {flat.numel()} bytes, " |
| 359 | f"which exceeds buffer_size_bytes={buffer_size_bytes}. " |
| 360 | f"Increase buffer_size_bytes to at least {flat.numel()}." |
| 361 | ) |
| 362 | |
| 363 | if total_bytes and total_bytes + flat.numel() > buffer_size_bytes: |
| 364 | # Drain queued copies so the consumer sees a fully-written buffer. |
| 365 | torch.cuda.current_stream().synchronize() |
| 366 | yield PackedIpcChunk( |
| 367 | names=names, |
| 368 | shapes=shapes, |
| 369 | dtype_names=[str(d).split(".")[-1] for d in dtypes], |
| 370 | tensor_sizes=tensor_sizes, |
| 371 | ipc_handle={gpu_uuid: ipc_args}, |
| 372 | ) |
| 373 | # Rebind to fresh lists so the yielded chunk's metadata is |
| 374 | # not mutated while the consumer is still reading. |
| 375 | names, shapes, dtypes, tensor_sizes = [], [], [], [] |
| 376 | total_bytes = 0 |
| 377 | |
| 378 | ipc_buffer[total_bytes : total_bytes + flat.numel()].copy_(flat) |
| 379 | names.append(name) |
| 380 | shapes.append(list(orig_tensor.shape)) |
| 381 | dtypes.append(orig_tensor.dtype) |
| 382 | tensor_sizes.append(flat.numel()) |
| 383 | total_bytes += flat.numel() |
| 384 | |
| 385 | if total_bytes: |
| 386 | torch.cuda.current_stream().synchronize() |
| 387 | yield PackedIpcChunk( |
| 388 | names=names, |
| 389 | shapes=shapes, |
| 390 | dtype_names=[str(d).split(".")[-1] for d in dtypes], |
| 391 | tensor_sizes=tensor_sizes, |
| 392 | ipc_handle={gpu_uuid: ipc_args}, |
| 393 | ) |
| 394 | |
| 395 | |
| 396 | class PackedBufferImporter: |
| 397 | """One-slot cache for the consumer-side mapping of a packed IPC buffer. |
| 398 | |
| 399 | torch's cross-process refcount pairs one ``reduce_tensor`` export with |
| 400 | one consumer rebuild-release cycle. The producer exports its buffer once |
| 401 | and ships the same args with every chunk, so rebuilding and releasing |
| 402 | per chunk decrements the counter once per chunk, drives it negative, |
| 403 | and the dropped buffer is then never reclaimable by ``ipc_collect`` on |
| 404 | the producer side (it leaks one buffer per transfer). Caching the |
| 405 | rebuilt buffer keyed by the rebuild args restores the pairing: chunks |
| 406 | of one transfer reuse the mapping (it aliases producer memory, so it |
| 407 | always reads the current chunk's bytes), and replacing the entry when |
| 408 | the next export arrives — or ``close()`` — releases the previous |
| 409 | mapping exactly once. |
| 410 | |
| 411 | The importer that consumes a multi-chunk transfer must be one object |
| 412 | reused across chunks; a fresh importer per chunk reintroduces the |
| 413 | per-chunk release. ``IPCWeightTransferEngine`` owns one instance per |
| 414 | engine and closes it on shutdown. |
| 415 | """ |
| 416 | |
| 417 | def __init__(self) -> None: |
| 418 | self._entry: tuple[tuple, torch.Tensor] | None = None |
| 419 | |
| 420 | def rebuild(self, list_args: list) -> torch.Tensor: |
| 421 | from torch.multiprocessing.reductions import rebuild_cuda_tensor |
| 422 | |
| 423 | key = tuple(list_args) |
| 424 | if self._entry is not None and self._entry[0] == key: |
| 425 | return self._entry[1] |
| 426 | packed = rebuild_cuda_tensor(*list_args) |
| 427 | self._entry = (key, packed) |
| 428 | return packed |
| 429 | |
| 430 | def close(self) -> None: |
| 431 | """Release the held mapping (one producer-side refcount decrement).""" |
| 432 | self._entry = None |
| 433 | |
| 434 | |
| 435 | def packed_ipc_consumer( |
| 436 | ipc_handle: dict[str, tuple], |
| 437 | names: list[str], |
| 438 | shapes: list[list[int]], |
| 439 | dtype_names: list[str], |
| 440 | tensor_sizes: list[int], |
| 441 | device_index: int, |
| 442 | importer: "PackedBufferImporter | None" = None, |
| 443 | ) -> list[tuple[str, torch.Tensor]]: |
| 444 | """Unpack a single packed IPC chunk into named tensors. |
| 445 | |
| 446 | Reconstructs the packed buffer via rebuild_cuda_tensor, unpacks |
| 447 | into individual tensors, and clones each into independent storage |
| 448 | before returning. |
| 449 | |
| 450 | The clone is intentional: the producer reuses one IPC buffer across |
| 451 | chunks, so any tensor view that aliases the buffer would observe the |
| 452 | *next* chunk's bytes as soon as the producer's generator is resumed. |
| 453 | Callers that retain references past their own update_weights call |
| 454 | (notably vLLM's layerwise reload, which buffers ``bound_args`` for |
| 455 | replay in ``_layerwise_process``) would otherwise replay against |
| 456 | stale data and silently corrupt multi-chunk weight transfers. |
| 457 | |
| 458 | Args: |
| 459 | ipc_handle: Mapping of GPU UUID to rebuild_cuda_tensor args tuple |
| 460 | names: Parameter names in the packed buffer |
| 461 | shapes: Parameter shapes |
| 462 | dtype_names: Parameter dtype name strings (e.g. "float16") |
| 463 | tensor_sizes: Size in bytes of each parameter in the packed buffer |
| 464 | device_index: Local CUDA device index |
| 465 | importer: Import cache that must be shared across all chunks of one |
| 466 | export so the buffer mapping is opened and released exactly once |
| 467 | (see ``PackedBufferImporter``). ``None`` builds an ephemeral one, |
| 468 | which is only safe for single-chunk consumption. |
| 469 | |
| 470 | """ |
| 471 | props = torch.cuda.get_device_properties(device_index) |
| 472 | physical_gpu_id = str(props.uuid) |
| 473 | |
| 474 | if physical_gpu_id not in ipc_handle: |
| 475 | raise ValueError( |
| 476 | f"IPC handle not found for GPU UUID {physical_gpu_id}. " |
| 477 | f"Available UUIDs: {list(ipc_handle.keys())}" |
| 478 | ) |
| 479 | |
| 480 | args = ipc_handle[physical_gpu_id] |
| 481 | list_args = list(args) |
| 482 | list_args[6] = device_index |
| 483 | if importer is None: |
| 484 | importer = PackedBufferImporter() |
| 485 | packed = importer.rebuild(list_args) |
| 486 | |
| 487 | content_size = sum(tensor_sizes) |
| 488 | packed = packed[:content_size] |
| 489 | |
| 490 | dtypes = [getattr(torch, dn) for dn in dtype_names] |
| 491 | return [ |
| 492 | (name, t.clone()) |
| 493 | for name, t in unpack_tensor(packed, names, shapes, dtypes, tensor_sizes) |
| 494 | ] |