| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs). |
| 4 | |
| 5 | 8-GPU layout: |
| 6 | Training — 4 GPUs, PyTorch FSDP2 (fully_shard), as Ray actors |
| 7 | Inference — 4 GPUs, a `vllm serve` HTTP server with expert parallelism + |
| 8 | data parallelism (TP=1, DP=4, enable_expert_parallel |
| 9 | → EP_SIZE = TP×DP = 4) |
| 10 | |
| 11 | The inference side is a standalone HTTP server (spawned by this script with |
| 12 | `vllm serve`), so both the weight-sync control plane (HTTP) and the NCCL data |
| 13 | plane run inside the rank-0 FSDP Ray actor. That lets the trainer use the |
| 14 | unified `TrainerWeightTransferEngine.send_weights()` with an |
| 15 | `HTTPVLLMWeightSyncClient` — one call drives start/update/finish on the server |
| 16 | concurrently with the NCCL broadcast. Every FSDP rank builds an engine and calls |
| 17 | `send_weights()`, so all 4 participate in the incremental `full_tensor()` |
| 18 | all-gather; only rank 0 holds a communicator and broadcasts (it is the only |
| 19 | trainer rank in the NCCL group). |
| 20 | |
| 21 | GPU split (single node): the server takes GPUs 0-3 (CUDA_VISIBLE_DEVICES), and |
| 22 | Ray (training) is restricted to GPUs 4-7. |
| 23 | |
| 24 | Steps: |
| 25 | 1. Launch the vLLM HTTP server (EP+DP, dummy weights) on GPUs 0-3. |
| 26 | 2. Launch 4 FSDP training workers (Ray) on GPUs 4-7. |
| 27 | 3. Generate from prompts over HTTP → gibberish (random weights). |
| 28 | 4. Pause generation, transfer weights FSDP → server over NCCL, resume. |
| 29 | 5. Generate from prompts → sensible output (synced weights). |
| 30 | |
| 31 | Assumes a single-node cluster with 8 GPUs. |
| 32 | """ |
| 33 | |
| 34 | import json |
| 35 | import os |
| 36 | import subprocess |
| 37 | import sys |
| 38 | import time |
| 39 | |
| 40 | import ray |
| 41 | import requests |
| 42 | import torch |
| 43 | import torch.distributed as dist |
| 44 | from huggingface_hub import snapshot_download |
| 45 | from openai import OpenAI |
| 46 | from torch.distributed.fsdp import fully_shard |
| 47 | from transformers import AutoModelForCausalLM |
| 48 | |
| 49 | from vllm.distributed.weight_transfer import ( |
| 50 | HTTPVLLMWeightSyncClient, |
| 51 | ModuleSource, |
| 52 | WeightTransferTrainerFactory, |
| 53 | ) |
| 54 | from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo |
| 55 | from vllm.utils.network_utils import get_ip, get_open_port |
| 56 | |
| 57 | MODEL_NAME = "Qwen/Qwen3-30B-A3B" |
| 58 | SERVED_MODEL_NAME = "policy" |
| 59 | |
| 60 | FSDP_WORLD_SIZE = 4 |
| 61 | INFERENCE_TP_SIZE = 1 |
| 62 | INFERENCE_DP_SIZE = 4 |
| 63 | |
| 64 | # Training (FSDP) GPUs are reserved through Ray; the inference server then runs |
| 65 | # on the complementary GPUs (see main()). We do NOT hard-code the split via |
| 66 | # CUDA_VISIBLE_DEVICES before ray.init(): that only restricts Ray when ray.init() |
| 67 | # *starts* a local cluster, and is silently ignored when it connects to an |
| 68 | # existing one (e.g. a shared/managed Ray cluster), causing training and the |
| 69 | # server to collide on the same physical GPUs. |
| 70 | SERVER_PORT = 8000 |
| 71 | BASE_URL = f"http://localhost:{SERVER_PORT}" |
| 72 | |
| 73 | |
| 74 | @ray.remote(num_gpus=1) |
| 75 | class FSDPTrainWorker: |
| 76 | """One FSDP2 training worker per GPU. Four of these form the FSDP group. |
| 77 | Rank 0 additionally drives weight transfer to the vLLM server. |
| 78 | """ |
| 79 | |
| 80 | def __init__( |
| 81 | self, |
| 82 | model_name: str, |
| 83 | rank: int, |
| 84 | fsdp_world_size: int, |
| 85 | fsdp_master_addr: str, |
| 86 | fsdp_master_port: int, |
| 87 | ): |
| 88 | self.rank = rank |
| 89 | self.engine = None |
| 90 | |
| 91 | os.environ["MASTER_ADDR"] = fsdp_master_addr |
| 92 | os.environ["MASTER_PORT"] = str(fsdp_master_port) |
| 93 | |
| 94 | dist.init_process_group(backend="nccl", rank=rank, world_size=fsdp_world_size) |
| 95 | torch.accelerator.set_device_index(0) |
| 96 | |
| 97 | model = AutoModelForCausalLM.from_pretrained( |
| 98 | model_name, torch_dtype=torch.bfloat16 |
| 99 | ) |
| 100 | |
| 101 | for layer in model.model.layers: |
| 102 | fully_shard(layer) |
| 103 | fully_shard(model) |
| 104 | |
| 105 | self.model = model |
| 106 | |
| 107 | self.transfer_port = None |
| 108 | self.transfer_master_address = None |
| 109 | |
| 110 | def get_rank(self): |
| 111 | return self.rank |
| 112 | |
| 113 | def get_gpu_ids(self): |
| 114 | """Physical GPU id(s) Ray assigned to this worker (for server/train split).""" |
| 115 | return ray.get_gpu_ids() |
| 116 | |
| 117 | # ---- weight-transfer setup (rank 0 only) ---- |
| 118 | |
| 119 | def setup_transfer_endpoint(self): |
| 120 | """Create the NCCL rendezvous endpoint for weight transfer.""" |
| 121 | assert self.rank == 0 |
| 122 | self.transfer_port = get_open_port() |
| 123 | self.transfer_master_address = get_ip() |
| 124 | return self.transfer_master_address, self.transfer_port |
| 125 | |
| 126 | def setup_engine( |
| 127 | self, |
| 128 | base_url: str, |
| 129 | transfer_master_address: str, |
| 130 | transfer_port: int, |
| 131 | transfer_world_size: int, |
| 132 | ): |
| 133 | """Build the trainer engine on every FSDP rank. |
| 134 | |
| 135 | Called on all ranks with the shared rendezvous endpoint. Rank 0 is the |
| 136 | sender: `trainer_init` opens its rank-0 NCCL endpoint and, on a worker |
| 137 | thread, calls the server's `init_weight_transfer_engine` over HTTP so |
| 138 | both ends rendezvous together. The other ranks skip the rendezvous and |
| 139 | only join the FSDP all-gather during send_weights. |
| 140 | """ |
| 141 | self.engine = WeightTransferTrainerFactory.trainer_init( |
| 142 | init_info=NCCLTrainerInitInfo( |
| 143 | master_address=transfer_master_address, |
| 144 | master_port=transfer_port, |
| 145 | world_size=transfer_world_size, |
| 146 | rank=self.rank, # FSDP rank; sender is rank 0 |
| 147 | packed=True, |
| 148 | ), |
| 149 | client=HTTPVLLMWeightSyncClient(base_url), |
| 150 | # Yields sharded DTensors; the engine reads global shape/dtype for |
| 151 | # metadata (no gather) and calls full_tensor() at broadcast time. |
| 152 | source=ModuleSource(self.model), |
| 153 | ) |
| 154 | |
| 155 | # ---- collective ops (ALL FSDP ranks must call concurrently) ---- |
| 156 | |
| 157 | def gather_and_broadcast_weights(self): |
| 158 | """All-gather full parameters and broadcast them to the vLLM server. |
| 159 | |
| 160 | Called on all FSDP ranks. `send_weights` gathers each param via |
| 161 | `full_tensor()` (a collective every rank must enter in the same order); |
| 162 | only rank 0 (the sender) drives the server-side update_weights |
| 163 | concurrently with the NCCL broadcast — the other ranks only gather. |
| 164 | """ |
| 165 | self.engine.send_weights() |
| 166 | |
| 167 | |
| 168 | def start_vllm_server(server_gpus: str) -> subprocess.Popen: |
| 169 | """Spawn a `vllm serve` HTTP server (EP+DP) on `server_gpus` and wait for it.""" |
| 170 | serve_args = [ |
| 171 | "vllm", |
| 172 | "serve", |
| 173 | MODEL_NAME, |
| 174 | "--served-model-name", |
| 175 | SERVED_MODEL_NAME, |
| 176 | "--tensor-parallel-size", |
| 177 | str(INFERENCE_TP_SIZE), |
| 178 | "--data-parallel-size", |
| 179 | str(INFERENCE_DP_SIZE), |
| 180 | "--enable-expert-parallel", |
| 181 | "--enforce-eager", |
| 182 | "--load-format", |
| 183 | "dummy", |
| 184 | "--gpu-memory-utilization", |
| 185 | "0.7", |
| 186 | "--port", |
| 187 | str(SERVER_PORT), |
| 188 | "--weight-transfer-config", |
| 189 | json.dumps({"backend": "nccl"}), |
| 190 | ] |
| 191 | env = os.environ.copy() |
| 192 | env["CUDA_VISIBLE_DEVICES"] = server_gpus |
| 193 | env["VLLM_SERVER_DEV_MODE"] = "1" # exposes the weight-transfer endpoints |
| 194 | env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" |
| 195 | print(f"[server] Launching: {' '.join(serve_args)} (GPUs {server_gpus})") |
| 196 | proc = subprocess.Popen( |
| 197 | serve_args, |
| 198 | env=env, |
| 199 | stdout=sys.stdout, |
| 200 | stderr=sys.stderr, |
| 201 | start_new_session=True, |
| 202 | ) |
| 203 | |
| 204 | # Wait for the server to come up (model load can take a while). |
| 205 | deadline = time.monotonic() + 1800 |
| 206 | while True: |
| 207 | if proc.poll() is not None: |
| 208 | raise RuntimeError("vLLM server exited before becoming ready.") |
| 209 | try: |
| 210 | if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200: |
| 211 | break |
| 212 | except requests.RequestException: |
| 213 | pass |
| 214 | if time.monotonic() > deadline: |
| 215 | raise RuntimeError("vLLM server failed to start in time.") |
| 216 | time.sleep(2) |
| 217 | print("[server] Ready.") |
| 218 | return proc |
| 219 | |
| 220 | |
| 221 | def generate_completions(client: OpenAI, prompts: list[str]) -> list[str]: |
| 222 | """Generate completions for a batch of prompts via the OpenAI HTTP API.""" |
| 223 | results = [] |
| 224 | for prompt in prompts: |
| 225 | response = client.completions.create( |
| 226 | model=SERVED_MODEL_NAME, |
| 227 | prompt=prompt, |
| 228 | max_tokens=32, |
| 229 | temperature=0, |
| 230 | ) |
| 231 | results.append(response.choices[0].text) |
| 232 | return results |
| 233 | |
| 234 | |
| 235 | def main(): |
| 236 | # Download model weights to local/shared disk once. |
| 237 | local_model_path = snapshot_download(MODEL_NAME) |
| 238 | print(f"[init] Model downloaded to {local_model_path}") |
| 239 | |
| 240 | ray.init() |
| 241 | |
| 242 | # FSDP rendezvous address (single-node). |
| 243 | fsdp_master_addr = get_ip() |
| 244 | fsdp_master_port = get_open_port() |
| 245 | |
| 246 | # Launch the FSDP training workers first so Ray reserves their GPUs, then |
| 247 | # place the inference server on the GPUs Ray did NOT use. This keeps the two |
| 248 | # on disjoint physical GPUs whether ray.init() started a fresh cluster or |
| 249 | # connected to an existing one. |
| 250 | fsdp_workers = [ |
| 251 | FSDPTrainWorker.remote( |
| 252 | local_model_path, |
| 253 | rank, |
| 254 | FSDP_WORLD_SIZE, |
| 255 | fsdp_master_addr, |
| 256 | fsdp_master_port, |
| 257 | ) |
| 258 | for rank in range(FSDP_WORLD_SIZE) |
| 259 | ] |
| 260 | ray.get([w.get_rank.remote() for w in fsdp_workers]) |
| 261 | print(f"[init] {FSDP_WORLD_SIZE} FSDP training workers ready.") |
| 262 | |
| 263 | # Discover the physical GPUs Ray assigned to training; run the server on the |
| 264 | # complementary GPUs. |
| 265 | training_gpus = { |
| 266 | int(g) |
| 267 | for ids in ray.get([w.get_gpu_ids.remote() for w in fsdp_workers]) |
| 268 | for g in ids |
| 269 | } |
| 270 | num_gpus = int(ray.cluster_resources().get("GPU", 0)) |
| 271 | num_server_gpus = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE |
| 272 | server_gpu_ids = [g for g in range(num_gpus) if g not in training_gpus][ |
| 273 | :num_server_gpus |
| 274 | ] |
| 275 | if len(server_gpu_ids) < num_server_gpus: |
| 276 | raise RuntimeError( |
| 277 | f"Need {num_server_gpus} free GPUs for the inference server but only " |
| 278 | f"found {server_gpu_ids} (training uses {sorted(training_gpus)} of " |
| 279 | f"{num_gpus} cluster GPUs)." |
| 280 | ) |
| 281 | server_gpus = ",".join(str(g) for g in server_gpu_ids) |
| 282 | print(f"[init] Training GPUs {sorted(training_gpus)}; server GPUs [{server_gpus}].") |
| 283 | |
| 284 | # Start the inference server on the complementary GPUs. |
| 285 | server_proc = start_vllm_server(server_gpus) |
| 286 | try: |
| 287 | client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY") |
| 288 | |
| 289 | prompts = [ |
| 290 | "Hello, my name is", |
| 291 | "The president of the United States is", |
| 292 | "The capital of France is", |
| 293 | "The future of AI is", |
| 294 | ] |
| 295 | |
| 296 | # Generate with dummy weights — expect gibberish. |
| 297 | print("[generate] Generating with dummy weights...") |
| 298 | outputs = generate_completions(client, prompts) |
| 299 | print("-" * 60) |
| 300 | print("BEFORE weight sync (dummy weights):") |
| 301 | print("-" * 60) |
| 302 | for prompt, text in zip(prompts, outputs): |
| 303 | print(f"Prompt: {prompt!r}") |
| 304 | print(f"Generated: {text!r}") |
| 305 | print("-" * 60) |
| 306 | |
| 307 | # --- Weight-transfer setup --- |
| 308 | print("[transfer] Setting up weight-transfer endpoint...") |
| 309 | transfer_addr, transfer_port = ray.get( |
| 310 | fsdp_workers[0].setup_transfer_endpoint.remote() |
| 311 | ) |
| 312 | print(f"[transfer] Endpoint ready at {transfer_addr}:{transfer_port}") |
| 313 | |
| 314 | transfer_world_size = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE + 1 |
| 315 | print( |
| 316 | f"[transfer] World size: {transfer_world_size} " |
| 317 | f"(1 trainer + {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} vLLM workers)" |
| 318 | ) |
| 319 | |
| 320 | # Build the trainer engine on all FSDP ranks (rank 0 is the sender). The |
| 321 | # sender drives the server's init_weight_transfer_engine (HTTP) while |
| 322 | # opening the trainer NCCL endpoint, so both ends rendezvous together; |
| 323 | # the other ranks build a null-client engine that only gathers. |
| 324 | print("[transfer] Initializing NCCL groups (all FSDP ranks)...") |
| 325 | ray.get( |
| 326 | [ |
| 327 | w.setup_engine.remote( |
| 328 | BASE_URL, transfer_addr, transfer_port, transfer_world_size |
| 329 | ) |
| 330 | for w in fsdp_workers |
| 331 | ] |
| 332 | ) |
| 333 | print("[transfer] NCCL groups initialized.") |
| 334 | |
| 335 | # --- Pause, transfer weights, resume --- |
| 336 | print("[sync] Pausing generation...") |
| 337 | requests.post(f"{BASE_URL}/pause", timeout=60).raise_for_status() |
| 338 | |
| 339 | # All ranks participate in the FSDP all-gather; rank 0 additionally |
| 340 | # drives start/update/finish on the server and the NCCL broadcast. |
| 341 | print("[sync] Broadcasting weights from FSDP → vLLM...") |
| 342 | ray.get([w.gather_and_broadcast_weights.remote() for w in fsdp_workers]) |
| 343 | print("[sync] Weight broadcast complete.") |
| 344 | |
| 345 | print("[sync] Resuming generation...") |
| 346 | requests.post(f"{BASE_URL}/resume", timeout=60).raise_for_status() |
| 347 | |
| 348 | # Generate with synced weights — expect sensible output. |
| 349 | print("[generate] Generating with synced weights...") |
| 350 | outputs_updated = generate_completions(client, prompts) |
| 351 | print("-" * 60) |
| 352 | print("AFTER weight sync (real weights):") |
| 353 | print("-" * 60) |
| 354 | for prompt, text in zip(prompts, outputs_updated): |
| 355 | print(f"Prompt: {prompt!r}") |
| 356 | print(f"Generated: {text!r}") |
| 357 | print("-" * 60) |
| 358 | finally: |
| 359 | print("[server] Shutting down...") |
| 360 | server_proc.terminate() |
| 361 | try: |
| 362 | server_proc.wait(timeout=30) |
| 363 | except subprocess.TimeoutExpired: |
| 364 | server_proc.kill() |
| 365 | |
| 366 | |
| 367 | if __name__ == "__main__": |
| 368 | main() |