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, 09:30 UTC
Comment fix for async rl example (#35244)
Signed-off-by: hao-aaron <ahao@anyscale.com>
Aaron Hao committed
Mar 19, 2026 at 12:46 UTC
4ee847e40094c66669f6095c034a4ff7ef8ad39f
1 file changed
+27
-14
examples/rl/rlhf_async_new_apis.py
+27
-14
@@ -2,25 +2,38 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""
4
Demonstrates async reinforcement learning using vLLM and Ray,
5
-with native weight syncing APIs at engine instance.
5
+with native weight syncing APIs and batch-invariant generation.
6
7
The script separates training and inference workloads onto distinct GPUs
8
so that Ray can manage process placement and inter-process communication.
9
-A Hugging Face Transformer model occupies one GPU for training, whereas a
10
-2x tensor-parallel vLLM inference engine occupies two GPUs.
9
+A Hugging Face Transformer model occupies one GPU for training, and a
10
+vLLM AsyncLLMEngine occupies another GPU for inference.
11
+
12
+Batch invariance is enabled so that generation output is deterministic
13
+regardless of how many requests are batched together. This is required
14
+for the validation phase to succeed. Batch invariance currently requires
15
+NVIDIA GPUs with compute capability 9.0 or higher:
16
+ - H-series: H100, H200
17
+ - B-series: B100, B200
18
19
The example performs the following steps:
13
-* Load the training model on one gpu (scheduled via ray)
14
-* Initialize the inference model with dummy weights across
15
- two gpus using vLLM's tensor parallelism and Ray placement groups.
16
-* Generate gibberish from a list of prompts using the randomly initialized
17
- inference engine.
18
-* Pause generation once generation completes for one sequence
19
-* Update the weights of the training model and broadcast the updated weights
20
- to the inference engine by using a Ray collective RPC group.
21
-* Resume generation and print out the results
22
-
23
-This example assumes a single-node cluster with three GPUs, but Ray
20
+* Load the training model (Qwen3-1.7B) on one GPU via a Ray actor.
21
+* Initialize the inference engine with a base model (Qwen3-1.7B-Base)
22
+ on a separate GPU using vLLM's AsyncLLMEngine with Ray as the
23
+ distributed executor backend.
24
+* Set up an NCCL-based weight transfer channel between the trainer
25
+ and the inference engine.
26
+* Submit generation requests for a batch of prompts.
27
+* Pause generation once any request reaches a token threshold.
28
+* Broadcast the training model's weights to the inference engine
29
+ via the NCCL weight transfer engine, replacing the base weights.
30
+* Resume generation and collect results, noting which tokens were
31
+ generated before vs. after the weight swap.
32
+* Validate correctness by launching a fresh vLLM instance loaded
33
+ directly with the training model and comparing its output to the
34
+ post-swap tokens from the weight-synced engine.
35
+
36
+This example assumes a single-node cluster with two GPUs, but Ray
37
supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
38
workloads. Residual GPU activity interferes with vLLM memory profiling and
39
causes unexpected behavior.