-
-
Notifications
You must be signed in to change notification settings - Fork 19.3k
[docs] Add docs for new RL flows #36188
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from 3 commits
Commits
Show all changes
13 commits
Select commit
Hold shift + click to select a range
76e9da5
docs
hao-aaron b2558fc
Merge remote-tracking branch 'upstream/main' into rl-docs
hao-aaron da84089
merge
hao-aaron dbcd3e4
x
hao-aaron a28f138
x
hao-aaron ba5909a
moved examples
hao-aaron b64d66b
Add missing title subs
hmellor f5b03a2
Use title method for naming examples subdirs
hmellor 6d02124
Merge remote-tracking branch 'upstream/main' into rl-docs
hao-aaron 888630b
removed legacy weight sync
hao-aaron 177136d
Update distributed.yaml
hmellor 713ef48
pre-commit
hmellor 9b59ed5
Merge branch 'main' into rl-docs
hmellor File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,63 @@ | ||
| # Async Reinforcement Learning | ||
|
|
||
| ## Overview | ||
|
|
||
| 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. | ||
|
|
||
| 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. | ||
|
|
||
| However, this overlap introduces a complication: weights must be updated in the inference engine mid-flight, while requests may still be in progress. | ||
|
|
||
| ## The Pause and Resume API | ||
|
|
||
| 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. | ||
|
|
||
| ### pause_generation | ||
|
|
||
| ```python | ||
| await engine.pause_generation(mode="keep", clear_cache=True) | ||
| ``` | ||
|
|
||
| The `mode` parameter controls how in-flight requests are handled: | ||
|
|
||
| | Mode | Behavior | | ||
| |------|----------| | ||
| | `"abort"` | Abort all in-flight requests immediately and return partial results (default) | | ||
| | `"wait"` | Wait for all in-flight requests to finish before pausing | | ||
| | `"keep"` | Freeze requests in the queue; they resume when `resume_generation` is called | | ||
|
|
||
| The `clear_cache` parameter controls whether to clear the KV cache and prefix cache after pausing. | ||
|
|
||
| ### resume_generation | ||
|
|
||
| ```python | ||
| await engine.resume_generation() | ||
| ``` | ||
|
|
||
| Resumes the scheduler after a pause. Any requests frozen with `mode="keep"` will continue generating. | ||
|
|
||
| ### HTTP Endpoints | ||
|
|
||
| When using the vLLM HTTP server, the same functionality is available via: | ||
|
|
||
| - `POST /pause?mode=keep` - Pause generation | ||
| - `POST /resume` - Resume generation | ||
|
|
||
| !!! note "Data Parallelism" | ||
| 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. | ||
|
|
||
| ## Typical Async RL Flow | ||
|
|
||
| A typical async RL loop with weight syncing looks like this: | ||
|
|
||
| 1. Start generating rollouts from the current policy | ||
| 2. Once enough tokens have been generated (or a batch is ready), pause generation with `mode="keep"` | ||
| 3. Sync the updated weights from the trainer to the inference engine (see [Weight Transfer](weight_transfer/README.md)) | ||
| 4. Resume generation -- in-flight requests continue with the new weights | ||
| 5. Repeat | ||
|
|
||
| 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. | ||
|
hao-aaron marked this conversation as resolved.
Outdated
|
||
|
|
||
| ## Example | ||
|
|
||
| The [async RLHF example](../examples/offline_inference/weight_syncing.md#rlhf-async-new-apis) demonstrates this pattern with `vllm.AsyncLLMEngine`, NCCL weight transfer, and mid-flight pause/resume with validation. | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
| @@ -0,0 +1,79 @@ | ||||||
| # Weight Transfer | ||||||
|
|
||||||
| 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. | ||||||
|
|
||||||
| ## Architecture | ||||||
|
|
||||||
| The weight transfer system follows a **two-phase protocol** with a pluggable backend design: | ||||||
|
|
||||||
| 1. **Initialization** (`init_weight_transfer_engine`): Establishes the communication channel between the trainer and inference workers. Called once before the training loop begins. | ||||||
| 2. **Weight Update** (`update_weights`): Transfers updated weights from the trainer to the inference engine. Called after each training step (or batch of steps). | ||||||
|
|
||||||
| ## Available Backends | ||||||
|
|
||||||
| | Backend | Transport | Use Case | | ||||||
| |---------|-----------|----------| | ||||||
| | [NCCL](nccl.md) | NCCL broadcast |Separate GPUs for training and inference | | ||||||
| | [IPC](ipc.md) | CUDA IPC handles |Colocated training and inference on same GPU | | ||||||
|
|
||||||
| ## Configuration | ||||||
|
|
||||||
| Specify the weight transfer backend through `WeightTransferConfig`. The backend determines which engine handles the weight synchronization. | ||||||
|
|
||||||
| ### Programmatic (Offline Inference) | ||||||
|
|
||||||
| ```python | ||||||
| from vllm import LLM | ||||||
| from vllm.config import WeightTransferConfig | ||||||
|
|
||||||
| llm = LLM( | ||||||
| model="my-model", | ||||||
| weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc" | ||||||
| ) | ||||||
| ``` | ||||||
|
|
||||||
| ### CLI (Online Serving) | ||||||
|
|
||||||
| ```bash | ||||||
| vllm serve my-model \ | ||||||
| --weight-transfer-config '{"backend": "nccl"}' | ||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @hao-aaron I just wanted to give you the option to recommend this, I think it looks nicer but it's up to you
Suggested change
|
||||||
| ``` | ||||||
|
|
||||||
| The `backend` field accepts `"nccl"` (default) or `"ipc"`. | ||||||
|
|
||||||
| ## API Endpoints | ||||||
|
|
||||||
| When running vLLM as an HTTP server, the following endpoints are available for weight transfer: | ||||||
|
|
||||||
| | Endpoint | Method | Description | | ||||||
| |----------|--------|-------------| | ||||||
| | `/init_weight_transfer_engine` | POST | Initialize the weight transfer engine with backend-specific info | | ||||||
| | `/update_weights` | POST | Trigger a weight update with backend-specific metadata | | ||||||
| | `/pause` | POST | Pause generation before weight sync to handle inflight requests | | ||||||
| | `/resume` | POST | Resume generation after weight sync | | ||||||
| | `/get_world_size` | GET | Get the number of inference workers (useful for NCCL world size calculation) | | ||||||
|
|
||||||
| !!! note | ||||||
| The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set. | ||||||
|
|
||||||
| ## Trainer-Side API | ||||||
|
|
||||||
| Both backends provide static methods that the trainer calls to send weights. The general pattern is: | ||||||
|
|
||||||
| ```python | ||||||
| # 1. Initialize the transfer engine (backend-specific) | ||||||
| EngineClass.trainer_init(init_info) | ||||||
|
|
||||||
| # 2. Send weights to inference workers | ||||||
| EngineClass.trainer_send_weights( | ||||||
| iterator=model.named_parameters(), | ||||||
| trainer_args=backend_specific_args, | ||||||
| ) | ||||||
| ``` | ||||||
|
|
||||||
| See the [NCCL](nccl.md) and [IPC](ipc.md) pages for backend-specific trainer APIs and full examples. | ||||||
|
|
||||||
| ## Extending the System | ||||||
|
|
||||||
| 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. | ||||||
|
|
||||||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,164 @@ | ||
| # Base Class and Custom Engines | ||
|
|
||
| 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`. | ||
|
|
||
| ## WeightTransferEngine | ||
|
|
||
| ::: vllm.distributed.weight_transfer.base.WeightTransferEngine | ||
|
|
||
| The `WeightTransferEngine` is a generic abstract class parameterized by two dataclass types: | ||
|
|
||
| - **`TInitInfo`** (extends `WeightTransferInitInfo`): Backend-specific initialization parameters. | ||
| - **`TUpdateInfo`** (extends `WeightTransferUpdateInfo`): Backend-specific weight update metadata. | ||
|
|
||
| ### Abstract Methods | ||
|
|
||
| Subclasses must implement these four methods: | ||
|
|
||
| | Method | Side | Description | | ||
| |--------|------|-------------| | ||
| | `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker | | ||
| | `receive_weights(update_info, load_weights)` | Inference | Receive weights and call `load_weights` incrementally | | ||
| | `shutdown()` | Inference | Clean up resources | | ||
| | `trainer_send_weights(iterator, trainer_args)` | Trainer | Static method to send weights from the trainer process | | ||
|
|
||
| ### Request Classes | ||
|
|
||
| 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. | ||
|
|
||
| ```python | ||
| from vllm.distributed.weight_transfer.base import ( | ||
| WeightTransferInitRequest, | ||
| WeightTransferUpdateRequest, | ||
| ) | ||
|
|
||
| # Init request (dict is converted to backend-specific TInitInfo) | ||
| init_request = WeightTransferInitRequest( | ||
| init_info={"master_address": "10.0.0.1", "master_port": 29500, ...} | ||
| ) | ||
|
|
||
| # Update request (dict is converted to backend-specific TUpdateInfo) | ||
| update_request = WeightTransferUpdateRequest( | ||
| update_info={"names": [...], "dtype_names": [...], "shapes": [...]} | ||
| ) | ||
| ``` | ||
|
|
||
| ### WeightTransferUpdateInfo | ||
|
|
||
| The base `WeightTransferUpdateInfo` includes an `is_checkpoint_format` flag: | ||
|
|
||
| ```python | ||
| @dataclass | ||
| class WeightTransferUpdateInfo(ABC): | ||
| is_checkpoint_format: bool = True | ||
| ``` | ||
|
|
||
| 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. | ||
|
|
||
| ## Implementing a Custom Engine | ||
|
|
||
| To create a custom weight transfer backend: | ||
|
|
||
| ### 1. Define Info Dataclasses | ||
|
|
||
| ```python | ||
| from dataclasses import dataclass | ||
| from vllm.distributed.weight_transfer.base import ( | ||
| WeightTransferEngine, | ||
| WeightTransferInitInfo, | ||
| WeightTransferUpdateInfo, | ||
| ) | ||
|
|
||
| @dataclass | ||
| class MyInitInfo(WeightTransferInitInfo): | ||
| endpoint: str | ||
| token: str | ||
|
|
||
| @dataclass | ||
| class MyUpdateInfo(WeightTransferUpdateInfo): | ||
| names: list[str] | ||
| dtype_names: list[str] | ||
| shapes: list[list[int]] | ||
| # Add custom fields as needed | ||
| ``` | ||
|
|
||
| ### 2. Implement the Engine | ||
|
|
||
| ```python | ||
| from collections.abc import Callable, Iterator | ||
| from typing import Any | ||
| import torch | ||
|
|
||
| class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]): | ||
| init_info_cls = MyInitInfo | ||
| update_info_cls = MyUpdateInfo | ||
|
|
||
| def init_transfer_engine(self, init_info: MyInitInfo) -> None: | ||
| # Set up connection to trainer using init_info.endpoint, etc. | ||
| ... | ||
|
|
||
| def receive_weights( | ||
| self, | ||
| update_info: MyUpdateInfo, | ||
| load_weights: Callable[[list[tuple[str, torch.Tensor]]], None], | ||
| ) -> None: | ||
| # Receive each weight and call load_weights incrementally | ||
| for name, dtype_name, shape in zip( | ||
| update_info.names, update_info.dtype_names, update_info.shapes | ||
| ): | ||
| dtype = getattr(torch, dtype_name) | ||
| weight = self._fetch_weight(name, shape, dtype) | ||
| load_weights([(name, weight)]) | ||
|
|
||
| def shutdown(self) -> None: | ||
| # Clean up resources | ||
| ... | ||
|
|
||
| @staticmethod | ||
| def trainer_send_weights( | ||
| iterator: Iterator[tuple[str, torch.Tensor]], | ||
| trainer_args: dict[str, Any], | ||
| ) -> None: | ||
| # Send weights from the trainer process | ||
| for name, tensor in iterator: | ||
| # Send tensor via custom transport | ||
| ... | ||
| ``` | ||
|
|
||
| !!! important | ||
| 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. | ||
|
|
||
| ### 3. Register with the Factory | ||
|
|
||
| ```python | ||
| from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory | ||
|
|
||
| # Option 1: Lazy loading (recommended for built-in engines) | ||
| WeightTransferEngineFactory.register_engine( | ||
| "my_backend", | ||
| "my_package.my_module", | ||
| "MyWeightTransferEngine", | ||
| ) | ||
|
|
||
| # Option 2: Direct class registration | ||
| WeightTransferEngineFactory.register_engine( | ||
| "my_backend", | ||
| MyWeightTransferEngine, | ||
| ) | ||
| ``` | ||
|
|
||
| Once registered, users can select your backend via `WeightTransferConfig(backend="my_backend")`. | ||
|
|
||
| ## WeightTransferEngineFactory | ||
|
|
||
| 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. | ||
|
|
||
| ```python | ||
| from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory | ||
|
|
||
| # Create an engine from config | ||
| engine = WeightTransferEngineFactory.create_engine( | ||
| config=weight_transfer_config, | ||
| parallel_config=parallel_config, | ||
| ) | ||
| ``` |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,72 @@ | ||
| # IPC Engine | ||
|
|
||
| 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. | ||
|
|
||
| ## When to Use IPC | ||
|
|
||
| - Training and inference on the **same GPU** (colocated) | ||
| - You want to minimize memory overhead by sharing tensors in-place | ||
|
|
||
| ## How It Works | ||
|
|
||
| 1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`. | ||
| 2. IPC handles are sent to the inference engine via **Ray RPC** or **HTTP POST**. | ||
|
hao-aaron marked this conversation as resolved.
Outdated
|
||
| 3. The inference worker reconstructs the tensors from the handles, reading directly from the trainer's GPU memory. | ||
|
|
||
| !!! warning | ||
| 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. | ||
|
|
||
| ## Initialization | ||
|
|
||
| The IPC backend requires no initialization on either side. The `init_transfer_engine` call is a no-op for IPC. | ||
|
|
||
| ## Sending Weights | ||
|
|
||
| IPC supports two transport modes for delivering the handles: | ||
|
|
||
| ### Ray Mode | ||
|
|
||
| Used when vLLM is running as a Ray actor: | ||
|
|
||
| ```python | ||
| from vllm.distributed.weight_transfer.ipc_engine import ( | ||
| IPCTrainerSendWeightsArgs, | ||
| IPCWeightTransferEngine, | ||
| ) | ||
|
|
||
| trainer_args = IPCTrainerSendWeightsArgs( | ||
| mode="ray", | ||
| llm_handle=llm_actor_handle, | ||
| ) | ||
|
|
||
| IPCWeightTransferEngine.trainer_send_weights( | ||
| iterator=model.named_parameters(), | ||
| trainer_args=trainer_args, | ||
| ) | ||
| ``` | ||
|
|
||
| In Ray mode, the engine calls `llm_handle.update_weights.remote(...)` directly, passing the IPC handles via Ray's serialization. | ||
|
|
||
| ### HTTP Mode | ||
|
|
||
| Used when vLLM is running as an HTTP server: | ||
|
|
||
| ```python | ||
| trainer_args = IPCTrainerSendWeightsArgs( | ||
| mode="http", | ||
| url="http://localhost:8000", | ||
| ) | ||
|
|
||
| IPCWeightTransferEngine.trainer_send_weights( | ||
| iterator=model.named_parameters(), | ||
| trainer_args=trainer_args, | ||
| ) | ||
| ``` | ||
|
|
||
| In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint. | ||
|
|
||
| See [`IPCTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/ipc_engine.py) for the full list of configurable fields. | ||
| ## Examples | ||
|
|
||
| - [RLHF with IPC weight syncing (offline, Ray)](../../examples/offline_inference/weight_syncing.md#rlhf-ipc) - Colocated training and inference on a single GPU using Ray placement groups and CUDA IPC handles | ||
| - [RLHF with IPC weight syncing (online serving, HTTP)](../../examples/online_serving/weight_syncing.md#rlhf-http-ipc) - Weight transfer with a vLLM HTTP server where both server and trainer share the same GPU | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.