main
md 122 lines 5.93 KB
Rendered Raw
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 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 share the **same GPU(s)** (colocated)
8
9 ## How It Works
10
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 materializes the full tensor for each parameter onto its own GPU before generating the handle — which `ModuleSource` does for you.
12 2. Every trainer rank contributes its handles to an all-gather; the sender merges them so each payload maps every GPU UUID to its args, and ships the merged handles to the inference engine through the client. Each worker reads only the handle for 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 Unlike NCCL, IPC transfer is straight-line: `update_weights` *is* the transfer,
16 and it rides the client, so there is no concurrent broadcast to overlap with.
17
18 !!! note
19 The handle all-gather in step 2 runs over the **default** process group, so it
20 assumes that group is exactly the set of colocated trainer ranks and that the
21 sender is a member. It is a no-op when no distributed group exists.
22
23 !!! warning
24 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.
25
26 ## Inference Side
27
28 ```python
29 from vllm import LLM
30 from vllm.config import WeightTransferConfig
31
32 llm = LLM(model="my-model", weight_transfer_config=WeightTransferConfig(backend="ipc"))
33 ```
34
35 ```bash
36 vllm serve my-model --weight-transfer-config '{"backend": "ipc"}'
37 ```
38
39 IPC needs no data-plane rendezvous, so `init_transfer_engine` opens no channel —
40 it only records the `packed` flag the trainer ships at the handshake, which
41 `receive_weights` then reads. Whether a transfer is packed is therefore never
42 something you configure on the inference side.
43
44 ## Trainer Side
45
46 ```python
47 from vllm.distributed.weight_transfer import (
48 ModuleSource,
49 HTTPVLLMWeightSyncClient,
50 WeightTransferTrainerFactory,
51 )
52 from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
53
54 engine = WeightTransferTrainerFactory.trainer_init(
55 init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 is the sender
56 client=HTTPVLLMWeightSyncClient("http://localhost:8000"),
57 source=ModuleSource(model),
58 )
59
60 engine.send_weights() # once per sync
61 ```
62
63 `send_weights()` drives `start_weight_update`, the transfer itself, and
64 `finish_weight_update`, and holds strong references to the IPC-shared copies
65 until after the post-send barrier — the consumer's views would otherwise dangle.
66
67 Any [`VLLMWeightSyncClient`](base.md#vllmweightsyncclient) works here — the
68 built-in HTTP and Ray clients, or an adapter for your own stack. The engine is
69 identical either way.
70
71 ### `IPCTrainerInitInfo`
72
73 | Field | Default | Description |
74 | ----- | ------- | ----------- |
75 | `rank` | — | Keyword-only. This trainer process's rank; **0 is the sender** |
76 | `packed` | `False` | Chunked, bounded-memory transfer (see below) |
77 | `packed_buffer_size_bytes` | 1 GiB | Chunk size when `packed=True` |
78
79 `packed` is a must-agree wire param: `trainer_init` ships it to the worker, which
80 records it and decodes accordingly. It is not a `WeightTransferConfig` field and
81 not a per-round `update_weights` field. `packed_buffer_size_bytes` is
82 producer-only — the consumer rebuilds from the IPC handle plus the per-chunk
83 `tensor_sizes`, so it never needs the buffer size.
84
85 ## Packed (Chunked) Transfer
86
87 By default, all weights are sent in a single `update_weights` call. For large
88 models, this requires the full model to reside in GPU memory on both sides
89 simultaneously. Setting `packed=True` enables **chunked transfer** with bounded
90 GPU memory:
91
92 - Weights are concatenated into fixed-size packed buffers (`packed_buffer_size_bytes`).
93 - 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.
94 - After each chunk is consumed, the GPU memory for that chunk can be reclaimed.
95
96 ```python
97 engine = WeightTransferTrainerFactory.trainer_init(
98 init_info=IPCTrainerInitInfo(
99 rank=0,
100 packed=True,
101 packed_buffer_size_bytes=256 * 1024 * 1024, # 256 MB chunks
102 ),
103 client=client,
104 source=ModuleSource(model),
105 )
106 ```
107
108 On a multi-rank trainer the producer reuses one buffer across chunks, so packed
109 mode carries a per-chunk barrier across ranks: without it, a rank could overwrite
110 its buffer while its colocated worker is still reading the current chunk. This is
111 handled inside `send_weights()`.
112
113 ## Rank-Local Updates
114
115 For different parameter subsets per worker, `update_info` may be a list indexed
116 by worker rank. This form is IPC-only because collective backends require every
117 worker to participate.
118
119 ## Examples
120
121 - [RLHF with IPC weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_ipc.py) - **Start here.** Server and training model share a single GPU; HTTP control plane, CUDA IPC data plane. Launches and tears down its own server
122 - [RLHF with IPC + FSDP2 and expert parallelism](../../../examples/rl/rlhf_ipc_fsdp_ep.py) - Multi-rank trainer colocated with a `--data-parallel-size 4` server on the same 4 GPUs: every FSDP rank builds an engine and joins the handle all-gather, with packed chunking and sleep/wake around the transfer