| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """RLHF weight syncing against a `vllm serve` HTTP server, using NCCL for the |
| 4 | data plane. |
| 5 | |
| 6 | * OpenAI-compatible API for inference requests |
| 7 | * HTTP endpoints for the weight-transfer control plane |
| 8 | * NCCL for the weight data plane |
| 9 | |
| 10 | 3-GPU layout (single node): |
| 11 | Inference — GPUs 0-1, `vllm serve` with TP=2 and fp8 quantization |
| 12 | Training — GPU 2, a bf16 Hugging Face model in this process |
| 13 | (the server quantizes to fp8 as it loads) |
| 14 | |
| 15 | The script starts the server itself, then: |
| 16 | |
| 17 | 1. Generate over HTTP → gibberish (server started with dummy weights). |
| 18 | 2. Pause generation, sync real weights trainer → server over NCCL, resume. |
| 19 | 3. Generate again → sensible output. |
| 20 | |
| 21 | Run: |
| 22 | $ python examples/rl/rlhf_http_nccl.py |
| 23 | """ |
| 24 | |
| 25 | import os |
| 26 | import subprocess |
| 27 | import sys |
| 28 | import time |
| 29 | |
| 30 | import requests |
| 31 | import torch |
| 32 | from openai import OpenAI |
| 33 | from transformers import AutoModelForCausalLM |
| 34 | |
| 35 | from vllm.distributed.weight_transfer import ( |
| 36 | HTTPVLLMWeightSyncClient, |
| 37 | ModuleSource, |
| 38 | WeightTransferTrainerFactory, |
| 39 | ) |
| 40 | from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo |
| 41 | from vllm.utils.network_utils import get_ip, get_open_port |
| 42 | |
| 43 | MODEL_NAME = "facebook/opt-125m" |
| 44 | |
| 45 | SERVER_PORT = 8000 |
| 46 | BASE_URL = f"http://localhost:{SERVER_PORT}" |
| 47 | |
| 48 | INFERENCE_TP_SIZE = 2 |
| 49 | # Physical GPUs for the server; the trainer takes the next one. |
| 50 | SERVER_DEVICE_IDS = "0,1" |
| 51 | TRAINER_DEVICE = "cuda:2" |
| 52 | |
| 53 | PROMPTS = [ |
| 54 | "Hello, my name is", |
| 55 | "The president of the United States is", |
| 56 | "The capital of France is", |
| 57 | "The future of AI is", |
| 58 | ] |
| 59 | |
| 60 | |
| 61 | def start_vllm_server() -> subprocess.Popen: |
| 62 | """Spawn `vllm serve` and block until it is healthy.""" |
| 63 | serve_args = [ |
| 64 | "vllm", |
| 65 | "serve", |
| 66 | MODEL_NAME, |
| 67 | "--tensor-parallel-size", |
| 68 | str(INFERENCE_TP_SIZE), |
| 69 | "--device-ids", |
| 70 | SERVER_DEVICE_IDS, |
| 71 | "--quantization", |
| 72 | "fp8", |
| 73 | "--enforce-eager", |
| 74 | "--load-format", |
| 75 | "dummy", |
| 76 | "--port", |
| 77 | str(SERVER_PORT), |
| 78 | "--weight-transfer-config", |
| 79 | '{"backend": "nccl"}', |
| 80 | ] |
| 81 | env = os.environ.copy() |
| 82 | # Exposes the weight-transfer and pause/resume endpoints. |
| 83 | env["VLLM_SERVER_DEV_MODE"] = "1" |
| 84 | print(f"[server] Launching: {' '.join(serve_args)}") |
| 85 | proc = subprocess.Popen( |
| 86 | serve_args, |
| 87 | env=env, |
| 88 | stdout=sys.stdout, |
| 89 | stderr=sys.stderr, |
| 90 | start_new_session=True, |
| 91 | ) |
| 92 | |
| 93 | deadline = time.monotonic() + 900 |
| 94 | while True: |
| 95 | if proc.poll() is not None: |
| 96 | raise RuntimeError("vLLM server exited before becoming ready.") |
| 97 | try: |
| 98 | if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200: |
| 99 | break |
| 100 | except requests.RequestException: |
| 101 | pass |
| 102 | if time.monotonic() > deadline: |
| 103 | raise RuntimeError("vLLM server failed to start in time.") |
| 104 | time.sleep(2) |
| 105 | print("[server] Ready.") |
| 106 | return proc |
| 107 | |
| 108 | |
| 109 | def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]: |
| 110 | """Generate completions using the OpenAI-compatible API.""" |
| 111 | results = [] |
| 112 | for prompt in prompts: |
| 113 | response = client.completions.create( |
| 114 | model=model, |
| 115 | prompt=prompt, |
| 116 | max_tokens=32, |
| 117 | temperature=0, |
| 118 | ) |
| 119 | results.append(response.choices[0].text) |
| 120 | return results |
| 121 | |
| 122 | |
| 123 | def pause_generation(base_url: str) -> None: |
| 124 | """Pause generation via HTTP endpoint.""" |
| 125 | requests.post(f"{base_url}/pause", timeout=60).raise_for_status() |
| 126 | |
| 127 | |
| 128 | def resume_generation(base_url: str) -> None: |
| 129 | """Resume generation via HTTP endpoint.""" |
| 130 | requests.post(f"{base_url}/resume", timeout=60).raise_for_status() |
| 131 | |
| 132 | |
| 133 | def get_world_size(base_url: str) -> int: |
| 134 | """Get the number of inference workers from the vLLM server.""" |
| 135 | response = requests.get(f"{base_url}/get_world_size", timeout=10) |
| 136 | response.raise_for_status() |
| 137 | return response.json()["world_size"] |
| 138 | |
| 139 | |
| 140 | def print_generations(label: str, prompts: list[str], outputs: list[str]) -> None: |
| 141 | print("-" * 50) |
| 142 | print(label) |
| 143 | print("-" * 50) |
| 144 | for prompt, generated_text in zip(prompts, outputs): |
| 145 | print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}") |
| 146 | print("-" * 50) |
| 147 | |
| 148 | |
| 149 | def main(): |
| 150 | server_proc = start_vllm_server() |
| 151 | try: |
| 152 | # The trainer sits on the GPU after the server's, and is NCCL rank 0. |
| 153 | torch.accelerator.set_device_index(TRAINER_DEVICE) |
| 154 | |
| 155 | print(f"[trainer] Loading training model: {MODEL_NAME} on {TRAINER_DEVICE}") |
| 156 | train_model = AutoModelForCausalLM.from_pretrained( |
| 157 | MODEL_NAME, dtype=torch.bfloat16 |
| 158 | ) |
| 159 | train_model.to(TRAINER_DEVICE) |
| 160 | |
| 161 | client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY") |
| 162 | |
| 163 | # Generate with dummy weights — expect nonsense. |
| 164 | outputs = generate_completions(client, MODEL_NAME, PROMPTS) |
| 165 | print_generations("BEFORE weight sync (dummy weights):", PROMPTS, outputs) |
| 166 | |
| 167 | # The transfer NCCL group is the trainer plus every inference worker. |
| 168 | world_size = get_world_size(BASE_URL) + 1 |
| 169 | master_address = get_ip() |
| 170 | master_port = get_open_port() |
| 171 | print( |
| 172 | f"[transfer] Rendezvous at {master_address}:{master_port}, " |
| 173 | f"world_size={world_size} (1 trainer + {world_size - 1} vLLM workers)" |
| 174 | ) |
| 175 | |
| 176 | # `trainer_init` drives the handshake: it initializes the server's |
| 177 | # transfer engine while opening the trainer's own NCCL endpoint, so both |
| 178 | # ends rendezvous together. |
| 179 | engine = WeightTransferTrainerFactory.trainer_init( |
| 180 | init_info=NCCLTrainerInitInfo( |
| 181 | master_address=master_address, |
| 182 | master_port=master_port, |
| 183 | world_size=world_size, |
| 184 | rank=0, # single-GPU trainer is the sole (sender) rank |
| 185 | packed=True, |
| 186 | ), |
| 187 | client=HTTPVLLMWeightSyncClient(BASE_URL), |
| 188 | source=ModuleSource(train_model), |
| 189 | ) |
| 190 | |
| 191 | pause_generation(BASE_URL) |
| 192 | |
| 193 | # Drives start_weight_update / update_weights / finish_weight_update, |
| 194 | # concurrent with the NCCL broadcast. |
| 195 | print("[sync] Broadcasting weights via NCCL...") |
| 196 | engine.send_weights() |
| 197 | print("[sync] Weight broadcast complete.") |
| 198 | |
| 199 | resume_generation(BASE_URL) |
| 200 | |
| 201 | # Generate with the synced weights — expect sensible output. |
| 202 | outputs_updated = generate_completions(client, MODEL_NAME, PROMPTS) |
| 203 | print_generations("AFTER weight sync (real weights):", PROMPTS, outputs_updated) |
| 204 | finally: |
| 205 | print("[server] Shutting down...") |
| 206 | server_proc.terminate() |
| 207 | try: |
| 208 | server_proc.wait(timeout=30) |
| 209 | except subprocess.TimeoutExpired: |
| 210 | server_proc.kill() |
| 211 | |
| 212 | |
| 213 | if __name__ == "__main__": |
| 214 | main() |