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 6, 10:45 UTC
[BUG] Exclude SKIP_TENSORS from get_layer_size() + new weight sync example for dpep (#37334)
Signed-off-by: ahao-anyscale <ahao@anyscale.com>
Aaron Hao committed
Mar 18, 2026 at 17:45 UTC
5f82706a21012df9508913969ab7c9932e127417
2 files changed
+351
-2
examples/rl/rlhf_nccl_fsdp_ep.py
new
+339
@@ -0,0 +1,339 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs).
5
+
6
+8-GPU layout:
7
+ Training — 4 GPUs, PyTorch FSDP2 (fully_shard)
8
+ Inference — 4 GPUs, vLLM AsyncLLMEngine with expert parallelism +
9
+ data parallelism (TP=1, DP=4, enable_expert_parallel
10
+ → EP_SIZE = TP×DP = 4)
11
+
12
+FSDP workers are Ray actors that form a single FSDP2 process group.
13
+Rank 0 gathers full parameters via DTensor.full_tensor() and broadcasts
14
+them to the vLLM inference engine through the NCCL weight-transfer API.
15
+
16
+The inference engine uses AsyncLLMEngine which automatically spawns
17
+DP worker processes (no manual placement group needed). Weight sync
18
+uses pause_generation / resume_generation.
19
+
20
+Steps:
21
+ 1. Launch 4 FSDP training workers.
22
+ 2. Launch AsyncLLMEngine with EP+DP (dummy weights).
23
+ 3. Generate from prompts → gibberish (random weights).
24
+ 4. Pause generation, transfer weights from FSDP, resume.
25
+ 5. Generate from prompts → sensible output (synced weights).
26
+
27
+Assumes a single-node cluster with 8 GPUs.
28
+"""
29
+
30
+import asyncio
31
+import os
32
+import uuid
33
+from dataclasses import asdict
34
+
35
+import ray
36
+import torch
37
+import torch.distributed as dist
38
+from huggingface_hub import snapshot_download
39
+from torch.distributed.fsdp import fully_shard
40
+from transformers import AutoModelForCausalLM
41
+
42
+import vllm
43
+from vllm import SamplingParams
44
+from vllm.config import WeightTransferConfig
45
+from vllm.distributed.weight_transfer.base import (
46
+ WeightTransferInitRequest,
47
+ WeightTransferUpdateRequest,
48
+)
49
+from vllm.distributed.weight_transfer.nccl_engine import (
50
+ NCCLTrainerSendWeightsArgs,
51
+ NCCLWeightTransferEngine,
52
+ NCCLWeightTransferInitInfo,
53
+ NCCLWeightTransferUpdateInfo,
54
+)
55
+from vllm.utils.network_utils import get_ip, get_open_port
56
+from vllm.v1.executor import Executor
57
+
58
+MODEL_NAME = "Qwen/Qwen3-30B-A3B"
59
+
60
+FSDP_WORLD_SIZE = 4
61
+INFERENCE_TP_SIZE = 1
62
+INFERENCE_DP_SIZE = 4
63
+
64
+
65
+@ray.remote(num_gpus=1)
66
+class FSDPTrainWorker:
67
+ """
68
+ One FSDP2 training worker per GPU. Four of these form the FSDP group.
69
+ Rank 0 additionally handles weight transfer to the vLLM engine.
70
+ """
71
+
72
+ def __init__(
73
+ self,
74
+ model_name: str,
75
+ rank: int,
76
+ fsdp_world_size: int,
77
+ fsdp_master_addr: str,
78
+ fsdp_master_port: int,
79
+ ):
80
+ self.rank = rank
81
+
82
+ os.environ["MASTER_ADDR"] = fsdp_master_addr
83
+ os.environ["MASTER_PORT"] = str(fsdp_master_port)
84
+
85
+ dist.init_process_group(backend="nccl", rank=rank, world_size=fsdp_world_size)
86
+ torch.accelerator.set_device_index(0)
87
+
88
+ model = AutoModelForCausalLM.from_pretrained(
89
+ model_name, torch_dtype=torch.bfloat16
90
+ )
91
+
92
+ self.weight_names = [n for n, _ in model.named_parameters()]
93
+ self.weight_dtype_names = [
94
+ str(p.dtype).split(".")[-1] for _, p in model.named_parameters()
95
+ ]
96
+ self.weight_shapes = [list(p.shape) for _, p in model.named_parameters()]
97
+
98
+ for layer in model.model.layers:
99
+ fully_shard(layer)
100
+ fully_shard(model)
101
+
102
+ self.model = model
103
+
104
+ self.transfer_port = None
105
+ self.transfer_master_address = None
106
+ self.model_update_group = None
107
+
108
+ def get_rank(self):
109
+ return self.rank
110
+
111
+ # ---- weight-transfer setup (rank 0 only) ----
112
+
113
+ def setup_transfer_endpoint(self):
114
+ """Create the NCCL rendezvous endpoint for weight transfer."""
115
+ assert self.rank == 0
116
+ self.transfer_port = get_open_port()
117
+ self.transfer_master_address = get_ip()
118
+ return self.transfer_master_address, self.transfer_port
119
+
120
+ def init_weight_transfer_group(self, transfer_world_size: int):
121
+ """Join the weight-transfer NCCL group as rank 0 (the source)."""
122
+ assert self.rank == 0
123
+ self.model_update_group = NCCLWeightTransferEngine.trainer_init(
124
+ dict(
125
+ master_address=self.transfer_master_address,
126
+ master_port=self.transfer_port,
127
+ world_size=transfer_world_size,
128
+ ),
129
+ )
130
+
131
+ def get_weight_metadata(self):
132
+ """Return weight names, dtypes, and shapes captured before FSDP wrapping."""
133
+ return self.weight_names, self.weight_dtype_names, self.weight_shapes
134
+
135
+ # ---- collective ops (ALL FSDP ranks must call concurrently) ----
136
+
137
+ def gather_and_broadcast_weights(self, packed: bool = True):
138
+ """
139
+ All-gather full parameters and broadcast them to vLLM.
140
+ Only rank 0 performs the actual NCCL broadcast; others just
141
+ participate in the FSDP all-gather.
142
+
143
+ full_tensor() is a collective — all FSDP ranks must call it
144
+ for each parameter in the same order. Rank 0 additionally
145
+ feeds each gathered tensor to the weight-transfer engine.
146
+ """
147
+ if self.rank == 0:
148
+
149
+ def _full_param_iter():
150
+ for name, param in self.model.named_parameters():
151
+ yield name, param.full_tensor()
152
+
153
+ trainer_args = NCCLTrainerSendWeightsArgs(
154
+ group=self.model_update_group,
155
+ packed=packed,
156
+ )
157
+ NCCLWeightTransferEngine.trainer_send_weights(
158
+ iterator=_full_param_iter(),
159
+ trainer_args=trainer_args,
160
+ )
161
+ else:
162
+ for _, param in self.model.named_parameters():
163
+ param.full_tensor()
164
+
165
+
166
+def create_async_engine(**kwargs):
167
+ """Create an AsyncLLMEngine directly (no subclass needed)."""
168
+ engine_args = vllm.AsyncEngineArgs(**kwargs)
169
+ vllm_config = engine_args.create_engine_config()
170
+ executor_class = Executor.get_class(vllm_config)
171
+ return vllm.AsyncLLMEngine(
172
+ vllm_config=vllm_config,
173
+ executor_class=executor_class,
174
+ log_requests=engine_args.enable_log_requests,
175
+ log_stats=not engine_args.disable_log_stats,
176
+ )
177
+
178
+
179
+async def generate_batch(engine, prompts, sampling_params):
180
+ """Generate completions for a batch of prompts."""
181
+
182
+ async def gen_one(prompt):
183
+ output = None
184
+ async for request_output in engine.generate(
185
+ {"prompt": prompt},
186
+ sampling_params,
187
+ request_id=str(uuid.uuid4()),
188
+ ):
189
+ output = request_output
190
+ return output
191
+
192
+ return await asyncio.gather(*[gen_one(p) for p in prompts])
193
+
194
+
195
+async def main():
196
+ ray.init()
197
+
198
+ # Download model weights to local/shared disk once.
199
+ local_model_path = snapshot_download(MODEL_NAME)
200
+ print(f"[init] Model downloaded to {local_model_path}")
201
+
202
+ # FSDP rendezvous address (single-node)
203
+ fsdp_master_addr = get_ip()
204
+ fsdp_master_port = get_open_port()
205
+
206
+ # Launch 4 FSDP training workers.
207
+ # Ray allocates 1 GPU per worker; AsyncLLMEngine's internal DP
208
+ # placement groups will land on the remaining 4 GPUs.
209
+ fsdp_workers = [
210
+ FSDPTrainWorker.remote(
211
+ local_model_path,
212
+ rank,
213
+ FSDP_WORLD_SIZE,
214
+ fsdp_master_addr,
215
+ fsdp_master_port,
216
+ )
217
+ for rank in range(FSDP_WORLD_SIZE)
218
+ ]
219
+ ray.get([w.get_rank.remote() for w in fsdp_workers])
220
+ print(f"[init] {FSDP_WORLD_SIZE} FSDP training workers ready.")
221
+
222
+ # Launch vLLM with expert parallelism + data parallelism.
223
+ # AsyncLLMEngine with data_parallel_backend="ray" creates its own
224
+ # placement groups internally — no manual placement group needed.
225
+ print("[engine] Creating AsyncLLMEngine...")
226
+ engine = create_async_engine(
227
+ model=local_model_path,
228
+ enforce_eager=True,
229
+ tensor_parallel_size=INFERENCE_TP_SIZE,
230
+ data_parallel_size=INFERENCE_DP_SIZE,
231
+ enable_expert_parallel=True,
232
+ distributed_executor_backend="ray",
233
+ data_parallel_backend="ray",
234
+ weight_transfer_config=WeightTransferConfig(backend="nccl"),
235
+ load_format="dummy",
236
+ gpu_memory_utilization=0.7,
237
+ )
238
+ print("[engine] AsyncLLMEngine created.")
239
+
240
+ prompts = [
241
+ "Hello, my name is",
242
+ "The president of the United States is",
243
+ "The capital of France is",
244
+ "The future of AI is",
245
+ ]
246
+ sampling_params = SamplingParams(temperature=0)
247
+
248
+ # Generate with dummy weights — expect gibberish.
249
+ print("[generate] Starting generation with dummy weights...")
250
+ outputs = await generate_batch(engine, prompts, sampling_params)
251
+ print("[generate] Generation complete.")
252
+
253
+ print("-" * 60)
254
+ print("BEFORE weight sync (dummy weights):")
255
+ print("-" * 60)
256
+ for output in outputs:
257
+ print(f"Prompt: {output.prompt!r}")
258
+ print(f"Generated: {output.outputs[0].text!r}")
259
+ print("-" * 60)
260
+
261
+ # --- Weight-transfer setup ---
262
+ print("[transfer] Setting up weight-transfer endpoint...")
263
+ transfer_addr, transfer_port = ray.get(
264
+ fsdp_workers[0].setup_transfer_endpoint.remote()
265
+ )
266
+ print(f"[transfer] Endpoint ready at {transfer_addr}:{transfer_port}")
267
+
268
+ transfer_world_size = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE + 1
269
+ print(
270
+ f"[transfer] World size: {transfer_world_size} "
271
+ f"(1 trainer + {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} vLLM workers)"
272
+ )
273
+
274
+ print("[transfer] Initializing NCCL groups...")
275
+ train_handle = fsdp_workers[0].init_weight_transfer_group.remote(
276
+ transfer_world_size
277
+ )
278
+ await engine.init_weight_transfer_engine(
279
+ WeightTransferInitRequest(
280
+ init_info=asdict(
281
+ NCCLWeightTransferInitInfo(
282
+ master_address=transfer_addr,
283
+ master_port=transfer_port,
284
+ rank_offset=1,
285
+ world_size=transfer_world_size,
286
+ )
287
+ )
288
+ )
289
+ )
290
+ ray.get(train_handle)
291
+ print("[transfer] NCCL groups initialized.")
292
+
293
+ # --- Pause, transfer weights, resume ---
294
+ print("[sync] Pausing generation...")
295
+ await engine.pause_generation(mode="abort")
296
+ print("[sync] Generation paused.")
297
+
298
+ names, dtype_names, shapes = ray.get(fsdp_workers[0].get_weight_metadata.remote())
299
+ print(f"[sync] Got metadata for {len(names)} parameters.")
300
+
301
+ print("[sync] Broadcasting weights from FSDP → vLLM...")
302
+ broadcast_handles = [
303
+ w.gather_and_broadcast_weights.remote(packed=True) for w in fsdp_workers
304
+ ]
305
+ await engine.update_weights(
306
+ WeightTransferUpdateRequest(
307
+ update_info=asdict(
308
+ NCCLWeightTransferUpdateInfo(
309
+ names=names,
310
+ dtype_names=dtype_names,
311
+ shapes=shapes,
312
+ packed=True,
313
+ )
314
+ )
315
+ )
316
+ )
317
+ ray.get(broadcast_handles)
318
+ print("[sync] Weight broadcast complete.")
319
+
320
+ print("[sync] Resuming generation...")
321
+ await engine.resume_generation()
322
+ print("[sync] Generation resumed.")
323
+
324
+ # Generate with synced weights — expect sensible output.
325
+ print("[generate] Starting generation with synced weights...")
326
+ outputs_updated = await generate_batch(engine, prompts, sampling_params)
327
+ print("[generate] Generation complete.")
328
+
329
+ print("-" * 60)
330
+ print("AFTER weight sync (real weights):")
331
+ print("-" * 60)
332
+ for output in outputs_updated:
333
+ print(f"Prompt: {output.prompt!r}")
334
+ print(f"Generated: {output.outputs[0].text!r}")
335
+ print("-" * 60)
336
+
337
+
338
+if __name__ == "__main__":
339
+ asyncio.run(main())
vllm/model_executor/model_loader/reload/utils.py
+12
-2
@@ -27,5 +27,15 @@ def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors:
27
28
29
def get_layer_size(layer: torch.nn.Module) -> int:
30
- """Calculate total number of elements across all tensors in a layer."""
31
- return sum(tensor.numel() for tensor in get_layer_tensors(layer).values())
30
+ """Calculate total number of elements across loadable tensors in a layer.
31
+
32
+ Excludes SKIP_TENSORS (e.g. _expert_map) which are never moved to meta
33
+ device and never loaded via weight_loader during layerwise reload.
34
+ """
35
+ from .meta import SKIP_TENSORS
36
+
37
+ return sum(
38
+ tensor.numel()
39
+ for name, tensor in get_layer_tensors(layer).items()
40
+ if name not in SKIP_TENSORS
41
+ )