| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Demonstrates async reinforcement learning using vLLM and Ray, |
| 4 | with native weight syncing APIs and batch-invariant generation. |
| 5 | |
| 6 | The script separates training and inference workloads onto distinct GPUs |
| 7 | so that Ray can manage process placement and inter-process communication. |
| 8 | A Hugging Face Transformer model occupies one GPU for training, and a |
| 9 | vLLM AsyncLLMEngine occupies another GPU for inference. |
| 10 | |
| 11 | Batch invariance is enabled so that generation output is deterministic |
| 12 | regardless of how many requests are batched together. This is required |
| 13 | for the validation phase to succeed. Batch invariance currently requires |
| 14 | NVIDIA GPUs with compute capability 9.0 or higher: |
| 15 | - H-series: H100, H200 |
| 16 | - B-series: B100, B200 |
| 17 | |
| 18 | The example performs the following steps: |
| 19 | * Load the training model (Qwen3-1.7B) on one GPU via a Ray actor. |
| 20 | * Initialize the inference engine with a base model (Qwen3-1.7B-Base) |
| 21 | on a separate GPU using vLLM's AsyncLLMEngine with Ray as the |
| 22 | distributed executor backend. |
| 23 | * Set up an NCCL-based weight transfer channel between the trainer |
| 24 | and the inference engine. |
| 25 | * Submit generation requests for a batch of prompts. |
| 26 | * Pause generation once any request reaches a token threshold. |
| 27 | * Broadcast the training model's weights to the inference engine |
| 28 | via the NCCL weight transfer engine, replacing the base weights. |
| 29 | * Resume generation and collect results, noting which tokens were |
| 30 | generated before vs. after the weight swap. |
| 31 | * Validate correctness by launching a fresh vLLM instance loaded |
| 32 | directly with the training model and comparing its output to the |
| 33 | post-swap tokens from the weight-synced engine. |
| 34 | |
| 35 | This example assumes a single-node cluster with two GPUs, but Ray |
| 36 | supports multi-node clusters. vLLM expects the GPUs are only used for vLLM |
| 37 | workloads. Residual GPU activity interferes with vLLM memory profiling and |
| 38 | causes unexpected behavior. |
| 39 | """ |
| 40 | |
| 41 | import asyncio |
| 42 | import uuid |
| 43 | |
| 44 | import ray |
| 45 | import torch |
| 46 | from transformers import AutoModelForCausalLM, AutoTokenizer |
| 47 | |
| 48 | import vllm |
| 49 | from vllm import SamplingParams |
| 50 | from vllm.config import WeightTransferConfig |
| 51 | from vllm.distributed.weight_transfer import ( |
| 52 | ModuleSource, |
| 53 | RayVLLMWeightSyncClient, |
| 54 | WeightTransferTrainerFactory, |
| 55 | ) |
| 56 | from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo |
| 57 | from vllm.platforms import current_platform |
| 58 | from vllm.utils.network_utils import get_ip, get_open_port |
| 59 | from vllm.v1.executor import Executor |
| 60 | |
| 61 | MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base" |
| 62 | MODEL_NAME_V2 = "Qwen/Qwen3-1.7B" |
| 63 | PAUSE_TOKEN_THRESHOLD = 10 |
| 64 | ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "FLASH_ATTN" |
| 65 | |
| 66 | |
| 67 | class MyLLM(vllm.AsyncLLMEngine): |
| 68 | """Configure the vLLM worker for Ray placement group execution.""" |
| 69 | |
| 70 | def __init__(self, **kwargs): |
| 71 | engine_args = vllm.AsyncEngineArgs(**kwargs) |
| 72 | vllm_config = engine_args.create_engine_config() |
| 73 | executor_class = Executor.get_class(vllm_config) |
| 74 | super().__init__( |
| 75 | vllm_config=vllm_config, |
| 76 | executor_class=executor_class, |
| 77 | log_requests=engine_args.enable_log_requests, |
| 78 | log_stats=not engine_args.disable_log_stats, |
| 79 | ) |
| 80 | self._generation_paused = False |
| 81 | self._request_pause_flag = False |
| 82 | |
| 83 | async def do_generate( |
| 84 | self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams |
| 85 | ) -> tuple[vllm.RequestOutput, int]: |
| 86 | """Generate a single request, setting the request pause flag once the |
| 87 | token count reaches the threshold. |
| 88 | |
| 89 | Returns (output, pause_token_index). pause_token_index is the number |
| 90 | of tokens generated before the weight change, or -1 if no pause. |
| 91 | """ |
| 92 | pause_token_index = -1 |
| 93 | prev_token_count = 0 |
| 94 | async for request_output in self.generate( |
| 95 | {"prompt_token_ids": prompt_token_ids}, |
| 96 | sampling_params, |
| 97 | request_id=str(uuid.uuid4()), |
| 98 | ): |
| 99 | output = request_output |
| 100 | cur_token_count = len(output.outputs[0].token_ids) |
| 101 | if ( |
| 102 | cur_token_count >= PAUSE_TOKEN_THRESHOLD |
| 103 | and not self._request_pause_flag |
| 104 | ): |
| 105 | self._request_pause_flag = True |
| 106 | if self._generation_paused and pause_token_index == -1: |
| 107 | pause_token_index = prev_token_count |
| 108 | prev_token_count = cur_token_count |
| 109 | return output, pause_token_index |
| 110 | |
| 111 | async def pause_after_n_tokens(self): |
| 112 | """Wait for any request to set the pause flag, then pause.""" |
| 113 | while not self._request_pause_flag: |
| 114 | await asyncio.sleep(0) |
| 115 | await super().pause_generation(mode="keep") |
| 116 | await asyncio.sleep(5) |
| 117 | self._generation_paused = True |
| 118 | |
| 119 | |
| 120 | @ray.remote(num_gpus=1) |
| 121 | class TrainModel: |
| 122 | """Ray actor that wraps the training model on a dedicated GPU.""" |
| 123 | |
| 124 | def __init__(self, model_name: str): |
| 125 | from vllm.model_executor.determinism.batch_invariant import ( |
| 126 | init_batch_invariance, |
| 127 | ) |
| 128 | |
| 129 | # need to init all env vars for batch invariance which affect nccl ops |
| 130 | init_batch_invariance() |
| 131 | |
| 132 | self.model = AutoModelForCausalLM.from_pretrained( |
| 133 | model_name, dtype=torch.bfloat16 |
| 134 | ).to("cuda:0") |
| 135 | self.port = get_open_port() |
| 136 | self.master_address = get_ip() |
| 137 | |
| 138 | def init_weight_transfer(self, world_size, llm_handle): |
| 139 | """Build the trainer-side weight-transfer engine and rendezvous.""" |
| 140 | self.engine = WeightTransferTrainerFactory.trainer_init( |
| 141 | init_info=NCCLTrainerInitInfo( |
| 142 | master_address=self.master_address, |
| 143 | master_port=self.port, |
| 144 | world_size=world_size, |
| 145 | rank=0, # single-GPU trainer is the sole (sender) rank |
| 146 | packed=True, |
| 147 | ), |
| 148 | client=RayVLLMWeightSyncClient(llm_handle), |
| 149 | source=ModuleSource(self.model), |
| 150 | ) |
| 151 | |
| 152 | def broadcast_weights(self): |
| 153 | """Push weights to the inference engine (drives start/update/finish).""" |
| 154 | self.engine.send_weights() |
| 155 | |
| 156 | @torch.inference_mode() |
| 157 | def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]: |
| 158 | """Greedy-decode max_new_tokens from the given context.""" |
| 159 | input_ids = torch.tensor([token_ids], device="cuda:0") |
| 160 | output = self.model.generate( |
| 161 | input_ids, |
| 162 | max_new_tokens=max_new_tokens, |
| 163 | do_sample=False, |
| 164 | ) |
| 165 | new_token_ids = output[0, len(token_ids) :].tolist() |
| 166 | return new_token_ids |
| 167 | |
| 168 | |
| 169 | # Build platform-specific env vars for Ray |
| 170 | ray_env_vars = {} |
| 171 | |
| 172 | if current_platform.is_rocm(): |
| 173 | # Workaround for RCCL bug. See https://github.com/ROCm/rocm-systems/issues/5756 |
| 174 | ray_env_vars["RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES"] = "1" |
| 175 | else: |
| 176 | # Enable batch invariance for deterministic outputs on NVIDIA |
| 177 | ray_env_vars["VLLM_BATCH_INVARIANT"] = "1" |
| 178 | |
| 179 | ray.init(runtime_env={"env_vars": ray_env_vars}) |
| 180 | |
| 181 | # Launch the training model actor. Ray's resource scheduler will allocate |
| 182 | # 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs. |
| 183 | train_model = TrainModel.remote(MODEL_NAME_V2) |
| 184 | |
| 185 | rocm_determinism_kwargs = {} |
| 186 | if current_platform.is_rocm(): |
| 187 | # ROCm: To minimize non-determinism, we set fixed seed, no prefix caching, and |
| 188 | # sequential request processing (max_num_seqs=1). |
| 189 | rocm_determinism_kwargs = { |
| 190 | "seed": 0, |
| 191 | "enable_prefix_caching": False, |
| 192 | "max_num_seqs": 1, |
| 193 | } |
| 194 | |
| 195 | # Build platform-specific LLM kwargs |
| 196 | llm_kwargs = dict( |
| 197 | model=MODEL_NAME_V1, |
| 198 | enforce_eager=True, |
| 199 | max_model_len=8192, |
| 200 | distributed_executor_backend="ray", |
| 201 | attention_backend=ATTN_BACKEND, |
| 202 | gpu_memory_utilization=0.75, |
| 203 | weight_transfer_config=WeightTransferConfig(backend="nccl"), |
| 204 | ) |
| 205 | llm_kwargs.update(rocm_determinism_kwargs) |
| 206 | |
| 207 | # Launch the vLLM inference engine. |
| 208 | # With data_parallel_backend="ray", vLLM's CoreEngineActorManager creates |
| 209 | # its own placement groups internally for each DP rank, so we must NOT |
| 210 | # create an outer placement group (it would reserve GPUs and hide them |
| 211 | # from the internal DP resource check). |
| 212 | llm = ray.remote( |
| 213 | num_cpus=0, |
| 214 | num_gpus=0, |
| 215 | )(MyLLM).remote(**llm_kwargs) |
| 216 | |
| 217 | PROMPTS = [ |
| 218 | "The president of the United States is", |
| 219 | "The capital of France is", |
| 220 | "The largest ocean on Earth is", |
| 221 | "The speed of light in a vacuum is", |
| 222 | "The chemical formula for water is", |
| 223 | "The tallest mountain in the world is", |
| 224 | "The first person to walk on the moon was", |
| 225 | "The Great Wall of China was built to", |
| 226 | "Photosynthesis is the process by which", |
| 227 | "The theory of general relativity was proposed by", |
| 228 | "The boiling point of water at sea level is", |
| 229 | "The largest planet in our solar system is", |
| 230 | "DNA stands for deoxyribonucleic acid and it", |
| 231 | ] |
| 232 | |
| 233 | tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_V1) |
| 234 | batch_prompt_token_ids = [ |
| 235 | tokenizer.encode(prompt, add_special_tokens=False) for prompt in PROMPTS |
| 236 | ] |
| 237 | |
| 238 | |
| 239 | # Set up the communication channel between the training process and the |
| 240 | # inference engine. |
| 241 | world_size = 2 # 1 trainer + 1 inference worker |
| 242 | ray.get(train_model.init_weight_transfer.remote(world_size, llm)) |
| 243 | |
| 244 | |
| 245 | N_NEW_TOKENS = 100 |
| 246 | |
| 247 | # ── Phase 1: concurrent requests with weight sync ─────────────────── |
| 248 | print(f"\n{'=' * 50}") |
| 249 | print(f"Prompts ({len(PROMPTS)}):") |
| 250 | for p in PROMPTS: |
| 251 | print(f" - {p!r}") |
| 252 | print(f"{'=' * 50}") |
| 253 | |
| 254 | sampling_params = SamplingParams( |
| 255 | temperature=0, max_tokens=PAUSE_TOKEN_THRESHOLD + N_NEW_TOKENS |
| 256 | ) |
| 257 | |
| 258 | gen_futures = [ |
| 259 | llm.do_generate.remote(ptids, sampling_params) for ptids in batch_prompt_token_ids |
| 260 | ] |
| 261 | |
| 262 | ray.get(llm.pause_after_n_tokens.remote()) |
| 263 | |
| 264 | ray.get(train_model.broadcast_weights.remote()) |
| 265 | |
| 266 | ray.get(llm.resume_generation.remote()) |
| 267 | results = ray.get(gen_futures) |
| 268 | |
| 269 | for i, (output, pause_idx) in enumerate(results): |
| 270 | all_token_ids = list(output.outputs[0].token_ids) |
| 271 | before_text = tokenizer.decode(all_token_ids[:pause_idx]) |
| 272 | after_text = tokenizer.decode(all_token_ids[pause_idx:]) |
| 273 | print(f"\n Request {i} ({PROMPTS[i]!r}):") |
| 274 | print(f" Old weights ({pause_idx} tokens): {before_text!r}") |
| 275 | n_after = len(all_token_ids) - pause_idx |
| 276 | print(f" New weights ({n_after} tokens): {after_text!r}") |
| 277 | |
| 278 | # ── Phase 2: validate with a fresh V2 vLLM instance ──────────────── |
| 279 | # This validation relies on batch-invariant (deterministic) generation to |
| 280 | # compare outputs from the weight-synced engine against a fresh V2 instance. |
| 281 | # On NVIDIA, batch invariance is fully supported, so we require 100% exact |
| 282 | # token match. On ROCm, batch invariance is not yet fully implemented |
| 283 | # (see https://github.com/vllm-project/vllm/issues/27433 and |
| 284 | # https://github.com/vllm-project/vllm/issues/33123), so residual |
| 285 | # non-determinism (e.g. GEMM accumulation order, missing kernel overrides) |
| 286 | # can cause single-token divergences that don't indicate a weight-sync |
| 287 | # failure. We relax the pass rate to 90% on ROCm to accommodate this; a |
| 288 | # real regression (broken weight transfer) would cause ~0% pass rate, not 90%+. |
| 289 | MIN_PASS_RATE = 1.0 if not current_platform.is_rocm() else 0.9 |
| 290 | |
| 291 | print(f"\n{'=' * 50}") |
| 292 | print("VALIDATION: comparing weight-synced vLLM with fresh V2 instance") |
| 293 | if current_platform.is_rocm(): |
| 294 | print(f" (ROCm mode: requiring >= {MIN_PASS_RATE:.0%} exact match rate)") |
| 295 | print(f"{'=' * 50}") |
| 296 | |
| 297 | ray.get(llm.shutdown.remote()) |
| 298 | ray.kill(llm) |
| 299 | ray.kill(train_model) |
| 300 | |
| 301 | llm_v2_kwargs = dict( |
| 302 | model=MODEL_NAME_V2, |
| 303 | enforce_eager=True, |
| 304 | max_model_len=8192, |
| 305 | gpu_memory_utilization=0.75, |
| 306 | distributed_executor_backend="ray", |
| 307 | attention_backend=ATTN_BACKEND, |
| 308 | ) |
| 309 | llm_v2_kwargs.update(rocm_determinism_kwargs) |
| 310 | |
| 311 | llm_v2 = ray.remote( |
| 312 | num_cpus=0, |
| 313 | num_gpus=0, |
| 314 | )(MyLLM).remote(**llm_v2_kwargs) |
| 315 | |
| 316 | val_futures = [ |
| 317 | llm_v2.do_generate.remote( |
| 318 | list(output.prompt_token_ids) + list(output.outputs[0].token_ids)[:pause_idx], |
| 319 | SamplingParams( |
| 320 | temperature=0, max_tokens=len(output.outputs[0].token_ids) - pause_idx |
| 321 | ), |
| 322 | ) |
| 323 | for output, pause_idx in results |
| 324 | ] |
| 325 | val_results = ray.get(val_futures) |
| 326 | |
| 327 | num_pass = 0 |
| 328 | num_total = len(results) |
| 329 | for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_results)): |
| 330 | expected = list(output.outputs[0].token_ids)[pause_idx:] |
| 331 | actual = list(val_output.outputs[0].token_ids) |
| 332 | match = actual == expected |
| 333 | |
| 334 | if match: |
| 335 | num_pass += 1 |
| 336 | print(f" [PASS] {PROMPTS[i]!r}") |
| 337 | else: |
| 338 | print(f" [FAIL] {PROMPTS[i]!r}") |
| 339 | print(f" weight-synced vLLM: {tokenizer.decode(expected)!r}") |
| 340 | print(f" V2 vLLM: {tokenizer.decode(actual)!r}") |
| 341 | for j, (e, a) in enumerate(zip(expected, actual)): |
| 342 | if e != a: |
| 343 | print( |
| 344 | f" first divergence at output token {j}: " |
| 345 | f"expected {e} ({tokenizer.decode([e])!r}) vs " |
| 346 | f"actual {a} ({tokenizer.decode([a])!r})" |
| 347 | ) |
| 348 | break |
| 349 | |
| 350 | ray.get(llm_v2.shutdown.remote()) |
| 351 | ray.kill(llm_v2) |
| 352 | |
| 353 | pass_rate = num_pass / num_total |
| 354 | print(f"\n Result: {num_pass}/{num_total} prompts passed ({pass_rate:.0%})") |
| 355 | print(f" Required: >= {MIN_PASS_RATE:.0%}") |
| 356 | |
| 357 | assert pass_rate >= MIN_PASS_RATE, ( |
| 358 | f"Validation pass rate {pass_rate:.0%} ({num_pass}/{num_total}) " |
| 359 | f"is below the required {MIN_PASS_RATE:.0%} threshold. " |
| 360 | f"See failures above for details." |
| 361 | ) |
| 362 | print("=" * 50) |