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, 13:15 UTC
[docs] Add docs for new RL flows (#36188)
Signed-off-by: ahao-anyscale <ahao@anyscale.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Aaron Hao committed
Mar 18, 2026 at 02:04 UTC
47a1f11bffdd12cd59d90d79ff9867b7b3ac5b69
18 files changed
+514
-760
.buildkite/test-amd.yaml
+5
-5
@@ -1573,7 +1573,7 @@ steps:
1573
- tests/compile/fullgraph/test_basic_correctness.py
1574
- examples/offline_inference/rlhf.py
1575
- examples/offline_inference/rlhf_colocate.py
1576
- - examples/offline_inference/new_weight_syncing/
1576
+ - examples/rl/
1577
- tests/examples/offline_inference/data_parallel.py
1578
- tests/v1/distributed
1579
- tests/v1/engine/test_engine_core_client.py
@@ -1615,7 +1615,7 @@ steps:
1615
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
1616
- popd
1617
# NEW rlhf examples
1618
- - pushd ../examples/offline_inference/new_weight_syncing
1618
+ - pushd ../examples/rl
1619
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
1620
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
1621
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
@@ -2660,7 +2660,7 @@ steps:
2660
- tests/v1/entrypoints/openai/test_multi_api_servers.py
2661
- tests/v1/shutdown
2662
- tests/v1/worker/test_worker_memory_snapshot.py
2663
- - examples/offline_inference/new_weight_syncing/
2663
+ - examples/rl/
2664
commands:
2665
# Work around HIP bug tracked here: https://github.com/ROCm/hip/issues/3876
2666
# TODO: Remove when the bug is fixed in a future ROCm release
@@ -3325,7 +3325,7 @@ steps:
3325
- tests/compile/fullgraph/test_basic_correctness.py
3326
- examples/offline_inference/rlhf.py
3327
- examples/offline_inference/rlhf_colocate.py
3328
- - examples/offline_inference/new_weight_syncing/
3328
+ - examples/rl/
3329
- tests/examples/offline_inference/data_parallel.py
3330
- tests/v1/distributed
3331
- tests/v1/engine/test_engine_core_client.py
@@ -3367,7 +3367,7 @@ steps:
3367
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
3368
- popd
3369
# NEW rlhf examples
3370
- - pushd ../examples/offline_inference/new_weight_syncing
3370
+ - pushd ../examples/rl
3371
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
3372
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
3373
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
.buildkite/test_areas/distributed.yaml
+12
-17
@@ -82,7 +82,7 @@ steps:
82
83
- label: Distributed Torchrun + Examples (4 GPUs)
84
timeout_in_minutes: 30
85
- working_dir: "/vllm-workspace/tests"
85
+ working_dir: "/vllm-workspace"
86
num_devices: 4
87
source_file_dependencies:
88
- vllm/distributed/
@@ -90,33 +90,28 @@ steps:
90
- tests/distributed/test_torchrun_example_moe.py
91
- examples/offline_inference/rlhf.py
92
- examples/offline_inference/rlhf_colocate.py
93
- - examples/offline_inference/new_weight_syncing/
93
+ - examples/rl/
94
- tests/examples/offline_inference/data_parallel.py
95
commands:
96
# https://github.com/NVIDIA/nccl/issues/1838
97
- export NCCL_CUMEM_HOST_ENABLE=0
98
# test with torchrun tp=2 and external_dp=2
99
- - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
99
+ - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
100
# test with torchrun tp=2 and pp=2
101
- - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
101
+ - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
102
# test with torchrun tp=4 and dp=1
103
- - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
103
+ - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
104
# test with torchrun tp=2, pp=2 and dp=1
105
- - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
105
+ - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
106
# test with torchrun tp=1 and dp=4 with ep
107
- - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
107
+ - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
108
# test with torchrun tp=2 and dp=2 with ep
109
- - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
109
+ - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
110
# test with internal dp
111
- - python3 ../examples/offline_inference/data_parallel.py --enforce-eager
112
- # OLD rlhf examples
113
- - cd ../examples/offline_inference
114
- - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
115
- - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
116
- # NEW rlhf examples
117
- - cd new_weight_syncing
118
- - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
119
- - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
111
+ - python3 examples/offline_inference/data_parallel.py --enforce-eager
112
+ # rlhf examples
113
+ - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
114
+ - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
115
116
- label: Distributed DP Tests (4 GPUs)
117
timeout_in_minutes: 30
docs/mkdocs/hooks/generate_examples.py
+9
-1
@@ -23,15 +23,18 @@ def title(text: str) -> str:
23
# Custom substitutions
24
subs = {
25
"io": "IO",
26
- "api": "API",
26
+ "rl": "RL",
27
+ "api(s?)": r"API\1",
28
"cli": "CLI",
29
"cpu": "CPU",
30
+ "ipc": "IPC",
31
"llm": "LLM",
32
"mae": "MAE",
33
"ner": "NER",
34
"tpu": "TPU",
35
"gguf": "GGUF",
36
"lora": "LoRA",
37
+ "nccl": "NCCL",
38
"rlhf": "RLHF",
39
"vllm": "vLLM",
40
"openai": "OpenAI",
@@ -196,6 +199,11 @@ class Example:
199
200
201
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
202
+ # Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are
203
+ # title-cased (e.g. "Offline Inference" instead of "Offline inference").
204
+ import mkdocs_awesome_nav.nav.directory as _nav_dir
205
+
206
+ _nav_dir.dirname_to_title = title
207
logger.info("Generating example documentation")
208
logger.debug("Root directory: %s", ROOT_DIR.resolve())
209
logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
docs/training/async_rl.md
new
+63
@@ -0,0 +1,63 @@
1
+# Async Reinforcement Learning
2
+
3
+## Overview
4
+
5
+In a standard RL training loop, generation and training happen sequentially: the policy generates rollouts, then training runs on those rollouts, and the cycle repeats. During generation the training accelerators sit idle, and vice versa.
6
+
7
+The **one-off pipelining** approach separates the generation and training phases into two parallel coroutines, allowing the model to generate new samples while simultaneously training on previously generated data. This can lead to better GPU utilization and greater training throughput.
8
+
9
+However, this overlap introduces a complication: weights must be updated in the inference engine mid-flight, while requests may still be in progress.
10
+
11
+## The Pause and Resume API
12
+
13
+To safely update weights while the inference engine is running, vLLM provides `pause_generation` and `resume_generation` methods. These let the trainer coordinate a clean window for weight synchronization without losing in-flight work.
14
+
15
+### pause_generation
16
+
17
+```python
18
+await engine.pause_generation(mode="keep", clear_cache=True)
19
+```
20
+
21
+The `mode` parameter controls how in-flight requests are handled:
22
+
23
+| Mode | Behavior |
24
+| ---- | -------- |
25
+| `"abort"` | Abort all in-flight requests immediately and return partial results (default) |
26
+| `"wait"` | Wait for all in-flight requests to finish before pausing |
27
+| `"keep"` | Freeze requests in the queue; they resume when `resume_generation` is called |
28
+
29
+The `clear_cache` parameter controls whether to clear the KV cache and prefix cache after pausing.
30
+
31
+### resume_generation
32
+
33
+```python
34
+await engine.resume_generation()
35
+```
36
+
37
+Resumes the scheduler after a pause. Any requests frozen with `mode="keep"` will continue generating.
38
+
39
+### HTTP Endpoints
40
+
41
+When using the vLLM HTTP server, the same functionality is available via:
42
+
43
+- `POST /pause?mode=keep` - Pause generation
44
+- `POST /resume` - Resume generation
45
+
46
+!!! note "Data Parallelism"
47
+ When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update.
48
+
49
+## Typical Async RL Flow
50
+
51
+A typical async RL loop with weight syncing looks like this:
52
+
53
+1. Start generating rollouts from the current policy
54
+2. Once trainer has new weights to update to, pause generation with `mode="keep"`
55
+3. Sync the updated weights from the trainer to the inference engine (see [Weight Transfer](weight_transfer/README.md))
56
+4. Resume generation -- in-flight requests continue with the new weights
57
+5. Repeat
58
+
59
+The key insight is that requests paused with `mode="keep"` will produce tokens from the **old** weights before the pause and tokens from the **new** weights after resume. The `clear_cache` parameter controls whether the KV cache is invalidated during the pause. When `clear_cache=True`, previously cached key-value entries are discarded, so all tokens generated after resume will be computed entirely with the new weights. When `clear_cache=False`, existing KV cache entries are retained, meaning some tokens in context may still reflect the old weights (stale KV cache).
60
+
61
+## Example
62
+
63
+The [async RLHF example](../examples/rl/rlhf_async_new_apis.md) demonstrates this pattern with `vllm.AsyncLLMEngine`, NCCL weight transfer, and mid-flight pause/resume with validation.
docs/training/rlhf.md
+2
-4
@@ -16,11 +16,9 @@ The following open-source RL libraries use vLLM for fast rollouts (sorted alphab
16
- [Unsloth](https://github.com/unslothai/unsloth)
17
- [verl](https://github.com/volcengine/verl)
18
19
-See the following basic examples to get started if you don't want to use an existing library:
19
+For weight synchronization between training and inference, see the [Weight Transfer](weight_transfer/README.md) documentation, which covers the pluggable backend system with [NCCL](weight_transfer/nccl.md) (multi-GPU) and [IPC](weight_transfer/ipc.md) (same-GPU) engines.
20
21
-- [Training and inference processes are located on separate GPUs (inspired by OpenRLHF)](../examples/offline_inference/rlhf.md)
22
-- [Training and inference processes are colocated on the same GPUs using Ray](../examples/offline_inference/rlhf_colocate.md)
23
-- [Utilities for performing RLHF with vLLM](../examples/offline_inference/rlhf_utils.md)
21
+For pipelining generation and training to improve GPU utilization and throughput, see the [Async Reinforcement Learning](async_rl.md) guide, which covers the pause/resume API for safely updating weights mid-flight.
22
23
See the following notebooks showing how to use vLLM for GRPO:
24
docs/training/weight_transfer/README.md
new
+78
@@ -0,0 +1,78 @@
1
+# Weight Transfer
2
+
3
+vLLM provides a pluggable weight transfer system for synchronizing model weights from a training process to the inference engine during reinforcement learning (RL) workflows. This is essential for RLHF, GRPO, and other online RL methods where the policy model is iteratively updated during training and the updated weights must be reflected in the inference engine for rollout generation.
4
+
5
+## Architecture
6
+
7
+The weight transfer system follows a **two-phase protocol** with a pluggable backend design:
8
+
9
+1. **Initialization** (`init_weight_transfer_engine`): Establishes the communication channel between the trainer and inference workers. Called once before the training loop begins.
10
+2. **Weight Update** (`update_weights`): Transfers updated weights from the trainer to the inference engine. Called after each training step (or batch of steps).
11
+
12
+## Available Backends
13
+
14
+| Backend | Transport | Use Case |
15
+| ------- | --------- | -------- |
16
+| [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
17
+| [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
18
+
19
+## Configuration
20
+
21
+Specify the weight transfer backend through `WeightTransferConfig`. The backend determines which engine handles the weight synchronization.
22
+
23
+### Programmatic (Offline Inference)
24
+
25
+```python
26
+from vllm import LLM
27
+from vllm.config import WeightTransferConfig
28
+
29
+llm = LLM(
30
+ model="my-model",
31
+ weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc"
32
+)
33
+```
34
+
35
+### CLI (Online Serving)
36
+
37
+```bash
38
+vllm serve my-model \
39
+ --weight-transfer-config '{"backend": "nccl"}'
40
+```
41
+
42
+The `backend` field accepts `"nccl"` (default) or `"ipc"`.
43
+
44
+## API Endpoints
45
+
46
+When running vLLM as an HTTP server, the following endpoints are available for weight transfer:
47
+
48
+| Endpoint | Method | Description |
49
+| -------- | ------ | ----------- |
50
+| `/init_weight_transfer_engine` | POST | Initialize the weight transfer engine with backend-specific info |
51
+| `/update_weights` | POST | Trigger a weight update with backend-specific metadata |
52
+| `/pause` | POST | Pause generation before weight sync to handle inflight requests |
53
+| `/resume` | POST | Resume generation after weight sync |
54
+| `/get_world_size` | GET | Get the number of inference workers (useful for NCCL world size calculation) |
55
+
56
+!!! note
57
+ The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set.
58
+
59
+## Trainer-Side API
60
+
61
+Both backends provide static methods that the trainer calls to send weights. The general pattern is:
62
+
63
+```python
64
+# 1. Initialize the transfer engine (backend-specific)
65
+EngineClass.trainer_init(init_info)
66
+
67
+# 2. Send weights to inference workers
68
+EngineClass.trainer_send_weights(
69
+ iterator=model.named_parameters(),
70
+ trainer_args=backend_specific_args,
71
+)
72
+```
73
+
74
+See the [NCCL](nccl.md) and [IPC](ipc.md) pages for backend-specific trainer APIs and full examples.
75
+
76
+## Extending the System
77
+
78
+The weight transfer system is designed to be extensible. You can implement custom backends by subclassing `WeightTransferEngine` and registering them with the factory. See the [Base Class](base.md) page for details.
docs/training/weight_transfer/base.md
new
+162
@@ -0,0 +1,162 @@
1
+# Base Class and Custom Engines
2
+
3
+The weight transfer system is built on an abstract base class that defines the contract between vLLM's worker infrastructure and the transport backend. You can implement custom backends by subclassing `WeightTransferEngine` and registering them with the `WeightTransferEngineFactory`.
4
+
5
+## WeightTransferEngine
6
+
7
+The `WeightTransferEngine` is a generic abstract class parameterized by two dataclass types:
8
+
9
+- **`TInitInfo`** (extends `WeightTransferInitInfo`): Backend-specific initialization parameters.
10
+- **`TUpdateInfo`** (extends `WeightTransferUpdateInfo`): Backend-specific weight update metadata.
11
+
12
+### Abstract Methods
13
+
14
+Subclasses must implement these four methods:
15
+
16
+| Method | Side | Description |
17
+| ------ | ---- | ----------- |
18
+| `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker |
19
+| `receive_weights(update_info, load_weights)` | Inference | Receive weights and call `load_weights` incrementally |
20
+| `shutdown()` | Inference | Clean up resources |
21
+| `trainer_send_weights(iterator, trainer_args)` | Trainer | Static method to send weights from the trainer process |
22
+
23
+### Request Classes
24
+
25
+The API-level request classes provide backend-agnostic serialization using plain dictionaries. The engine's `parse_init_info` and `parse_update_info` methods convert these dictionaries into typed dataclasses.
26
+
27
+```python
28
+from vllm.distributed.weight_transfer.base import (
29
+ WeightTransferInitRequest,
30
+ WeightTransferUpdateRequest,
31
+)
32
+
33
+# Init request (dict is converted to backend-specific TInitInfo)
34
+init_request = WeightTransferInitRequest(
35
+ init_info={"master_address": "10.0.0.1", "master_port": 29500, ...}
36
+)
37
+
38
+# Update request (dict is converted to backend-specific TUpdateInfo)
39
+update_request = WeightTransferUpdateRequest(
40
+ update_info={"names": [...], "dtype_names": [...], "shapes": [...]}
41
+)
42
+```
43
+
44
+### WeightTransferUpdateInfo
45
+
46
+The base `WeightTransferUpdateInfo` includes an `is_checkpoint_format` flag:
47
+
48
+```python
49
+@dataclass
50
+class WeightTransferUpdateInfo(ABC):
51
+ is_checkpoint_format: bool = True
52
+```
53
+
54
+When `is_checkpoint_format=True` (the default), vLLM applies layerwise weight processing (repacking, renaming, etc.) on the received weights before loading them. Set to `False` if the trainer has already converted weights to the kernel format expected by the model.
55
+
56
+## Implementing a Custom Engine
57
+
58
+To create a custom weight transfer backend:
59
+
60
+### 1. Define Info Dataclasses
61
+
62
+```python
63
+from dataclasses import dataclass
64
+from vllm.distributed.weight_transfer.base import (
65
+ WeightTransferEngine,
66
+ WeightTransferInitInfo,
67
+ WeightTransferUpdateInfo,
68
+)
69
+
70
+@dataclass
71
+class MyInitInfo(WeightTransferInitInfo):
72
+ endpoint: str
73
+ token: str
74
+
75
+@dataclass
76
+class MyUpdateInfo(WeightTransferUpdateInfo):
77
+ names: list[str]
78
+ dtype_names: list[str]
79
+ shapes: list[list[int]]
80
+ # Add custom fields as needed
81
+```
82
+
83
+### 2. Implement the Engine
84
+
85
+```python
86
+from collections.abc import Callable, Iterator
87
+from typing import Any
88
+import torch
89
+
90
+class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
91
+ init_info_cls = MyInitInfo
92
+ update_info_cls = MyUpdateInfo
93
+
94
+ def init_transfer_engine(self, init_info: MyInitInfo) -> None:
95
+ # Set up connection to trainer using init_info.endpoint, etc.
96
+ ...
97
+
98
+ def receive_weights(
99
+ self,
100
+ update_info: MyUpdateInfo,
101
+ load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
102
+ ) -> None:
103
+ # Receive each weight and call load_weights incrementally
104
+ for name, dtype_name, shape in zip(
105
+ update_info.names, update_info.dtype_names, update_info.shapes
106
+ ):
107
+ dtype = getattr(torch, dtype_name)
108
+ weight = self._fetch_weight(name, shape, dtype)
109
+ load_weights([(name, weight)])
110
+
111
+ def shutdown(self) -> None:
112
+ # Clean up resources
113
+ ...
114
+
115
+ @staticmethod
116
+ def trainer_send_weights(
117
+ iterator: Iterator[tuple[str, torch.Tensor]],
118
+ trainer_args: dict[str, Any],
119
+ ) -> None:
120
+ # Send weights from the trainer process
121
+ for name, tensor in iterator:
122
+ # Send tensor via custom transport
123
+ ...
124
+```
125
+
126
+!!! important
127
+ The `load_weights` callable passed to `receive_weights` should be called **incrementally** (one or a few weights at a time) rather than accumulating all weights first. This avoids GPU out-of-memory errors with large models.
128
+
129
+### 3. Register with the Factory
130
+
131
+```python
132
+from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
133
+
134
+# Option 1: Lazy loading (recommended for built-in engines)
135
+WeightTransferEngineFactory.register_engine(
136
+ "my_backend",
137
+ "my_package.my_module",
138
+ "MyWeightTransferEngine",
139
+)
140
+
141
+# Option 2: Direct class registration
142
+WeightTransferEngineFactory.register_engine(
143
+ "my_backend",
144
+ MyWeightTransferEngine,
145
+)
146
+```
147
+
148
+Once registered, users can select your backend via `WeightTransferConfig(backend="my_backend")`.
149
+
150
+## WeightTransferEngineFactory
151
+
152
+The factory uses a registry pattern with lazy loading. Built-in engines (`nccl` and `ipc`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
153
+
154
+```python
155
+from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
156
+
157
+# Create an engine from config
158
+engine = WeightTransferEngineFactory.create_engine(
159
+ config=weight_transfer_config,
160
+ parallel_config=parallel_config,
161
+)
162
+```
docs/training/weight_transfer/ipc.md
new
+73
@@ -0,0 +1,73 @@
1
+# IPC Engine
2
+
3
+The IPC weight transfer engine uses **CUDA IPC** (Inter-Process Communication) handles to share GPU memory directly between the trainer and inference workers on the **same node and same GPU**. This avoids any data copying, making it a efficient option when colocating training and inference.
4
+
5
+## When to Use IPC
6
+
7
+- Training and inference on the **same GPU** (colocated)
8
+- You want to minimize memory overhead by sharing tensors in-place
9
+
10
+## How It Works
11
+
12
+1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`.
13
+2. IPC handles are sent to the inference engine via **Ray.remote()** or **HTTP POST**.
14
+3. The inference worker reconstructs the tensors from the handles, reading directly from the trainer's GPU memory.
15
+
16
+!!! warning
17
+ IPC handles involve sending serialized Python objects. When using HTTP transport, you must set `VLLM_ALLOW_INSECURE_SERIALIZATION=1` on both the server and client. This is because IPC handles are pickled and base64-encoded for HTTP transmission.
18
+
19
+## Initialization
20
+
21
+The IPC backend requires no initialization on either side. The `init_transfer_engine` call is a no-op for IPC.
22
+
23
+## Sending Weights
24
+
25
+IPC supports two transport modes for delivering the handles:
26
+
27
+### Ray Mode
28
+
29
+Used when vLLM is running as a Ray actor:
30
+
31
+```python
32
+from vllm.distributed.weight_transfer.ipc_engine import (
33
+ IPCTrainerSendWeightsArgs,
34
+ IPCWeightTransferEngine,
35
+)
36
+
37
+trainer_args = IPCTrainerSendWeightsArgs(
38
+ mode="ray",
39
+ llm_handle=llm_actor_handle,
40
+)
41
+
42
+IPCWeightTransferEngine.trainer_send_weights(
43
+ iterator=model.named_parameters(),
44
+ trainer_args=trainer_args,
45
+)
46
+```
47
+
48
+In Ray mode, the engine calls `llm_handle.update_weights.remote(...)` directly, passing the IPC handles via Ray's serialization.
49
+
50
+### HTTP Mode
51
+
52
+Used when vLLM is running as an HTTP server:
53
+
54
+```python
55
+trainer_args = IPCTrainerSendWeightsArgs(
56
+ mode="http",
57
+ url="http://localhost:8000",
58
+)
59
+
60
+IPCWeightTransferEngine.trainer_send_weights(
61
+ iterator=model.named_parameters(),
62
+ trainer_args=trainer_args,
63
+)
64
+```
65
+
66
+In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint.
67
+
68
+See [`IPCTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/ipc_engine.py) for the full list of configurable fields.
69
+
70
+## Examples
71
+
72
+- [RLHF with IPC weight syncing (offline, Ray)](../../examples/rl/rlhf_ipc.md) - Colocated training and inference on a single GPU using Ray placement groups and CUDA IPC handles
73
+- [RLHF with IPC weight syncing (online serving, HTTP)](../../examples/rl/rlhf_http_ipc.md) - Weight transfer with a vLLM HTTP server where both server and trainer share the same GPU
docs/training/weight_transfer/nccl.md
new
+110
@@ -0,0 +1,110 @@
1
+# NCCL Engine
2
+
3
+The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) broadcast operations to transfer weights from the trainer to inference workers. It supports **multi-node** and **multi-GPU** setups where the trainer and inference engine run on separate GPUs.
4
+
5
+## When to Use NCCL
6
+
7
+- Training and inference on **separate GPUs** (possibly across nodes)
8
+- **Tensor-parallel** inference with multiple workers that all need the updated weights
9
+- You need high-bandwidth, low-latency weight transfer over NVLink or InfiniBand
10
+
11
+## How It Works
12
+
13
+1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction).
14
+2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads weights incrementally.
15
+3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py).
16
+
17
+## Initialization
18
+
19
+NCCL requires explicit process group setup. The trainer and inference workers must agree on a master address, port, and world size.
20
+
21
+### Inference Side
22
+
23
+```python
24
+from vllm.distributed.weight_transfer.base import WeightTransferInitRequest
25
+
26
+# rank_offset accounts for the trainer occupying rank 0
27
+llm.init_weight_transfer_engine(
28
+ WeightTransferInitRequest(
29
+ init_info=dict(
30
+ master_address=master_address,
31
+ master_port=master_port,
32
+ rank_offset=1,
33
+ world_size=world_size, # trainer + all inference workers
34
+ )
35
+ )
36
+)
37
+```
38
+
39
+### Trainer Side
40
+
41
+```python
42
+from vllm.distributed.weight_transfer.nccl_engine import (
43
+ NCCLWeightTransferEngine,
44
+)
45
+
46
+group = NCCLWeightTransferEngine.trainer_init(
47
+ dict(
48
+ master_address=master_address,
49
+ master_port=master_port,
50
+ world_size=world_size,
51
+ )
52
+)
53
+```
54
+
55
+!!! note
56
+ `trainer_init` always assigns the trainer to rank 0. Inference workers start at `rank_offset` (typically 1).
57
+
58
+## Sending Weights
59
+
60
+```python
61
+from vllm.distributed.weight_transfer.nccl_engine import (
62
+ NCCLTrainerSendWeightsArgs,
63
+ NCCLWeightTransferEngine,
64
+)
65
+
66
+trainer_args = NCCLTrainerSendWeightsArgs(
67
+ group=group,
68
+ packed=True, # use packed broadcasting for efficiency
69
+)
70
+
71
+NCCLWeightTransferEngine.trainer_send_weights(
72
+ iterator=model.named_parameters(),
73
+ trainer_args=trainer_args,
74
+)
75
+```
76
+
77
+See [`NCCLTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/nccl_engine.py) for the full list of configurable fields.
78
+
79
+### Packed Tensor Broadcasting
80
+
81
+When `packed=True`, multiple weight tensors are packed into large contiguous buffers before broadcasting. This reduces the number of NCCL operations and uses double/triple buffering with dedicated CUDA streams for overlap between packing, broadcasting, and unpacking.
82
+
83
+Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightTransferUpdateInfo`) must use matching `packed_buffer_size_bytes` and `packed_num_buffers` values.
84
+
85
+## Receiving Weights (Inference Side)
86
+
87
+The inference side triggers weight reception by calling `update_weights`:
88
+
89
+```python
90
+from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
91
+
92
+llm.update_weights(
93
+ WeightTransferUpdateRequest(
94
+ update_info=dict(
95
+ names=names,
96
+ dtype_names=dtype_names,
97
+ shapes=shapes,
98
+ packed=True,
99
+ )
100
+ )
101
+)
102
+```
103
+
104
+The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters.
105
+
106
+## Examples
107
+
108
+- [RLHF with NCCL weight syncing (offline, Ray)](../../examples/rl/rlhf_nccl.md) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
109
+- [RLHF with async weight syncing (offline, Ray)](../../examples/rl/rlhf_async_new_apis.md) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
110
+- [RLHF with NCCL weight syncing (online serving, HTTP)](../../examples/rl/rlhf_http_nccl.md) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
examples/offline_inference/rlhf.py
deleted
-147
@@ -1,147 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray.
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 GPU 0 for training, whereas a
9
-tensor-parallel vLLM inference engine occupies GPU 1–2.
10
-
11
-The example performs the following steps:
12
-
13
-* Load the training model on GPU 0.
14
-* Split the inference model across GPUs 1–2 using vLLM's tensor parallelism
15
- and Ray placement groups.
16
-* Generate text from a list of prompts using the inference engine.
17
-* Update the weights of the training model and broadcast the updated weights
18
- to the inference engine by using a Ray collective RPC group. Note that
19
- for demonstration purposes we simply zero out the weights.
20
-
21
-For a production-ready implementation that supports multiple training and
22
-inference replicas, see the OpenRLHF framework:
23
-https://github.com/OpenRLHF/OpenRLHF
24
-
25
-This example assumes a single-node cluster with three GPUs, but Ray
26
-supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
27
-workloads. Residual GPU activity interferes with vLLM memory profiling and
28
-causes unexpected behavior.
29
-"""
30
-
31
-import os
32
-
33
-import ray
34
-import torch
35
-from ray.util.placement_group import placement_group
36
-from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
37
-from rlhf_utils import stateless_init_process_group
38
-from transformers import AutoModelForCausalLM
39
-
40
-from vllm import LLM, SamplingParams
41
-from vllm.utils.network_utils import get_ip, get_open_port
42
-
43
-
44
-class MyLLM(LLM):
45
- """Configure the vLLM worker for Ray placement group execution."""
46
-
47
- def __init__(self, *args, **kwargs):
48
- # Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
49
- # so that vLLM can manage its own device placement within the worker.
50
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
51
- super().__init__(*args, **kwargs)
52
-
53
-
54
-# Load the OPT-125M model onto GPU 0 for the training workload.
55
-train_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
56
-train_model.to("cuda:0")
57
-
58
-# Initialize Ray and set the visible devices. The vLLM engine will
59
-# be placed on GPUs 1 and 2.
60
-os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
61
-ray.init()
62
-
63
-# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
64
-# Learn more about Ray placement groups:
65
-# https://docs.ray.io/en/latest/ray-core/scheduling/placement-group.html
66
-pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
67
-ray.get(pg_inference.ready())
68
-scheduling_inference = PlacementGroupSchedulingStrategy(
69
- placement_group=pg_inference,
70
- placement_group_capture_child_tasks=True,
71
- placement_group_bundle_index=0,
72
-)
73
-
74
-# Launch the vLLM inference engine. The `enforce_eager` flag reduces
75
-# start-up latency.
76
-llm = ray.remote(
77
- num_cpus=0,
78
- num_gpus=0,
79
- scheduling_strategy=scheduling_inference,
80
-)(MyLLM).remote(
81
- model="facebook/opt-125m",
82
- enforce_eager=True,
83
- worker_extension_cls="rlhf_utils.WorkerExtension",
84
- tensor_parallel_size=2,
85
- distributed_executor_backend="ray",
86
-)
87
-
88
-# Generate text from the prompts.
89
-prompts = [
90
- "Hello, my name is",
91
- "The president of the United States is",
92
- "The capital of France is",
93
- "The future of AI is",
94
-]
95
-
96
-sampling_params = SamplingParams(temperature=0)
97
-
98
-outputs = ray.get(llm.generate.remote(prompts, sampling_params))
99
-
100
-print("-" * 50)
101
-for output in outputs:
102
- prompt = output.prompt
103
- generated_text = output.outputs[0].text
104
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
105
- print("-" * 50)
106
-
107
-# Set up the communication channel between the training process and the
108
-# inference engine.
109
-master_address = get_ip()
110
-master_port = get_open_port()
111
-
112
-handle = llm.collective_rpc.remote(
113
- "init_weight_update_group", args=(master_address, master_port, 1, 3)
114
-)
115
-
116
-model_update_group = stateless_init_process_group(
117
- master_address, master_port, 0, 3, torch.device("cuda:0")
118
-)
119
-ray.get(handle)
120
-
121
-# Simulate a training step by zeroing out all model weights.
122
-# In a real RLHF training loop the weights would be updated using the gradient
123
-# from an RL objective such as PPO on a reward model.
124
-for name, p in train_model.named_parameters():
125
- p.data.zero_()
126
-
127
-# Synchronize the updated weights to the inference engine.
128
-for name, p in train_model.named_parameters():
129
- dtype_name = str(p.dtype).split(".")[-1]
130
- handle = llm.collective_rpc.remote(
131
- "update_weight", args=(name, dtype_name, p.shape)
132
- )
133
- model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream())
134
- ray.get(handle)
135
-
136
-# Verify that the inference weights have been updated.
137
-assert all(ray.get(llm.collective_rpc.remote("check_weights_changed")))
138
-
139
-# Generate text with the updated model. The output is expected to be nonsense
140
-# because the weights are zero.
141
-outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
142
-print("-" * 50)
143
-for output in outputs_updated:
144
- prompt = output.prompt
145
- generated_text = output.outputs[0].text
146
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
147
- print("-" * 50)
examples/offline_inference/rlhf_colocate.py
deleted
-256
@@ -1,256 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Demonstrates how to co-locate a vLLM inference worker and training
5
-actors on the same set of GPUs for reinforcement learning from human feedback
6
-(RLHF) workloads.
7
-
8
-Ray serves as the distributed execution framework in this example. Ray
9
-placement groups allocate both training actors and vLLM workers to the
10
-same GPU bundles, enabling fast, in-GPU communication between the two
11
-components.
12
-
13
-The script shows how to do the following:
14
-
15
-* Configure environment variables (`VLLM_RAY_PER_WORKER_GPUS` and
16
- `VLLM_RAY_BUNDLE_INDICES`) so that vLLM workers land on the desired
17
- devices.
18
-* Exchange tensors between processes by means of CUDA inter-process
19
- communication (IPC). CUDA IPC sidesteps NCCL limitations that occur
20
- when multiple processes share a single GPU.
21
-
22
-Note that this example assumes a single-node cluster with four GPUs, but Ray
23
-supports multi-node clusters. vLLM expects exclusive use of the GPUs during
24
-its initialization for memory profiling. Residual GPU activity interferes
25
-with vLLM memory profiling and causes unexpected behavior.
26
-
27
-Learn more about Ray placement groups:
28
-https://docs.ray.io/en/latest/placement-groups.html
29
-"""
30
-
31
-import gc
32
-import os
33
-import sys
34
-
35
-import ray
36
-import torch
37
-import zmq
38
-from ray.util.placement_group import placement_group
39
-from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
40
-from torch.multiprocessing.reductions import reduce_tensor
41
-
42
-from vllm import LLM
43
-
44
-if torch.version.hip is not None:
45
- print("Skipping test for ROCm. Ray is unsupported on vLLM ROCm.")
46
- sys.exit(0)
47
-
48
-
49
-class MyLLM(LLM):
50
- """Configure the vLLM worker for Ray placement group execution.
51
-
52
- The constructor sets environment variables that allow multiple vLLM
53
- workers to share a single physical GPU and that encode the bundle
54
- indices assigned by the placement group.
55
-
56
- Args:
57
- *args: Positional arguments forwarded to `vllm.LLM`.
58
- bundle_indices (list[int]): Placement-group bundle indices
59
- assigned to this worker.
60
- **kwargs: Keyword arguments forwarded to `vllm.LLM`.
61
- """
62
-
63
- def __init__(self, *args, bundle_indices: list[int], **kwargs):
64
- # Prevent Ray from manipulating the top-level CUDA_VISIBLE_DEVICES variable
65
- # so that vLLM can its own device placement inside the worker.
66
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
67
- # Each worker uses 0.4 GPU so that two instances fit on the same GPUs.
68
- os.environ["VLLM_RAY_PER_WORKER_GPUS"] = "0.4"
69
- os.environ["VLLM_RAY_BUNDLE_INDICES"] = ",".join(map(str, bundle_indices))
70
- print(f"creating LLM with bundle_indices={bundle_indices}")
71
- super().__init__(*args, **kwargs)
72
-
73
-
74
-class RayTrainingActor:
75
- """Training actor that hosts a Facebook OPT-125M model from Hugging Face.
76
-
77
- The model is loaded onto the first GPU assigned to this actor, and expose
78
- the CUDA IPC handles so that colocated vLLM workers can map tensors
79
- directly.
80
- """
81
-
82
- def __init__(self):
83
- # Ray sets CUDA_VISIBLE_DEVICES to the GPUs assigned to this actor.
84
- from transformers import AutoModelForCausalLM
85
-
86
- self.model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
87
- self.model.to("cuda:0")
88
- # Zero out all the parameters.
89
- for name, p in self.model.named_parameters():
90
- p.data.zero_()
91
- torch.accelerator.synchronize()
92
- # The argument for `get_device_uuid` is the index of the GPU in the
93
- # list of visible devices.
94
- from vllm.platforms import current_platform
95
-
96
- self.device_uuid = current_platform.get_device_uuid(0)
97
- self.zmq_context = zmq.Context()
98
- self.zmq_address_counter = 0
99
- self.zmq_handle = None
100
-
101
- def report_device_id(self) -> str:
102
- return self.device_uuid
103
-
104
- def get_zmq_handles(self) -> dict[str, str]:
105
- suffix = f"{self.device_uuid}-{self.zmq_address_counter}"
106
- self.zmq_handle = f"ipc:///tmp/rl-colocate-zmq-{suffix}.sock"
107
- self.zmq_address_counter += 1
108
- return {self.device_uuid: self.zmq_handle}
109
-
110
- def update_weights(self):
111
- # align size to avoid misaligned address
112
- align_size = 256
113
-
114
- def get_size(p: torch.Tensor) -> int:
115
- return (p.nbytes + align_size - 1) // align_size * align_size
116
-
117
- named_parameters: dict[str, torch.nn.Parameter] = dict(
118
- self.model.named_parameters()
119
- )
120
- max_tensor_size = max(get_size(p) for p in named_parameters.values())
121
- # use max_tensor_size * 2 as buffer size
122
- buffer = torch.empty(max_tensor_size * 2, dtype=torch.uint8, device="cuda:0")
123
- s = self.zmq_context.socket(zmq.REQ)
124
- s.bind(self.zmq_handle)
125
- handle = reduce_tensor(buffer)
126
-
127
- offset = 0
128
- buckets: list[tuple[list[dict], list[torch.Tensor]]] = []
129
- named_tensors: list[dict] = []
130
- real_tensors: list[torch.Tensor] = []
131
- for name, p in named_parameters.items():
132
- size = get_size(p)
133
- if offset + size > buffer.numel():
134
- buckets.append((named_tensors, real_tensors))
135
- named_tensors, real_tensors = [], []
136
- offset = 0
137
- # assume tensors are contiguous
138
- named_tensors.append(
139
- {"name": name, "dtype": p.dtype, "shape": p.shape, "offset": offset}
140
- )
141
- real_tensors.append(p)
142
- offset += size
143
- if named_tensors:
144
- buckets.append((named_tensors, real_tensors))
145
- s.send_pyobj(handle)
146
- s.recv()
147
- for named_tensors, real_tensors in buckets:
148
- offset = 0
149
- for p in real_tensors:
150
- buffer[offset : offset + p.nbytes].data.copy_(
151
- p.data.view(-1).view(dtype=torch.uint8), non_blocking=True
152
- )
153
- offset += get_size(p)
154
- torch.accelerator.synchronize()
155
- s.send_pyobj(named_tensors)
156
- s.recv()
157
- s.send_pyobj(None)
158
- s.recv()
159
- s.close()
160
- del buffer
161
- gc.collect()
162
- torch.accelerator.empty_cache()
163
-
164
-
165
-# Ray manages four GPUs.
166
-
167
-os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
168
-ray.init()
169
-
170
-# Co-locate vLLM instances and training actors on the same set of GPUs:
171
-# * GPU 0 and 1: training actor 0, training actor 1, and vLLM instance 0
172
-# (tensor parallelism = 2).
173
-# * GPU 2 and 3: training actor 2, training actor 3, and vLLM instance 1
174
-# (tensor parallelism = 2).
175
-
176
-pg = placement_group([{"GPU": 1, "CPU": 0}] * 4)
177
-ray.get(pg.ready())
178
-print(f"placement group has bundles {pg.bundle_specs=}")
179
-
180
-training_actors = []
181
-training_actor_device_ids = []
182
-inference_engines = []
183
-inference_engine_device_ids = []
184
-
185
-for bundle_index in [0, 1, 2, 3]:
186
- training_actor = ray.remote(
187
- num_cpus=0,
188
- num_gpus=0.4,
189
- scheduling_strategy=PlacementGroupSchedulingStrategy(
190
- placement_group=pg,
191
- placement_group_capture_child_tasks=True,
192
- placement_group_bundle_index=bundle_index,
193
- ),
194
- )(RayTrainingActor).remote()
195
- training_actors.append(training_actor)
196
-
197
-for bundle_index, training_actor in enumerate(training_actors):
198
- device_id = ray.get(training_actor.report_device_id.remote())
199
- print(f"training actor {bundle_index} is on {device_id}")
200
- training_actor_device_ids.append(device_id)
201
-
202
-for i, bundle_indices in enumerate([[0, 1], [2, 3]]):
203
- # Use the following syntax instead of the @ray.remote decorator so that
204
- # the placement group is customized for each bundle.
205
- llm = ray.remote(
206
- num_cpus=0,
207
- num_gpus=0,
208
- scheduling_strategy=PlacementGroupSchedulingStrategy(
209
- placement_group=pg,
210
- placement_group_capture_child_tasks=True,
211
- ),
212
- )(MyLLM).remote(
213
- model="facebook/opt-125m",
214
- enforce_eager=True,
215
- worker_extension_cls="rlhf_utils.ColocateWorkerExtension",
216
- tensor_parallel_size=2,
217
- distributed_executor_backend="ray",
218
- gpu_memory_utilization=0.4,
219
- bundle_indices=bundle_indices,
220
- )
221
- inference_engines.append(llm)
222
- # Do not call any method on the inference engine at this point; the call
223
- # blocks until the vLLM instance finishes initialization.
224
-
225
-for i, llm in enumerate(inference_engines):
226
- inference_engine_device_ids.append(
227
- ray.get(llm.collective_rpc.remote("report_device_id", args=tuple()))
228
- )
229
- print(f"inference engine {i} is on {inference_engine_device_ids[-1]}")
230
-
231
-# Verify placement: the first two training actors share the same GPUs as
232
-# the first inference engine.
233
-assert training_actor_device_ids[:2] == inference_engine_device_ids[0]
234
-# Verify placement: the last two training actors share the same GPUs as
235
-# the second inference engine.
236
-assert training_actor_device_ids[2:] == inference_engine_device_ids[1]
237
-
238
-print("Gather all the ZMQ handles from the training actors.")
239
-zmq_handles = {}
240
-for actor in training_actors:
241
- zmq_handles.update(ray.get(actor.get_zmq_handles.remote()))
242
-
243
-print(f"ZMQ handles: {zmq_handles}")
244
-
245
-print("Update the weights of the inference engines.")
246
-ray.get(
247
- [actor.update_weights.remote() for actor in training_actors]
248
- + [
249
- llm.collective_rpc.remote("update_weights_from_ipc", args=(zmq_handles,))
250
- for llm in inference_engines
251
- ]
252
-)
253
-
254
-print("Check if the weights are updated.")
255
-for llm in inference_engines:
256
- assert ray.get(llm.collective_rpc.remote("check_weights_changed", args=tuple()))
examples/offline_inference/rlhf_online_quant.py
deleted
-162
@@ -1,162 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray.
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 GPU 0 for training, whereas a
9
-tensor-parallel vLLM inference engine occupies GPU 1–2.
10
-
11
-The example performs the following steps:
12
-
13
-* Load the training model on GPU 0.
14
-* Split the inference model across GPUs 1–2 using vLLM's tensor parallelism
15
- and Ray placement groups.
16
-* Generate text from a list of prompts using the inference engine.
17
-* Update the weights of the training model and broadcast the updated weights
18
- to the inference engine by using a Ray collective RPC group. Note that
19
- for demonstration purposes we simply zero out the weights.
20
-
21
-For a production-ready implementation that supports multiple training and
22
-inference replicas, see the OpenRLHF framework:
23
-https://github.com/OpenRLHF/OpenRLHF
24
-
25
-This example assumes a single-node cluster with three GPUs, but Ray
26
-supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
27
-workloads. Residual GPU activity interferes with vLLM memory profiling and
28
-causes unexpected behavior.
29
-"""
30
-
31
-import json
32
-import os
33
-
34
-import ray
35
-import torch
36
-from ray.util.placement_group import placement_group
37
-from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
38
-from rlhf_utils import stateless_init_process_group
39
-from torchao.core.config import config_to_dict
40
-from torchao.quantization import (
41
- Float8DynamicActivationFloat8WeightConfig,
42
- PerRow,
43
-)
44
-from transformers import AutoModelForCausalLM
45
-
46
-from vllm import LLM, SamplingParams
47
-from vllm.utils.network_utils import get_ip, get_open_port
48
-
49
-
50
-class MyLLM(LLM):
51
- """Configure the vLLM worker for Ray placement group execution."""
52
-
53
- def __init__(self, *args, **kwargs):
54
- # Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
55
- # so that vLLM can manage its own device placement within the worker.
56
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
57
- super().__init__(*args, **kwargs)
58
-
59
-
60
-# Load the OPT-125M model onto GPU 0 for the training workload.
61
-train_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
62
-train_model.to("cuda:0")
63
-
64
-# Initialize Ray and set the visible devices. The vLLM engine will
65
-# be placed on GPUs 1 and 2.
66
-os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
67
-ray.init()
68
-
69
-# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
70
-# Learn more about Ray placement groups:
71
-# https://docs.ray.io/en/latest/ray-core/scheduling/placement-group.html
72
-pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
73
-ray.get(pg_inference.ready())
74
-scheduling_inference = PlacementGroupSchedulingStrategy(
75
- placement_group=pg_inference,
76
- placement_group_capture_child_tasks=True,
77
- placement_group_bundle_index=0,
78
-)
79
-
80
-# Launch the vLLM inference engine. The `enforce_eager` flag reduces
81
-# start-up latency.
82
-
83
-# generate torchao quantization config for RL rollout
84
-# see https://github.com/vllm-project/vllm/pull/23014 for instructions to
85
-# use serialized config files instead of passing around json string
86
-config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
87
-
88
-json_str = json.dumps(config_to_dict(config))
89
-
90
-llm = ray.remote(
91
- num_cpus=0,
92
- num_gpus=0,
93
- scheduling_strategy=scheduling_inference,
94
-)(MyLLM).remote(
95
- model="facebook/opt-125m",
96
- hf_overrides={"quantization_config_dict_json": json_str},
97
- enforce_eager=True,
98
- worker_extension_cls="rlhf_utils.WorkerExtension",
99
- tensor_parallel_size=2,
100
- distributed_executor_backend="ray",
101
-)
102
-
103
-# Generate text from the prompts.
104
-prompts = [
105
- "Hello, my name is",
106
- "The president of the United States is",
107
- "The capital of France is",
108
- "The future of AI is",
109
-]
110
-
111
-sampling_params = SamplingParams(temperature=0)
112
-
113
-outputs = ray.get(llm.generate.remote(prompts, sampling_params))
114
-
115
-print("-" * 50)
116
-for output in outputs:
117
- prompt = output.prompt
118
- generated_text = output.outputs[0].text
119
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
120
- print("-" * 50)
121
-
122
-# Set up the communication channel between the training process and the
123
-# inference engine.
124
-master_address = get_ip()
125
-master_port = get_open_port()
126
-
127
-handle = llm.collective_rpc.remote(
128
- "init_weight_update_group", args=(master_address, master_port, 1, 3)
129
-)
130
-
131
-model_update_group = stateless_init_process_group(
132
- master_address, master_port, 0, 3, torch.device("cuda:0")
133
-)
134
-ray.get(handle)
135
-
136
-# Simulate a training step by zeroing out all model weights.
137
-# In a real RLHF training loop the weights would be updated using the gradient
138
-# from an RL objective such as PPO on a reward model.
139
-for name, p in train_model.named_parameters():
140
- p.data.zero_()
141
-
142
-# Synchronize the updated weights to the inference engine.
143
-for name, p in train_model.named_parameters():
144
- dtype_name = str(p.dtype).split(".")[-1]
145
- handle = llm.collective_rpc.remote(
146
- "update_weight", args=(name, dtype_name, p.shape)
147
- )
148
- model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream())
149
- ray.get(handle)
150
-
151
-# Verify that the inference weights have been updated.
152
-assert all(ray.get(llm.collective_rpc.remote("check_weights_changed")))
153
-
154
-# Generate text with the updated model. The output is expected to be nonsense
155
-# because the weights are zero.
156
-outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
157
-print("-" * 50)
158
-for output in outputs_updated:
159
- prompt = output.prompt
160
- generated_text = output.outputs[0].text
161
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
162
- print("-" * 50)
examples/offline_inference/rlhf_utils.py
deleted
-168
@@ -1,168 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import gc
4
-from collections.abc import Callable
5
-from typing import TypedDict
6
-
7
-import torch
8
-import zmq
9
-
10
-
11
-def stateless_init_process_group(master_address, master_port, rank, world_size, device):
12
- """
13
- vLLM provides `StatelessProcessGroup` to create a process group
14
- without considering the global process group in torch.distributed.
15
- It is recommended to create `StatelessProcessGroup`, and then initialize
16
- the data-plane communication (NCCL) between external (train processes)
17
- and vLLM workers.
18
- """
19
- from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
20
- from vllm.distributed.utils import StatelessProcessGroup
21
-
22
- pg = StatelessProcessGroup.create(
23
- host=master_address, port=master_port, rank=rank, world_size=world_size
24
- )
25
- pynccl = PyNcclCommunicator(pg, device=device)
26
- return pynccl
27
-
28
-
29
-class WorkerExtension:
30
- """
31
- The class for vLLM's worker to inherit from.
32
- By defining an extension class, the code can work no matter what is
33
- the underlying worker class.
34
-
35
- NOTE: we define this class in a separate module, and the main module
36
- should pass the full qualified name as `worker_extension_cls` argument.
37
- """
38
-
39
- def init_weight_update_group(
40
- self, master_address, master_port, rank_offset, world_size
41
- ):
42
- from vllm.distributed.parallel_state import get_world_group
43
-
44
- rank = get_world_group().rank + rank_offset
45
- self.model_update_group = stateless_init_process_group(
46
- master_address,
47
- master_port,
48
- rank,
49
- world_size,
50
- self.device,
51
- )
52
-
53
- def update_weight(self, name, dtype_name, shape):
54
- dtype = getattr(torch, dtype_name)
55
- weight = torch.empty(shape, dtype=dtype, device="cuda")
56
- self.model_update_group.broadcast(
57
- weight, src=0, stream=torch.cuda.current_stream()
58
- )
59
-
60
- self.model_runner.model.load_weights(weights=[(name, weight)])
61
-
62
- del weight
63
-
64
- def check_weights_changed(self):
65
- """
66
- Check if the weights are updated to 0.
67
- """
68
- weights_updated = True
69
- for name, p in self.model_runner.model.named_parameters():
70
- weights_updated = weights_updated and torch.allclose(p, torch.zeros_like(p))
71
- return weights_updated
72
-
73
-
74
-def rebuild_ipc(
75
- handle: tuple[Callable, tuple], device_id: int | None = None
76
-) -> torch.Tensor:
77
- func, args = handle
78
- list_args = list(args)
79
- if device_id is not None:
80
- # the key is to change device id to the current device id
81
- # in case two processes have different CUDA_VISIBLE_DEVICES
82
- list_args[6] = device_id
83
- buffer = func(*list_args)
84
- return buffer
85
-
86
-
87
-class FlattenedTensorMetadata(TypedDict):
88
- name: str
89
- shape: torch.Size
90
- dtype: torch.dtype
91
- # specify the start offset of this tensor in shared ipc_buffer tensor
92
- offset: int
93
-
94
-
95
-class ColocateWorkerExtension:
96
- """
97
- The class for vLLM's worker to inherit from, in the colocate setting.
98
- By defining an extension class, the code can work no matter what is
99
- the underlying worker class.
100
-
101
- NOTE: we define this class in a separate module, and the main module
102
- should pass the full qualified name as `worker_extension_cls` argument.
103
- """
104
-
105
- def update_weights_from_ipc(self, zmq_handles: dict[str, str]):
106
- from vllm.model_executor.model_loader.utils import process_weights_after_loading
107
-
108
- assert self.device is not None
109
- if not hasattr(self, "_zmq_ctx") or self._zmq_ctx is None:
110
- self._zmq_ctx = zmq.Context()
111
- socket = self._zmq_ctx.socket(zmq.REP)
112
- socket.connect(zmq_handles[self.report_device_id()])
113
- buffer: torch.Tensor | None = None
114
- while True:
115
- payload: tuple[Callable, tuple] | list[FlattenedTensorMetadata] | None = (
116
- socket.recv_pyobj()
117
- )
118
- if payload is None:
119
- # means the update is done
120
- process_weights_after_loading(
121
- self.model_runner.model, self.model_config, self.device
122
- )
123
- torch.accelerator.synchronize()
124
- socket.send(b"")
125
- break
126
- if isinstance(payload, tuple):
127
- # an ipc handle that vLLM can use `func, args = handle`
128
- # and `func(*args)` to rebuild GPU tensor.
129
- buffer = rebuild_ipc(payload, self.device.index)
130
- assert buffer.dtype == torch.uint8
131
- socket.send(b"")
132
- continue
133
- assert isinstance(payload, list)
134
- assert buffer is not None
135
- weights = []
136
- for item in payload:
137
- shape = item["shape"]
138
- if isinstance(shape, (list, tuple)):
139
- shape = torch.Size(shape)
140
- assert isinstance(shape, torch.Size)
141
- dtype, offset = item["dtype"], item["offset"]
142
- size = dtype.itemsize * shape.numel()
143
- tensor = buffer[offset : offset + size].view(dtype=dtype).view(shape)
144
- weights.append((item["name"], tensor))
145
- self.model_runner.model.load_weights(weights=weights)
146
- del weights
147
- torch.accelerator.synchronize()
148
- socket.send(b"")
149
-
150
- socket.close()
151
- del buffer
152
- gc.collect()
153
- torch.accelerator.empty_cache()
154
-
155
- def report_device_id(self) -> str:
156
- from vllm.platforms import current_platform
157
-
158
- self.device_uuid = current_platform.get_device_uuid(self.device.index)
159
- return self.device_uuid
160
-
161
- def check_weights_changed(self):
162
- """
163
- Check if the weights are updated to 0.
164
- """
165
- weights_updated = True
166
- for name, p in self.model_runner.model.named_parameters():
167
- weights_updated = weights_updated and torch.allclose(p, torch.zeros_like(p))
168
- return weights_updated