| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC** |
| 4 | weight transfer and **packed** tensors. |
| 5 | |
| 6 | Multi-rank version of `rlhf_http_ipc.py`: the trainer is 4 FSDP2 Ray actors |
| 7 | colocated with a data-parallel `vllm serve` on the same 4 physical GPUs. |
| 8 | |
| 9 | 4-GPU layout (single node), all colocated: |
| 10 | Training — 4 GPUs, PyTorch FSDP2 (fully_shard), as Ray actors |
| 11 | Inference — the same 4 GPUs, `vllm serve --data-parallel-size 4 -tp 1 |
| 12 | --enable-expert-parallel` (EP_SIZE = TP x DP = 4) |
| 13 | |
| 14 | IPC requires the trainer and the server to sit on the same GPUs, so the script |
| 15 | reserves the training GPUs through Ray first, asks Ray which ones it got, and |
| 16 | pins the server to exactly those with `--device-ids`. |
| 17 | |
| 18 | Both sides share each GPU, so the server is capped with |
| 19 | `--gpu-memory-utilization` and its weights are moved aside for the transfer: |
| 20 | |
| 21 | 1. `/sleep?level=1` — offload server weights to CPU, drop KV cache |
| 22 | 2. `/wake_up?tags=weights` — weights back on GPU, KV cache still free |
| 23 | 3. packed IPC transfer — overwrite weights with room to spare |
| 24 | 4. `/wake_up?tags=kv_cache&tags=scheduling` — re-allocate KV cache, resume |
| 25 | |
| 26 | Every FSDP rank builds an ``IPCTrainerWeightTransferEngine`` (via ``trainer_init``) |
| 27 | and calls ``send_weights()``; all ranks join the IPC handle all-gather, and only |
| 28 | rank 0 (the sender) ships the merged handles and drives the server. |
| 29 | |
| 30 | This example was run on 4xH100. |
| 31 | |
| 32 | Run: |
| 33 | $ python examples/rl/rlhf_ipc_fsdp_ep.py |
| 34 | """ |
| 35 | |
| 36 | from __future__ import annotations |
| 37 | |
| 38 | import os |
| 39 | import subprocess |
| 40 | import sys |
| 41 | import time |
| 42 | |
| 43 | import ray |
| 44 | import requests |
| 45 | import torch |
| 46 | import torch.distributed as dist |
| 47 | from huggingface_hub import snapshot_download |
| 48 | from openai import OpenAI |
| 49 | from torch.distributed.fsdp import fully_shard |
| 50 | from transformers import AutoModelForCausalLM |
| 51 | |
| 52 | from vllm.distributed.weight_transfer import ( |
| 53 | HTTPVLLMWeightSyncClient, |
| 54 | ModuleSource, |
| 55 | WeightTransferTrainerFactory, |
| 56 | ) |
| 57 | from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo |
| 58 | from vllm.utils.network_utils import get_ip, get_open_port |
| 59 | |
| 60 | MODEL_NAME = "Qwen/Qwen3-30B-A3B" |
| 61 | SERVED_MODEL_NAME = "policy" |
| 62 | |
| 63 | FSDP_WORLD_SIZE = 4 |
| 64 | INFERENCE_TP_SIZE = 1 |
| 65 | INFERENCE_DP_SIZE = 4 |
| 66 | |
| 67 | # Packed IPC transfer with a 1 GB chunk buffer. |
| 68 | PACKED = True |
| 69 | PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 |
| 70 | |
| 71 | # The server shares each GPU with a training rank, so cap what it reserves. |
| 72 | SERVER_GPU_MEMORY_UTILIZATION = 0.35 |
| 73 | |
| 74 | SERVER_PORT = 8000 |
| 75 | BASE_URL = f"http://localhost:{SERVER_PORT}" |
| 76 | |
| 77 | PROMPTS = [ |
| 78 | "Hello, my name is", |
| 79 | "The president of the United States is", |
| 80 | "The capital of France is", |
| 81 | "The future of AI is", |
| 82 | ] |
| 83 | |
| 84 | |
| 85 | @ray.remote(num_gpus=1) |
| 86 | class FSDPTrainWorker: |
| 87 | """One FSDP2 worker per GPU; colocated with one vLLM DP rank.""" |
| 88 | |
| 89 | def __init__( |
| 90 | self, |
| 91 | model_name: str, |
| 92 | rank: int, |
| 93 | fsdp_world_size: int, |
| 94 | fsdp_master_addr: str, |
| 95 | fsdp_master_port: int, |
| 96 | ): |
| 97 | self.rank = rank |
| 98 | |
| 99 | os.environ["MASTER_ADDR"] = fsdp_master_addr |
| 100 | os.environ["MASTER_PORT"] = str(fsdp_master_port) |
| 101 | |
| 102 | dist.init_process_group(backend="nccl", rank=rank, world_size=fsdp_world_size) |
| 103 | torch.accelerator.set_device_index(0) |
| 104 | |
| 105 | model = AutoModelForCausalLM.from_pretrained( |
| 106 | model_name, torch_dtype=torch.bfloat16 |
| 107 | ) |
| 108 | |
| 109 | for layer in model.model.layers: |
| 110 | fully_shard(layer) |
| 111 | fully_shard(model) |
| 112 | |
| 113 | self.model = model |
| 114 | |
| 115 | def get_rank(self): |
| 116 | return self.rank |
| 117 | |
| 118 | def get_gpu_ids(self): |
| 119 | """Physical GPU id(s) Ray assigned to this worker.""" |
| 120 | return ray.get_gpu_ids() |
| 121 | |
| 122 | def setup_engine(self, base_url: str): |
| 123 | """Build the trainer IPC engine. Called on every FSDP rank.""" |
| 124 | self.engine = WeightTransferTrainerFactory.trainer_init( |
| 125 | init_info=IPCTrainerInitInfo( |
| 126 | rank=self.rank, # FSDP rank; sender is 0 |
| 127 | packed=PACKED, |
| 128 | packed_buffer_size_bytes=PACKED_BUFFER_SIZE_BYTES, |
| 129 | ), |
| 130 | client=HTTPVLLMWeightSyncClient(base_url), |
| 131 | source=ModuleSource(self.model), |
| 132 | ) |
| 133 | |
| 134 | def gather_and_broadcast_weights_ipc(self): |
| 135 | """Send the current weights to vLLM. Called on every FSDP rank.""" |
| 136 | self.engine.send_weights() |
| 137 | |
| 138 | |
| 139 | def start_vllm_server(model_path: str, device_ids: str) -> subprocess.Popen: |
| 140 | """Spawn a `vllm serve` HTTP server (DP+EP) pinned to `device_ids`.""" |
| 141 | serve_args = [ |
| 142 | "vllm", |
| 143 | "serve", |
| 144 | model_path, |
| 145 | "--served-model-name", |
| 146 | SERVED_MODEL_NAME, |
| 147 | "--tensor-parallel-size", |
| 148 | str(INFERENCE_TP_SIZE), |
| 149 | "--data-parallel-size", |
| 150 | str(INFERENCE_DP_SIZE), |
| 151 | "--enable-expert-parallel", |
| 152 | # Pins the server to the same physical GPUs as the training ranks. |
| 153 | "--device-ids", |
| 154 | device_ids, |
| 155 | "--enable-sleep-mode", |
| 156 | "--enforce-eager", |
| 157 | "--load-format", |
| 158 | "dummy", |
| 159 | "--gpu-memory-utilization", |
| 160 | str(SERVER_GPU_MEMORY_UTILIZATION), |
| 161 | "--port", |
| 162 | str(SERVER_PORT), |
| 163 | "--weight-transfer-config", |
| 164 | '{"backend": "ipc"}', |
| 165 | ] |
| 166 | env = os.environ.copy() |
| 167 | env["VLLM_SERVER_DEV_MODE"] = "1" # exposes weight-transfer + sleep endpoints |
| 168 | env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" # IPC handles over HTTP |
| 169 | env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" |
| 170 | print(f"[server] Launching: {' '.join(serve_args)} (GPUs {device_ids})") |
| 171 | proc = subprocess.Popen( |
| 172 | serve_args, |
| 173 | env=env, |
| 174 | stdout=sys.stdout, |
| 175 | stderr=sys.stderr, |
| 176 | start_new_session=True, |
| 177 | ) |
| 178 | |
| 179 | # Wait for the server to come up (model load can take a while). |
| 180 | deadline = time.monotonic() + 1800 |
| 181 | while True: |
| 182 | if proc.poll() is not None: |
| 183 | raise RuntimeError("vLLM server exited before becoming ready.") |
| 184 | try: |
| 185 | if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200: |
| 186 | break |
| 187 | except requests.RequestException: |
| 188 | pass |
| 189 | if time.monotonic() > deadline: |
| 190 | raise RuntimeError("vLLM server failed to start in time.") |
| 191 | time.sleep(2) |
| 192 | print("[server] Ready.") |
| 193 | return proc |
| 194 | |
| 195 | |
| 196 | def generate_completions(client: OpenAI, prompts: list[str]) -> list[str]: |
| 197 | """Generate completions via the OpenAI HTTP API.""" |
| 198 | results = [] |
| 199 | for prompt in prompts: |
| 200 | response = client.completions.create( |
| 201 | model=SERVED_MODEL_NAME, |
| 202 | prompt=prompt, |
| 203 | max_tokens=32, |
| 204 | temperature=0, |
| 205 | ) |
| 206 | results.append(response.choices[0].text) |
| 207 | return results |
| 208 | |
| 209 | |
| 210 | def sleep_engine(base_url: str, level: int) -> None: |
| 211 | """Put the engine to sleep (level 1 offloads weights, drops KV cache).""" |
| 212 | response = requests.post(f"{base_url}/sleep", params={"level": level}, timeout=600) |
| 213 | response.raise_for_status() |
| 214 | |
| 215 | |
| 216 | def wake_up_engine(base_url: str, tags: list[str] | None = None) -> None: |
| 217 | """Wake the engine, optionally only for specific memory tags.""" |
| 218 | params = [("tags", tag) for tag in tags] if tags else None |
| 219 | response = requests.post(f"{base_url}/wake_up", params=params, timeout=600) |
| 220 | response.raise_for_status() |
| 221 | |
| 222 | |
| 223 | def print_generations(label: str, prompts: list[str], outputs: list[str]) -> None: |
| 224 | print("-" * 60) |
| 225 | print(label) |
| 226 | print("-" * 60) |
| 227 | for prompt, text in zip(prompts, outputs): |
| 228 | print(f"Prompt: {prompt!r}") |
| 229 | print(f"Generated: {text!r}") |
| 230 | print("-" * 60) |
| 231 | |
| 232 | |
| 233 | def main(): |
| 234 | ray.init( |
| 235 | runtime_env={ |
| 236 | "env_vars": { |
| 237 | # The trainer pickles IPC handles for the HTTP client. |
| 238 | "VLLM_ALLOW_INSECURE_SERIALIZATION": "1", |
| 239 | } |
| 240 | } |
| 241 | ) |
| 242 | |
| 243 | local_model_path = snapshot_download(MODEL_NAME) |
| 244 | print(f"[init] Model downloaded to {local_model_path}") |
| 245 | |
| 246 | fsdp_master_addr = get_ip() |
| 247 | fsdp_master_port = get_open_port() |
| 248 | |
| 249 | # Launch the training workers first so Ray reserves their GPUs; the server |
| 250 | # is then pinned to those same physical GPUs. |
| 251 | fsdp_workers = [ |
| 252 | FSDPTrainWorker.remote( |
| 253 | local_model_path, |
| 254 | rank, |
| 255 | FSDP_WORLD_SIZE, |
| 256 | fsdp_master_addr, |
| 257 | fsdp_master_port, |
| 258 | ) |
| 259 | for rank in range(FSDP_WORLD_SIZE) |
| 260 | ] |
| 261 | ray.get([w.get_rank.remote() for w in fsdp_workers]) |
| 262 | print(f"[init] {FSDP_WORLD_SIZE} FSDP training workers ready.") |
| 263 | |
| 264 | training_gpus = sorted( |
| 265 | int(g) |
| 266 | for ids in ray.get([w.get_gpu_ids.remote() for w in fsdp_workers]) |
| 267 | for g in ids |
| 268 | ) |
| 269 | if len(training_gpus) != INFERENCE_TP_SIZE * INFERENCE_DP_SIZE: |
| 270 | raise RuntimeError( |
| 271 | f"Need {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} colocated GPUs but " |
| 272 | f"Ray assigned training to {training_gpus}." |
| 273 | ) |
| 274 | device_ids = ",".join(str(g) for g in training_gpus) |
| 275 | print(f"[init] Colocating training and inference on GPUs [{device_ids}].") |
| 276 | |
| 277 | server_proc = start_vllm_server(local_model_path, device_ids) |
| 278 | try: |
| 279 | client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY") |
| 280 | |
| 281 | print("[generate] Generating with dummy weights...") |
| 282 | outputs = generate_completions(client, PROMPTS) |
| 283 | print_generations("BEFORE weight sync (dummy weights):", PROMPTS, outputs) |
| 284 | |
| 285 | # --- Weight transfer --- |
| 286 | print("[transfer] Initializing IPC weight transfer (all FSDP ranks)...") |
| 287 | ray.get([w.setup_engine.remote(BASE_URL) for w in fsdp_workers]) |
| 288 | |
| 289 | print("[sync] Sleeping engine (offload weights + free KV cache)...") |
| 290 | sleep_engine(BASE_URL, level=1) |
| 291 | |
| 292 | print("[sync] Waking weights (KV cache stays free)...") |
| 293 | wake_up_engine(BASE_URL, tags=["weights"]) |
| 294 | |
| 295 | print("[sync] Packed IPC transfer FSDP -> vLLM...") |
| 296 | ray.get([w.gather_and_broadcast_weights_ipc.remote() for w in fsdp_workers]) |
| 297 | print("[sync] Weight transfer complete.") |
| 298 | |
| 299 | print("[sync] Waking KV cache + scheduling...") |
| 300 | wake_up_engine(BASE_URL, tags=["kv_cache", "scheduling"]) |
| 301 | |
| 302 | print("[generate] Generating with synced weights...") |
| 303 | outputs_updated = generate_completions(client, PROMPTS) |
| 304 | print_generations("AFTER weight sync (real weights):", PROMPTS, outputs_updated) |
| 305 | finally: |
| 306 | print("[server] Shutting down...") |
| 307 | server_proc.terminate() |
| 308 | try: |
| 309 | server_proc.wait(timeout=30) |
| 310 | except subprocess.TimeoutExpired: |
| 311 | server_proc.kill() |
| 312 | |
| 313 | |
| 314 | if __name__ == "__main__": |
| 315 | main() |