[Bugfix] Import each packed IPC export once on the consumer side (#51259)
Signed-off-by: acmore <hezj07+github@gmail.com>
acmore committed
Aug 11, 2026 at 10:05 UTC
f37dff22fc77e21ac857f2023faa7d5b8f9c9591
2 files changed
+59
-4
vllm/distributed/weight_transfer/ipc_engine.py
+12
-1
@@ -27,6 +27,7 @@ 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
)
@@ -143,6 +144,9 @@ class IPCWeightTransferEngine(
144
# Set from the trainer-supplied init info at the handshake; defaults are
145
# only for the (unreachable) receive-before-init case.
146
self.packed = False
147
+ # Shared across all chunks of every packed transfer this engine
148
+ # receives; see PackedBufferImporter for the refcount contract.
149
+ self._packed_importer = PackedBufferImporter()
150
151
def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None:
152
"""
@@ -170,6 +174,12 @@ class IPCWeightTransferEngine(
174
)
175
176
finalize_layerwise_reload(self.model, self.model_config)
177
+ # Every reduce_tensor call is a fresh export with its own refcount
178
+ # slot, so releasing once per update always balances this update's
179
+ # export and lets the trainer reclaim its staging buffer. Callers
180
+ # that skip finish are still covered by the replace-on-next-export
181
+ # path inside the importer.
182
+ self._packed_importer.close()
183
184
def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None:
185
"""
@@ -202,6 +212,7 @@ class IPCWeightTransferEngine(
212
dtype_names=update_info.dtype_names,
213
tensor_sizes=update_info.tensor_sizes,
214
device_index=device_index,
215
+ importer=self._packed_importer,
216
)
217
else:
218
assert isinstance(update_info.ipc_handles, list)
@@ -236,7 +247,7 @@ class IPCWeightTransferEngine(
247
self.model.load_weights(weights)
248
249
def shutdown(self) -> None:
239
- pass
250
+ self._packed_importer.close()
251
252
@staticmethod
253
def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
vllm/distributed/weight_transfer/packed_tensor.py
+47
-3
@@ -381,6 +381,45 @@ def packed_ipc_producer(
381
)
382
383
384
+class PackedBufferImporter:
385
+ """One-slot cache for the consumer-side mapping of a packed IPC buffer.
386
+
387
+ torch's cross-process refcount pairs one ``reduce_tensor`` export with
388
+ one consumer rebuild-release cycle. The producer exports its buffer once
389
+ and ships the same args with every chunk, so rebuilding and releasing
390
+ per chunk decrements the counter once per chunk, drives it negative,
391
+ and the dropped buffer is then never reclaimable by ``ipc_collect`` on
392
+ the producer side (it leaks one buffer per transfer). Caching the
393
+ rebuilt buffer keyed by the rebuild args restores the pairing: chunks
394
+ of one transfer reuse the mapping (it aliases producer memory, so it
395
+ always reads the current chunk's bytes), and replacing the entry when
396
+ the next export arrives — or ``close()`` — releases the previous
397
+ mapping exactly once.
398
+
399
+ The importer that consumes a multi-chunk transfer must be one object
400
+ reused across chunks; a fresh importer per chunk reintroduces the
401
+ per-chunk release. ``IPCWeightTransferEngine`` owns one instance per
402
+ engine and closes it on shutdown.
403
+ """
404
+
405
+ def __init__(self) -> None:
406
+ self._entry: tuple[tuple, torch.Tensor] | None = None
407
+
408
+ def rebuild(self, list_args: list) -> torch.Tensor:
409
+ from torch.multiprocessing.reductions import rebuild_cuda_tensor
410
+
411
+ key = tuple(list_args)
412
+ if self._entry is not None and self._entry[0] == key:
413
+ return self._entry[1]
414
+ packed = rebuild_cuda_tensor(*list_args)
415
+ self._entry = (key, packed)
416
+ return packed
417
+
418
+ def close(self) -> None:
419
+ """Release the held mapping (one producer-side refcount decrement)."""
420
+ self._entry = None
421
+
422
+
423
def packed_ipc_consumer(
424
ipc_handle: dict[str, tuple],
425
names: list[str],
@@ -388,6 +427,7 @@ def packed_ipc_consumer(
427
dtype_names: list[str],
428
tensor_sizes: list[int],
429
device_index: int,
430
+ importer: "PackedBufferImporter | None" = None,
431
) -> list[tuple[str, torch.Tensor]]:
432
"""Unpack a single packed IPC chunk into named tensors.
433
@@ -410,9 +450,11 @@ def packed_ipc_consumer(
450
dtype_names: Parameter dtype name strings (e.g. "float16")
451
tensor_sizes: Size in bytes of each parameter in the packed buffer
452
device_index: Local CUDA device index
453
+ importer: Import cache that must be shared across all chunks of one
454
+ export so the buffer mapping is opened and released exactly once
455
+ (see ``PackedBufferImporter``). ``None`` builds an ephemeral one,
456
+ which is only safe for single-chunk consumption.
457
"""
414
- from torch.multiprocessing.reductions import rebuild_cuda_tensor
415
-
458
props = torch.cuda.get_device_properties(device_index)
459
physical_gpu_id = str(props.uuid)
460
@@ -425,7 +467,9 @@ def packed_ipc_consumer(
467
args = ipc_handle[physical_gpu_id]
468
list_args = list(args)
469
list_args[6] = device_index
428
- packed = rebuild_cuda_tensor(*list_args)
470
+ if importer is None:
471
+ importer = PackedBufferImporter()
472
+ packed = importer.rebuild(list_args)
473
474
content_size = sum(tensor_sizes)
475
packed = packed[:content_size]