diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 17957924051..0b413965032 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import io +import pickle import random import threading import time @@ -10,12 +12,15 @@ from unittest import mock import multiprocess as mp import numpy as np import pytest +import torch import torch.distributed as dist from vllm.distributed.device_communicators import shm_broadcast from vllm.distributed.device_communicators.shm_broadcast import ( MessageQueue, ShmRingBuffer, + _rebuild_tensor, + _reduce_tensor, check_shm_free_space, ) from vllm.distributed.utils import StatelessProcessGroup @@ -354,6 +359,150 @@ def test_message_queue_busy_to_idle(): distributed_run(worker_fn_test_busy_to_idle, 4) +@worker_fn_wrapper +def worker_fn_tensor_broadcast(): + rank = dist.get_rank() + writer_rank = 0 + message_queue = MessageQueue.create_from_process_group( + dist.group.WORLD, 8 * 1024 * 1024, 4, writer_rank + ) + + # Both ranks construct the identical reference payload. + torch.manual_seed(42) + payload = { + # 2MiB: rides the shm ring as an out-of-band buffer (the receiving + # side must copy out of the reusable ring chunk). + "mid": torch.randn(1024, 512), + # 16MiB > max_chunk_bytes: overflows to the zmq socket (the + # receiving side aliases the zmq.Frame zero-copy). + "big": torch.randn(4096, 2048, dtype=torch.bfloat16), + "nested": ["plain", 123, {"inner": torch.arange(5)}], + } + + if rank == writer_rank: + with mock.patch( + "vllm.distributed.device_communicators.shm_broadcast._reduce_tensor", + wraps=_reduce_tensor, + ) as wrapped_reduce: + message_queue.enqueue(payload) + assert wrapped_reduce.call_count == 3 + # Cycle the ring (max_chunks=4) several times over so that aliased + # ring chunks would be overwritten. + for i in range(16): + message_queue.enqueue({"junk": torch.full((1024, 512), float(i))}) + else: + received = message_queue.dequeue(timeout=30) + for key in ("mid", "big"): + assert torch.equal(received[key], payload[key]), key + assert received[key].dtype == payload[key].dtype, key + assert torch.equal(received["nested"][2]["inner"], torch.arange(5)) + + snapshot = received["mid"].clone() + for i in range(16): + junk = message_queue.dequeue(timeout=30) + assert torch.equal(junk["junk"], torch.full((1024, 512), float(i))) + # Tensors received via the shm ring must not alias chunk memory + # that the writer has reused for subsequent messages. + assert torch.equal(received["mid"], snapshot) + # Rebuilt tensors must be writable, like regular tensors. + received["mid"] += 1.0 + received["big"][0, 0] = 1.0 + + dist.barrier() + print(f"tensor broadcast passed the test! Rank {rank}") + + +def test_tensor_broadcast(): + distributed_run(worker_fn_tensor_broadcast, 2) + + +def _dumps_oob(obj) -> tuple[bytes, list]: + """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor + dispatch table + out-of-band buffers >= 1MiB.""" + buffers = [] + + def callback(buf: pickle.PickleBuffer) -> bool: + raw = buf.raw() + if raw.nbytes < 1024 * 1024: + return True + buffers.append(raw) + return False + + bio = io.BytesIO() + pickler = pickle.Pickler( + bio, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=callback + ) + pickler.dispatch_table = {torch.Tensor: _reduce_tensor} + pickler.dump(obj) + return bio.getvalue(), buffers + + +@pytest.mark.parametrize( + "case", + [ + "small", + "mid", + "bf16", + "fp8", + "empty", + "scalar", + "noncontig", + "requires_grad", + "conj", + "param", + ], +) +def test_tensor_pickle_roundtrip(case: str): + tensor = { + # Inlined in-band (< 1MiB) and out-of-band (>= 1MiB) buffers. + "small": lambda: torch.randn(100, 10), + "mid": lambda: torch.randn(1024, 512), + # Dtypes numpy doesn't recognize. + "bf16": lambda: torch.randn(512, 512, dtype=torch.bfloat16), + "fp8": lambda: torch.randn(32, 32).to(torch.float8_e4m3fn), + # Shape edge cases. + "empty": lambda: torch.empty(0, 8), + "scalar": lambda: torch.tensor(3.14), + "noncontig": lambda: torch.randn(64, 64).t(), + # These fall back to torch's default reducer. + "requires_grad": lambda: torch.randn(8, 8, requires_grad=True), + "conj": lambda: torch.randn(4, dtype=torch.complex64).conj(), + "param": lambda: torch.nn.Parameter(torch.randn(4), requires_grad=False), + }[case]() + + data, buffers = _dumps_oob({"tensor": tensor, "meta": list(range(10))}) + received = pickle.loads(data, buffers=buffers)["tensor"] + + assert received.shape == tensor.shape + assert received.dtype == tensor.dtype + if tensor.dtype == torch.float8_e4m3fn: + assert torch.equal(received.view(torch.uint8), tensor.view(torch.uint8)) + else: + assert torch.equal(received, tensor) + assert received.requires_grad == tensor.requires_grad + assert isinstance(received, type(tensor)) + if tensor.numel() and not tensor.requires_grad: + # Rebuilt tensors must be writable, like regular tensors. + received.view(-1)[0] = 1.0 + + +@pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"]) +def test_reduce_tensor_fallback(case: str): + """Tensors the zero-copy reducer can't safely alias must fall back to + torch's default reduction.""" + if case == "cuda": + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + tensor = torch.randn(4, device="cuda") + elif case == "requires_grad": + tensor = torch.randn(8, requires_grad=True) + else: + tensor = torch.randn(4, dtype=torch.complex64).conj() + + reduced = _reduce_tensor(tensor) + assert reduced[0] is not _rebuild_tensor + + @pytest.mark.parametrize("should_warn", [False, True]) def test_reader_timeout_caps_indefinite_waits(should_warn): with ( diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index a788cecc6ce..52a95ce1d8d 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -48,6 +48,7 @@ CHECK_IMPORTS = { "vllm/distributed/device_communicators/shm_object_storage.py", "vllm/distributed/weight_transfer/ipc_engine.py", "vllm/distributed/weight_transfer/clients.py", + "tests/distributed/test_shm_broadcast.py", "tests/distributed/test_weight_transfer.py", "vllm/utils/hashing.py", "tests/multimodal/media/test_base.py", diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index afabdf18c80..e59b14f7a6b 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copyreg import functools +import io import os import pickle import shutil @@ -385,6 +387,70 @@ class ShmRingBuffer: yield buf +def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.Tensor: + """Rebuild a tensor from an out-of-band pickle buffer. + + Counterpart of `_reduce_tensor`. Note that pickle passes the original + buffer-providing object from `loads(buffers=...)` straight to this + function (no `PickleBuffer` wrapper on the receiving side), so `buf` is + a `zmq.Frame`, a `memoryview` of a shared-memory ring chunk, or `bytes` + if the buffer was serialized in-band. + """ + dtype = getattr(torch, dtype_str) + assert isinstance(dtype, torch.dtype) + if isinstance(buf, zmq.Frame): + # ZMQ frames own their message memory independently of any context, + # so the tensor can safely alias it with zero copies. The tensor's + # storage keeps the frame (and thus its bytes) alive via a strong + # reference for as long as the tensor is. + try: + return torch.frombuffer(buf, dtype=torch.uint8).view(dtype).view(shape) + except ValueError: + # Empty or read-only frame buffer; fall through to the copy path. + pass + # Shared-memory ring buffer chunks are reused by the writer once all + # readers have marked them read, so we must copy out of them. bytearray + # (vs bytes) keeps the resulting tensor writable, matching normal tensor + # semantics. + raw = bytearray(buf) + if not raw: + assert 0 in shape + return torch.empty(shape, dtype=dtype) + return torch.frombuffer(raw, dtype=torch.uint8).view(dtype).view(shape) + + +def _reduce_tensor(tensor: torch.Tensor): + """Reduce a CPU tensor to a `PickleBuffer` for out-of-band pickling. + + `torch.Tensor.__reduce_ex__` copies the tensor bytes into the pickle + byte stream via `torch.serialization` and never emits a `PickleBuffer`, + which defeats the out-of-band buffer handling in `MessageQueue.enqueue`. + This reducer instead exposes the tensor's memory directly, so large + tensors (e.g. `prompt_embeds` in `SchedulerOutput`) traverse the queue + without being copied into and back out of the pickled message. + """ + if ( + tensor.device.type == "cpu" + and tensor.layout == torch.strided + and not tensor.requires_grad + ): + try: + # The uint8 view exposes the raw bytes via the buffer protocol, + # including for dtypes numpy doesn't recognize (bfloat16, fp8, ...). + # reshape(-1) first so that 0-dim tensors can be viewed as well. + raw = tensor.contiguous().reshape(-1).view(torch.uint8).numpy() + except RuntimeError: + # Exotic tensors (e.g. with the conjugate bit set) that don't + # support aliasing views; let torch handle them. + pass + else: + dtype_str = str(tensor.dtype).removeprefix("torch.") + return _rebuild_tensor, (PickleBuffer(raw), tuple(tensor.shape), dtype_str) + + # Fall back to torch's default (copying) reduction. + return tensor.__reduce_ex__(pickle.HIGHEST_PROTOCOL) + + @dataclass class Handle: local_reader_ranks: list[int] = field(default_factory=list) @@ -771,9 +837,23 @@ class MessageQueue: total_bytes += len(raw_buf) + 4 return False - all_buffers[0] = pickle.dumps( - obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback - ) + # CPU tensors are routed through `_reduce_tensor` so that their + # bytes are emitted as out-of-band buffers instead of being + # copied into the pickle stream by torch's default reducer. + # Start from `copyreg.dispatch_table` to preserve globally + # registered reducers (e.g. `re.Pattern`); the per-pickler + # dispatch table would otherwise shadow them. + dispatch_table = dict(copyreg.dispatch_table) + dispatch_table[torch.Tensor] = _reduce_tensor + with io.BytesIO() as bio: + pickler = pickle.Pickler( + bio, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=oob_callback, + ) + pickler.dispatch_table = dispatch_table + pickler.dump(obj) + all_buffers[0] = bio.getvalue() if self.n_local_reader > 0: if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes: with self.acquire_write(timeout) as buf: