@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 201 lines 9.17 KB
Rendered Raw
1 # NCCL Engine
2
3 The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) broadcast operations to transfer weights from the trainer to inference workers. It supports **multi-node** and **multi-GPU** setups where the trainer and inference engine run on separate GPUs.
4
5 ## When to Use NCCL
6
7 - Training and inference on **separate GPUs** (possibly across nodes)
8 - **Tensor-parallel** inference with multiple workers that all need the updated weights
9 - You need high-bandwidth, low-latency weight transfer over NVLink or InfiniBand
10
11 ## How It Works
12
13 1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction). The trainer is rank 0; workers start at `rank_offset` (1).
14 2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads the weights.
15 3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py).
16
17 The workers' `update_weights` and the trainer's broadcast run at the same time —
18 both sides rendezvous inside the same NCCL calls. The trainer engine
19 owns that concurrency internally.
20
21 ## Inference Side
22
23 The inference side takes a plain backend selector. The rendezvous parameters and
24 the packing wire params arrive from the trainer at the init handshake.
25
26 ```python
27 from vllm import LLM
28 from vllm.config import WeightTransferConfig
29
30 llm = LLM(model="my-model", weight_transfer_config=WeightTransferConfig(backend="nccl"))
31 ```
32
33 ```bash
34 vllm serve my-model --weight-transfer-config '{"backend": "nccl"}'
35 ```
36
37 Nothing else is required: `init_weight_transfer_engine`, `start_weight_update`,
38 `update_weights`, and `finish_weight_update` are all driven remotely by the
39 trainer engine.
40
41 ## Trainer Side
42
43 ```python
44 from vllm.distributed.weight_transfer import (
45 ModuleSource,
46 HTTPVLLMWeightSyncClient,
47 WeightTransferTrainerFactory,
48 )
49 from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
50
51 engine = WeightTransferTrainerFactory.trainer_init(
52 init_info=NCCLTrainerInitInfo(
53 master_address=master_address,
54 master_port=master_port,
55 world_size=world_size, # trainer + all inference workers
56 rank=0, # this trainer rank; rank 0 is the sender
57 packed=True,
58 ),
59 client=HTTPVLLMWeightSyncClient("http://localhost:8000"), # or RayVLLMWeightSyncClient(llm)
60 source=ModuleSource(model),
61 )
62
63 engine.send_weights() # once per sync
64 ```
65
66 `trainer_init` drives the whole handshake: it kicks off the inference side's
67 `init_weight_transfer_engine` (through the client) on a side thread while opening
68 the trainer's own rank-0 endpoint, since both ends must rendezvous together. It
69 builds the worker's init info itself, with `rank_offset=1` and the same packed
70 params, so the two sides cannot disagree.
71
72 `send_weights()` then drives `start_weight_update`, `update_weights` — run
73 concurrently with the broadcast — and `finish_weight_update`, and returns only
74 once every transfer has drained.
75
76 ### `NCCLTrainerInitInfo`
77
78 | Field | Default | Description |
79 | ----- | ------- | ----------- |
80 | `master_address` | — | Rendezvous host |
81 | `master_port` | — | Rendezvous port |
82 | `world_size` | — | Full trainer + worker NCCL group size |
83 | `rank` | — | Keyword-only. This trainer process's rank; **0 is the sender** |
84 | `packed` | `True` | Use packed broadcasting |
85 | `packed_buffer_size_bytes` | 1 GiB | Packed buffer size |
86 | `packed_num_buffers` | 2 | Number of rotating buffers (double/triple buffering) |
87
88 `packed` defaults to `True` here. The worker-side default is `False`, but it only
89 applies when no trainer ships a value — which, on this path, never happens.
90
91 ### Packed Tensor Broadcasting
92
93 When `packed=True`, weight tensors are packed into large contiguous buffers
94 before broadcasting. This cuts the number of NCCL operations and uses
95 double/triple buffering with dedicated CUDA streams to overlap packing,
96 broadcasting, and unpacking.
97
98 You set `packed`, `packed_buffer_size_bytes`, and `packed_num_buffers` **only on
99 `NCCLTrainerInitInfo`**. The trainer propagates them to the worker inside
100 `trainer_init`, the worker records them at the handshake, and `receive_weights`
101 decodes with exactly the values the trainer encoded with. They are not per-round
102 `update_weights` fields.
103
104 !!! note "Memory"
105 The rotating buffers are live for the whole transfer:
106 `packed_buffer_size_bytes * packed_num_buffers` on each side (2 GiB at the
107 defaults). Lower `packed_buffer_size_bytes` if that is too much headroom.
108
109 ### The two `WeightSource` channels must agree
110
111 Dense NCCL is the backend that reads *both*
112 [`WeightSource`](base.md#weightsource) channels, so it is the one where a
113 disagreement between them is fatal. The engine builds the per-round update info
114 from `metadata()` and ships it ahead of the bytes; the worker sizes its receive
115 buffers from that info, and in packed mode cuts its chunk boundaries from it. The
116 bytes themselves come from iterating the source.
117
118 If iteration disagrees with what `metadata()` declared — a reordered, omitted, or
119 re-dtyped parameter — the two sides split the same byte stream differently. The
120 transfer then either hangs in NCCL waiting for a length that never arrives, or
121 loads garbage into the model.
122
123 The sender therefore checks each pair against the declared metadata as it goes,
124 one comparison per parameter, and raises naming the first divergent parameter
125 rather than letting it reach the wire. `ModuleSource` satisfies this by
126 construction. If you write a custom source — a Megatron export, an MoE re-fusing
127 pass — this is the invariant to test first.
128
129 !!! note
130 IPC does not read `metadata()` at all: it derives the update info from
131 iteration as it goes, so it cannot observe a divergence.
132
133 ## Sparse NCCL
134
135 Sparse, flat-index weight patches use
136 `WeightTransferConfig(backend="sparse_nccl")`. Names, full shapes, and flat
137 indices are interpreted in **checkpoint/Hugging Face coordinates**. Every
138 inference rank receives the same checkpoint-global patch, then the model's native
139 `load_weights()` maps it to rank-local TP/EP and packed runtime parameters.
140
141 Tensor parallelism is supported and does not have to match the trainer's layout.
142 Pipeline-parallel behavior is not covered by the current sparse NCCL GPU test.
143
144 Sparse is a **delta** backend: each call supplies replacement patches rather
145 than a stable stream of the model's parameters. The engine therefore takes no
146 `WeightSource`; patches go straight to `send_weights(patches)`, and an empty
147 patch list is a no-op.
148
149 ```python
150 from vllm.distributed.weight_transfer import (
151 RayVLLMWeightSyncClient,
152 WeightTransferTrainerFactory,
153 )
154 from vllm.distributed.weight_transfer.sparse_nccl_engine import (
155 SparseNCCLTrainerInitInfo,
156 SparseWeightPatch,
157 )
158
159 client = RayVLLMWeightSyncClient(llm)
160 engine = WeightTransferTrainerFactory.trainer_init(
161 init_info=SparseNCCLTrainerInitInfo(
162 master_address=master_address,
163 master_port=master_port,
164 world_size=world_size,
165 rank=0,
166 ),
167 client=client,
168 )
169
170 patches = [
171 SparseWeightPatch(
172 name="model.layers.0.mlp.down_proj.weight",
173 indices=flat_indices, # int32, 1-D
174 values=new_values, # same length as indices
175 full_shape=tuple(param.shape), # required when sending via the engine
176 )
177 ]
178 engine.send_weights(patches)
179 ```
180
181 Each `send_weights(patches)` call owns a complete one-shot `start` / `update` / `finish` lifecycle.
182
183 RL infrastructure that owns the generic worker lifecycle can keep one logical update open across bounded chunks without adding trainer-side session state:
184
185 ```python
186 client.start_weight_update()
187 for patches in patch_chunks:
188 engine.send_weight_chunk(patches)
189 client.finish_weight_update()
190 ```
191
192 Sparse NCCL sends only `O(nnz)` indices and values. Checkpoint application still uses `O(N)` staging through the native loader, and the caller owns export/diff state and restart or reseed after a partial failure.
193
194 [`rlhf_sparse_nccl.py`](../../../examples/rl/rlhf_sparse_nccl.py) demonstrates per-expert checkpoint patches with a Qwen3 MoE TP2/EP2 engine.
195
196 ## Examples
197
198 - [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server
199 - [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire
200 - [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Qwen3 MoE per-expert checkpoint updates with one trainer GPU and two TP2/EP2 inference GPUs
201 - [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model; uses an in-process `AsyncLLMEngine` with `RayVLLMWeightSyncClient`