@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e0a45f1455
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 14:04 UTC

[Feat][RL] IPC weight sync optimizations: multigpu support and chunked packed tensors (#37476)

Signed-off-by: ahao-anyscale <ahao@anyscale.com> Signed-off-by: hao-aaron <ahao@anyscale.com>

Aaron Hao committed May 15, 2026 at 07:53 UTC e0a45f1455ce27693f7cbcac262900589e77e89f
10 files changed +1595 -459
docs/training/weight_transfer/ipc.md
+40 -9
@@ -1,21 +1,37 @@
1 # IPC Engine
2
3 -The IPC weight transfer engine uses **CUDA IPC** (Inter-Process Communication) handles to share GPU memory directly between the trainer and inference workers on the **same node and same GPU**. This avoids any data copying, making it a efficient option when colocating training and inference.
3 +The IPC weight transfer engine uses **CUDA IPC** (Inter-Process Communication) handles to share GPU memory directly between the trainer and inference workers on the **same GPU**. This avoids any data copying, making it the most efficient option when colocating training and inference. Multi-GPU setups are supported — weights are all gathered by each GPU and are extracted by the correct colocated process.
4
5 ## When to Use IPC
6
7 -- Training and inference on the **same GPU** (colocated)
8 -- You want to minimize memory overhead by sharing tensors in-place
7 +- Training and inference share the **same GPU(s)** (colocated)
8
9 ## How It Works
10
12 -1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`.
13 -2. IPC handles are sent to the inference engine via **Ray.remote()** or **HTTP POST**.
14 -3. The inference worker reconstructs the tensors from the handles, reading directly from the trainer's GPU memory.
11 +1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`. In multi-GPU setups (e.g. FSDP), each trainer rank must all-gather the full tensor for each layer onto its own GPU before generating the IPC handle.
12 +2. IPC handles for each gpu are sent to the inference engine via **Ray**, **HTTP**, or a **custom callable**. Each rank only reads the handle corresponding to its own GPU.
13 +3. The inference worker reconstructs the tensors from the handles using `rebuild_cuda_tensor`, reading directly from the trainer's GPU memory.
14
15 !!! warning
16 IPC handles involve sending serialized Python objects. When using HTTP transport, you must set `VLLM_ALLOW_INSECURE_SERIALIZATION=1` on both the server and client. This is because IPC handles are pickled and base64-encoded for HTTP transmission.
17
18 +## Packed (Chunked) Transfer
19 +
20 +By default, all weights are sent in a single API call. For large models, this requires the full model to reside in GPU memory on both sides simultaneously. Setting `packed=True` enables **chunked transfer** with bounded GPU memory:
21 +
22 +- Weights are concatenated into fixed-size packed buffers (controlled by `packed_buffer_size_bytes`).
23 +- Each chunk is sent as a separate `update_weights` call within a single `start_weight_update` / `finish_weight_update` bracket, so the layerwise reload pass is initialized once at the start and finalized once at the end regardless of chunk count.
24 +- After each chunk is consumed, the GPU memory for that chunk can be reclaimed.
25 +
26 +```python
27 +trainer_args = IPCTrainerSendWeightsArgs(
28 + send_mode="ray",
29 + llm_handle=llm_actor_handle,
30 + packed=True,
31 + packed_buffer_size_bytes=256 * 1024 * 1024, # 256 MB chunks
32 +)
33 +```
34 +
35 ## Initialization
36
37 The IPC backend requires no initialization on either side. The `init_transfer_engine` call is a no-op for IPC.
@@ -35,7 +51,7 @@ from vllm.distributed.weight_transfer.ipc_engine import (
51 )
52
53 trainer_args = IPCTrainerSendWeightsArgs(
38 - mode="ray",
54 + send_mode="ray",
55 llm_handle=llm_actor_handle,
56 )
57 # start
@@ -57,7 +73,7 @@ Used when vLLM is running as an HTTP server:
73
74 ```python
75 trainer_args = IPCTrainerSendWeightsArgs(
60 - mode="http",
76 + send_mode="http",
77 url="http://localhost:8000",
78 )
79
@@ -77,7 +93,22 @@ response = requests.post(url, json={}, timeout=60)
93 response.raise_for_status()
94 ```
95
80 -In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint. As with Ray mode, you must call `start_weight_update` before and `finish_weight_update` after.
96 +In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint. Because the worker deserializes the payload via `pickle.loads`, the vLLM server must be started with `VLLM_ALLOW_INSECURE_SERIALIZATION=1`.
97 +
98 +```python
99 +def my_custom_sender(update_info: IPCWeightTransferUpdateInfo):
100 + # Custom logic to deliver update_info to vLLM
101 + ...
102 +
103 +trainer_args = IPCTrainerSendWeightsArgs(
104 + send_mode=my_custom_sender,
105 +)
106 +
107 +IPCWeightTransferEngine.trainer_send_weights(
108 + iterator=model.named_parameters(),
109 + trainer_args=trainer_args,
110 +)
111 +```
112
113 See [`IPCTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/ipc_engine.py) for the full list of configurable fields.
114
examples/rl/rlhf_http_ipc.py
+1 -1
@@ -173,7 +173,7 @@ def main():
173 start_weight_update(BASE_URL, is_checkpoint_format=False)
174
175 print("Broadcasting weights via CUDA IPC (HTTP)...")
176 - trainer_args = IPCTrainerSendWeightsArgs(mode="http", url=BASE_URL)
176 + trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)
177 IPCWeightTransferEngine.trainer_send_weights(
178 iterator=train_model.named_parameters(),
179 trainer_args=trainer_args,
examples/rl/rlhf_ipc.py
+9 -5
@@ -70,10 +70,14 @@ class TrainModel:
70 self.llm_handle.init_weight_transfer_engine.remote(dict(init_info=dict()))
71 )
72
73 - def broadcast_weights(self, llm_handle: ray.actor.ActorHandle):
73 + def broadcast_weights(
74 + self, llm_handle: ray.actor.ActorHandle, packed: bool = False
75 + ):
76 """Broadcast weights to the inference engine using IPC."""
77 self.llm_handle = llm_handle
76 - trainer_args = IPCTrainerSendWeightsArgs(mode="ray", llm_handle=llm_handle)
78 + trainer_args = IPCTrainerSendWeightsArgs(
79 + send_mode="ray", llm_handle=llm_handle, packed=packed
80 + )
81 IPCWeightTransferEngine.trainer_send_weights(
82 iterator=self.train_model.named_parameters(),
83 trainer_args=trainer_args,
@@ -141,10 +145,10 @@ ray.get(llm.finish_weight_update.remote())
145
146 ray.get(llm.wake_up.remote(tags=["scheduling"]))
147
144 -# Generate text with the updated model.
145 -outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
148 +outputs_packed = ray.get(llm.generate.remote(prompts, sampling_params))
149 print("-" * 50)
147 -for output in outputs_updated:
150 +print("Results after packed/chunked IPC weight sync:")
151 +for output in outputs_packed:
152 prompt = output.prompt
153 generated_text = output.outputs[0].text
154 print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
examples/rl/rlhf_ipc_fsdp_ep.py new
+425
@@ -0,0 +1,425 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC**
5 +weight transfer and **packed** tensors.
6 +
7 +Layout (4 GPUs, TP=1, DP=4, EP):
8 + * One Ray placement group per GPU.
9 + * Each PG holds one FSDP training worker and one vLLM ``LLM`` instance
10 + (sync API) using fractional GPUs so both fit on the same device.
11 + * The 4 ``LLM`` instances form a DP group via env-var-based SPMD
12 + coordination (``VLLM_DP_RANK``, ``VLLM_DP_SIZE``, etc.), the same
13 + mechanism used by ``examples/offline_inference/data_parallel.py``.
14 + * A ``DataParallelInferenceEngine`` actor spawns all 4 LLM actors,
15 + waits for initialization, and orchestrates generation / weight-sync.
16 +
17 +Uses the built-in ``ray`` send_mode: each FSDP worker calls
18 +``trainer_send_weights`` targeting its colocated LLM actor.
19 +
20 +This example was run on 4xH100.
21 +"""
22 +
23 +from __future__ import annotations
24 +
25 +import os
26 +from dataclasses import asdict
27 +
28 +import ray
29 +import torch
30 +import torch.distributed as dist
31 +from huggingface_hub import snapshot_download
32 +from ray.util.placement_group import placement_group
33 +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
34 +from torch.distributed._tensor import DTensor
35 +from torch.distributed.fsdp import fully_shard
36 +from transformers import AutoModelForCausalLM
37 +
38 +from vllm import LLM, SamplingParams
39 +from vllm.config import WeightTransferConfig
40 +from vllm.distributed.weight_transfer.ipc_engine import (
41 + IPCTrainerSendWeightsArgs,
42 + IPCWeightTransferEngine,
43 + IPCWeightTransferInitInfo,
44 +)
45 +from vllm.utils.network_utils import get_ip, get_open_port
46 +
47 +TRAIN_GPU_FRACTION = float(os.environ.get("RLHF_IPC_TRAIN_GPU_FRACTION", "0.42"))
48 +VLLM_GPU_FRACTION = float(os.environ.get("RLHF_IPC_VLLM_GPU_FRACTION", "0.42"))
49 +
50 +MODEL_NAME = "Qwen/Qwen3-30B-A3B"
51 +
52 +FSDP_WORLD_SIZE = 4
53 +INFERENCE_TP_SIZE = 1
54 +INFERENCE_DP_SIZE = 4
55 +
56 +
57 +class MyLLM(LLM):
58 + """LLM subclass that configures DP env vars for SPMD coordination."""
59 +
60 + def __init__(
61 + self,
62 + *args,
63 + dp_rank: int = 0,
64 + dp_size: int = 1,
65 + dp_master_ip: str = "127.0.0.1",
66 + dp_master_port: int = 0,
67 + **kwargs,
68 + ):
69 + os.environ.pop("CUDA_VISIBLE_DEVICES", None)
70 + os.environ["VLLM_RAY_PER_WORKER_GPUS"] = str(VLLM_GPU_FRACTION)
71 + os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
72 + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
73 +
74 + os.environ["VLLM_DP_RANK"] = str(dp_rank)
75 + os.environ["VLLM_DP_RANK_LOCAL"] = str(dp_rank)
76 + os.environ["VLLM_DP_SIZE"] = str(dp_size)
77 + os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
78 + os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
79 +
80 + super().__init__(*args, **kwargs)
81 +
82 + def ready(self):
83 + return True
84 +
85 +
86 +@ray.remote(num_cpus=0, num_gpus=TRAIN_GPU_FRACTION)
87 +class FSDPTrainWorker:
88 + """One FSDP2 worker per GPU; colocated with vLLM DP rank via placement group."""
89 +
90 + def __init__(
91 + self,
92 + model_name: str,
93 + rank: int,
94 + fsdp_world_size: int,
95 + fsdp_master_addr: str,
96 + fsdp_master_port: int,
97 + ):
98 + self.rank = rank
99 +
100 + os.environ["MASTER_ADDR"] = fsdp_master_addr
101 + os.environ["MASTER_PORT"] = str(fsdp_master_port)
102 +
103 + dist.init_process_group(backend="nccl", rank=rank, world_size=fsdp_world_size)
104 + torch.accelerator.set_device_index(0)
105 +
106 + model = AutoModelForCausalLM.from_pretrained(
107 + model_name, torch_dtype=torch.bfloat16
108 + )
109 +
110 + self.weight_names = [n for n, _ in model.named_parameters()]
111 + self.weight_dtype_names = [
112 + str(p.dtype).split(".")[-1] for _, p in model.named_parameters()
113 + ]
114 + self.weight_shapes = [list(p.shape) for _, p in model.named_parameters()]
115 +
116 + for layer in model.model.layers:
117 + fully_shard(layer)
118 + fully_shard(model)
119 +
120 + self.model = model
121 +
122 + def get_rank(self):
123 + return self.rank
124 +
125 + def get_weight_metadata(self):
126 + return self.weight_names, self.weight_dtype_names, self.weight_shapes
127 +
128 + def gather_and_broadcast_weights_ipc(self, llm_handle, packed: bool = True):
129 + """All-gather full params; all ranks create IPC handles, rank 0 sends.
130 +
131 + All ranks must call trainer_send_weights so they participate in the
132 + all_gather_object collective inside _all_gather_and_merge_handles.
133 + Only rank 0 actually sends the payload to vLLM (gated by _is_rank_zero).
134 + """
135 +
136 + def _full_param_iter():
137 + # HF's Qwen3MoeExperts (and other recent HF MoE impls) packs
138 + # all experts into two fused 3-D tensors per layer:
139 + # experts.gate_up_proj shape (E, 2*I, H)
140 + # experts.down_proj shape (E, H, I)
141 + # vLLM's Qwen3MoE load_weights still expects the older
142 + # per-expert HF layout (experts.<i>.gate_proj.weight,
143 + # experts.<i>.up_proj.weight, experts.<i>.down_proj.weight),
144 + # so we un-fuse on the fly. Split order matches HF's forward:
145 + # gate, up = linear(x, gate_up_proj[i]).chunk(2, dim=-1)
146 + # → rows [:I] of gate_up_proj[i] are gate, rows [I:] are up.
147 + params = self.model.state_dict()
148 + for name in list(params.keys()):
149 + param = params.pop(name)
150 + if isinstance(param, DTensor):
151 + tensor = param.full_tensor().detach().contiguous()
152 + else:
153 + tensor = param.detach().contiguous()
154 + del param
155 +
156 + if name.endswith(".experts.gate_up_proj") and tensor.dim() == 3:
157 + prefix = name[: -len(".gate_up_proj")]
158 + num_experts, two_inter, _ = tensor.shape
159 + inter = two_inter // 2
160 + for i in range(num_experts):
161 + expert = tensor[i]
162 + yield (
163 + f"{prefix}.{i}.gate_proj.weight",
164 + expert[:inter].contiguous(),
165 + )
166 + yield (
167 + f"{prefix}.{i}.up_proj.weight",
168 + expert[inter:].contiguous(),
169 + )
170 + del tensor
171 + elif name.endswith(".experts.down_proj") and tensor.dim() == 3:
172 + prefix = name[: -len(".down_proj")]
173 + num_experts = tensor.shape[0]
174 + for i in range(num_experts):
175 + yield (
176 + f"{prefix}.{i}.down_proj.weight",
177 + tensor[i].contiguous(),
178 + )
179 + del tensor
180 + else:
181 + yield name, tensor
182 +
183 + trainer_args = IPCTrainerSendWeightsArgs(
184 + send_mode="ray",
185 + llm_handle=llm_handle,
186 + packed=packed,
187 + packed_buffer_size_bytes=1024 * 1024 * 1024, # 1 GB
188 + )
189 + IPCWeightTransferEngine.trainer_send_weights(
190 + iterator=_full_param_iter(),
191 + trainer_args=trainer_args,
192 + )
193 +
194 +
195 +@ray.remote(num_cpus=1)
196 +class DataParallelInferenceEngine:
197 + """Manages a pool of DP-sharded vLLM LLM actors.
198 +
199 + Spawns one MyLLM actor per placement group, waits for all engines to
200 + finish initializing, and exposes generation / weight-sync helpers.
201 + """
202 +
203 + def __init__(
204 + self,
205 + model: str,
206 + pgs: list,
207 + dp_master_ip: str,
208 + dp_master_port: int,
209 + ):
210 + dp_size = len(pgs)
211 + self.llm_actors = []
212 + for r in range(dp_size):
213 + sched = PlacementGroupSchedulingStrategy(
214 + placement_group=pgs[r],
215 + placement_group_capture_child_tasks=True,
216 + )
217 + actor = (
218 + ray.remote(num_cpus=0, num_gpus=0)(MyLLM)
219 + .options(scheduling_strategy=sched)
220 + .remote(
221 + model=model,
222 + enforce_eager=True,
223 + tensor_parallel_size=INFERENCE_TP_SIZE,
224 + distributed_executor_backend="ray",
225 + enable_expert_parallel=True,
226 + gpu_memory_utilization=0.35,
227 + weight_transfer_config=WeightTransferConfig(backend="ipc"),
228 + enable_sleep_mode=True,
229 + load_format="dummy",
230 + dp_rank=r,
231 + dp_size=dp_size,
232 + dp_master_ip=dp_master_ip,
233 + dp_master_port=dp_master_port,
234 + )
235 + )
236 + self.llm_actors.append(actor)
237 +
238 + ray.get([actor.ready.remote() for actor in self.llm_actors])
239 +
240 + def get_llm_actors(self):
241 + return self.llm_actors
242 +
243 + def generate(self, prompts: list[str], sampling_params):
244 + """Distribute prompts round-robin across DP ranks and collect results."""
245 + dp_size = len(self.llm_actors)
246 + per_rank: list[list[str]] = [[] for _ in range(dp_size)]
247 + indices: list[list[int]] = [[] for _ in range(dp_size)]
248 +
249 + for i, prompt in enumerate(prompts):
250 + rank = i % dp_size
251 + per_rank[rank].append(prompt)
252 + indices[rank].append(i)
253 +
254 + refs = [
255 + actor.generate.remote(per_rank[r], sampling_params)
256 + for r, actor in enumerate(self.llm_actors)
257 + if per_rank[r]
258 + ]
259 + all_outputs = ray.get(refs)
260 +
261 + ordered = [None] * len(prompts)
262 + rank_idx = 0
263 + for r in range(dp_size):
264 + if per_rank[r]:
265 + for local_i, orig_i in enumerate(indices[r]):
266 + ordered[orig_i] = all_outputs[rank_idx][local_i]
267 + rank_idx += 1
268 + return ordered
269 +
270 + def init_weight_transfer(self):
271 + ray.get(
272 + [
273 + actor.init_weight_transfer_engine.remote(
274 + dict(init_info=asdict(IPCWeightTransferInitInfo()))
275 + )
276 + for actor in self.llm_actors
277 + ]
278 + )
279 +
280 + def start_weight_update(self, is_checkpoint_format: bool = True):
281 + ray.get(
282 + [
283 + actor.start_weight_update.remote(
284 + is_checkpoint_format=is_checkpoint_format
285 + )
286 + for actor in self.llm_actors
287 + ]
288 + )
289 +
290 + def finish_weight_update(self):
291 + ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])
292 +
293 + def sleep(self, level: int = 0):
294 + ray.get([actor.sleep.remote(level=level) for actor in self.llm_actors])
295 +
296 + def wake_up(self, tags: list[str] | None = None):
297 + ray.get([actor.wake_up.remote(tags=tags) for actor in self.llm_actors])
298 +
299 +
300 +def main():
301 + ray.init(
302 + runtime_env={
303 + "env_vars": {
304 + "VLLM_ALLOW_INSECURE_SERIALIZATION": "1",
305 + }
306 + }
307 + )
308 +
309 + assert TRAIN_GPU_FRACTION + VLLM_GPU_FRACTION <= 1.0, (
310 + "Train + vLLM GPU fractions must sum to at most 1.0 per bundle."
311 + )
312 +
313 + local_model_path = snapshot_download(MODEL_NAME)
314 + print(f"[init] Model downloaded to {local_model_path}")
315 +
316 + fsdp_master_addr = get_ip()
317 + fsdp_master_port = get_open_port()
318 + dp_master_port = get_open_port()
319 + dp_master_ip = get_ip()
320 +
321 + # Create one placement group per DP rank (one GPU each).
322 + pgs = []
323 + for _ in range(INFERENCE_DP_SIZE):
324 + pg = placement_group([{"GPU": 1, "CPU": 1}])
325 + pgs.append(pg)
326 + ray.get([pg.ready() for pg in pgs])
327 + print(f"[init] {len(pgs)} placement groups ready.")
328 +
329 + # Launch FSDP training workers, one per PG.
330 + scheduling = [
331 + PlacementGroupSchedulingStrategy(
332 + placement_group=pgs[r],
333 + placement_group_capture_child_tasks=True,
334 + )
335 + for r in range(FSDP_WORLD_SIZE)
336 + ]
337 +
338 + fsdp_workers = [
339 + FSDPTrainWorker.options(scheduling_strategy=scheduling[r]).remote(
340 + local_model_path,
341 + r,
342 + FSDP_WORLD_SIZE,
343 + fsdp_master_addr,
344 + fsdp_master_port,
345 + )
346 + for r in range(FSDP_WORLD_SIZE)
347 + ]
348 + ray.get([w.get_rank.remote() for w in fsdp_workers])
349 + print(f"[init] {FSDP_WORLD_SIZE} FSDP workers ready.")
350 +
351 + # Launch DP inference engine (spawns and initializes all LLM actors).
352 + inference_engine = DataParallelInferenceEngine.remote(
353 + model=local_model_path,
354 + pgs=pgs,
355 + dp_master_ip=dp_master_ip,
356 + dp_master_port=dp_master_port,
357 + )
358 + llm_actors = ray.get(inference_engine.get_llm_actors.remote())
359 + print(f"[init] {INFERENCE_DP_SIZE} LLM actors ready.")
360 +
361 + prompts = [
362 + "Hello, my name is",
363 + "The president of the United States is",
364 + "The capital of France is",
365 + "The future of AI is",
366 + ]
367 + sampling_params = SamplingParams(temperature=0)
368 +
369 + print("[generate] Generating with dummy weights...")
370 + outputs = ray.get(inference_engine.generate.remote(prompts, sampling_params))
371 + print("-" * 60)
372 + print("BEFORE weight sync (dummy weights):")
373 + print("-" * 60)
374 + for output in outputs:
375 + print(f"Prompt: {output.prompt!r}")
376 + print(f"Generated: {output.outputs[0].text!r}")
377 + print("-" * 60)
378 +
379 + # --- Weight transfer ---
380 + print("[transfer] Initializing IPC weight transfer...")
381 + ray.get(inference_engine.init_weight_transfer.remote())
382 +
383 + # Two-phase sleep/wake pattern:
384 + # 1. sleep(level=1) — offload weights to CPU, discard KV cache
385 + # 2. wake_up(tags=["weights"]) — bring weights back to GPU (KV cache still free)
386 + # 3. IPC weight transfer — overwrite weights, plenty of room without KV cache
387 + # 4. wake_up(tags=["kv_cache"]) — re-allocate KV cache for inference
388 + print("[sync] Sleeping engines (offload weights + free KV cache)...")
389 + ray.get(inference_engine.sleep.remote(level=1))
390 +
391 + print("[sync] Waking weights (KV cache stays free)...")
392 + ray.get(inference_engine.wake_up.remote(tags=["weights"]))
393 +
394 + print("[sync] Starting weight update...")
395 + ray.get(inference_engine.start_weight_update.remote(is_checkpoint_format=True))
396 +
397 + print("[sync] Packed IPC transfer FSDP → vLLM...")
398 + ray.get(
399 + [
400 + w.gather_and_broadcast_weights_ipc.remote(llm_actors, packed=True)
401 + for w in fsdp_workers
402 + ]
403 + )
404 +
405 + ray.get(inference_engine.finish_weight_update.remote())
406 + print("[sync] Weight transfer complete.")
407 +
408 + print("[sync] Waking KV cache + scheduling...")
409 + ray.get(inference_engine.wake_up.remote(tags=["kv_cache", "scheduling"]))
410 +
411 + print("[generate] Generating with synced weights...")
412 + outputs_updated = ray.get(
413 + inference_engine.generate.remote(prompts, sampling_params)
414 + )
415 + print("-" * 60)
416 + print("AFTER weight sync (real weights):")
417 + print("-" * 60)
418 + for output in outputs_updated:
419 + print(f"Prompt: {output.prompt!r}")
420 + print(f"Generated: {output.outputs[0].text!r}")
421 + print("-" * 60)
422 +
423 +
424 +if __name__ == "__main__":
425 + main()
tests/distributed/test_packed_tensor.py
+478 -146
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for packed tensor broadcasting functionality.
4
5 -Unit tests for packed_broadcast_producer and packed_broadcast_consumer.
5 +Unit tests for packed_nccl_broadcast_producer and packed_nccl_broadcast_consumer.
6 These utilities enable efficient batched tensor transfer over NCCL.
7 """
8
@@ -11,8 +11,12 @@ import torch
11
12 from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferUpdateInfo
13 from vllm.distributed.weight_transfer.packed_tensor import (
14 - packed_broadcast_consumer,
15 - packed_broadcast_producer,
14 + pack_tensors,
15 + packed_ipc_consumer,
16 + packed_ipc_producer,
17 + packed_nccl_broadcast_consumer,
18 + packed_nccl_broadcast_producer,
19 + unpack_tensor,
20 )
21
22
@@ -90,91 +94,18 @@ class TestNCCLWeightTransferUpdateInfoPacked:
94 assert info.packed is True
95
96
93 -# --- Unit Tests: packed_broadcast_producer ---
97 +# --- Unit Tests: packed_nccl_broadcast_producer ---
98
99
100 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
101 class TestPackedBroadcastProducer:
98 - """Test packed_broadcast_producer function."""
99 -
100 - def test_producer_broadcasts_tensors(self):
101 - """Test that producer broadcasts all tensors."""
102 - params = create_mock_model_params()
103 - params_cuda = [(name, tensor.cuda()) for name, tensor in params]
104 -
105 - mock_group = MockCommunicationGroup()
106 -
107 - # Use a small target size to force multiple batches
108 - packed_broadcast_producer(
109 - iterator=iter(params_cuda),
110 - group=mock_group,
111 - src=0,
112 - post_iter_func=lambda x: x[1],
113 - buffer_size_bytes=500,
114 - )
115 -
116 - # Should have broadcasted some tensors
117 - assert mock_group.broadcast_count > 0
118 - assert len(mock_group.broadcasted_tensors) > 0
119 -
120 - def test_producer_single_large_tensor(self):
121 - """Test with a single tensor larger than target size."""
122 - # Create a large tensor
123 - large_tensor = torch.randn(1000, 1000, dtype=torch.float32).cuda()
124 - params = [("large_weight", large_tensor)]
125 -
126 - mock_group = MockCommunicationGroup()
127 -
128 - # Small target size to force the tensor to exceed it
129 - packed_broadcast_producer(
130 - iterator=iter(params),
131 - group=mock_group,
132 - src=0,
133 - post_iter_func=lambda x: x[1],
134 - buffer_size_bytes=100,
135 - )
136 -
137 - # Should still broadcast the tensor (at least 1 broadcast)
138 - assert mock_group.broadcast_count >= 1
139 - assert len(mock_group.broadcasted_tensors) >= 1
140 -
141 - # Verify the total broadcasted size matches the tensor
142 - expected_size = large_tensor.numel() * large_tensor.element_size()
143 - actual_size = sum(t.numel() for t in mock_group.broadcasted_tensors)
144 - assert actual_size == expected_size
145 -
146 - def test_producer_multiple_batches(self):
147 - """Test that tensors are properly batched when exceeding target size."""
148 - # Create many small tensors
149 - params = [
150 - (f"weight_{i}", torch.randn(10, 10, dtype=torch.float32).cuda())
151 - for i in range(20)
152 - ]
153 -
154 - mock_group = MockCommunicationGroup()
155 -
156 - # Small target size to force multiple batches
157 - packed_broadcast_producer(
158 - iterator=iter(params),
159 - group=mock_group,
160 - src=0,
161 - post_iter_func=lambda x: x[1],
162 - buffer_size_bytes=2000,
163 - )
164 -
165 - # Should have multiple broadcasts
166 - assert mock_group.broadcast_count > 1
167 -
168 - # Total size should match sum of all tensors
169 - expected_total = sum(t.numel() * t.element_size() for _, t in params)
170 - actual_total = sum(t.numel() for t in mock_group.broadcasted_tensors)
171 - assert actual_total == expected_total
102 + """Test packed_nccl_broadcast_producer function."""
103
104 def test_producer_empty_iterator(self):
105 """Test producer handles empty iterator gracefully."""
106 mock_group = MockCommunicationGroup()
107
177 - packed_broadcast_producer(
108 + packed_nccl_broadcast_producer(
109 iterator=iter([]),
110 group=mock_group,
111 src=0,
@@ -186,64 +117,6 @@ class TestPackedBroadcastProducer:
117 assert mock_group.broadcast_count == 0
118
119
189 -# --- Unit Tests: packed_broadcast_consumer ---
190 -
191 -
192 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
193 -class TestPackedBroadcastConsumer:
194 - """Test packed_broadcast_consumer function."""
195 -
196 - def test_consumer_receives_tensors(self):
197 - """Test that consumer receives and unpacks tensors."""
198 - params = create_mock_model_params()
199 - params_cuda = [(name, tensor.cuda()) for name, tensor in params]
200 -
201 - buffer_size = 2000
202 -
203 - # First, run producer to get the broadcasted tensors
204 - producer_group = MockCommunicationGroup()
205 -
206 - packed_broadcast_producer(
207 - iterator=iter(params_cuda),
208 - group=producer_group,
209 - src=0,
210 - post_iter_func=lambda x: x[1],
211 - buffer_size_bytes=buffer_size,
212 - )
213 -
214 - # Now run consumer with the broadcasted tensors
215 - consumer_group = MockConsumerCommunicationGroup(
216 - producer_group.broadcasted_tensors
217 - )
218 -
219 - state_dict_info = create_state_dict_info(params_cuda)
220 -
221 - unpacked_tensors = {}
222 -
223 - def post_unpack_func(tensor_list):
224 - for name, tensor in tensor_list:
225 - unpacked_tensors[name] = tensor.clone()
226 -
227 - packed_broadcast_consumer(
228 - iterator=iter(state_dict_info.items()),
229 - group=consumer_group,
230 - src=0,
231 - post_unpack_func=post_unpack_func,
232 - buffer_size_bytes=buffer_size,
233 - )
234 -
235 - # Verify all parameters were unpacked
236 - assert len(unpacked_tensors) == len(params)
237 -
238 - # Verify each tensor matches the original
239 - for name, original_tensor in params_cuda:
240 - assert name in unpacked_tensors
241 - unpacked = unpacked_tensors[name]
242 - assert unpacked.shape == original_tensor.shape
243 - assert unpacked.dtype == original_tensor.dtype
244 - assert torch.allclose(unpacked, original_tensor, rtol=1e-5, atol=1e-7)
245 -
246 -
120 # --- Integration Tests: Producer-Consumer Roundtrip ---
121
122
@@ -260,7 +133,7 @@ class TestPackedBroadcastRoundtrip:
133 buffer_size = 1000
134 producer_group = MockCommunicationGroup()
135
263 - packed_broadcast_producer(
136 + packed_nccl_broadcast_producer(
137 iterator=iter(params_cuda),
138 group=producer_group,
139 src=0,
@@ -279,7 +152,7 @@ class TestPackedBroadcastRoundtrip:
152 for name, tensor in tensor_list:
153 unpacked_tensors[name] = tensor.clone()
154
282 - packed_broadcast_consumer(
155 + packed_nccl_broadcast_consumer(
156 iterator=iter(state_dict_info.items()),
157 group=consumer_group,
158 src=0,
@@ -306,7 +179,7 @@ class TestPackedBroadcastRoundtrip:
179 buffer_size = 500
180 producer_group = MockCommunicationGroup()
181
309 - packed_broadcast_producer(
182 + packed_nccl_broadcast_producer(
183 iterator=iter(params),
184 group=producer_group,
185 src=0,
@@ -325,7 +198,7 @@ class TestPackedBroadcastRoundtrip:
198 for name, tensor in tensor_list:
199 unpacked_tensors[name] = tensor.clone()
200
328 - packed_broadcast_consumer(
201 + packed_nccl_broadcast_consumer(
202 iterator=iter(state_dict_info.items()),
203 group=consumer_group,
204 src=0,
@@ -341,7 +214,7 @@ class TestPackedBroadcastRoundtrip:
214 assert unpacked.dtype == original_tensor.dtype
215 assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
216
344 - @pytest.mark.parametrize("target_size", [100, 1000, 10000, 100000])
217 + @pytest.mark.parametrize("target_size", [100, 100000])
218 def test_roundtrip_different_batch_sizes(self, target_size):
219 """Test roundtrip with different target batch sizes."""
220 params = create_mock_model_params(num_layers=5)
@@ -349,7 +222,7 @@ class TestPackedBroadcastRoundtrip:
222
223 producer_group = MockCommunicationGroup()
224
352 - packed_broadcast_producer(
225 + packed_nccl_broadcast_producer(
226 iterator=iter(params_cuda),
227 group=producer_group,
228 src=0,
@@ -368,7 +241,7 @@ class TestPackedBroadcastRoundtrip:
241 for name, tensor in tensor_list:
242 unpacked_tensors[name] = tensor.clone()
243
371 - packed_broadcast_consumer(
244 + packed_nccl_broadcast_consumer(
245 iterator=iter(state_dict_info.items()),
246 group=consumer_group,
247 src=0,
@@ -407,7 +280,7 @@ class TestPackedBroadcastRoundtrip:
280 buffer_size = 500
281 producer_group = MockCommunicationGroup()
282
410 - packed_broadcast_producer(
283 + packed_nccl_broadcast_producer(
284 iterator=iter(params),
285 group=producer_group,
286 src=0,
@@ -426,7 +299,7 @@ class TestPackedBroadcastRoundtrip:
299 for name, tensor in tensor_list:
300 unpacked_tensors[name] = tensor.clone()
301
429 - packed_broadcast_consumer(
302 + packed_nccl_broadcast_consumer(
303 iterator=iter(state_dict_info.items()),
304 group=consumer_group,
305 src=0,
@@ -441,3 +314,462 @@ class TestPackedBroadcastRoundtrip:
314 assert unpacked.shape == original_tensor.shape
315 assert unpacked.dtype == original_tensor.dtype
316 assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
317 +
318 +
319 +# --- Unit Tests: unpack_tensor ---
320 +
321 +
322 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
323 +class TestUnpackTensor:
324 + """Test the shared unpack_tensor function."""
325 +
326 + def test_unpack_produces_independent_copies(self):
327 + """Verify unpacked tensors don't share memory with packed buffer."""
328 + original = torch.randn(10, dtype=torch.float32).cuda()
329 + packed = original.contiguous().view(torch.uint8).view(-1)
330 +
331 + result = unpack_tensor(
332 + packed,
333 + names=["w"],
334 + shapes=[[10]],
335 + dtypes=[torch.float32],
336 + tensor_sizes=[packed.numel()],
337 + )
338 +
339 + # Mutate the packed buffer
340 + packed.zero_()
341 +
342 + # Unpacked tensor should be unaffected
343 + assert torch.allclose(result[0][1], original)
344 +
345 +
346 +# --- Unit Tests: pack_tensors ---
347 +
348 +
349 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
350 +class TestPackTensors:
351 + """Test the shared pack_tensors function."""
352 +
353 + def test_pack_basic(self):
354 + """Test packing a few tensors into one buffer."""
355 + params = [
356 + ("w1", torch.randn(10, 20, dtype=torch.float32).cuda()),
357 + ("w2", torch.randn(5, dtype=torch.float16).cuda()),
358 + ]
359 +
360 + chunk = pack_tensors(
361 + iterator=iter(params),
362 + post_iter_func=lambda x: x[1],
363 + buffer_size_bytes=10_000_000,
364 + )
365 +
366 + assert chunk is not None
367 + assert len(chunk.names) == 2
368 + assert chunk.names == ["w1", "w2"]
369 + assert chunk.shapes == [[10, 20], [5]]
370 + assert chunk.dtypes == [torch.float32, torch.float16]
371 + assert chunk.packed_tensor.dtype == torch.uint8
372 +
373 + def test_pack_respects_buffer_limit(self):
374 + """Test that packing stops when buffer_size_bytes is exceeded."""
375 + params = [
376 + (f"w{i}", torch.randn(100, 100, dtype=torch.float32).cuda())
377 + for i in range(10)
378 + ]
379 +
380 + chunk = pack_tensors(
381 + iterator=iter(params),
382 + post_iter_func=lambda x: x[1],
383 + buffer_size_bytes=50_000,
384 + )
385 +
386 + assert chunk is not None
387 + assert len(chunk.names) < 10
388 +
389 + def test_pack_empty_iterator(self):
390 + """Test that an empty iterator returns None."""
391 + chunk = pack_tensors(
392 + iterator=iter([]),
393 + post_iter_func=lambda x: x[1],
394 + buffer_size_bytes=1000,
395 + )
396 + assert chunk is None
397 +
398 + def test_pack_single_tensor_larger_than_buffer_warns(self):
399 + """Test that a tensor exceeding buffer_size_bytes emits a warning."""
400 + big = torch.randn(1000, 1000, dtype=torch.float32).cuda()
401 + params = [("big", big)]
402 +
403 + import warnings
404 +
405 + with warnings.catch_warnings(record=True) as w:
406 + warnings.simplefilter("always")
407 + chunk = pack_tensors(
408 + iterator=iter(params),
409 + post_iter_func=lambda x: x[1],
410 + buffer_size_bytes=100,
411 + )
412 + assert chunk is not None
413 + assert len(chunk.names) == 1
414 + assert any("exceeds buffer_size_bytes" in str(wi.message) for wi in w)
415 +
416 + def test_pack_unpack_roundtrip(self):
417 + """Test pack then unpack produces identical tensors."""
418 + params = [
419 + ("a", torch.randn(8, 16, dtype=torch.float32).cuda()),
420 + ("b", torch.randn(4, dtype=torch.float16).cuda()),
421 + ("c", torch.randn(3, 5, 7, dtype=torch.bfloat16).cuda()),
422 + ]
423 +
424 + chunk = pack_tensors(
425 + iterator=iter(params),
426 + post_iter_func=lambda x: x[1],
427 + buffer_size_bytes=10_000_000,
428 + )
429 +
430 + assert chunk is not None
431 + result = unpack_tensor(
432 + chunk.packed_tensor,
433 + chunk.names,
434 + chunk.shapes,
435 + chunk.dtypes,
436 + chunk.tensor_sizes,
437 + )
438 +
439 + assert len(result) == len(params)
440 + for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
441 + assert orig_name == res_name
442 + assert res_tensor.shape == orig_tensor.shape
443 + assert res_tensor.dtype == orig_tensor.dtype
444 + assert torch.allclose(res_tensor, orig_tensor, rtol=1e-4, atol=1e-6)
445 +
446 + def test_pack_multiple_chunks(self):
447 + """Test consuming an iterator across multiple pack_tensors calls."""
448 + params = [
449 + (f"w{i}", torch.randn(50, 50, dtype=torch.float32).cuda()) for i in range(6)
450 + ]
451 + it = iter(params)
452 +
453 + all_names = []
454 + chunks = []
455 + while True:
456 + chunk = pack_tensors(it, lambda x: x[1], buffer_size_bytes=12_000)
457 + if chunk is None:
458 + break
459 + chunks.append(chunk)
460 + all_names.extend(chunk.names)
461 +
462 + assert len(chunks) > 1
463 + assert all_names == [f"w{i}" for i in range(6)]
464 +
465 +
466 +# --- Unit Tests: packed_ipc_producer ---
467 +
468 +
469 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
470 +class TestPackedIpcProducer:
471 + """Test the packed_ipc_producer generator."""
472 +
473 + def test_producer_yields_chunks(self):
474 + """Test that the producer yields PackedIpcChunk objects."""
475 + params = [
476 + (f"w{i}", torch.randn(50, 50, dtype=torch.float32).cuda()) for i in range(6)
477 + ]
478 +
479 + chunks = list(
480 + packed_ipc_producer(
481 + iterator=iter(params),
482 + gpu_uuid="test-uuid",
483 + post_iter_func=lambda x: x[1],
484 + buffer_size_bytes=12_000,
485 + )
486 + )
487 +
488 + assert len(chunks) > 1
489 +
490 + def test_producer_ipc_handle_has_uuid(self):
491 + """Test that each chunk's ipc_handle is keyed by the given UUID."""
492 + params = [("w", torch.randn(10, dtype=torch.float32).cuda())]
493 +
494 + chunks = list(
495 + packed_ipc_producer(
496 + iterator=iter(params),
497 + gpu_uuid="my-gpu-uuid",
498 + post_iter_func=lambda x: x[1],
499 + buffer_size_bytes=10_000_000,
500 + )
501 + )
502 +
503 + assert "my-gpu-uuid" in chunks[0].ipc_handle
504 +
505 + def test_producer_dtype_names_are_strings(self):
506 + """Test that dtype_names are string representations."""
507 + params = [
508 + ("a", torch.randn(10, dtype=torch.float32).cuda()),
509 + ("b", torch.randn(10, dtype=torch.float16).cuda()),
510 + ]
511 +
512 + chunks = list(
513 + packed_ipc_producer(
514 + iterator=iter(params),
515 + gpu_uuid="uuid",
516 + post_iter_func=lambda x: x[1],
517 + buffer_size_bytes=10_000_000,
518 + )
519 + )
520 +
521 + assert chunks[0].dtype_names == ["float32", "float16"]
522 +
523 + def test_producer_empty_iterator(self):
524 + """Test producer with empty iterator yields nothing."""
525 + chunks = list(
526 + packed_ipc_producer(
527 + iterator=iter([]),
528 + gpu_uuid="uuid",
529 + post_iter_func=lambda x: x[1],
530 + buffer_size_bytes=1000,
531 + )
532 + )
533 + assert len(chunks) == 0
534 +
535 +
536 +# --- Integration Tests: IPC Producer-Consumer Roundtrip ---
537 +
538 +
539 +def _ipc_consumer_worker(cmd_q, ack_q, result_q, done_event, device_index):
540 + """Worker that consumes chunks streamed one at a time from the parent.
541 +
542 + CUDA IPC requires the consumer to be in a separate process from the
543 + producer. The producer reuses a single IPC buffer between chunks, so
544 + the parent must wait for our ack (sent after we copy the chunk to
545 + CPU) before advancing the producer.
546 + """
547 + try:
548 + torch.accelerator.set_device_index(device_index)
549 + all_results = []
550 + while True:
551 + cd = cmd_q.get()
552 + if cd is None:
553 + break
554 + result = packed_ipc_consumer(
555 + ipc_handle=cd["ipc_handle"],
556 + names=cd["names"],
557 + shapes=cd["shapes"],
558 + dtype_names=cd["dtype_names"],
559 + tensor_sizes=cd["tensor_sizes"],
560 + device_index=device_index,
561 + )
562 + # .cpu() forces a GPU→CPU copy off the shared IPC buffer, so
563 + # the producer is free to overwrite it once we ack.
564 + all_results.extend([(name, tensor.cpu()) for name, tensor in result])
565 + del result
566 + ack_q.put("ack")
567 + result_q.put(("ok", all_results))
568 + except Exception as e:
569 + result_q.put(("error", str(e)))
570 + # Keep the process alive until the parent has finished reading from
571 + # the result queue — torch serializes CPU tensors via fd sharing,
572 + # which requires this process's resource-sharer server to be running.
573 + done_event.wait(timeout=60)
574 +
575 +
576 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
577 +class TestPackedIpcRoundtrip:
578 + """Test IPC producer-consumer roundtrip using real CUDA IPC.
579 +
580 + These tests spawn a child process for the consumer because
581 + rebuild_cuda_tensor requires a separate process from the one that
582 + called reduce_tensor.
583 + """
584 +
585 + def _get_gpu_uuid(self) -> str:
586 + device_index = torch.cuda.current_device()
587 + props = torch.cuda.get_device_properties(device_index)
588 + return str(props.uuid)
589 +
590 + def _run_roundtrip(self, chunk_iter, device_index, timeout=30):
591 + """Stream chunks through a child consumer one at a time.
592 +
593 + ``packed_ipc_producer`` reuses a single IPC buffer for every
594 + chunk, so the producer must not be advanced until the consumer
595 + has finished reading the current chunk. We enforce that with an
596 + ack queue: the consumer puts ``"ack"`` after it has copied the
597 + chunk to CPU, and only then do we pull the next chunk from the
598 + generator.
599 +
600 + Returns ``(num_chunks, results)``.
601 + """
602 + import multiprocessing as mp
603 +
604 + ctx = mp.get_context("spawn")
605 + cmd_q = ctx.Queue()
606 + ack_q = ctx.Queue()
607 + result_q = ctx.Queue()
608 + done_event = ctx.Event()
609 + proc = ctx.Process(
610 + target=_ipc_consumer_worker,
611 + args=(cmd_q, ack_q, result_q, done_event, device_index),
612 + )
613 + proc.start()
614 +
615 + num_chunks = 0
616 + try:
617 + for chunk in chunk_iter:
618 + cmd_q.put(
619 + {
620 + "ipc_handle": chunk.ipc_handle,
621 + "names": chunk.names,
622 + "shapes": chunk.shapes,
623 + "dtype_names": chunk.dtype_names,
624 + "tensor_sizes": chunk.tensor_sizes,
625 + }
626 + )
627 + if ack_q.get(timeout=timeout) != "ack":
628 + raise RuntimeError("Consumer did not ack chunk")
629 + num_chunks += 1
630 + cmd_q.put(None)
631 + status, payload = result_q.get(timeout=timeout)
632 + finally:
633 + done_event.set()
634 + proc.join(timeout=10)
635 + if proc.is_alive():
636 + proc.kill()
637 +
638 + if status == "error":
639 + raise RuntimeError(f"Consumer process failed: {payload}")
640 + # Reclaim IPC-shared memory now that the child has released it
641 + torch.cuda.ipc_collect()
642 + return num_chunks, payload
643 +
644 + def test_roundtrip_basic(self):
645 + """Test basic IPC producer -> consumer roundtrip."""
646 + params = [
647 + ("w1", torch.randn(10, 20, dtype=torch.float32).cuda()),
648 + ("w2", torch.randn(5, dtype=torch.float16).cuda()),
649 + ]
650 + gpu_uuid = self._get_gpu_uuid()
651 + device_index = torch.cuda.current_device()
652 +
653 + num_chunks, result = self._run_roundtrip(
654 + packed_ipc_producer(
655 + iterator=iter(params),
656 + gpu_uuid=gpu_uuid,
657 + post_iter_func=lambda x: x[1],
658 + buffer_size_bytes=10_000_000,
659 + ),
660 + device_index,
661 + )
662 +
663 + assert num_chunks == 1
664 + assert len(result) == 2
665 + for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
666 + assert orig_name == res_name
667 + assert res_tensor.shape == orig_tensor.shape
668 + assert res_tensor.dtype == orig_tensor.dtype
669 + assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-4, atol=1e-6)
670 +
671 + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
672 + def test_roundtrip_dtypes(self, dtype):
673 + """Test IPC roundtrip with different dtypes."""
674 + params = create_mock_model_params(num_layers=2, dtype=dtype)
675 + params_cuda = [(n, t.cuda()) for n, t in params]
676 + gpu_uuid = self._get_gpu_uuid()
677 + device_index = torch.cuda.current_device()
678 +
679 + _, result = self._run_roundtrip(
680 + packed_ipc_producer(
681 + iterator=iter(params_cuda),
682 + gpu_uuid=gpu_uuid,
683 + post_iter_func=lambda x: x[1],
684 + buffer_size_bytes=10_000_000,
685 + ),
686 + device_index,
687 + )
688 +
689 + assert len(result) == len(params_cuda)
690 + for (orig_name, orig_tensor), (res_name, res_tensor) in zip(
691 + params_cuda, result
692 + ):
693 + assert orig_name == res_name
694 + assert res_tensor.dtype == dtype
695 + assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-4, atol=1e-6)
696 +
697 + def test_roundtrip_multiple_chunks(self):
698 + """Test IPC roundtrip across multiple chunks."""
699 + params = [
700 + (f"layer{i}.weight", torch.randn(100, 100, dtype=torch.float32).cuda())
701 + for i in range(8)
702 + ]
703 + gpu_uuid = self._get_gpu_uuid()
704 + device_index = torch.cuda.current_device()
705 +
706 + num_chunks, result = self._run_roundtrip(
707 + packed_ipc_producer(
708 + iterator=iter(params),
709 + gpu_uuid=gpu_uuid,
710 + post_iter_func=lambda x: x[1],
711 + buffer_size_bytes=50_000,
712 + ),
713 + device_index,
714 + )
715 +
716 + assert num_chunks > 1
717 + assert len(result) == len(params)
718 + for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
719 + assert orig_name == res_name
720 + assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-5, atol=1e-7)
721 +
722 + def test_roundtrip_non_contiguous(self):
723 + """Test IPC roundtrip with non-contiguous tensors."""
724 + params = [
725 + ("transposed", torch.randn(20, 10, dtype=torch.float32).cuda().T),
726 + ("sliced", torch.randn(40, 30, dtype=torch.float16).cuda()[::2, ::2]),
727 + ]
728 + gpu_uuid = self._get_gpu_uuid()
729 + device_index = torch.cuda.current_device()
730 +
731 + for _, t in params:
732 + assert not t.is_contiguous()
733 +
734 + _, result = self._run_roundtrip(
735 + packed_ipc_producer(
736 + iterator=iter(params),
737 + gpu_uuid=gpu_uuid,
738 + post_iter_func=lambda x: x[1],
739 + buffer_size_bytes=10_000_000,
740 + ),
741 + device_index,
742 + )
743 +
744 + for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
745 + assert orig_name == res_name
746 + assert res_tensor.shape == orig_tensor.shape
747 + assert res_tensor.dtype == orig_tensor.dtype
748 + assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-4, atol=1e-6)
749 +
750 + def test_consumer_wrong_uuid_raises(self):
751 + """Test that consumer raises ValueError for unknown GPU UUID."""
752 + params = [("w", torch.randn(10, dtype=torch.float32).cuda())]
753 + gpu_uuid = self._get_gpu_uuid()
754 +
755 + chunks = list(
756 + packed_ipc_producer(
757 + iterator=iter(params),
758 + gpu_uuid=gpu_uuid,
759 + post_iter_func=lambda x: x[1],
760 + buffer_size_bytes=10_000_000,
761 + )
762 + )
763 +
764 + c = chunks[0]
765 + fake_handle = {"fake-uuid-12345": c.ipc_handle[gpu_uuid]}
766 +
767 + with pytest.raises(ValueError, match="IPC handle not found"):
768 + packed_ipc_consumer(
769 + ipc_handle=fake_handle,
770 + names=c.names,
771 + shapes=c.shapes,
772 + dtype_names=c.dtype_names,
773 + tensor_sizes=c.tensor_sizes,
774 + device_index=torch.cuda.current_device(),
775 + )
tests/distributed/test_weight_transfer.py
+54 -75
@@ -389,7 +389,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
389
390 # Create a dummy tensor and IPC handle
391 dummy_tensor = torch.ones(10, 10, device="cuda:0")
392 - ipc_handle = reduce_tensor(dummy_tensor)
392 + _, ipc_handle = reduce_tensor(dummy_tensor)
393 gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
394 ipc_handles = [{gpu_uuid: ipc_handle}]
395
@@ -410,7 +410,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
410 pytest.skip("Need at least 1 GPU for this test")
411
412 dummy_tensor = torch.ones(10, 10, device="cuda:0")
413 - ipc_handle = reduce_tensor(dummy_tensor)
413 + _, ipc_handle = reduce_tensor(dummy_tensor)
414 gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
415 ipc_handles = [{gpu_uuid: ipc_handle}, {gpu_uuid: ipc_handle}]
416
@@ -428,7 +428,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
428 pytest.skip("Need at least 1 GPU for this test")
429
430 dummy_tensor = torch.ones(10, 10, device="cuda:0")
431 - ipc_handle = reduce_tensor(dummy_tensor)
431 + _, ipc_handle = reduce_tensor(dummy_tensor)
432 gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
433 ipc_handles = [{gpu_uuid: ipc_handle}, {gpu_uuid: ipc_handle}]
434
@@ -446,7 +446,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
446 pytest.skip("Need at least 1 GPU for this test")
447
448 dummy_tensor = torch.ones(10, 10, device="cuda:0")
449 - ipc_handle = reduce_tensor(dummy_tensor)
449 + _, ipc_handle = reduce_tensor(dummy_tensor)
450 gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
451 ipc_handles = [{gpu_uuid: ipc_handle}] # Only one handle
452
@@ -458,65 +458,9 @@ class TestIPCWeightTransferUpdateInfoValidation:
458 ipc_handles=ipc_handles,
459 )
460
461 - def test_valid_update_info_from_pickled(self, monkeypatch):
462 - """Test creating IPCWeightTransferUpdateInfo from pickled handles."""
463 - if torch.accelerator.device_count() < 1:
464 - pytest.skip("Need at least 1 GPU for this test")
465 -
466 - monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
467 -
468 - dummy_tensor = torch.ones(10, 10, device="cuda:0")
469 - ipc_handle = reduce_tensor(dummy_tensor)
470 - gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
471 - ipc_handles = [{gpu_uuid: ipc_handle}]
472 -
473 - pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
474 -
475 - info = IPCWeightTransferUpdateInfo(
476 - names=["layer.weight"],
477 - dtype_names=["float32"],
478 - shapes=[[10, 10]],
479 - ipc_handles_pickled=pickled,
480 - )
481 - assert info.ipc_handles == ipc_handles
482 - assert info.ipc_handles_pickled is None
483 -
484 - def test_pickled_requires_insecure_serialization_flag(self, monkeypatch):
485 - """Test that pickled handles are rejected unless env flag is enabled."""
486 - monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")
487 -
488 - with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"):
489 - IPCWeightTransferUpdateInfo(
490 - names=[],
491 - dtype_names=[],
492 - shapes=[],
493 - ipc_handles_pickled=base64.b64encode(pickle.dumps([])).decode("utf-8"),
494 - )
495 -
496 - def test_both_handles_and_pickled_raises(self):
497 - """Test that providing both ipc_handles and ipc_handles_pickled raises."""
498 - if torch.accelerator.device_count() < 1:
499 - pytest.skip("Need at least 1 GPU for this test")
500 -
501 - dummy_tensor = torch.ones(10, 10, device="cuda:0")
502 - ipc_handle = reduce_tensor(dummy_tensor)
503 - gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
504 - ipc_handles = [{gpu_uuid: ipc_handle}]
505 -
506 - pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
507 -
508 - with pytest.raises(ValueError, match="Cannot specify both"):
509 - IPCWeightTransferUpdateInfo(
510 - names=["layer.weight"],
511 - dtype_names=["float32"],
512 - shapes=[[10, 10]],
513 - ipc_handles=ipc_handles,
514 - ipc_handles_pickled=pickled,
515 - )
516 -
517 - def test_neither_handles_nor_pickled_raises(self):
518 - """Test that providing neither ipc_handles nor ipc_handles_pickled raises."""
519 - with pytest.raises(ValueError, match="must be provided"):
461 + def test_missing_ipc_handles_raises(self):
462 + """Test that omitting ipc_handles raises TypeError."""
463 + with pytest.raises(TypeError):
464 IPCWeightTransferUpdateInfo(
465 names=["layer.weight"],
466 dtype_names=["float32"],
@@ -552,10 +496,10 @@ class TestIPCEngineParsing:
496 # Create dummy IPC handles
497 dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
498 dummy_tensor2 = torch.ones(50, device="cuda:0")
555 - ipc_handle1 = reduce_tensor(dummy_tensor1)
556 - ipc_handle2 = reduce_tensor(dummy_tensor2)
499 + _, ipc_args1 = reduce_tensor(dummy_tensor1)
500 + _, ipc_args2 = reduce_tensor(dummy_tensor2)
501 gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
558 - ipc_handles = [{gpu_uuid: ipc_handle1}, {gpu_uuid: ipc_handle2}]
502 + ipc_handles = [{gpu_uuid: ipc_args1}, {gpu_uuid: ipc_args2}]
503
504 update_info = engine.parse_update_info(
505 {
@@ -585,10 +529,10 @@ class TestIPCEngineParsing:
529
530 dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
531 dummy_tensor2 = torch.ones(50, device="cuda:0")
588 - ipc_handle1 = reduce_tensor(dummy_tensor1)
589 - ipc_handle2 = reduce_tensor(dummy_tensor2)
532 + _, ipc_args1 = reduce_tensor(dummy_tensor1)
533 + _, ipc_args2 = reduce_tensor(dummy_tensor2)
534 gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
591 - ipc_handles = [{gpu_uuid: ipc_handle1}, {gpu_uuid: ipc_handle2}]
535 + ipc_handles = [{gpu_uuid: ipc_args1}, {gpu_uuid: ipc_args2}]
536
537 pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
538
@@ -604,10 +548,36 @@ class TestIPCEngineParsing:
548 assert isinstance(update_info, IPCWeightTransferUpdateInfo)
549 assert update_info.names == ["w1", "w2"]
550 assert len(update_info.ipc_handles) == 2
607 - assert update_info.ipc_handles_pickled is None
551 assert gpu_uuid in update_info.ipc_handles[0]
552 assert gpu_uuid in update_info.ipc_handles[1]
553
554 + def test_parse_update_info_both_handles_and_pickled_raises(self):
555 + """Test that providing both ipc_handles and ipc_handles_pickled raises."""
556 + if torch.accelerator.device_count() < 1:
557 + pytest.skip("Need at least 1 GPU for this test")
558 +
559 + config = WeightTransferConfig(backend="ipc")
560 + parallel_config = create_mock_parallel_config()
561 + engine = IPCWeightTransferEngine(config, parallel_config)
562 +
563 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
564 + _, ipc_handle = reduce_tensor(dummy_tensor)
565 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
566 + ipc_handles = [{gpu_uuid: ipc_handle}]
567 +
568 + pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
569 +
570 + with pytest.raises(ValueError, match="Cannot specify both"):
571 + engine.parse_update_info(
572 + {
573 + "names": ["layer.weight"],
574 + "dtype_names": ["float32"],
575 + "shapes": [[10, 10]],
576 + "ipc_handles": ipc_handles,
577 + "ipc_handles_pickled": pickled,
578 + }
579 + )
580 +
581
582 # --- Integration Test: IPC Weight Transfer Between Ray Tasks ---
583
@@ -629,13 +599,15 @@ class TrainerActor:
599 self.tensor.fill_(42.0) # Fill with 42 to verify correct transfer
600
601 # Create IPC handle (tensor must stay alive for IPC to work)
632 - ipc_handle = reduce_tensor(self.tensor)
602 + # reduce_tensor returns (rebuild_func, args); we only send args
603 + # since the receiver imports rebuild_cuda_tensor directly.
604 + _, ipc_args = reduce_tensor(self.tensor)
605 gpu_uuid = get_physical_gpu_id(0)
606
607 torch.accelerator.synchronize()
608
609 self.ipc_handle_dict = {
638 - "ipc_handle": ipc_handle,
610 + "ipc_handle": ipc_args,
611 "gpu_uuid": gpu_uuid,
612 "shape": tensor_shape,
613 "dtype": tensor_dtype,
@@ -652,6 +624,12 @@ def inference_receive_ipc_tensor(
624 mode: str = "ray",
625 ) -> dict:
626 """Inference task that receives tensor via IPCWeightTransferEngine."""
627 + import os
628 +
629 + # Worker-side: ipc_handles_pickled is deserialized via pickle.
630 + if mode == "http":
631 + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
632 +
633 from unittest.mock import MagicMock
634
635 import torch
@@ -684,7 +662,6 @@ def inference_receive_ipc_tensor(
662 # Clone tensor to keep it after engine cleans up
663 received_tensors.append((name, tensor.clone()))
664
687 - # Build update dict and go through parse_update_info (exercises __post_init__)
665 ipc_handles = [{ipc_handle_dict["gpu_uuid"]: ipc_handle_dict["ipc_handle"]}]
666
667 if mode == "ray":
@@ -695,6 +672,7 @@ def inference_receive_ipc_tensor(
672 "ipc_handles": ipc_handles,
673 }
674 elif mode == "http":
675 + # Simulate HTTP transport: pickle + base64 encode handles
676 pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
677 update_dict = {
678 "names": ["test.weight"],
@@ -743,7 +721,8 @@ def test_ipc_weight_transfer_between_processes(mode: str):
721
722 Parametrized over transport modes:
723 - 'ray': ipc_handles passed directly.
746 - - 'http': ipc_handles pickled + base64-encoded, unpickled via __post_init__.
724 + - 'http': ipc_handles pickled + base64-encoded, deserialized in
725 + parse_update_info before constructing the dataclass.
726
727 IPC requires same-GPU access, so we use a placement group to co-locate
728 the trainer actor and inference task on the same GPU.
@@ -801,7 +780,7 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():
780
781 # Create IPC handle with wrong GPU UUID
782 dummy_tensor = torch.ones(10, 10, device="cuda:0")
804 - ipc_handle = reduce_tensor(dummy_tensor)
783 + _, ipc_handle = reduce_tensor(dummy_tensor)
784 wrong_uuid = "wrong-uuid-12345"
785 ipc_handles = [{wrong_uuid: ipc_handle}]
786
tools/pre_commit/check_torch_cuda.py
+6 -1
@@ -15,7 +15,12 @@ _TORCH_CUDA_PATTERNS = [
15 r"\bcuda_device_count_stateless\(\)\b",
16 ]
17
18 -ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}
18 +ALLOWED_FILES = {
19 + "vllm/platforms/",
20 + "vllm/device_allocator/",
21 + "vllm/distributed/weight_transfer/ipc_engine.py",
22 + "tests/distributed/test_packed_tensor.py",
23 +}
24
25
26 def scan_file(path: str) -> int:
vllm/distributed/weight_transfer/ipc_engine.py
+295 -146
@@ -8,9 +8,10 @@ from dataclasses import asdict, dataclass
8 from typing import Any
9
10 import pybase64 as base64
11 +import ray
12 import requests
13 import torch
13 -from torch.multiprocessing.reductions import reduce_tensor
14 +from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor
15
16 from vllm import envs
17 from vllm.config.parallel import ParallelConfig
@@ -20,27 +21,43 @@ from vllm.distributed.weight_transfer.base import (
21 WeightTransferInitInfo,
22 WeightTransferUpdateInfo,
23 )
24 +from vllm.distributed.weight_transfer.packed_tensor import (
25 + DEFAULT_PACKED_BUFFER_SIZE_BYTES,
26 + packed_ipc_consumer,
27 + packed_ipc_producer,
28 +)
29
30
31 @dataclass
32 class IPCTrainerSendWeightsArgs:
33 """Arguments for IPC trainer_send_weights method."""
34
29 - mode: str
30 - """Transport mode: 'http' or 'ray'."""
35 + send_mode: str | Callable[["IPCWeightTransferUpdateInfo"], None]
36 + """How to send updates to vLLM. Either a string ('ray' or 'http') for
37 + built-in transports, or a callable that receives an
38 + IPCWeightTransferUpdateInfo and performs the send."""
39 llm_handle: Any = None
32 - """Ray ObjectRef to LLM handle (required for 'ray' mode)."""
40 + """Ray actor handle or list of handles (required for 'ray' send_mode)."""
41 url: str | None = None
34 - """Base URL for HTTP endpoint (required for 'http' mode)."""
42 + """Base URL for HTTP endpoint (required for 'http' send_mode)."""
43 + packed: bool = False
44 + """Whether to use packed tensor transfer for bounded-memory chunking."""
45 + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
46 + """Size in bytes for each packed tensor buffer when packed=True."""
47
48 def __post_init__(self):
49 """Validate that required arguments are provided for the selected mode."""
38 - if self.mode == "ray" and self.llm_handle is None:
39 - raise ValueError("llm_handle is required for 'ray' mode")
40 - if self.mode == "http" and self.url is None:
41 - raise ValueError("url is required for 'http' mode")
42 - if self.mode not in ("ray", "http"):
43 - raise ValueError(f"mode must be 'ray' or 'http', got {self.mode}")
50 + if callable(self.send_mode):
51 + return
52 + if self.send_mode == "ray" and self.llm_handle is None:
53 + raise ValueError("llm_handle is required for 'ray' send_mode")
54 + if self.send_mode == "http" and self.url is None:
55 + raise ValueError("url is required for 'http' send_mode")
56 + if self.send_mode not in ("ray", "http"):
57 + raise ValueError(
58 + f"send_mode must be 'ray', 'http', or a callable, "
59 + f"got {self.send_mode!r}"
60 + )
61
62
63 @dataclass
@@ -52,44 +69,22 @@ class IPCWeightTransferInitInfo(WeightTransferInitInfo):
69
70 @dataclass
71 class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
55 - """Update info for IPC weight transfer backend.
56 -
57 - Accepts IPC handles either directly via ``ipc_handles`` (Ray transport)
58 - or as a base64-encoded pickle via ``ipc_handles_pickled`` (HTTP transport).
59 - Exactly one of the two must be provided; if ``ipc_handles_pickled`` is set
60 - it is unpickled into ``ipc_handles`` during ``__post_init__``.
61 - """
72 + """Update info for IPC weight transfer backend."""
73
74 names: list[str]
75 dtype_names: list[str]
76 shapes: list[list[int]]
66 - ipc_handles: list[dict[str, tuple[Callable, tuple]]] | None = None
67 - """IPC handles mapping physical GPU UUID to (func, args) tuple.
68 - Each handle is a dictionary mapping GPU UUID strings to IPC handle tuples."""
69 - ipc_handles_pickled: str | None = None
70 - """Base64-encoded pickled IPC handles, used for HTTP transport."""
77 + ipc_handles: list[dict[str, tuple]] | dict[str, tuple]
78 + """IPC handles mapping physical GPU UUID to rebuild_cuda_tensor args.
79 + For non-packed mode: list of per-parameter handle dicts.
80 + For packed mode: single handle dict for the packed buffer."""
81 + tensor_sizes: list[int] | None = None
82 + """Per-parameter sizes in bytes within the packed buffer.
83 + Required when packed=True, unused otherwise."""
84 + packed: bool = False
85 + """Whether this update uses packed tensor format."""
86
87 def __post_init__(self):
73 - if self.ipc_handles_pickled is not None:
74 - if self.ipc_handles is not None:
75 - raise ValueError(
76 - "Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
77 - )
78 -
79 - if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
80 - raise ValueError(
81 - "Refusing to deserialize `ipc_handles_pickled` without "
82 - "VLLM_ALLOW_INSECURE_SERIALIZATION=1"
83 - )
84 -
85 - self.ipc_handles = pickle.loads(base64.b64decode(self.ipc_handles_pickled))
86 - self.ipc_handles_pickled = None
87 -
88 - if self.ipc_handles is None:
89 - raise ValueError(
90 - "Either `ipc_handles` or `ipc_handles_pickled` must be provided"
91 - )
92 -
88 num_params = len(self.names)
89 if len(self.dtype_names) != num_params:
90 raise ValueError(
@@ -101,11 +96,17 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
96 f"`shapes` should be of the same size as `names`: "
97 f"got {len(self.shapes)} and {len(self.names)}"
98 )
104 - if len(self.ipc_handles) != num_params:
99 + if (
100 + not self.packed
101 + and isinstance(self.ipc_handles, list)
102 + and len(self.ipc_handles) != num_params
103 + ):
104 raise ValueError(
105 f"`ipc_handles` should be of the same size as `names`: "
106 f"got {len(self.ipc_handles)} and {len(self.names)}"
107 )
108 + if self.packed and self.tensor_sizes is None:
109 + raise ValueError("`tensor_sizes` is required when packed=True")
110
111
112 class IPCWeightTransferEngine(
@@ -135,6 +136,36 @@ class IPCWeightTransferEngine(
136 """
137 super().__init__(config, parallel_config)
138
139 + def parse_update_info(
140 + self, update_dict: dict[str, Any]
141 + ) -> IPCWeightTransferUpdateInfo:
142 + """Parse update dict, deserializing pickled IPC handles if present.
143 +
144 + HTTP transport sends IPC handles as a base64-encoded pickle under the
145 + key ``ipc_handles_pickled``. This method deserializes them back into
146 + ``ipc_handles`` before constructing the typed dataclass, keeping
147 + serialization concerns out of the dataclass itself.
148 +
149 + Requires ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` because the
150 + payload is deserialized via ``pickle.loads``.
151 + """
152 + if "ipc_handles_pickled" in update_dict:
153 + if "ipc_handles" in update_dict:
154 + raise ValueError(
155 + "Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
156 + )
157 +
158 + if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
159 + raise ValueError(
160 + "Refusing to deserialize `ipc_handles_pickled` without "
161 + "VLLM_ALLOW_INSECURE_SERIALIZATION=1"
162 + )
163 +
164 + pickled = update_dict.pop("ipc_handles_pickled")
165 + update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled))
166 +
167 + return super().parse_update_info(update_dict)
168 +
169 def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None:
170 """
171 Initialize the weight transfer mechanism.
@@ -157,46 +188,52 @@ class IPCWeightTransferEngine(
188 Args:
189 update_info: IPC update info containing parameter names, dtypes, shapes,
190 and IPC handles. Each IPC handle is a mapping between physical
160 - GPU UUID and the IPC handle tuple (func, args).
191 + GPU UUID and the rebuild_cuda_tensor args tuple.
192 load_weights: Callable that loads weights into the model. Called
193 incrementally for each weight to avoid OOM.
194 """
164 - assert update_info.ipc_handles is not None
165 - weights = []
166 - for name, _dtype_name, _shape, ipc_handle in zip(
167 - update_info.names,
168 - update_info.dtype_names,
169 - update_info.shapes,
170 - update_info.ipc_handles,
171 - ):
172 - device_index = torch.accelerator.current_device_index()
173 - props = torch.cuda.get_device_properties(device_index)
174 - physical_gpu_id = str(props.uuid)
175 -
176 - if physical_gpu_id not in ipc_handle:
177 - raise ValueError(
178 - f"IPC handle not found for GPU UUID {physical_gpu_id}. "
179 - f"Available UUIDs: {list(ipc_handle.keys())}"
180 - )
181 -
182 - handle = ipc_handle[physical_gpu_id]
183 -
184 - func, args = handle
185 - list_args = list(args) # type: ignore
186 - # Index 6 is the device_index parameter in torch's
187 - # IPC handle tuple (rebuild_cuda_tensor). Update it
188 - # to the current device since the logical index can
189 - # differ between sender and receiver.
190 - list_args[6] = device_index
191 - weight = func(*list_args) # type: ignore
192 - weights.append((name, weight))
195 + device_index = torch.accelerator.current_device_index()
196
194 - load_weights(weights)
197 + if update_info.packed:
198 + assert update_info.tensor_sizes is not None
199 + assert isinstance(update_info.ipc_handles, dict)
200 + weights = packed_ipc_consumer(
201 + ipc_handle=update_info.ipc_handles,
202 + names=update_info.names,
203 + shapes=update_info.shapes,
204 + dtype_names=update_info.dtype_names,
205 + tensor_sizes=update_info.tensor_sizes,
206 + device_index=device_index,
207 + )
208 + load_weights(weights)
209 + else:
210 + assert isinstance(update_info.ipc_handles, list)
211 + weights = []
212 + for name, ipc_handle in zip(
213 + update_info.names,
214 + update_info.ipc_handles,
215 + ):
216 + props = torch.cuda.get_device_properties(device_index)
217 + physical_gpu_id = str(props.uuid)
218 +
219 + if physical_gpu_id not in ipc_handle:
220 + raise ValueError(
221 + f"IPC handle not found for GPU UUID "
222 + f"{physical_gpu_id}. "
223 + f"Available UUIDs: {list(ipc_handle.keys())}"
224 + )
225 +
226 + args = ipc_handle[physical_gpu_id]
227 + list_args = list(args)
228 + # Index 6 of the args from reduce_tensor is the device_index.
229 + # We need to overwrite it with the receiver's device index.
230 + list_args[6] = device_index
231 + weight = rebuild_cuda_tensor(*list_args)
232 + weights.append((name, weight))
233 +
234 + load_weights(weights)
235
236 def shutdown(self) -> None:
197 - """
198 - Shutdown the weight transfer engine.
199 - """
237 pass
238
239 @staticmethod
@@ -204,12 +241,17 @@ class IPCWeightTransferEngine(
241 iterator: Iterator[tuple[str, torch.Tensor]],
242 trainer_args: dict[str, Any] | IPCTrainerSendWeightsArgs,
243 ) -> None:
207 - """
208 - Send weights from trainer to inference workers via CUDA IPC.
244 + """Send weights from trainer to inference workers via CUDA IPC.
245 +
246 + Supports two transport modes ('ray' and 'http') and two transfer
247 + strategies:
248 + - Non-packed (default): all weights in a single API call.
249 + - Packed (packed=True): chunked transfer with bounded GPU memory.
250
210 - Supports two modes:
211 - - 'ray': Sends weights via Ray RPC to a Ray-based LLM handle
212 - - 'http': Sends weights via HTTP POST to a vLLM HTTP server
251 + For multi-GPU training, all ranks must call this method in
252 + parallel. IPC handles are all-gathered across ranks and merged
253 + so that each vLLM worker can find its own GPU UUID. Only rank 0
254 + sends the payload to vLLM.
255
256 .. note::
257 This method calls ``update_weights`` internally. The caller must
@@ -217,88 +259,195 @@ class IPCWeightTransferEngine(
259 after this method.
260
261 Args:
220 - iterator: Iterator of model parameters. Returns (name, tensor) tuples.
221 - Tensors should be on the same GPU as the inference workers.
222 - trainer_args: Dictionary containing IPC-specific arguments.
223 - Should contain keys from IPCTrainerSendWeightsArgs:
224 - - mode: 'ray' or 'http'
225 - - llm_handle: Ray ObjectRef (for 'ray' mode)
226 - - url: Base URL string (for 'http' mode)
227 -
228 - Example (Ray mode):
229 - >>> from vllm.distributed.weight_transfer.ipc_engine import (
230 - ... IPCWeightTransferEngine,
231 - ... IPCTrainerSendWeightsArgs,
232 - ... )
233 - >>> param_iter = ((n, p) for n, p in model.named_parameters())
234 - >>> args = IPCTrainerSendWeightsArgs(mode="ray", llm_handle=llm_handle)
235 - >>> IPCWeightTransferEngine.trainer_send_weights(param_iter, asdict(args))
236 -
237 - Example (HTTP mode):
238 - >>> args = IPCTrainerSendWeightsArgs(
239 - ... mode="http", url="http://localhost:8000"
240 - ... )
241 - >>> IPCWeightTransferEngine.trainer_send_weights(param_iter, asdict(args))
262 + iterator: Iterator of (name, tensor) pairs. For multi-GPU,
263 + each rank should yield the full tensor on its own GPU
264 + (e.g. via FSDP full_tensor()).
265 + trainer_args: IPCTrainerSendWeightsArgs or equivalent dict.
266 """
243 - # Parse trainer args - accept either dict or dataclass instance
244 - if isinstance(trainer_args, dict):
245 - args = IPCTrainerSendWeightsArgs(**trainer_args)
267 + args = (
268 + IPCTrainerSendWeightsArgs(**trainer_args)
269 + if isinstance(trainer_args, dict)
270 + else trainer_args
271 + )
272 + device_index = torch.accelerator.current_device_index()
273 + gpu_uuid = str(torch.cuda.get_device_properties(device_index).uuid)
274 + if args.packed:
275 + IPCWeightTransferEngine._send_packed(iterator, args, gpu_uuid)
276 else:
247 - args = trainer_args
277 + IPCWeightTransferEngine._send_unpacked(iterator, args, gpu_uuid)
278
249 - # Get physical GPU UUID
250 - device_index = torch.accelerator.current_device_index()
251 - props = torch.cuda.get_device_properties(device_index)
252 - gpu_uuid = str(props.uuid)
279 + @staticmethod
280 + def _is_rank_zero() -> bool:
281 + """Return True if this is rank 0 or no distributed group exists."""
282 + if not torch.distributed.is_initialized():
283 + return True
284 + return torch.distributed.get_rank() == 0
285 +
286 + @staticmethod
287 + def _all_gather_and_merge_handles(
288 + handles: list[dict[str, tuple]],
289 + ) -> list[dict[str, tuple]]:
290 + """All-gather and merge IPC handle dicts across ranks in one call.
291 +
292 + Each rank contributes a list of {gpu_uuid: ipc_args} dicts (one
293 + per parameter or one per chunk). A single all_gather_object
294 + collects every rank's full list, then rank 0 merges per-index so
295 + each dict maps every GPU UUID to its args.
296 +
297 + Non-rank-0 returns a list of empty dicts.
298 + No-op (returns handles unchanged) when no distributed group exists.
299 + """
300 + if (
301 + not torch.distributed.is_initialized()
302 + or torch.distributed.get_world_size() == 1
303 + ):
304 + return handles
305 +
306 + world_size = torch.distributed.get_world_size()
307 + gathered: list[list[dict[str, tuple]] | None] = [None] * world_size
308 + torch.distributed.all_gather_object(gathered, handles)
309 + torch.distributed.barrier()
310 + torch.cuda.synchronize()
311 +
312 + if torch.distributed.get_rank() == 0:
313 + merged: list[dict[str, tuple]] = []
314 + for param_idx in range(len(handles)):
315 + m: dict[str, tuple] = {}
316 + for rank_handles in gathered:
317 + if rank_handles is not None:
318 + m.update(rank_handles[param_idx])
319 + merged.append(m)
320 + return merged
321 + return [{} for _ in handles]
322
254 - # Collect weight metadata and create IPC handles
255 - names = []
256 - dtype_names = []
257 - shapes = []
258 - ipc_handles = []
323 + @staticmethod
324 + def _post_send_sync() -> None:
325 + """Barrier + ipc_collect after a send; no-op if single-GPU."""
326 + if (
327 + torch.distributed.is_initialized()
328 + and torch.distributed.get_world_size() > 1
329 + ):
330 + torch.distributed.barrier()
331 + torch.cuda.ipc_collect()
332 +
333 + @staticmethod
334 + def _send_unpacked(
335 + iterator: Iterator[tuple[str, torch.Tensor]],
336 + args: IPCTrainerSendWeightsArgs,
337 + gpu_uuid: str,
338 + ) -> None:
339 + """Send all weights in a single API call (non-packed mode)."""
340 + names: list[str] = []
341 + dtype_names: list[str] = []
342 + shapes: list[list[int]] = []
343 + ipc_handles: list[dict[str, tuple]] = []
344 + # Hold strong refs to every contiguous copy until the send + post-send
345 + # sync completes. reduce_tensor's returned args do NOT keep storage
346 + # alive, and non-contiguous inputs allocate fresh storage in
347 + # .contiguous() that would otherwise be GC'd before the consumer opens
348 + # the IPC handle.
349 + weight_refs: list[torch.Tensor] = []
350
351 for name, tensor in iterator:
352 names.append(name)
353 dtype_names.append(str(tensor.dtype).split(".")[-1])
354 shapes.append(list(tensor.shape))
355
265 - # Create IPC handle for this weight tensor
266 - # The tensor must remain in memory for IPC to work
356 weight = tensor.detach().contiguous()
268 - ipc_handle = reduce_tensor(weight)
269 - ipc_handles.append({gpu_uuid: ipc_handle})
270 -
271 - # Send weights based on mode
272 - if args.mode == "ray":
273 - # Ray mode: send via Ray RPC
274 - import ray
275 -
276 - update_info = asdict(
277 - IPCWeightTransferUpdateInfo(
278 - names=names,
279 - dtype_names=dtype_names,
280 - shapes=shapes,
281 - ipc_handles=ipc_handles,
357 + weight_refs.append(weight)
358 + _, ipc_args = reduce_tensor(weight)
359 + ipc_handles.append({gpu_uuid: ipc_args})
360 +
361 + ipc_handles = IPCWeightTransferEngine._all_gather_and_merge_handles(ipc_handles)
362 +
363 + if IPCWeightTransferEngine._is_rank_zero():
364 + IPCWeightTransferEngine._do_send(
365 + args=args,
366 + names=names,
367 + dtype_names=dtype_names,
368 + shapes=shapes,
369 + ipc_handles=ipc_handles,
370 + )
371 +
372 + IPCWeightTransferEngine._post_send_sync()
373 +
374 + @staticmethod
375 + def _send_packed(
376 + iterator: Iterator[tuple[str, torch.Tensor]],
377 + args: IPCTrainerSendWeightsArgs,
378 + gpu_uuid: str,
379 + ) -> None:
380 + """Send weights in bounded-memory chunks (packed mode)."""
381 + post_iter_func: Callable = lambda item: item[1]
382 +
383 + for chunk in packed_ipc_producer(
384 + iterator=iterator,
385 + gpu_uuid=gpu_uuid,
386 + post_iter_func=post_iter_func,
387 + buffer_size_bytes=args.packed_buffer_size_bytes,
388 + ):
389 + ipc_handle = IPCWeightTransferEngine._all_gather_and_merge_handles(
390 + [chunk.ipc_handle]
391 + )[0]
392 +
393 + if IPCWeightTransferEngine._is_rank_zero():
394 + IPCWeightTransferEngine._do_send(
395 + args=args,
396 + names=chunk.names,
397 + dtype_names=chunk.dtype_names,
398 + shapes=chunk.shapes,
399 + ipc_handles=ipc_handle,
400 + tensor_sizes=chunk.tensor_sizes,
401 + packed=True,
402 )
403 +
404 + IPCWeightTransferEngine._post_send_sync()
405 +
406 + @staticmethod
407 + def _do_send(
408 + args: IPCTrainerSendWeightsArgs,
409 + names: list[str],
410 + dtype_names: list[str],
411 + shapes: list[list[int]],
412 + ipc_handles: list[dict[str, tuple]] | dict[str, tuple],
413 + tensor_sizes: list[int] | None = None,
414 + packed: bool = False,
415 + ) -> None:
416 + """Send a single update payload via the configured transport."""
417 + update_fields: dict[str, Any] = {
418 + "names": names,
419 + "dtype_names": dtype_names,
420 + "shapes": shapes,
421 + "packed": packed,
422 + }
423 + if tensor_sizes is not None:
424 + update_fields["tensor_sizes"] = tensor_sizes
425 +
426 + update_fields["ipc_handles"] = ipc_handles
427 + update_info = IPCWeightTransferUpdateInfo(**update_fields)
428 +
429 + if callable(args.send_mode):
430 + args.send_mode(update_info)
431 + elif args.send_mode == "ray":
432 + handles = (
433 + args.llm_handle
434 + if isinstance(args.llm_handle, list)
435 + else [args.llm_handle]
436 )
437 ray.get(
285 - args.llm_handle.update_weights.remote(dict(update_info=update_info))
438 + [
439 + h.update_weights.remote(dict(update_info=asdict(update_info)))
440 + for h in handles
441 + ]
442 )
287 - elif args.mode == "http":
288 - # HTTP mode: send via HTTP POST with pickled handles
289 - # Pickle and base64 encode IPC handles for HTTP transmission
443 + elif args.send_mode == "http":
444 pickled_handles = base64.b64encode(pickle.dumps(ipc_handles)).decode(
445 "utf-8"
446 )
447 + http_fields = {k: v for k, v in update_fields.items() if k != "ipc_handles"}
448 + http_fields["ipc_handles_pickled"] = pickled_handles
449
450 url = f"{args.url}/update_weights"
295 - payload = {
296 - "update_info": {
297 - "names": names,
298 - "dtype_names": dtype_names,
299 - "shapes": shapes,
300 - "ipc_handles_pickled": pickled_handles,
301 - }
302 - }
451 + payload = {"update_info": http_fields}
452 response = requests.post(url, json=payload, timeout=300)
453 response.raise_for_status()
vllm/distributed/weight_transfer/nccl_engine.py
+4 -4
@@ -21,7 +21,7 @@ from vllm.distributed.weight_transfer.base import (
21 from vllm.distributed.weight_transfer.packed_tensor import (
22 DEFAULT_PACKED_BUFFER_SIZE_BYTES,
23 DEFAULT_PACKED_NUM_BUFFERS,
24 - packed_broadcast_consumer,
24 + packed_nccl_broadcast_consumer,
25 )
26
27
@@ -184,7 +184,7 @@ class NCCLWeightTransferEngine(
184 dtype = getattr(torch, dtype_name)
185 yield (name, (shape, dtype))
186
187 - packed_broadcast_consumer(
187 + packed_nccl_broadcast_consumer(
188 iterator=state_dict_info_iterator(),
189 group=self.model_update_group,
190 src=0,
@@ -247,10 +247,10 @@ class NCCLWeightTransferEngine(
247 if args.packed:
248 # Use packed tensor broadcasting for efficiency
249 from vllm.distributed.weight_transfer.packed_tensor import (
250 - packed_broadcast_producer,
250 + packed_nccl_broadcast_producer,
251 )
252
253 - packed_broadcast_producer(
253 + packed_nccl_broadcast_producer(
254 iterator=iterator,
255 group=args.group,
256 src=args.src,
vllm/distributed/weight_transfer/packed_tensor.py
+283 -72
@@ -4,9 +4,11 @@
4
5 import math
6 from collections.abc import Callable, Iterator
7 +from dataclasses import dataclass
8 from typing import Any
9
10 import torch
11 +from torch.multiprocessing.reductions import reduce_tensor
12
13 # Default values for packed tensor configuration.
14 # These are imported by NCCLWeightTransferUpdateInfo and trainer_send_weights.
@@ -14,7 +16,124 @@ DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 # 1GB
16 DEFAULT_PACKED_NUM_BUFFERS = 2
17
18
17 -def packed_broadcast_producer(
19 +def unpack_tensor(
20 + packed_tensor: torch.Tensor,
21 + names: list[str],
22 + shapes: list[list[int]],
23 + dtypes: list[torch.dtype],
24 + tensor_sizes: list[int],
25 +) -> list[tuple[str, torch.Tensor]]:
26 + """Unpack a packed uint8 tensor into a list of named tensors.
27 +
28 + The returned tensors are **views** of ``packed_tensor`` (the
29 + ``.contiguous()`` call is a no-op on already-contiguous row-slices).
30 + If ``packed_tensor`` lives in storage that may be reused — e.g. a
31 + reused CUDA IPC buffer — callers must clone the results before the
32 + underlying storage is overwritten.
33 +
34 + Args:
35 + packed_tensor: The packed torch.uint8 tensor to unpack
36 + names: List of tensor names
37 + shapes: List of tensor shapes
38 + dtypes: List of tensor dtypes
39 + tensor_sizes: List of tensor sizes in bytes
40 + """
41 + unpacked_tensors = packed_tensor.split(tensor_sizes)
42 +
43 + return [
44 + (name, tensor.contiguous().view(dtype).view(*shape))
45 + for name, shape, dtype, tensor in zip(names, shapes, dtypes, unpacked_tensors)
46 + ]
47 +
48 +
49 +@dataclass
50 +class PackedChunk:
51 + """Result of packing tensors into a single contiguous uint8 buffer."""
52 +
53 + packed_tensor: torch.Tensor
54 + names: list[str]
55 + shapes: list[list[int]]
56 + dtypes: list[torch.dtype]
57 + tensor_sizes: list[int]
58 +
59 +
60 +def pack_tensors(
61 + iterator: Iterator[tuple[str, torch.Tensor]],
62 + post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor],
63 + buffer_size_bytes: int,
64 + tensor_list: list[torch.Tensor] | None = None,
65 + current_size: int = 0,
66 +) -> PackedChunk | None:
67 + """Pack tensors from an iterator into a single contiguous uint8 buffer.
68 +
69 + Consumes from the iterator until the accumulated size exceeds
70 + buffer_size_bytes or the iterator is exhausted, then returns a
71 + PackedChunk. Returns None if no tensors were consumed.
72 +
73 + Args:
74 + iterator: Iterator of (name, tensor) pairs
75 + post_iter_func: Applied to each item before linearizing to uint8
76 + buffer_size_bytes: Max bytes before flushing
77 + tensor_list: Pre-existing tensor list to append to (for NCCL
78 + multi-buffer reuse). If None, a fresh list is created.
79 + current_size: Byte count already accumulated in tensor_list
80 + """
81 + if tensor_list is None:
82 + tensor_list = []
83 +
84 + names: list[str] = []
85 + shapes: list[list[int]] = []
86 + dtypes: list[torch.dtype] = []
87 + tensor_sizes: list[int] = []
88 + total_bytes = current_size
89 +
90 + while True:
91 + try:
92 + item = next(iterator)
93 + except StopIteration:
94 + break
95 +
96 + name, orig_tensor = item
97 + # Apply post processing and convert to linearized uint8 tensor
98 + tensor = post_iter_func(item).contiguous().view(torch.uint8).view(-1)
99 +
100 + if tensor.numel() > buffer_size_bytes:
101 + import warnings
102 +
103 + warnings.warn(
104 + f"Tensor '{name}' has size {tensor.numel()} bytes, which "
105 + f"exceeds buffer_size_bytes={buffer_size_bytes}.",
106 + stacklevel=2,
107 + )
108 +
109 + tensor_list.append(tensor)
110 + names.append(name)
111 + shapes.append(list(orig_tensor.shape))
112 + dtypes.append(orig_tensor.dtype)
113 + tensor_sizes.append(tensor.numel())
114 + total_bytes += tensor.numel()
115 +
116 + if total_bytes > buffer_size_bytes:
117 + break
118 +
119 + if not tensor_list:
120 + return None
121 +
122 + packed = torch.cat(tensor_list, dim=0)
123 + del tensor_list
124 + return PackedChunk(
125 + packed_tensor=packed,
126 + names=names,
127 + shapes=shapes,
128 + dtypes=dtypes,
129 + tensor_sizes=tensor_sizes,
130 + )
131 +
132 +
133 +# ── NCCL packed broadcast ──────────────────────────────────────────────
134 +
135 +
136 +def packed_nccl_broadcast_producer(
137 iterator: Iterator[tuple[str, torch.Tensor]],
138 group: Any,
139 src: int,
@@ -36,57 +155,31 @@ def packed_broadcast_producer(
155 Both producer and consumer must use the same value.
156
157 """
39 - target_packed_tensor_size = buffer_size_bytes
40 -
158 streams = [torch.cuda.Stream() for _ in range(num_buffers)]
159 + # Keep references to in-flight chunks so their packed_tensors
160 + # aren't freed while an async broadcast is still reading them.
161 + in_flight: list[PackedChunk | None] = [None] * num_buffers
162 buffer_idx = 0
163
44 - packing_tensor_list: list[list[torch.Tensor]] = [[] for _ in range(num_buffers)]
45 - packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)]
46 - packed_tensors: list[torch.Tensor] = [
47 - torch.empty(0, dtype=torch.uint8, device="cuda") for _ in range(num_buffers)
48 - ]
49 -
164 while True:
165 # Synchronize the current stream
166 streams[buffer_idx].synchronize()
167 + # Previous chunk on this buffer slot is now safe to free
168 + in_flight[buffer_idx] = None
169 # Start tasks for the new buffer in a new stream
170 with torch.cuda.stream(streams[buffer_idx]):
55 - try:
56 - # Initialize the packing tensor list and sizes
57 - packing_tensor_list[buffer_idx] = []
58 - packing_tensor_sizes[buffer_idx] = 0
59 - # Pack the tensors
60 - while True:
61 - # Apply post processing and convert to linearized uint8 tensor
62 - tensor = (
63 - post_iter_func(next(iterator))
64 - .contiguous()
65 - .view(torch.uint8)
66 - .view(-1)
67 - )
68 - packing_tensor_list[buffer_idx].append(tensor)
69 - packing_tensor_sizes[buffer_idx] += tensor.numel()
70 - if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size:
71 - break
72 - # Pack the tensors and call broadcast collective
73 - packed_tensors[buffer_idx] = torch.cat(
74 - packing_tensor_list[buffer_idx], dim=0
75 - )
76 - group.broadcast(packed_tensors[buffer_idx], src=src)
77 - # Move to the next buffer
78 - buffer_idx = (buffer_idx + 1) % num_buffers
79 - except StopIteration:
80 - # Do the last broadcast if there are remaining tensors
81 - if len(packing_tensor_list[buffer_idx]) > 0:
82 - packed_tensors[buffer_idx] = torch.cat(
83 - packing_tensor_list[buffer_idx], dim=0
84 - )
85 - group.broadcast(packed_tensors[buffer_idx], src=src)
171 + chunk = pack_tensors(iterator, post_iter_func, buffer_size_bytes)
172 + if chunk is None:
173 break
174 + # Pack the tensors and call broadcast collective
175 + group.broadcast(chunk.packed_tensor, src=src)
176 + # Hold reference until this stream is synchronized
177 + in_flight[buffer_idx] = chunk
178 + # Move to the next buffer
179 + buffer_idx = (buffer_idx + 1) % num_buffers
180
181
89 -def packed_broadcast_consumer(
182 +def packed_nccl_broadcast_consumer(
183 iterator: Iterator[tuple[str, tuple[list[int], torch.dtype]]],
184 group: Any,
185 src: int,
@@ -108,37 +201,6 @@ def packed_broadcast_consumer(
201 Both producer and consumer must use the same value.
202
203 """
111 -
112 - def unpack_tensor(
113 - packed_tensor: torch.Tensor,
114 - names: list[str],
115 - shapes: list[list[int]],
116 - dtypes: list[torch.dtype],
117 - tensor_sizes: list[int],
118 - ) -> list[tuple[str, torch.Tensor]]:
119 - """Unpack a single tensor into a list of tensors.
120 -
121 - Args:
122 - packed_tensor: The packed torch.uint8 tensor to unpack
123 - names: List of tensor names
124 - shapes: List of tensor shapes
125 - dtypes: List of tensor dtypes
126 - tensor_sizes: List of tensor sizes in bytes
127 -
128 - Returns:
129 - unpacked List[(name, tensor)]
130 - """
131 - unpacked_tensors = packed_tensor.split(tensor_sizes)
132 -
133 - unpacked_list = [
134 - (name, tensor.contiguous().view(dtype).view(*shape))
135 - for name, shape, dtype, tensor in zip(
136 - names, shapes, dtypes, unpacked_tensors
137 - )
138 - ]
139 -
140 - return unpacked_list
141 -
204 target_packed_tensor_size = buffer_size_bytes
205
206 streams = [torch.cuda.Stream() for _ in range(num_buffers)]
@@ -214,3 +276,152 @@ def packed_broadcast_consumer(
276 )
277 )
278 break
279 +
280 +
281 +# ── IPC packed transfer ────────────────────────────────────────────────
282 +
283 +
284 +@dataclass
285 +class PackedIpcChunk:
286 + """Metadata and IPC handle for a single packed chunk."""
287 +
288 + names: list[str]
289 + shapes: list[list[int]]
290 + dtype_names: list[str]
291 + tensor_sizes: list[int]
292 + ipc_handle: dict[str, tuple]
293 +
294 +
295 +def packed_ipc_producer(
296 + iterator: Iterator[tuple[str, torch.Tensor]],
297 + gpu_uuid: str,
298 + post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor],
299 + buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
300 +) -> Iterator[PackedIpcChunk]:
301 + """Pack tensors into a reusable IPC buffer and yield handles.
302 +
303 + Allocates a single GPU buffer of ``buffer_size_bytes`` and registers
304 + it for IPC once via ``reduce_tensor``. Each chunk's packed data is
305 + copied into this buffer before yielding, so only one IPC-shared
306 + allocation exists for the lifetime of the transfer.
307 +
308 + Callers **must** ensure the consumer has finished reading the buffer
309 + (e.g. ``ray.get`` returned) before resuming the generator for the
310 + next chunk.
311 +
312 + Args:
313 + iterator: Iterator of (name, tensor) pairs.
314 + gpu_uuid: Physical GPU UUID string for this rank.
315 + post_iter_func: Applied to each (name, tensor) before packing.
316 + buffer_size_bytes: Exact capacity of the reusable IPC buffer.
317 + Every chunk is guaranteed to fit within this size. A
318 + ``ValueError`` is raised if any single tensor exceeds it.
319 + """
320 + ipc_buffer = torch.empty(buffer_size_bytes, dtype=torch.uint8, device="cuda")
321 + _, ipc_args = reduce_tensor(ipc_buffer)
322 +
323 + names: list[str] = []
324 + shapes: list[list[int]] = []
325 + dtypes: list[torch.dtype] = []
326 + tensor_sizes: list[int] = []
327 + total_bytes = 0
328 +
329 + for name, orig_tensor in iterator:
330 + flat = (
331 + post_iter_func((name, orig_tensor)).contiguous().view(torch.uint8).view(-1)
332 + )
333 +
334 + if flat.numel() > buffer_size_bytes:
335 + raise ValueError(
336 + f"Tensor '{name}' has size {flat.numel()} bytes, "
337 + f"which exceeds buffer_size_bytes={buffer_size_bytes}. "
338 + f"Increase buffer_size_bytes to at least {flat.numel()}."
339 + )
340 +
341 + if total_bytes and total_bytes + flat.numel() > buffer_size_bytes:
342 + # Drain queued copies so the consumer sees a fully-written buffer.
343 + torch.cuda.current_stream().synchronize()
344 + yield PackedIpcChunk(
345 + names=names,
346 + shapes=shapes,
347 + dtype_names=[str(d).split(".")[-1] for d in dtypes],
348 + tensor_sizes=tensor_sizes,
349 + ipc_handle={gpu_uuid: ipc_args},
350 + )
351 + # Rebind to fresh lists so the yielded chunk's metadata is
352 + # not mutated while the consumer is still reading.
353 + names, shapes, dtypes, tensor_sizes = [], [], [], []
354 + total_bytes = 0
355 +
356 + ipc_buffer[total_bytes : total_bytes + flat.numel()].copy_(flat)
357 + names.append(name)
358 + shapes.append(list(orig_tensor.shape))
359 + dtypes.append(orig_tensor.dtype)
360 + tensor_sizes.append(flat.numel())
361 + total_bytes += flat.numel()
362 +
363 + if total_bytes:
364 + torch.cuda.current_stream().synchronize()
365 + yield PackedIpcChunk(
366 + names=names,
367 + shapes=shapes,
368 + dtype_names=[str(d).split(".")[-1] for d in dtypes],
369 + tensor_sizes=tensor_sizes,
370 + ipc_handle={gpu_uuid: ipc_args},
371 + )
372 +
373 +
374 +def packed_ipc_consumer(
375 + ipc_handle: dict[str, tuple],
376 + names: list[str],
377 + shapes: list[list[int]],
378 + dtype_names: list[str],
379 + tensor_sizes: list[int],
380 + device_index: int,
381 +) -> list[tuple[str, torch.Tensor]]:
382 + """Unpack a single packed IPC chunk into named tensors.
383 +
384 + Reconstructs the packed buffer via rebuild_cuda_tensor, unpacks
385 + into individual tensors, and clones each into independent storage
386 + before returning.
387 +
388 + The clone is intentional: the producer reuses one IPC buffer across
389 + chunks, so any tensor view that aliases the buffer would observe the
390 + *next* chunk's bytes as soon as the producer's generator is resumed.
391 + Callers that retain references past their own update_weights call
392 + (notably vLLM's layerwise reload, which buffers ``bound_args`` for
393 + replay in ``_layerwise_process``) would otherwise replay against
394 + stale data and silently corrupt multi-chunk weight transfers.
395 +
396 + Args:
397 + ipc_handle: Mapping of GPU UUID to rebuild_cuda_tensor args tuple
398 + names: Parameter names in the packed buffer
399 + shapes: Parameter shapes
400 + dtype_names: Parameter dtype name strings (e.g. "float16")
401 + tensor_sizes: Size in bytes of each parameter in the packed buffer
402 + device_index: Local CUDA device index
403 + """
404 + from torch.multiprocessing.reductions import rebuild_cuda_tensor
405 +
406 + props = torch.cuda.get_device_properties(device_index)
407 + physical_gpu_id = str(props.uuid)
408 +
409 + if physical_gpu_id not in ipc_handle:
410 + raise ValueError(
411 + f"IPC handle not found for GPU UUID {physical_gpu_id}. "
412 + f"Available UUIDs: {list(ipc_handle.keys())}"
413 + )
414 +
415 + args = ipc_handle[physical_gpu_id]
416 + list_args = list(args)
417 + list_args[6] = device_index
418 + packed = rebuild_cuda_tensor(*list_args)
419 +
420 + content_size = sum(tensor_sizes)
421 + packed = packed[:content_size]
422 +
423 + dtypes = [getattr(torch, dn) for dn in dtype_names]
424 + return [
425 + (name, t.clone())
426 + for name, t in unpack_tensor(packed, names, shapes, dtypes, tensor_sizes)
427 + ]