[Docs][RL] Rewrite weight-transfer docs; standardize examples (#51729)
Signed-off-by: haoaaron <ahao@anyscale.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Aaron Hao
co-authored by
Claude Opus 5
Claude Opus 5
parent
9b3382685e
commit
793ca6998a
@@ -847,8 +847,8 @@ steps:
|
||||
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
|
||||
- python3 ../examples/features/data_parallel/data_parallel_offline.py --enforce-eager
|
||||
# rlhf examples
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_nccl.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_ipc.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_http_nccl.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_http_ipc.py
|
||||
|
||||
- label: Elastic EP Scaling Test # TBD
|
||||
timeout_in_minutes: 180
|
||||
|
||||
@@ -133,9 +133,9 @@ steps:
|
||||
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
# test with internal dp
|
||||
- python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
|
||||
# rlhf examples
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
|
||||
# rlhf examples (each launches its own `vllm serve` and tears it down)
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py
|
||||
|
||||
- label: Distributed DP Tests (4 GPUs)
|
||||
key: distributed-dp-tests-4-gpus
|
||||
|
||||
@@ -4,12 +4,28 @@ vLLM provides a pluggable weight transfer system for synchronizing model weights
|
||||
|
||||
## Architecture
|
||||
|
||||
The weight transfer system follows a **four-phase protocol** with a pluggable backend design:
|
||||
Weight transfer has **two engines, one per process**, and they are symmetric:
|
||||
|
||||
1. **Initialization** (`init_weight_transfer_engine`): Establishes the communication channel between the trainer and inference workers. Called once before the training loop begins.
|
||||
2. **Start** (`start_weight_update`): Prepares the inference engine for a weight update.
|
||||
3. **Weight Update** (`update_weights`): Transfers updated weights from the trainer to the inference engine. May be called one or more times (e.g., for chunked transfers).
|
||||
4. **Finish** (`finish_weight_update`): Finalizes the weight update (e.g., runs post-processing for checkpoint-format weights). Called once after all weights have been transferred.
|
||||
| | Trainer process | Inference workers |
|
||||
| --- | --- | --- |
|
||||
| Class | `TrainerWeightTransferEngine` | `WeightTransferEngine` |
|
||||
| Built by | `WeightTransferTrainerFactory.trainer_init(...)` | vLLM, from `WeightTransferConfig` |
|
||||
| Driven by | `send_weights()` | the four-phase protocol below |
|
||||
| Holds | communicators, transfer plans, wire params | communicators, the target model |
|
||||
|
||||
The trainer engine is stateful: it owns its communicators and wire params, pulls
|
||||
weights from a [`WeightSource`](base.md#weightsource), and drives the inference
|
||||
side through a [`VLLMWeightSyncClient`](base.md#vllmweightsyncclient). Trainer
|
||||
code never has to know the transport, and never has to thread transfer state
|
||||
back in on every round — one `send_weights()` call per sync.
|
||||
|
||||
Under the hood every round is the same **four-phase protocol**, which the trainer
|
||||
engine drives on your behalf:
|
||||
|
||||
1. **Initialization** (`init_weight_transfer_engine`): establishes the communication channel between the trainer and inference workers. Called once, from `trainer_init`, before the training loop begins.
|
||||
2. **Start** (`start_weight_update`): prepares the inference engine for a weight update.
|
||||
3. **Weight Update** (`update_weights`): transfers updated weights. May be called one or more times (e.g. for chunked transfers).
|
||||
4. **Finish** (`finish_weight_update`): finalizes the update (e.g. runs post-processing for checkpoint-format weights). Called once after all weights have been transferred.
|
||||
|
||||
## Available Backends
|
||||
|
||||
@@ -19,11 +35,12 @@ The weight transfer system follows a **four-phase protocol** with a pluggable ba
|
||||
| [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
|
||||
| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
|
||||
|
||||
## Configuration
|
||||
## Quickstart
|
||||
|
||||
Specify the weight transfer backend through `WeightTransferConfig`. The backend determines which engine handles the weight synchronization.
|
||||
### Inference Side
|
||||
|
||||
### Programmatic (Offline Inference)
|
||||
The inference side takes only a backend name. Everything else about the transfer
|
||||
is decided by the trainer and shipped over at the init handshake.
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
@@ -31,22 +48,100 @@ from vllm.config import WeightTransferConfig
|
||||
|
||||
llm = LLM(
|
||||
model="my-model",
|
||||
weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc"
|
||||
weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc", "sparse_nccl"
|
||||
)
|
||||
```
|
||||
|
||||
### CLI (Online Serving)
|
||||
Or, for online serving:
|
||||
|
||||
```bash
|
||||
vllm serve my-model \
|
||||
--weight-transfer-config '{"backend": "nccl"}'
|
||||
```
|
||||
|
||||
The `backend` field accepts `"nccl"` (default), `"ipc"`, or `"sparse_nccl"`.
|
||||
### Trainer Side
|
||||
|
||||
Build the engine once, then call `send_weights()` once per sync:
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer import (
|
||||
ModuleSource,
|
||||
HTTPVLLMWeightSyncClient,
|
||||
WeightTransferTrainerFactory,
|
||||
)
|
||||
from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
|
||||
|
||||
# Once, before the training loop.
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=NCCLTrainerInitInfo(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
world_size=world_size, # trainer + all inference workers
|
||||
rank=0, # this trainer rank; rank 0 is the sender
|
||||
packed=True,
|
||||
),
|
||||
client=HTTPVLLMWeightSyncClient("http://localhost:8000"),
|
||||
source=ModuleSource(model),
|
||||
)
|
||||
|
||||
# Once per weight sync.
|
||||
for step in range(num_steps):
|
||||
train_one_step(model)
|
||||
engine.send_weights()
|
||||
```
|
||||
|
||||
`send_weights()` drives start → update → finish on the inference side *and* the
|
||||
data-plane transfer, including any concurrency the backend needs (NCCL, for
|
||||
example, must run the workers' `update_weights` at the same time as the
|
||||
trainer's broadcast, since both rendezvous inside the same NCCL calls).
|
||||
|
||||
There is no `backend=` argument on the trainer side: each `TrainerInitInfo`
|
||||
subclass declares its own `backend`, and the factory dispatches on it.
|
||||
|
||||
## Where Each Setting Lives
|
||||
|
||||
The trainer is the single source of truth for anything the two sides must agree
|
||||
on. This is the main rule to internalize:
|
||||
|
||||
| Setting | Lives on | Notes |
|
||||
| ------- | -------- | ----- |
|
||||
| Backend, inference side | `WeightTransferConfig(backend=...)` | A plain string selector — that is all this config holds |
|
||||
| Backend, trainer side | `init_info.backend` | A `ClassVar` on each init info; never passed by hand |
|
||||
| Wire params (`packed`, buffer sizes) | the backend's `TrainerInitInfo` | The trainer propagates them to the worker during `trainer_init`, so the two sides *cannot* disagree |
|
||||
| Weights to send | a [`WeightSource`](base.md#weightsource) | `ModuleSource(model)` covers plain and FSDP-sharded modules |
|
||||
| Per-round payload | the backend's update info | Built by the engine (from `source.metadata()`, or from a delta backend's per-round payload); you never construct it |
|
||||
|
||||
!!! note
|
||||
Wire params are deliberately *not* on `WeightTransferConfig` and *not* on the
|
||||
per-round update info. Because the trainer ships them at init and the worker
|
||||
reads what it was told, a mismatched `packed` flag between the two sides is
|
||||
unrepresentable rather than merely discouraged.
|
||||
|
||||
## Multi-Rank Trainers
|
||||
|
||||
For sharded trainers (FSDP, TP/PP/EP), **every trainer rank builds an engine and
|
||||
calls `send_weights()`**. Rank 0 is the sender: only it holds a communicator,
|
||||
talks to the client, and puts bytes on the wire. Non-sender ranks still iterate
|
||||
the `WeightSource`, because materializing a parameter is usually itself a
|
||||
collective (an FSDP `full_tensor()` all-gather, a Megatron export) that would
|
||||
deadlock if some ranks skipped it.
|
||||
|
||||
Pass each process's own `rank` on the init info. It is explicit rather than read
|
||||
from a global process group, which is ambiguous once several groups (FSDP / TP /
|
||||
PP / EP) exist.
|
||||
|
||||
```python
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=NCCLTrainerInitInfo(..., rank=torch.distributed.get_rank()),
|
||||
client=client,
|
||||
source=ModuleSource(model),
|
||||
)
|
||||
engine.send_weights() # called on every rank
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
When running vLLM as an HTTP server, the following endpoints are available for weight transfer:
|
||||
When running vLLM as an HTTP server, the following endpoints are available for weight transfer. `HTTPVLLMWeightSyncClient` speaks the first four for you.
|
||||
|
||||
| Endpoint | Method | Description |
|
||||
| -------- | ------ | ----------- |
|
||||
@@ -63,29 +158,9 @@ When running vLLM as an HTTP server, the following endpoints are available for w
|
||||
!!! 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. Start weight update on inference side
|
||||
llm.start_weight_update()
|
||||
|
||||
# 3. Send weights to inference workers
|
||||
EngineClass.trainer_send_weights(
|
||||
iterator=model.named_parameters(),
|
||||
trainer_args=backend_specific_args,
|
||||
)
|
||||
|
||||
# 4. Finish weight update on inference side
|
||||
llm.finish_weight_update(weight_version="step-42")
|
||||
```
|
||||
|
||||
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.
|
||||
Every piece of the system is replaceable: the weights you send (`WeightSource`),
|
||||
the control-plane transport (`VLLMWeightSyncClient`), and the transports
|
||||
themselves (both engine ABCs, each with its own factory registry). See
|
||||
[Base Classes](base.md).
|
||||
|
||||
@@ -1,36 +1,391 @@
|
||||
# Base Class and Custom Engines
|
||||
# Base Classes 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`.
|
||||
The weight transfer system is built from four abstractions, each independently
|
||||
replaceable:
|
||||
|
||||
## WeightTransferEngine
|
||||
| Abstraction | Side | Answers |
|
||||
| ----------- | ---- | ------- |
|
||||
| [`WeightSource`](#weightsource) | Trainer | *What* weights to send |
|
||||
| [`VLLMWeightSyncClient`](#vllmweightsyncclient) | Trainer | *How to reach* the inference engine — the adapter for your RL stack's own vLLM wrapper |
|
||||
| [`TrainerWeightTransferEngine`](#trainerweighttransferengine) | Trainer | *How to transmit* the bytes |
|
||||
| [`WeightTransferEngine`](#weighttransferengine) | Inference | *How to receive* them and load them |
|
||||
|
||||
The `WeightTransferEngine` is a generic abstract class parameterized by two dataclass types:
|
||||
The two engines are registered in two separate factories,
|
||||
[`WeightTransferTrainerFactory`](#weighttransfertrainerfactory) and
|
||||
[`WeightTransferEngineFactory`](#weighttransferenginefactory). They share backend
|
||||
names by convention, but a trainer process never instantiates a worker engine or
|
||||
vice versa, so the registries stay independent.
|
||||
|
||||
- **`TInitInfo`** (extends `WeightTransferInitInfo`): Backend-specific initialization parameters.
|
||||
- **`TUpdateInfo`** (extends `WeightTransferUpdateInfo`): Backend-specific weight update metadata.
|
||||
## Trainer Side
|
||||
|
||||
### Abstract Methods
|
||||
### WeightSource
|
||||
|
||||
Subclasses must implement these methods:
|
||||
A `WeightSource` is a **re-iterable** source of the trainer's weights with two
|
||||
channels:
|
||||
|
||||
| Method | Side | Description |
|
||||
| ------ | ---- | ----------- |
|
||||
| `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker |
|
||||
| `start_weight_update()` | Inference | Prepare for an update (e.g. begin layerwise reload); no-op for in-place engines |
|
||||
| `finish_weight_update()` | Inference | Finalize the update (e.g. finalize layerwise reload); no-op for in-place engines |
|
||||
| `receive_weights(update_info)` | Inference | Receive weights and load them into `self.model` |
|
||||
| `shutdown()` | Inference | Clean up resources |
|
||||
| `trainer_send_weights(iterator, trainer_args)` | Trainer | Static method to send weights from the trainer process |
|
||||
- **`metadata() -> list[ParamMeta]`** — the name, wire dtype, and full shape of
|
||||
every parameter, *without transferring anything*. Cheap when shapes are known
|
||||
locally (an FSDP `DTensor` knows its global shape); may be expensive on the
|
||||
first call for producers that must materialize to learn shapes (a
|
||||
Megatron-Bridge export), in which case it should cache.
|
||||
- **iteration** — yields fully-materialized `(name, tensor)` pairs, one at a
|
||||
time.
|
||||
|
||||
The base class provides two methods:
|
||||
```python
|
||||
@dataclass(frozen=True)
|
||||
class ParamMeta:
|
||||
name: str
|
||||
dtype: torch.dtype
|
||||
shape: tuple[int, ...]
|
||||
```
|
||||
|
||||
1. `__init__` : Engines receive `config` (`WeightTransferConfig`), `vllm_config` (`VllmConfig`), `device` (`torch.device`) and `model` (`nn.Module`)
|
||||
2. `update_weights(update_info_dict)`: Thin wrapper for `receive_weights`: parses
|
||||
the dict into user-specified data type, calls `receive_weights`, and synchronizes the device. Subclasses implement `receive_weights`.
|
||||
!!! warning "The two channels must agree, element for element"
|
||||
`metadata()` must declare exactly what iteration will yield: the same
|
||||
parameters, in the same order, with the same dtypes and shapes. This is an
|
||||
invariant of the ABC, not of any one backend — a source that reorders, omits,
|
||||
or re-dtypes a parameter between the two channels is broken even if the
|
||||
backend you happen to test against never notices. Backends are free to read
|
||||
both channels and to trust that they match; dense NCCL does, and
|
||||
[enforces it](nccl.md#the-two-weightsource-channels-must-agree).
|
||||
|
||||
Materializing is typically a collective, so **every trainer rank must iterate the
|
||||
same source in the same order, in lockstep**, or ranks deadlock. `metadata()` can
|
||||
itself be a collective for custom producers, so it too runs on every rank — only
|
||||
the sender ships the result. Under pipeline parallelism a rank may not own a
|
||||
parameter at all; iterating still drives the collective, and the yielded tensor
|
||||
is only meaningful on the sender.
|
||||
|
||||
`iter(source)` must yield a *fresh* pass each round.
|
||||
|
||||
#### ModuleSource
|
||||
|
||||
`ModuleSource(module)` is the common case, over `module.named_parameters()`. It
|
||||
handles plain and FSDP-sharded modules with no special casing: iteration
|
||||
all-gathers each `DTensor` via `full_tensor()`, while `metadata()` reads the
|
||||
*global* `.shape` / `.dtype` and so never triggers a gather.
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer import ModuleSource
|
||||
|
||||
source = ModuleSource(model)
|
||||
```
|
||||
|
||||
#### Custom Sources
|
||||
|
||||
Subclass `WeightSource` when the weights you want to send require additional processing to convert to a HF compatible format.
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer import ParamMeta, WeightSource
|
||||
from vllm.distributed.weight_transfer.base import materialize_full_tensor
|
||||
|
||||
|
||||
class MyExportSource(WeightSource):
|
||||
def __init__(self, model):
|
||||
self._model = model
|
||||
self._meta: list[ParamMeta] | None = None
|
||||
|
||||
def metadata(self) -> list[ParamMeta]:
|
||||
# Cache: for producers that must materialize to learn shapes, this is
|
||||
# the expensive channel. Runs on every rank (it may be a collective).
|
||||
if self._meta is None:
|
||||
self._meta = [
|
||||
ParamMeta(name, t.dtype, tuple(t.shape)) for name, t in self._export()
|
||||
]
|
||||
return self._meta
|
||||
|
||||
def __iter__(self):
|
||||
# Must yield exactly what metadata() declared, in the same order.
|
||||
for name, tensor in self._export():
|
||||
yield name, materialize_full_tensor(tensor)
|
||||
```
|
||||
|
||||
### VLLMWeightSyncClient
|
||||
|
||||
**This is the adapter for however your RL stack reaches vLLM.** Many RL frameworks wrap
|
||||
inference engines in their own abstractions, and each reaches vLLM its own way.
|
||||
`VLLMWeightSyncClient` is the single seam where that bespoke shape is adapted, so
|
||||
weight sync engines remain control plane agnostic.
|
||||
|
||||
The contract is only this: **however the wrapper is shaped, it must bottom out in
|
||||
the same four calls** — `init_weight_transfer_engine` once at setup, then
|
||||
`start_weight_update` → one or more `update_weights` → `finish_weight_update` per
|
||||
round. Everything a trainer engine needs from the inference side goes through them.
|
||||
|
||||
```python
|
||||
class VLLMWeightSyncClient(Protocol):
|
||||
def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: ...
|
||||
def start_weight_update(self) -> None: ...
|
||||
def update_weights(self, update_info: dict[str, Any]) -> None: ...
|
||||
def finish_weight_update(self, weight_version: str | None = None) -> None: ...
|
||||
```
|
||||
|
||||
It is a `@runtime_checkable` structural `Protocol` (PEP 544), which is what makes
|
||||
adapting cheap: **any object with those four methods already satisfies it**. An existing
|
||||
wrapper in your framework can usually become a client by gaining four forwarding
|
||||
methods.
|
||||
|
||||
Two implementations ship with vLLM:
|
||||
|
||||
| Client | Talks to |
|
||||
| ------ | -------- |
|
||||
| `RayVLLMWeightSyncClient(handle)` | One or more `AsyncLLM`/`LLM` Ray actors. Accepts a list and fans each call out to every handle, blocking on all of them, so a multi-actor (e.g. multi-DP) deployment is driven as one unit |
|
||||
| `HTTPVLLMWeightSyncClient(base_url, timeout=300)` | A vLLM server over the RLHF HTTP routes |
|
||||
|
||||
Custom weight sync clients can be implement like so:
|
||||
|
||||
```python
|
||||
class MyFrameworkWeightSyncClient:
|
||||
"""Adapts one RL framework's rollout pool to the four weight-sync calls."""
|
||||
|
||||
def __init__(self, rollout_pool):
|
||||
self.pool = rollout_pool # whatever your stack already has
|
||||
|
||||
def init_weight_transfer_engine(self, init_info):
|
||||
# Fan out to every replica and block: all of them receive weights.
|
||||
self.pool.broadcast_rpc("init_weight_transfer_engine", init_info=init_info)
|
||||
|
||||
def start_weight_update(self):
|
||||
self.pool.broadcast_rpc("start_weight_update")
|
||||
|
||||
def update_weights(self, update_info):
|
||||
self.pool.broadcast_rpc("update_weights", update_info=update_info)
|
||||
|
||||
def finish_weight_update(self, weight_version=None):
|
||||
self.pool.broadcast_rpc("finish_weight_update")
|
||||
if weight_version is not None:
|
||||
self.pool.broadcast_rpc("update_weight_version", weight_version)
|
||||
```
|
||||
|
||||
Two things to get right in any adapter:
|
||||
|
||||
- **Reach every replica, and block until all of them are done.** A weight update
|
||||
is not a load-balanced request: every worker holding a copy of the model must
|
||||
receive it. Returning before they all finish lets the trainer race ahead of
|
||||
workers still loading. (Both built-in clients do this — Ray by fanning out over
|
||||
its handles, HTTP because the server's DP client broadcasts internally.)
|
||||
- **Raise on failure.** Trainer engines rely on exceptions to surface
|
||||
inference-side errors; a client that swallows them turns a failed sync into
|
||||
silently stale weights, or into a hang for backends whose transfer rendezvouses
|
||||
with the worker.
|
||||
|
||||
!!! note
|
||||
HTTP cannot carry raw CUDA IPC handles, so `HTTPVLLMWeightSyncClient` pickles
|
||||
and base64-encodes them into an `ipc_handles_pickled` field. The worker
|
||||
deserializes it only when `VLLM_ALLOW_INSECURE_SERIALIZATION=1`. Backends
|
||||
whose payloads are JSON-native (NCCL) pass through untouched.
|
||||
|
||||
### TrainerWeightTransferEngine
|
||||
|
||||
The trainer-side engine: it holds the transport state (NCCL communicators, IPC
|
||||
device info, transfer plans), pulls weights from a `WeightSource`, and drives the
|
||||
inference side through a `VLLMWeightSyncClient`. It is generic over its init info
|
||||
type, constructed by the `trainer_init` classmethod factory, and driven by `send_weights()`.
|
||||
|
||||
| Method | Description |
|
||||
| ------ | ----------- |
|
||||
| `trainer_init(init_info, *, client, source=None)` | Classmethod. Rendezvous with the inference side and return a ready instance |
|
||||
| `send_weights()` | Push weights and drive the full update round trip |
|
||||
| `shutdown()` | Tear down communicators / process groups. Default no-op |
|
||||
|
||||
Both `trainer_init` and `send_weights` are called on **every** trainer rank.
|
||||
`is_sender` is resolved once, at `trainer_init`, from `init_info.rank`. Each
|
||||
engine holds the real client on every rank but guards the control-plane RPCs and
|
||||
the transmit on `self.is_sender`, so only the sender touches the wire; non-sender
|
||||
ranks still run every collective so the group stays aligned.
|
||||
|
||||
The trainer side takes **no `WeightTransferConfig`**. The backend comes from the
|
||||
init info's `backend` `ClassVar`, and the wire params ride the init info too.
|
||||
|
||||
#### TrainerInitInfo
|
||||
|
||||
The `init_info` passed to `trainer_init` above. It is how a caller configures a
|
||||
transfer: it selects the backend, says which rank this process is, and carries the
|
||||
wire params. Each backend subclasses it; the base class holds the one field every
|
||||
backend needs.
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class TrainerInitInfo:
|
||||
backend: ClassVar[str] # factory dispatch key
|
||||
rank: int = field(kw_only=True)
|
||||
|
||||
@property
|
||||
def is_sender(self) -> bool:
|
||||
return self.rank == 0
|
||||
```
|
||||
|
||||
- **`rank`** is this trainer process's rank, supplied **explicitly**. The engine
|
||||
does not read it from a global process group, which is ambiguous once several
|
||||
groups (FSDP / TP / PP / EP) exist. **Rank 0 is always the sender** — this is
|
||||
what `trainer_init` resolves into `is_sender`. It is keyword-only, so backend
|
||||
subclasses can add positional fields freely.
|
||||
- **`backend`** is a `ClassVar`, not an `__init__` field: it is a fixed
|
||||
per-backend constant that the factory reads to dispatch, which is why callers
|
||||
never pass a `backend=` argument. Every subclass must set it —
|
||||
`__init_subclass__` raises otherwise.
|
||||
|
||||
Subclasses also carry the transfer's **wire params** (`packed`, buffer sizes).
|
||||
The sender propagates them to the worker inside `trainer_init`, so the two sides
|
||||
cannot disagree. See [`NCCLTrainerInitInfo`](nccl.md#nccltrainerinitinfo) and
|
||||
[`IPCTrainerInitInfo`](ipc.md#ipctrainerinitinfo) for the concrete fields.
|
||||
|
||||
#### Full-Resync vs. Delta Backends
|
||||
|
||||
`source` is optional, which splits the backends into two shapes:
|
||||
|
||||
- **Full resync** (NCCL, IPC) — a stable `WeightSource` is fixed at
|
||||
`trainer_init` and re-iterated each round; `send_weights()` takes no
|
||||
arguments. These backends validate that `source` is non-null themselves.
|
||||
- **Delta** (sparse NCCL) — the payload differs every round, so there is no
|
||||
stable source. The engine takes no `source` and each round's payload is passed
|
||||
straight to `send_weights(patches)`.
|
||||
|
||||
#### Implementing a Custom Trainer Engine
|
||||
|
||||
```python
|
||||
from dataclasses import dataclass
|
||||
from typing import ClassVar
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
from vllm.distributed.weight_transfer.base import (
|
||||
TrainerInitInfo,
|
||||
TrainerWeightTransferEngine,
|
||||
VLLMWeightSyncClient,
|
||||
WeightSource,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MyTrainerInitInfo(TrainerInitInfo):
|
||||
backend: ClassVar[str] = "my_backend"
|
||||
|
||||
endpoint: str
|
||||
chunk_size_bytes: int = 256 * 1024 * 1024 # a wire param: shipped to the worker
|
||||
|
||||
|
||||
class MyTrainerWeightTransferEngine(TrainerWeightTransferEngine[MyTrainerInitInfo]):
|
||||
init_info_cls = MyTrainerInitInfo
|
||||
|
||||
def __init__(self, *, client, source, is_sender=True, chunk_size_bytes=0):
|
||||
super().__init__(client=client, source=source, is_sender=is_sender)
|
||||
self.chunk_size_bytes = chunk_size_bytes
|
||||
|
||||
@classmethod
|
||||
def trainer_init(
|
||||
cls,
|
||||
init_info: MyTrainerInitInfo,
|
||||
*,
|
||||
client: VLLMWeightSyncClient,
|
||||
source: WeightSource | None = None,
|
||||
) -> Self:
|
||||
if source is None:
|
||||
raise ValueError("my_backend requires a WeightSource.")
|
||||
engine = cls(
|
||||
client=client,
|
||||
source=source,
|
||||
is_sender=init_info.is_sender,
|
||||
chunk_size_bytes=init_info.chunk_size_bytes,
|
||||
)
|
||||
if engine.is_sender:
|
||||
# Ship the must-agree wire params so the worker decodes exactly as
|
||||
# this trainer encodes, then open the trainer-side endpoint.
|
||||
engine.client.init_weight_transfer_engine(
|
||||
{"chunk_size_bytes": init_info.chunk_size_bytes}
|
||||
)
|
||||
return engine
|
||||
|
||||
def send_weights(self) -> None:
|
||||
assert self.source is not None
|
||||
meta = self.source.metadata() # every rank: may be a collective
|
||||
if not self.is_sender:
|
||||
for _ in self.source: # stay in the trainer-side collective
|
||||
pass
|
||||
return
|
||||
|
||||
self.client.start_weight_update()
|
||||
self.client.update_weights(
|
||||
{
|
||||
"names": [m.name for m in meta],
|
||||
"dtype_names": [str(m.dtype).split(".")[-1] for m in meta],
|
||||
"shapes": [list(m.shape) for m in meta],
|
||||
}
|
||||
)
|
||||
for name, tensor in self.source:
|
||||
... # transmit
|
||||
self.client.finish_weight_update()
|
||||
```
|
||||
|
||||
Two things to get right, both of which have bitten the built-in backends:
|
||||
|
||||
- **Drain before returning.** `send_weights` must not return with transfers
|
||||
still in flight. Anything keeping a send buffer alive dies with the frame, and
|
||||
the inference side's `finish_weight_update` post-processing can otherwise
|
||||
finalize weights that have not landed.
|
||||
- **Never join a control-plane thread on the error path.** If you run
|
||||
`update_weights` on a side thread concurrently with a transmit (as NCCL does)
|
||||
and the transmit raises, the worker is still blocked in the matching
|
||||
collective and will never return. Shut the executor down without waiting, so
|
||||
the real exception surfaces instead of hanging.
|
||||
|
||||
### WeightTransferTrainerFactory
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer import WeightTransferTrainerFactory
|
||||
|
||||
# Lazy loading (recommended): the module is imported only when the backend is used
|
||||
WeightTransferTrainerFactory.register_engine(
|
||||
"my_backend",
|
||||
"my_package.my_module",
|
||||
"MyTrainerWeightTransferEngine",
|
||||
)
|
||||
|
||||
# Or register the class directly
|
||||
WeightTransferTrainerFactory.register_engine("my_backend", MyTrainerWeightTransferEngine)
|
||||
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=MyTrainerInitInfo(rank=0, endpoint="..."), # `backend` selects the engine
|
||||
client=client,
|
||||
source=source,
|
||||
)
|
||||
```
|
||||
|
||||
## Inference Side
|
||||
|
||||
### WeightTransferEngine
|
||||
|
||||
A generic abstract class parameterized by two dataclass types:
|
||||
|
||||
- **`TInitInfo`** (extends `WeightTransferInitInfo`): backend-specific initialization parameters.
|
||||
- **`TUpdateInfo`** (extends `WeightTransferUpdateInfo`): backend-specific weight update metadata.
|
||||
|
||||
Subclasses must implement five methods:
|
||||
|
||||
| Method | Description |
|
||||
| ------ | ----------- |
|
||||
| `init_transfer_engine(init_info)` | Initialize the communication channel on each inference worker, and record the trainer-supplied wire params |
|
||||
| `start_weight_update()` | Prepare for an update (e.g. begin layerwise reload); no-op for in-place engines |
|
||||
| `finish_weight_update()` | Finalize the update (e.g. finalize layerwise reload); no-op for in-place engines |
|
||||
| `receive_weights(update_info)` | Receive weights and load them into `self.model` |
|
||||
| `shutdown()` | Clean up resources |
|
||||
|
||||
The base class provides:
|
||||
|
||||
1. `__init__`, taking `config` (`WeightTransferConfig`), `vllm_config` (`VllmConfig`), `device` (`torch.device`), and `model` (`nn.Module`).
|
||||
2. `update_weights(update_info_dict)`, a thin wrapper for `receive_weights`: it parses the dict into the typed dataclass, calls `receive_weights`, and synchronizes the device.
|
||||
3. `parse_init_info` / `parse_update_info`, which convert API-level dicts into the typed dataclasses and raise `ValueError` on a bad payload.
|
||||
4. `set_weight_update_target` / `reset_weight_update_target`, used to retarget an update at the speculative draft model.
|
||||
|
||||
!!! note "Read wire params from the handshake, not the payload"
|
||||
Anything the two sides must agree on — `packed`, buffer geometry — arrives
|
||||
on the **init info** and should be stored on `self` in
|
||||
`init_transfer_engine`, then read from `self` in `receive_weights`. Per-round
|
||||
update info carries only per-round metadata. This is what makes a
|
||||
trainer/worker mismatch unrepresentable.
|
||||
|
||||
### 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.
|
||||
The API-level request classes provide backend-agnostic serialization using plain dictionaries.
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer.base import (
|
||||
@@ -49,25 +404,17 @@ update_request = WeightTransferUpdateRequest(
|
||||
)
|
||||
```
|
||||
|
||||
Using a built-in client, you never construct these by hand — `RayVLLMWeightSyncClient`
|
||||
wraps the dicts for you, and `HTTPVLLMWeightSyncClient` posts them as JSON.
|
||||
|
||||
At the LLM/API layer, call `start_draft_weight_update()` instead of
|
||||
`start_weight_update()` to target the speculative draft model;
|
||||
`update_weights` / `finish_weight_update` are unchanged.
|
||||
`update_weights` / `finish_weight_update` are unchanged. Engines that cannot
|
||||
support this set `supports_draft_weight_update = False`.
|
||||
|
||||
### WeightTransferUpdateInfo
|
||||
### Implementing a Custom Engine
|
||||
|
||||
The base `WeightTransferUpdateInfo` is a marker class for backend-specific update info:
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class WeightTransferUpdateInfo(ABC):
|
||||
pass
|
||||
```
|
||||
|
||||
## Implementing a Custom Engine
|
||||
|
||||
To create a custom weight transfer backend:
|
||||
|
||||
### 1. Define Info Dataclasses
|
||||
#### 1. Define Info Dataclasses
|
||||
|
||||
```python
|
||||
from dataclasses import dataclass
|
||||
@@ -80,29 +427,26 @@ from vllm.distributed.weight_transfer.base import (
|
||||
@dataclass
|
||||
class MyInitInfo(WeightTransferInitInfo):
|
||||
endpoint: str
|
||||
token: str
|
||||
chunk_size_bytes: int = 256 * 1024 * 1024 # must-agree wire param
|
||||
|
||||
@dataclass
|
||||
class MyUpdateInfo(WeightTransferUpdateInfo):
|
||||
names: list[str]
|
||||
dtype_names: list[str]
|
||||
shapes: list[list[int]]
|
||||
# Add custom fields as needed
|
||||
# Per-round metadata only.
|
||||
```
|
||||
|
||||
### 2. Implement the Engine
|
||||
#### 2. Implement the Engine
|
||||
|
||||
```python
|
||||
from collections.abc import 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.
|
||||
# Record the trainer's wire params, then set up the connection.
|
||||
self.chunk_size_bytes = init_info.chunk_size_bytes
|
||||
...
|
||||
|
||||
def start_weight_update(self) -> None:
|
||||
@@ -128,22 +472,12 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
|
||||
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
|
||||
...
|
||||
```
|
||||
|
||||
### 3. Register with the Factory
|
||||
#### 3. Register with the Factory
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
|
||||
from vllm.distributed.weight_transfer import WeightTransferEngineFactory
|
||||
|
||||
# Option 1: Lazy loading (recommended for built-in engines)
|
||||
WeightTransferEngineFactory.register_engine(
|
||||
@@ -159,14 +493,14 @@ WeightTransferEngineFactory.register_engine(
|
||||
)
|
||||
```
|
||||
|
||||
Once registered, users can select your backend via `WeightTransferConfig(backend="my_backend")`.
|
||||
Once registered, users select your backend via `WeightTransferConfig(backend="my_backend")`.
|
||||
|
||||
## WeightTransferEngineFactory
|
||||
### WeightTransferEngineFactory
|
||||
|
||||
The factory uses a registry pattern with lazy loading. Built-in engines (`nccl`, `ipc`, and `sparse_nccl`) 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
|
||||
from vllm.distributed.weight_transfer import WeightTransferEngineFactory
|
||||
|
||||
# Create an engine from config
|
||||
engine = WeightTransferEngineFactory.create_engine(
|
||||
@@ -176,3 +510,6 @@ engine = WeightTransferEngineFactory.create_engine(
|
||||
model=model,
|
||||
)
|
||||
```
|
||||
|
||||
vLLM calls this for you during worker startup; you only need it directly when
|
||||
embedding the engine in your own worker.
|
||||
|
||||
@@ -8,111 +8,109 @@ The IPC weight transfer engine uses **CUDA IPC** (Inter-Process Communication) h
|
||||
|
||||
## How It Works
|
||||
|
||||
1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`. In multi-GPU setups (e.g. FSDP), each trainer rank must all-gather the full tensor for each layer onto its own GPU before generating the IPC handle.
|
||||
2. IPC handles for each gpu are sent to the inference engine via **Ray**, **HTTP**, or a **custom callable**. Each rank only reads the handle corresponding to its own GPU.
|
||||
1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`. In multi-GPU setups (e.g. FSDP), each trainer rank materializes the full tensor for each parameter onto its own GPU before generating the handle — which `ModuleSource` does for you.
|
||||
2. Every trainer rank contributes its handles to an all-gather; the sender merges them so each payload maps every GPU UUID to its args, and ships the merged handles to the inference engine through the client. Each worker reads only the handle for its own GPU.
|
||||
3. The inference worker reconstructs the tensors from the handles using `rebuild_cuda_tensor`, reading directly from the trainer's GPU memory.
|
||||
|
||||
Unlike NCCL, IPC transfer is straight-line: `update_weights` *is* the transfer,
|
||||
and it rides the client, so there is no concurrent broadcast to overlap with.
|
||||
|
||||
!!! note
|
||||
The handle all-gather in step 2 runs over the **default** process group, so it
|
||||
assumes that group is exactly the set of colocated trainer ranks and that the
|
||||
sender is a member. It is a no-op when no distributed group exists.
|
||||
|
||||
!!! 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.
|
||||
|
||||
## Inference Side
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from vllm.config import WeightTransferConfig
|
||||
|
||||
llm = LLM(model="my-model", weight_transfer_config=WeightTransferConfig(backend="ipc"))
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve my-model --weight-transfer-config '{"backend": "ipc"}'
|
||||
```
|
||||
|
||||
IPC needs no data-plane rendezvous, so `init_transfer_engine` opens no channel —
|
||||
it only records the `packed` flag the trainer ships at the handshake, which
|
||||
`receive_weights` then reads. Whether a transfer is packed is therefore never
|
||||
something you configure on the inference side.
|
||||
|
||||
## Trainer Side
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer import (
|
||||
ModuleSource,
|
||||
HTTPVLLMWeightSyncClient,
|
||||
WeightTransferTrainerFactory,
|
||||
)
|
||||
from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
|
||||
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 is the sender
|
||||
client=HTTPVLLMWeightSyncClient("http://localhost:8000"),
|
||||
source=ModuleSource(model),
|
||||
)
|
||||
|
||||
engine.send_weights() # once per sync
|
||||
```
|
||||
|
||||
`send_weights()` drives `start_weight_update`, the transfer itself, and
|
||||
`finish_weight_update`, and holds strong references to the IPC-shared copies
|
||||
until after the post-send barrier — the consumer's views would otherwise dangle.
|
||||
|
||||
Any [`VLLMWeightSyncClient`](base.md#vllmweightsyncclient) works here — the
|
||||
built-in HTTP and Ray clients, or an adapter for your own stack. The engine is
|
||||
identical either way.
|
||||
|
||||
### `IPCTrainerInitInfo`
|
||||
|
||||
| Field | Default | Description |
|
||||
| ----- | ------- | ----------- |
|
||||
| `rank` | — | Keyword-only. This trainer process's rank; **0 is the sender** |
|
||||
| `packed` | `False` | Chunked, bounded-memory transfer (see below) |
|
||||
| `packed_buffer_size_bytes` | 1 GiB | Chunk size when `packed=True` |
|
||||
|
||||
`packed` is a must-agree wire param: `trainer_init` ships it to the worker, which
|
||||
records it and decodes accordingly. It is not a `WeightTransferConfig` field and
|
||||
not a per-round `update_weights` field. `packed_buffer_size_bytes` is
|
||||
producer-only — the consumer rebuilds from the IPC handle plus the per-chunk
|
||||
`tensor_sizes`, so it never needs the buffer size.
|
||||
|
||||
## Packed (Chunked) Transfer
|
||||
|
||||
By default, all weights are sent in a single API call. For large models, this requires the full model to reside in GPU memory on both sides simultaneously. Setting `packed=True` enables **chunked transfer** with bounded GPU memory:
|
||||
By default, all weights are sent in a single `update_weights` call. For large
|
||||
models, this requires the full model to reside in GPU memory on both sides
|
||||
simultaneously. Setting `packed=True` enables **chunked transfer** with bounded
|
||||
GPU memory:
|
||||
|
||||
- Weights are concatenated into fixed-size packed buffers (controlled by `packed_buffer_size_bytes`).
|
||||
- Weights are concatenated into fixed-size packed buffers (`packed_buffer_size_bytes`).
|
||||
- Each chunk is sent as a separate `update_weights` call within a single `start_weight_update` / `finish_weight_update` bracket, so the layerwise reload pass is initialized once at the start and finalized once at the end regardless of chunk count.
|
||||
- After each chunk is consumed, the GPU memory for that chunk can be reclaimed.
|
||||
|
||||
```python
|
||||
trainer_args = IPCTrainerSendWeightsArgs(
|
||||
send_mode="ray",
|
||||
llm_handle=llm_actor_handle,
|
||||
packed=True,
|
||||
packed_buffer_size_bytes=256 * 1024 * 1024, # 256 MB chunks
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=IPCTrainerInitInfo(
|
||||
rank=0,
|
||||
packed=True,
|
||||
packed_buffer_size_bytes=256 * 1024 * 1024, # 256 MB chunks
|
||||
),
|
||||
client=client,
|
||||
source=ModuleSource(model),
|
||||
)
|
||||
```
|
||||
|
||||
## 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(
|
||||
send_mode="ray",
|
||||
llm_handle=llm_actor_handle,
|
||||
)
|
||||
# start
|
||||
ray.get(llm_actor_handle.start_weight_update.remote())
|
||||
# send weights
|
||||
IPCWeightTransferEngine.trainer_send_weights(
|
||||
iterator=model.named_parameters(),
|
||||
trainer_args=trainer_args,
|
||||
)
|
||||
# finish
|
||||
ray.get(llm_actor_handle.finish_weight_update.remote())
|
||||
```
|
||||
|
||||
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(
|
||||
send_mode="http",
|
||||
url="http://localhost:8000",
|
||||
)
|
||||
|
||||
# start
|
||||
base_url = "http://localhost:8000"
|
||||
url = f"{base_url}/start_weight_update"
|
||||
response = requests.post(url, json={}, timeout=60)
|
||||
response.raise_for_status()
|
||||
# send weights
|
||||
IPCWeightTransferEngine.trainer_send_weights(
|
||||
iterator=model.named_parameters(),
|
||||
trainer_args=trainer_args,
|
||||
)
|
||||
# finish
|
||||
url = f"{base_url}/finish_weight_update"
|
||||
response = requests.post(url, json={}, timeout=60)
|
||||
response.raise_for_status()
|
||||
```
|
||||
|
||||
In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint. Because the worker deserializes the payload via `pickle.loads`, the vLLM server must be started with `VLLM_ALLOW_INSECURE_SERIALIZATION=1`.
|
||||
|
||||
```python
|
||||
def my_custom_sender(update_info: IPCWeightTransferUpdateInfo):
|
||||
# Custom logic to deliver update_info to vLLM
|
||||
...
|
||||
|
||||
trainer_args = IPCTrainerSendWeightsArgs(
|
||||
send_mode=my_custom_sender,
|
||||
)
|
||||
|
||||
IPCWeightTransferEngine.trainer_send_weights(
|
||||
iterator=model.named_parameters(),
|
||||
trainer_args=trainer_args,
|
||||
)
|
||||
```
|
||||
|
||||
See [`IPCTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/ipc_engine.py) for the full list of configurable fields.
|
||||
On a multi-rank trainer the producer reuses one buffer across chunks, so packed
|
||||
mode carries a per-chunk barrier across ranks: without it, a rank could overwrite
|
||||
its buffer while its colocated worker is still reading the current chunk. This is
|
||||
handled inside `send_weights()`.
|
||||
|
||||
## Examples
|
||||
|
||||
- [RLHF with IPC weight syncing (offline, Ray)](../../../examples/rl/rlhf_ipc.py) - 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/rl/rlhf_http_ipc.py) - Weight transfer with a vLLM HTTP server where both server and trainer share the same GPU
|
||||
- [RLHF with IPC weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_ipc.py) - **Start here.** Server and training model share a single GPU; HTTP control plane, CUDA IPC data plane. Launches and tears down its own server
|
||||
- [RLHF with IPC + FSDP2 and expert parallelism](../../../examples/rl/rlhf_ipc_fsdp_ep.py) - Multi-rank trainer colocated with a `--data-parallel-size 4` server on the same 4 GPUs: every FSDP rank builds an engine and joins the handle all-gather, with packed chunking and sleep/wake around the transfer
|
||||
|
||||
@@ -10,128 +10,186 @@ The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) b
|
||||
|
||||
## How It Works
|
||||
|
||||
1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction).
|
||||
1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction). The trainer is rank 0; workers start at `rank_offset` (1).
|
||||
2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads the weights.
|
||||
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).
|
||||
|
||||
## Initialization
|
||||
The workers' `update_weights` and the trainer's broadcast run at the same time —
|
||||
both sides rendezvous inside the same NCCL calls. The trainer engine
|
||||
owns that concurrency internally.
|
||||
|
||||
NCCL requires explicit process group setup. The trainer and inference workers must agree on a master address, port, and world size.
|
||||
## Inference Side
|
||||
|
||||
### Inference Side
|
||||
The inference side takes a plain backend selector. The rendezvous parameters and
|
||||
the packing wire params arrive from the trainer at the init handshake.
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer.base import WeightTransferInitRequest
|
||||
from vllm import LLM
|
||||
from vllm.config import WeightTransferConfig
|
||||
|
||||
# rank_offset accounts for the trainer occupying rank 0
|
||||
llm.init_weight_transfer_engine(
|
||||
WeightTransferInitRequest(
|
||||
init_info=dict(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
rank_offset=1,
|
||||
world_size=world_size, # trainer + all inference workers
|
||||
)
|
||||
)
|
||||
)
|
||||
llm = LLM(model="my-model", weight_transfer_config=WeightTransferConfig(backend="nccl"))
|
||||
```
|
||||
|
||||
### Trainer Side
|
||||
```bash
|
||||
vllm serve my-model --weight-transfer-config '{"backend": "nccl"}'
|
||||
```
|
||||
|
||||
Nothing else is required: `init_weight_transfer_engine`, `start_weight_update`,
|
||||
`update_weights`, and `finish_weight_update` are all driven remotely by the
|
||||
trainer engine.
|
||||
|
||||
## Trainer Side
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer.nccl_engine import (
|
||||
NCCLWeightTransferEngine,
|
||||
from vllm.distributed.weight_transfer import (
|
||||
ModuleSource,
|
||||
HTTPVLLMWeightSyncClient,
|
||||
WeightTransferTrainerFactory,
|
||||
)
|
||||
from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
|
||||
|
||||
group = NCCLWeightTransferEngine.trainer_init(
|
||||
dict(
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=NCCLTrainerInitInfo(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
world_size=world_size,
|
||||
)
|
||||
world_size=world_size, # trainer + all inference workers
|
||||
rank=0, # this trainer rank; rank 0 is the sender
|
||||
packed=True,
|
||||
),
|
||||
client=HTTPVLLMWeightSyncClient("http://localhost:8000"), # or RayVLLMWeightSyncClient(llm)
|
||||
source=ModuleSource(model),
|
||||
)
|
||||
|
||||
engine.send_weights() # once per sync
|
||||
```
|
||||
|
||||
!!! note
|
||||
`trainer_init` always assigns the trainer to rank 0. Inference workers start at `rank_offset` (typically 1).
|
||||
`trainer_init` drives the whole handshake: it kicks off the inference side's
|
||||
`init_weight_transfer_engine` (through the client) on a side thread while opening
|
||||
the trainer's own rank-0 endpoint, since both ends must rendezvous together. It
|
||||
builds the worker's init info itself, with `rank_offset=1` and the same packed
|
||||
params, so the two sides cannot disagree.
|
||||
|
||||
## Sending Weights
|
||||
`send_weights()` then drives `start_weight_update`, `update_weights` — run
|
||||
concurrently with the broadcast — and `finish_weight_update`, and returns only
|
||||
once every transfer has drained.
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer.nccl_engine import (
|
||||
NCCLTrainerSendWeightsArgs,
|
||||
NCCLWeightTransferEngine,
|
||||
)
|
||||
### `NCCLTrainerInitInfo`
|
||||
|
||||
trainer_args = NCCLTrainerSendWeightsArgs(
|
||||
group=group,
|
||||
packed=True, # use packed broadcasting for efficiency
|
||||
)
|
||||
| Field | Default | Description |
|
||||
| ----- | ------- | ----------- |
|
||||
| `master_address` | — | Rendezvous host |
|
||||
| `master_port` | — | Rendezvous port |
|
||||
| `world_size` | — | Full trainer + worker NCCL group size |
|
||||
| `rank` | — | Keyword-only. This trainer process's rank; **0 is the sender** |
|
||||
| `packed` | `True` | Use packed broadcasting |
|
||||
| `packed_buffer_size_bytes` | 1 GiB | Packed buffer size |
|
||||
| `packed_num_buffers` | 2 | Number of rotating buffers (double/triple buffering) |
|
||||
|
||||
NCCLWeightTransferEngine.trainer_send_weights(
|
||||
iterator=model.named_parameters(),
|
||||
trainer_args=trainer_args,
|
||||
)
|
||||
```
|
||||
|
||||
See [`NCCLTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/nccl_engine.py) for the full list of configurable fields.
|
||||
`packed` defaults to `True` here. The worker-side default is `False`, but it only
|
||||
applies when no trainer ships a value — which, on this path, never happens.
|
||||
|
||||
### Packed Tensor Broadcasting
|
||||
|
||||
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.
|
||||
When `packed=True`, weight tensors are packed into large contiguous buffers
|
||||
before broadcasting. This cuts the number of NCCL operations and uses
|
||||
double/triple buffering with dedicated CUDA streams to overlap packing,
|
||||
broadcasting, and unpacking.
|
||||
|
||||
Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightTransferUpdateInfo`) must use matching `packed_buffer_size_bytes` and `packed_num_buffers` values.
|
||||
You set `packed`, `packed_buffer_size_bytes`, and `packed_num_buffers` **only on
|
||||
`NCCLTrainerInitInfo`**. The trainer propagates them to the worker inside
|
||||
`trainer_init`, the worker records them at the handshake, and `receive_weights`
|
||||
decodes with exactly the values the trainer encoded with. They are not per-round
|
||||
`update_weights` fields.
|
||||
|
||||
## Receiving Weights (Inference Side)
|
||||
!!! note "Memory"
|
||||
The rotating buffers are live for the whole transfer:
|
||||
`packed_buffer_size_bytes * packed_num_buffers` on each side (2 GiB at the
|
||||
defaults). Lower `packed_buffer_size_bytes` if that is too much headroom.
|
||||
|
||||
The inference side triggers weight reception using the four-phase protocol:
|
||||
`init_weight_transfer_engine`, `start_weight_update`, `update_weights`,
|
||||
`finish_weight_update`. The init phase is shown [above](#initialization). The
|
||||
remaining three steps are:
|
||||
### The two `WeightSource` channels must agree
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
|
||||
Dense NCCL is the backend that reads *both*
|
||||
[`WeightSource`](base.md#weightsource) channels, so it is the one where a
|
||||
disagreement between them is fatal. The engine builds the per-round update info
|
||||
from `metadata()` and ships it ahead of the bytes; the worker sizes its receive
|
||||
buffers from that info, and in packed mode cuts its chunk boundaries from it. The
|
||||
bytes themselves come from iterating the source.
|
||||
|
||||
# 1. Start the weight update
|
||||
llm.start_weight_update()
|
||||
If iteration disagrees with what `metadata()` declared — a reordered, omitted, or
|
||||
re-dtyped parameter — the two sides split the same byte stream differently. The
|
||||
transfer then either hangs in NCCL waiting for a length that never arrives, or
|
||||
loads garbage into the model.
|
||||
|
||||
# 2. Receive weights (can be called multiple times for chunked transfers)
|
||||
llm.update_weights(
|
||||
WeightTransferUpdateRequest(
|
||||
update_info=dict(
|
||||
names=names,
|
||||
dtype_names=dtype_names,
|
||||
shapes=shapes,
|
||||
packed=True,
|
||||
)
|
||||
)
|
||||
)
|
||||
The sender therefore checks each pair against the declared metadata as it goes,
|
||||
one comparison per parameter, and raises naming the first divergent parameter
|
||||
rather than letting it reach the wire. `ModuleSource` satisfies this by
|
||||
construction. If you write a custom source — a Megatron export, an MoE re-fusing
|
||||
pass — this is the invariant to test first.
|
||||
|
||||
# 3. Finish the weight update
|
||||
llm.finish_weight_update()
|
||||
```
|
||||
|
||||
The `names`, `dtype_names`, and `shapes` lists describe each parameter. These
|
||||
must match the order in which the trainer iterates over its parameters.
|
||||
|
||||
`start_weight_update` must be called before `update_weights`, and
|
||||
`finish_weight_update` must be called after all weight chunks have been
|
||||
transferred. The NCCL engine receives checkpoint-format weights and applies
|
||||
layerwise reload processing automatically inside `start_weight_update` /
|
||||
`finish_weight_update`.
|
||||
!!! note
|
||||
IPC does not read `metadata()` at all: it derives the update info from
|
||||
iteration as it goes, so it cannot observe a divergence.
|
||||
|
||||
## Sparse NCCL
|
||||
|
||||
Sparse, flat-index weight patches use a separate backend,
|
||||
`WeightTransferConfig(backend="sparse_nccl")`, implemented by
|
||||
`SparseNCCLWeightTransferEngine`. It shares only NCCL process-group
|
||||
`WeightTransferConfig(backend="sparse_nccl")`. It shares only NCCL process-group
|
||||
initialization with the dense engine; patches are applied directly in place to
|
||||
existing parameters (no layerwise reload). The current sparse MVP requires
|
||||
`TP=1` and `PP=1`. See the example below.
|
||||
existing parameters, with no layerwise reload. The current sparse MVP requires
|
||||
`TP=1` and `PP=1`, and assumes a single-rank trainer.
|
||||
|
||||
Sparse is a **delta** backend: each round's patches are a fresh set of deltas
|
||||
from the latest optimizer step, not a stable stream of the model's parameters. So
|
||||
the engine takes **no `WeightSource`** — passing one raises — and each round's
|
||||
patches go straight to `send_weights(patches)`. A round with no patches is a
|
||||
no-op.
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer import (
|
||||
RayVLLMWeightSyncClient,
|
||||
WeightTransferTrainerFactory,
|
||||
)
|
||||
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
|
||||
SparseNCCLTrainerInitInfo,
|
||||
SparseWeightPatch,
|
||||
)
|
||||
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=SparseNCCLTrainerInitInfo(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
world_size=world_size,
|
||||
rank=0,
|
||||
),
|
||||
client=RayVLLMWeightSyncClient(llm),
|
||||
)
|
||||
|
||||
patches = [
|
||||
SparseWeightPatch(
|
||||
name="model.layers.0.mlp.down_proj.weight",
|
||||
indices=flat_indices, # int32, 1-D
|
||||
values=new_values, # same length as indices
|
||||
full_shape=tuple(param.shape), # required when sending via the engine
|
||||
)
|
||||
]
|
||||
engine.send_weights(patches)
|
||||
```
|
||||
|
||||
`SparseNCCLTrainerInitInfo` takes the same rendezvous fields as the dense
|
||||
backend and no packed params — sparse transfers are never packed.
|
||||
|
||||
Patches are validated on the trainer *before* `start_weight_update`: int32
|
||||
indices, 1-D, and matching `indices`/`values` lengths, with `full_shape` set. The
|
||||
worker checks the same invariants when applying, but by then the broadcasts are
|
||||
under way, where a size mismatch wedges both sides instead of raising.
|
||||
|
||||
A trainer can hold a dense engine and a sparse engine at once — a full resync
|
||||
through the dense engine, cheap sparse deltas in between — as
|
||||
[`rlhf_sparse_nccl.py`](../../../examples/rl/rlhf_sparse_nccl.py) demonstrates.
|
||||
|
||||
## Examples
|
||||
|
||||
- [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
|
||||
- [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server
|
||||
- [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire
|
||||
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `backend="sparse_nccl"` and currently require `TP=1` and `PP=1`
|
||||
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
|
||||
- [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
|
||||
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model; uses an in-process `AsyncLLMEngine` with `RayVLLMWeightSyncClient`
|
||||
|
||||
+132
-108
@@ -1,45 +1,35 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
|
||||
via HTTP API, with IPC-based weight syncing APIs.
|
||||
RLHF weight syncing against a `vllm serve` HTTP server, using CUDA IPC for the
|
||||
data plane.
|
||||
|
||||
Unlike rlhf_nccl.py which uses NCCL and can use separate GPUs, this script
|
||||
uses CUDA IPC which requires the training model and vLLM server to be on the
|
||||
same GPU. Memory must be carefully managed to fit both models.
|
||||
* OpenAI-compatible API for inference requests
|
||||
* HTTP endpoints for the weight-transfer control plane
|
||||
* CUDA IPC handles for the weight data plane
|
||||
|
||||
Unlike rlhf.py which creates a vLLM instance programmatically, this script
|
||||
assumes you have already started a vLLM server using `vllm serve`. It uses:
|
||||
- OpenAI-compatible API for inference requests
|
||||
- HTTP endpoints for weight transfer control plane
|
||||
- CUDA IPC for actual weight data transfer
|
||||
1-GPU layout (single node): IPC shares GPU memory directly, so the server (TP=1)
|
||||
and the training model both live on GPU 0. The server is started with
|
||||
`--gpu-memory-utilization 0.5` to leave room for the training model.
|
||||
|
||||
Prerequisites:
|
||||
Start a vLLM server with weight transfer enabled and reduced GPU memory
|
||||
utilization to leave room for the training model:
|
||||
The script starts the server itself, then:
|
||||
|
||||
$ VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_INSECURE_SERIALIZATION=1 \
|
||||
vllm serve facebook/opt-125m --enforce-eager \
|
||||
--weight-transfer-config '{"backend": "ipc"}' \
|
||||
--load-format dummy \
|
||||
--gpu-memory-utilization 0.5
|
||||
1. Generate over HTTP → gibberish (server started with dummy weights).
|
||||
2. Pause generation, sync real weights trainer → server over IPC, resume.
|
||||
3. Generate again → sensible output.
|
||||
|
||||
Then run this script:
|
||||
IPC handles are pickled for HTTP transport, so both sides need
|
||||
`VLLM_ALLOW_INSECURE_SERIALIZATION=1`; this script sets it for itself and for
|
||||
the server it spawns.
|
||||
|
||||
$ python rlhf_http_ipc.py
|
||||
|
||||
The example performs the following steps:
|
||||
|
||||
* Load the training model on GPU 0 (same GPU as the vLLM server).
|
||||
* Generate text using the vLLM server via OpenAI-compatible API. The output
|
||||
is expected to be nonsense because the server is initialized with dummy weights.
|
||||
* Initialize weight transfer via HTTP endpoint (no-op for IPC).
|
||||
* Broadcast the real weights from the training model to the vLLM server
|
||||
using CUDA IPC handles.
|
||||
* Generate text again to show normal output after the weight update.
|
||||
Run:
|
||||
$ python examples/rl/rlhf_http_ipc.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
import requests
|
||||
import torch
|
||||
@@ -53,12 +43,76 @@ from vllm.distributed.weight_transfer import (
|
||||
)
|
||||
from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
|
||||
|
||||
BASE_URL = "http://localhost:8000"
|
||||
MODEL_NAME = "facebook/opt-125m"
|
||||
|
||||
# Enable insecure serialization for IPC handle serialization
|
||||
SERVER_PORT = 8000
|
||||
BASE_URL = f"http://localhost:{SERVER_PORT}"
|
||||
|
||||
# IPC requires colocation: the server and the training model share this GPU.
|
||||
SERVER_DEVICE_IDS = "0"
|
||||
TRAINER_DEVICE = "cuda:0"
|
||||
# Leave room on the shared GPU for the training model.
|
||||
SERVER_GPU_MEMORY_UTILIZATION = 0.5
|
||||
|
||||
# Needed to (de)serialize IPC handles across the HTTP boundary.
|
||||
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
|
||||
def start_vllm_server() -> subprocess.Popen:
|
||||
"""Spawn `vllm serve` and block until it is healthy."""
|
||||
serve_args = [
|
||||
"vllm",
|
||||
"serve",
|
||||
MODEL_NAME,
|
||||
"--tensor-parallel-size",
|
||||
"1",
|
||||
"--device-ids",
|
||||
SERVER_DEVICE_IDS,
|
||||
"--enforce-eager",
|
||||
"--load-format",
|
||||
"dummy",
|
||||
"--gpu-memory-utilization",
|
||||
str(SERVER_GPU_MEMORY_UTILIZATION),
|
||||
"--port",
|
||||
str(SERVER_PORT),
|
||||
"--weight-transfer-config",
|
||||
'{"backend": "ipc"}',
|
||||
]
|
||||
env = os.environ.copy()
|
||||
# Exposes the weight-transfer and pause/resume endpoints.
|
||||
env["VLLM_SERVER_DEV_MODE"] = "1"
|
||||
env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
||||
print(f"[server] Launching: {' '.join(serve_args)}")
|
||||
proc = subprocess.Popen(
|
||||
serve_args,
|
||||
env=env,
|
||||
stdout=sys.stdout,
|
||||
stderr=sys.stderr,
|
||||
start_new_session=True,
|
||||
)
|
||||
|
||||
deadline = time.monotonic() + 900
|
||||
while True:
|
||||
if proc.poll() is not None:
|
||||
raise RuntimeError("vLLM server exited before becoming ready.")
|
||||
try:
|
||||
if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200:
|
||||
break
|
||||
except requests.RequestException:
|
||||
pass
|
||||
if time.monotonic() > deadline:
|
||||
raise RuntimeError("vLLM server failed to start in time.")
|
||||
time.sleep(2)
|
||||
print("[server] Ready.")
|
||||
return proc
|
||||
|
||||
|
||||
def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
|
||||
"""Generate completions using the OpenAI-compatible API."""
|
||||
@@ -76,100 +130,70 @@ def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list
|
||||
|
||||
def pause_generation(base_url: str) -> None:
|
||||
"""Pause generation via HTTP endpoint."""
|
||||
url = f"{base_url}/pause"
|
||||
response = requests.post(url, timeout=60)
|
||||
response.raise_for_status()
|
||||
requests.post(f"{base_url}/pause", timeout=60).raise_for_status()
|
||||
|
||||
|
||||
def resume_generation(base_url: str) -> None:
|
||||
"""Resume generation via HTTP endpoint."""
|
||||
url = f"{base_url}/resume"
|
||||
response = requests.post(url, timeout=60)
|
||||
response.raise_for_status()
|
||||
requests.post(f"{base_url}/resume", timeout=60).raise_for_status()
|
||||
|
||||
|
||||
def get_world_size(base_url: str) -> int:
|
||||
"""Get world size from the vLLM server."""
|
||||
url = f"{base_url}/get_world_size"
|
||||
response = requests.get(url, timeout=10)
|
||||
response.raise_for_status()
|
||||
return response.json()["world_size"]
|
||||
|
||||
|
||||
def main():
|
||||
# IPC requires the training model to be on the same GPU as the vLLM server
|
||||
# The server should be started on GPU 0 with reduced memory utilization
|
||||
device = "cuda:0"
|
||||
torch.accelerator.set_device_index(device)
|
||||
|
||||
# Load the training model on the same GPU as the server
|
||||
# Use bfloat16 to reduce memory footprint
|
||||
print(f"Loading training model: {MODEL_NAME} on {device}")
|
||||
print(
|
||||
"Note: Ensure the vLLM server was started with --gpu-memory-utilization 0.5 "
|
||||
"or lower to leave room for the training model."
|
||||
)
|
||||
train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
|
||||
train_model.to(device)
|
||||
train_model.eval() # Set to eval mode to save memory
|
||||
|
||||
# Create OpenAI client pointing to the vLLM server
|
||||
client = OpenAI(
|
||||
base_url=f"{BASE_URL}/v1",
|
||||
api_key="EMPTY", # vLLM doesn't require an API key by default
|
||||
)
|
||||
|
||||
# Test prompts
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
# Generate text before weight update. The output is expected to be nonsense
|
||||
# because the server is initialized with dummy weights.
|
||||
def print_generations(label: str, prompts: list[str], outputs: list[str]) -> None:
|
||||
print("-" * 50)
|
||||
print("Generating text BEFORE weight update (expect nonsense):")
|
||||
print(label)
|
||||
print("-" * 50)
|
||||
outputs = generate_completions(client, MODEL_NAME, prompts)
|
||||
for prompt, generated_text in zip(prompts, outputs):
|
||||
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
|
||||
print("-" * 50)
|
||||
|
||||
print("Initializing weight transfer (IPC backend)...")
|
||||
|
||||
# The trainer engine drives the inference side over HTTP. init for IPC is a
|
||||
# no-op rendezvous; the same client carries start/update/finish.
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender
|
||||
client=HTTPVLLMWeightSyncClient(BASE_URL),
|
||||
source=ModuleSource(train_model),
|
||||
)
|
||||
def main():
|
||||
server_proc = start_vllm_server()
|
||||
try:
|
||||
# The training model must sit on the same physical GPU as the server.
|
||||
torch.accelerator.set_device_index(TRAINER_DEVICE)
|
||||
|
||||
# Pause generation before weight sync
|
||||
pause_generation(BASE_URL)
|
||||
print(f"[trainer] Loading training model: {MODEL_NAME} on {TRAINER_DEVICE}")
|
||||
train_model = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL_NAME, dtype=torch.bfloat16
|
||||
)
|
||||
train_model.to(TRAINER_DEVICE)
|
||||
train_model.eval() # eval mode to save memory on the shared GPU
|
||||
|
||||
print("Broadcasting weights via CUDA IPC (HTTP)...")
|
||||
# One call drives start_weight_update / update_weights / finish_weight_update.
|
||||
engine.send_weights()
|
||||
client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY")
|
||||
|
||||
# Resume generation after weight sync
|
||||
resume_generation(BASE_URL)
|
||||
# Generate with dummy weights — expect nonsense.
|
||||
outputs = generate_completions(client, MODEL_NAME, PROMPTS)
|
||||
print_generations("BEFORE weight sync (dummy weights):", PROMPTS, outputs)
|
||||
|
||||
# Generate text after weight update. The output is expected to be normal
|
||||
# because the real weights are now loaded.
|
||||
print("-" * 50)
|
||||
print("Generating text AFTER weight update:")
|
||||
print("-" * 50)
|
||||
outputs_updated = generate_completions(client, MODEL_NAME, prompts)
|
||||
for prompt, generated_text in zip(prompts, outputs_updated):
|
||||
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
|
||||
print("-" * 50)
|
||||
# IPC needs no data-plane rendezvous; `trainer_init` only ships the
|
||||
# `packed` flag, which the server must decode with.
|
||||
print("[transfer] Initializing IPC weight transfer...")
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender
|
||||
client=HTTPVLLMWeightSyncClient(BASE_URL),
|
||||
source=ModuleSource(train_model),
|
||||
)
|
||||
|
||||
# Note: The training model and IPC handles remain in memory.
|
||||
# In a real RLHF training loop, you would update the training model
|
||||
# and create new IPC handles for each weight update.
|
||||
pause_generation(BASE_URL)
|
||||
|
||||
# Drives start_weight_update / update_weights / finish_weight_update.
|
||||
print("[sync] Sharing weights via CUDA IPC...")
|
||||
engine.send_weights()
|
||||
print("[sync] Weight transfer complete.")
|
||||
|
||||
resume_generation(BASE_URL)
|
||||
|
||||
# Generate with the synced weights — expect sensible output.
|
||||
outputs_updated = generate_completions(client, MODEL_NAME, PROMPTS)
|
||||
print_generations("AFTER weight sync (real weights):", PROMPTS, outputs_updated)
|
||||
finally:
|
||||
print("[server] Shutting down...")
|
||||
server_proc.terminate()
|
||||
try:
|
||||
server_proc.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
server_proc.kill()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+146
-102
@@ -1,38 +1,33 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
|
||||
via HTTP API, with native weight syncing APIs.
|
||||
RLHF weight syncing against a `vllm serve` HTTP server, using NCCL for the
|
||||
data plane.
|
||||
|
||||
Unlike rlhf.py which creates a vLLM instance programmatically, this script
|
||||
assumes you have already started a vLLM server using `vllm serve`. It uses:
|
||||
- OpenAI-compatible API for inference requests
|
||||
- HTTP endpoints for weight transfer control plane
|
||||
- NCCL for actual weight data transfer
|
||||
* OpenAI-compatible API for inference requests
|
||||
* HTTP endpoints for the weight-transfer control plane
|
||||
* NCCL for the weight data plane
|
||||
|
||||
Prerequisites:
|
||||
Start a vLLM server with weight transfer enabled:
|
||||
3-GPU layout (single node):
|
||||
Inference — GPUs 0-1, `vllm serve` with TP=2 and fp8 quantization
|
||||
Training — GPU 2, a bf16 Hugging Face model in this process
|
||||
(the server quantizes to fp8 as it loads)
|
||||
|
||||
$ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
|
||||
--enforce-eager \
|
||||
--weight-transfer-config '{"backend": "nccl"}' \
|
||||
--load-format dummy
|
||||
The script starts the server itself, then:
|
||||
|
||||
Then run this script:
|
||||
1. Generate over HTTP → gibberish (server started with dummy weights).
|
||||
2. Pause generation, sync real weights trainer → server over NCCL, resume.
|
||||
3. Generate again → sensible output.
|
||||
|
||||
$ python rlhf_http.py
|
||||
|
||||
The example performs the following steps:
|
||||
|
||||
* Load the training model on GPU 0.
|
||||
* Generate text using the vLLM server via OpenAI-compatible API. The output
|
||||
is expected to be nonsense because the server is initialized with dummy weights.
|
||||
* Initialize weight transfer via HTTP endpoint.
|
||||
* Broadcast the real weights from the training model to the vLLM server
|
||||
using NCCL.
|
||||
* Generate text again to show normal output after the weight update.
|
||||
Run:
|
||||
$ python examples/rl/rlhf_http_nccl.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
import requests
|
||||
import torch
|
||||
from openai import OpenAI
|
||||
@@ -46,9 +41,71 @@ from vllm.distributed.weight_transfer import (
|
||||
from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
|
||||
from vllm.utils.network_utils import get_ip, get_open_port
|
||||
|
||||
BASE_URL = "http://localhost:8000"
|
||||
MODEL_NAME = "facebook/opt-125m"
|
||||
|
||||
SERVER_PORT = 8000
|
||||
BASE_URL = f"http://localhost:{SERVER_PORT}"
|
||||
|
||||
INFERENCE_TP_SIZE = 2
|
||||
# Physical GPUs for the server; the trainer takes the next one.
|
||||
SERVER_DEVICE_IDS = "0,1"
|
||||
TRAINER_DEVICE = "cuda:2"
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
|
||||
def start_vllm_server() -> subprocess.Popen:
|
||||
"""Spawn `vllm serve` and block until it is healthy."""
|
||||
serve_args = [
|
||||
"vllm",
|
||||
"serve",
|
||||
MODEL_NAME,
|
||||
"--tensor-parallel-size",
|
||||
str(INFERENCE_TP_SIZE),
|
||||
"--device-ids",
|
||||
SERVER_DEVICE_IDS,
|
||||
"--quantization",
|
||||
"fp8",
|
||||
"--enforce-eager",
|
||||
"--load-format",
|
||||
"dummy",
|
||||
"--port",
|
||||
str(SERVER_PORT),
|
||||
"--weight-transfer-config",
|
||||
'{"backend": "nccl"}',
|
||||
]
|
||||
env = os.environ.copy()
|
||||
# Exposes the weight-transfer and pause/resume endpoints.
|
||||
env["VLLM_SERVER_DEV_MODE"] = "1"
|
||||
print(f"[server] Launching: {' '.join(serve_args)}")
|
||||
proc = subprocess.Popen(
|
||||
serve_args,
|
||||
env=env,
|
||||
stdout=sys.stdout,
|
||||
stderr=sys.stderr,
|
||||
start_new_session=True,
|
||||
)
|
||||
|
||||
deadline = time.monotonic() + 900
|
||||
while True:
|
||||
if proc.poll() is not None:
|
||||
raise RuntimeError("vLLM server exited before becoming ready.")
|
||||
try:
|
||||
if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200:
|
||||
break
|
||||
except requests.RequestException:
|
||||
pass
|
||||
if time.monotonic() > deadline:
|
||||
raise RuntimeError("vLLM server failed to start in time.")
|
||||
time.sleep(2)
|
||||
print("[server] Ready.")
|
||||
return proc
|
||||
|
||||
|
||||
def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
|
||||
"""Generate completions using the OpenAI-compatible API."""
|
||||
@@ -66,105 +123,92 @@ def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list
|
||||
|
||||
def pause_generation(base_url: str) -> None:
|
||||
"""Pause generation via HTTP endpoint."""
|
||||
url = f"{base_url}/pause"
|
||||
response = requests.post(url, timeout=60)
|
||||
response.raise_for_status()
|
||||
requests.post(f"{base_url}/pause", timeout=60).raise_for_status()
|
||||
|
||||
|
||||
def resume_generation(base_url: str) -> None:
|
||||
"""Resume generation via HTTP endpoint."""
|
||||
url = f"{base_url}/resume"
|
||||
response = requests.post(url, timeout=60)
|
||||
response.raise_for_status()
|
||||
requests.post(f"{base_url}/resume", timeout=60).raise_for_status()
|
||||
|
||||
|
||||
def get_world_size(base_url: str) -> int:
|
||||
"""Get world size from the vLLM server."""
|
||||
url = f"{base_url}/get_world_size"
|
||||
response = requests.get(url, timeout=10)
|
||||
"""Get the number of inference workers from the vLLM server."""
|
||||
response = requests.get(f"{base_url}/get_world_size", timeout=10)
|
||||
response.raise_for_status()
|
||||
return response.json()["world_size"]
|
||||
|
||||
|
||||
def main():
|
||||
# Get the inference world size from the vLLM server
|
||||
inference_world_size = get_world_size(BASE_URL)
|
||||
world_size = inference_world_size + 1 # +1 for the trainer
|
||||
device = f"cuda:{inference_world_size}"
|
||||
torch.accelerator.set_device_index(device)
|
||||
|
||||
# Load the training model
|
||||
print(f"Loading training model: {MODEL_NAME}")
|
||||
train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
|
||||
train_model.to(device)
|
||||
|
||||
# Create OpenAI client pointing to the vLLM server
|
||||
client = OpenAI(
|
||||
base_url=f"{BASE_URL}/v1",
|
||||
api_key="EMPTY", # vLLM doesn't require an API key by default
|
||||
)
|
||||
|
||||
# Test prompts
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
# Generate text before weight update. The output is expected to be nonsense
|
||||
# because the server is initialized with dummy weights.
|
||||
def print_generations(label: str, prompts: list[str], outputs: list[str]) -> None:
|
||||
print("-" * 50)
|
||||
print("Generating text BEFORE weight update (expect nonsense):")
|
||||
print(label)
|
||||
print("-" * 50)
|
||||
outputs = generate_completions(client, MODEL_NAME, prompts)
|
||||
for prompt, generated_text in zip(prompts, outputs):
|
||||
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
|
||||
print("-" * 50)
|
||||
|
||||
# Set up the communication channel between the training process and the
|
||||
# vLLM server. The trainer is rank 0, vLLM worker(s) start after it.
|
||||
master_address = get_ip()
|
||||
master_port = get_open_port()
|
||||
|
||||
print(f"Initializing weight transfer: master={master_address}:{master_port}")
|
||||
def main():
|
||||
server_proc = start_vllm_server()
|
||||
try:
|
||||
# The trainer sits on the GPU after the server's, and is NCCL rank 0.
|
||||
torch.accelerator.set_device_index(TRAINER_DEVICE)
|
||||
|
||||
# The trainer engine owns the full handshake. `trainer_init` kicks off the
|
||||
# server's init_weight_transfer_engine (via the HTTP client) on a worker
|
||||
# thread while opening the trainer's NCCL endpoint, so both ends rendezvous
|
||||
# together — no manual threading needed here.
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=NCCLTrainerInitInfo(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
world_size=world_size,
|
||||
rank=0, # single-GPU trainer is the sole (sender) rank
|
||||
packed=True,
|
||||
),
|
||||
client=HTTPVLLMWeightSyncClient(BASE_URL),
|
||||
source=ModuleSource(train_model),
|
||||
)
|
||||
print(f"[trainer] Loading training model: {MODEL_NAME} on {TRAINER_DEVICE}")
|
||||
train_model = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL_NAME, dtype=torch.bfloat16
|
||||
)
|
||||
train_model.to(TRAINER_DEVICE)
|
||||
|
||||
# Pause generation before weight sync
|
||||
pause_generation(BASE_URL)
|
||||
client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY")
|
||||
|
||||
# One call drives start_weight_update / update_weights / finish_weight_update
|
||||
# over HTTP, concurrent with the NCCL broadcast.
|
||||
print("Broadcasting weights via NCCL...")
|
||||
engine.send_weights()
|
||||
# Generate with dummy weights — expect nonsense.
|
||||
outputs = generate_completions(client, MODEL_NAME, PROMPTS)
|
||||
print_generations("BEFORE weight sync (dummy weights):", PROMPTS, outputs)
|
||||
|
||||
# Resume generation after weight sync
|
||||
resume_generation(BASE_URL)
|
||||
# The transfer NCCL group is the trainer plus every inference worker.
|
||||
world_size = get_world_size(BASE_URL) + 1
|
||||
master_address = get_ip()
|
||||
master_port = get_open_port()
|
||||
print(
|
||||
f"[transfer] Rendezvous at {master_address}:{master_port}, "
|
||||
f"world_size={world_size} (1 trainer + {world_size - 1} vLLM workers)"
|
||||
)
|
||||
|
||||
# Generate text after weight update. The output is expected to be normal
|
||||
# because the real weights are now loaded.
|
||||
print("-" * 50)
|
||||
print("Generating text AFTER weight update:")
|
||||
print("-" * 50)
|
||||
outputs_updated = generate_completions(client, MODEL_NAME, prompts)
|
||||
for prompt, generated_text in zip(prompts, outputs_updated):
|
||||
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
|
||||
print("-" * 50)
|
||||
# `trainer_init` drives the handshake: it initializes the server's
|
||||
# transfer engine while opening the trainer's own NCCL endpoint, so both
|
||||
# ends rendezvous together.
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=NCCLTrainerInitInfo(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
world_size=world_size,
|
||||
rank=0, # single-GPU trainer is the sole (sender) rank
|
||||
packed=True,
|
||||
),
|
||||
client=HTTPVLLMWeightSyncClient(BASE_URL),
|
||||
source=ModuleSource(train_model),
|
||||
)
|
||||
|
||||
pause_generation(BASE_URL)
|
||||
|
||||
# Drives start_weight_update / update_weights / finish_weight_update,
|
||||
# concurrent with the NCCL broadcast.
|
||||
print("[sync] Broadcasting weights via NCCL...")
|
||||
engine.send_weights()
|
||||
print("[sync] Weight broadcast complete.")
|
||||
|
||||
resume_generation(BASE_URL)
|
||||
|
||||
# Generate with the synced weights — expect sensible output.
|
||||
outputs_updated = generate_completions(client, MODEL_NAME, PROMPTS)
|
||||
print_generations("AFTER weight sync (real weights):", PROMPTS, outputs_updated)
|
||||
finally:
|
||||
print("[server] Shutting down...")
|
||||
server_proc.terminate()
|
||||
try:
|
||||
server_proc.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
server_proc.kill()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,151 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray,
|
||||
with IPC-based weight syncing APIs
|
||||
|
||||
The script colocates the training and inference workloads onto the same GPU using Ray.
|
||||
|
||||
The example performs the following steps:
|
||||
|
||||
* Request a placement group of 1 GPU.
|
||||
* Place the inference model on the above GPU using the placement group.
|
||||
* Place and load the training model on the same GPU using the placement group.
|
||||
* Generate text from a list of prompts using the inference engine.
|
||||
* Update the weights of the training model and broadcast the updated weights
|
||||
to the inference engine by using CUDA IPC handles. Note that
|
||||
for demonstration purposes we simply zero out the weights.
|
||||
|
||||
This example assumes a single-node cluster with a single GPU,
|
||||
but can be extended to multiple GPUs.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import ray
|
||||
from ray.util.placement_group import placement_group
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import WeightTransferConfig
|
||||
from vllm.distributed.weight_transfer import (
|
||||
ModuleSource,
|
||||
RayVLLMWeightSyncClient,
|
||||
WeightTransferTrainerFactory,
|
||||
)
|
||||
from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
|
||||
|
||||
|
||||
class MyLLM(LLM):
|
||||
"""Configure the vLLM worker for Ray placement group execution."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
# Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
|
||||
# so that vLLM can manage its own device placement within the worker.
|
||||
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
|
||||
# Each worker uses 0.4 GPU so that two instances fit on the same GPU.
|
||||
os.environ["VLLM_RAY_PER_WORKER_GPUS"] = "0.4"
|
||||
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
|
||||
# needed for ipc handle serialization
|
||||
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
# Load the OPT-125M model onto GPU 0 for the training workload.
|
||||
|
||||
MODEL_NAME = "facebook/opt-125m"
|
||||
|
||||
|
||||
@ray.remote
|
||||
class TrainModel:
|
||||
def __init__(self, llm_handle: ray.actor.ActorHandle):
|
||||
self.train_model = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL_NAME,
|
||||
)
|
||||
self.train_model.to("cuda:0")
|
||||
self.llm_handle = llm_handle
|
||||
|
||||
def init_weight_transfer(self):
|
||||
"""Build the trainer-side IPC engine (no rendezvous needed for IPC)."""
|
||||
self.engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender
|
||||
client=RayVLLMWeightSyncClient(self.llm_handle),
|
||||
source=ModuleSource(self.train_model),
|
||||
)
|
||||
|
||||
def broadcast_weights(self):
|
||||
"""Broadcast weights to the inference engine using IPC.
|
||||
|
||||
Drives start/update/finish on the inference side internally.
|
||||
"""
|
||||
self.engine.send_weights()
|
||||
|
||||
|
||||
ray.init()
|
||||
|
||||
pg_colocate = placement_group([{"GPU": 1, "CPU": 0}])
|
||||
ray.get(pg_colocate.ready())
|
||||
|
||||
|
||||
llm = ray.remote(
|
||||
num_cpus=0,
|
||||
num_gpus=0,
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
||||
placement_group=pg_colocate,
|
||||
placement_group_capture_child_tasks=True,
|
||||
),
|
||||
)(MyLLM).remote(
|
||||
model=MODEL_NAME,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=1,
|
||||
distributed_executor_backend="ray",
|
||||
gpu_memory_utilization=0.7,
|
||||
weight_transfer_config=WeightTransferConfig(backend="ipc"),
|
||||
load_format="dummy",
|
||||
)
|
||||
|
||||
train_model = TrainModel.options(
|
||||
num_gpus=0.1,
|
||||
num_cpus=0,
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
||||
placement_group=pg_colocate, placement_group_capture_child_tasks=True
|
||||
),
|
||||
).remote(llm)
|
||||
|
||||
|
||||
# Generate text from the prompts.
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
sampling_params = SamplingParams(temperature=0)
|
||||
|
||||
outputs = ray.get(llm.generate.remote(prompts, sampling_params))
|
||||
|
||||
print("-" * 50)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
|
||||
print("-" * 50)
|
||||
|
||||
ray.get(llm.sleep.remote(level=0))
|
||||
|
||||
ray.get(train_model.init_weight_transfer.remote())
|
||||
# One call drives start_weight_update / update_weights / finish_weight_update.
|
||||
ray.get(train_model.broadcast_weights.remote())
|
||||
|
||||
ray.get(llm.wake_up.remote(tags=["scheduling"]))
|
||||
|
||||
outputs_packed = ray.get(llm.generate.remote(prompts, sampling_params))
|
||||
print("-" * 50)
|
||||
print("Results after packed/chunked IPC weight sync:")
|
||||
for output in outputs_packed:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
|
||||
print("-" * 50)
|
||||
+181
-232
@@ -4,96 +4,88 @@
|
||||
RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC**
|
||||
weight transfer and **packed** tensors.
|
||||
|
||||
Layout (4 GPUs, TP=1, DP=4, EP):
|
||||
* One Ray placement group per GPU.
|
||||
* Each PG holds one FSDP training worker and one vLLM ``LLM`` instance
|
||||
(sync API) using fractional GPUs so both fit on the same device.
|
||||
* The 4 ``LLM`` instances form a DP group via env-var-based SPMD
|
||||
coordination (``VLLM_DP_RANK``, ``VLLM_DP_SIZE``, etc.), the same
|
||||
mechanism used by ``examples/offline_inference/data_parallel.py``.
|
||||
* A ``DataParallelInferenceEngine`` actor spawns all 4 LLM actors,
|
||||
waits for initialization, and orchestrates generation / weight-sync.
|
||||
Multi-rank version of `rlhf_http_ipc.py`: the trainer is 4 FSDP2 Ray actors
|
||||
colocated with a data-parallel `vllm serve` on the same 4 physical GPUs.
|
||||
|
||||
4-GPU layout (single node), all colocated:
|
||||
Training — 4 GPUs, PyTorch FSDP2 (fully_shard), as Ray actors
|
||||
Inference — the same 4 GPUs, `vllm serve --data-parallel-size 4 -tp 1
|
||||
--enable-expert-parallel` (EP_SIZE = TP x DP = 4)
|
||||
|
||||
IPC requires the trainer and the server to sit on the same GPUs, so the script
|
||||
reserves the training GPUs through Ray first, asks Ray which ones it got, and
|
||||
pins the server to exactly those with `--device-ids`.
|
||||
|
||||
Both sides share each GPU, so the server is capped with
|
||||
`--gpu-memory-utilization` and its weights are moved aside for the transfer:
|
||||
|
||||
1. `/sleep?level=1` — offload server weights to CPU, drop KV cache
|
||||
2. `/wake_up?tags=weights` — weights back on GPU, KV cache still free
|
||||
3. packed IPC transfer — overwrite weights with room to spare
|
||||
4. `/wake_up?tags=kv_cache&tags=scheduling` — re-allocate KV cache, resume
|
||||
|
||||
Every FSDP rank builds an ``IPCTrainerWeightTransferEngine`` (via ``trainer_init``)
|
||||
and calls ``send_weights()``; all ranks join the IPC handle all-gather, and only
|
||||
rank 0 (the sender) ships the merged handles and drives the DP LLM actors through
|
||||
its ``RayVLLMWeightSyncClient``.
|
||||
rank 0 (the sender) ships the merged handles and drives the server.
|
||||
|
||||
This example was run on 4xH100.
|
||||
|
||||
Run:
|
||||
$ python examples/rl/rlhf_ipc_fsdp_ep.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
import ray
|
||||
import requests
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from huggingface_hub import snapshot_download
|
||||
from ray.util.placement_group import placement_group
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
from openai import OpenAI
|
||||
from torch.distributed.fsdp import fully_shard
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import WeightTransferConfig
|
||||
from vllm.distributed.weight_transfer import (
|
||||
HTTPVLLMWeightSyncClient,
|
||||
ModuleSource,
|
||||
RayVLLMWeightSyncClient,
|
||||
WeightTransferTrainerFactory,
|
||||
)
|
||||
from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
|
||||
from vllm.utils.network_utils import get_ip, get_open_port
|
||||
|
||||
# The inference side only needs the backend; the packed wire params are
|
||||
# trainer-side and get propagated to the workers at the init handshake.
|
||||
WEIGHT_TRANSFER_CONFIG = WeightTransferConfig(backend="ipc")
|
||||
# Packed IPC transfer with a 1 GB buffer (matches the per-chunk buffer size).
|
||||
PACKED = True
|
||||
PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024
|
||||
|
||||
TRAIN_GPU_FRACTION = float(os.environ.get("RLHF_IPC_TRAIN_GPU_FRACTION", "0.42"))
|
||||
VLLM_GPU_FRACTION = float(os.environ.get("RLHF_IPC_VLLM_GPU_FRACTION", "0.42"))
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen3-30B-A3B"
|
||||
SERVED_MODEL_NAME = "policy"
|
||||
|
||||
FSDP_WORLD_SIZE = 4
|
||||
INFERENCE_TP_SIZE = 1
|
||||
INFERENCE_DP_SIZE = 4
|
||||
|
||||
# Packed IPC transfer with a 1 GB chunk buffer.
|
||||
PACKED = True
|
||||
PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024
|
||||
|
||||
class MyLLM(LLM):
|
||||
"""LLM subclass that configures DP env vars for SPMD coordination."""
|
||||
# The server shares each GPU with a training rank, so cap what it reserves.
|
||||
SERVER_GPU_MEMORY_UTILIZATION = 0.35
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
dp_rank: int = 0,
|
||||
dp_size: int = 1,
|
||||
dp_master_ip: str = "127.0.0.1",
|
||||
dp_master_port: int = 0,
|
||||
**kwargs,
|
||||
):
|
||||
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
|
||||
os.environ["VLLM_RAY_PER_WORKER_GPUS"] = str(VLLM_GPU_FRACTION)
|
||||
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
|
||||
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
||||
SERVER_PORT = 8000
|
||||
BASE_URL = f"http://localhost:{SERVER_PORT}"
|
||||
|
||||
os.environ["VLLM_DP_RANK"] = str(dp_rank)
|
||||
os.environ["VLLM_DP_RANK_LOCAL"] = str(dp_rank)
|
||||
os.environ["VLLM_DP_SIZE"] = str(dp_size)
|
||||
os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
|
||||
os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
|
||||
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def ready(self):
|
||||
return True
|
||||
PROMPTS = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
|
||||
@ray.remote(num_cpus=0, num_gpus=TRAIN_GPU_FRACTION)
|
||||
@ray.remote(num_gpus=1)
|
||||
class FSDPTrainWorker:
|
||||
"""One FSDP2 worker per GPU; colocated with vLLM DP rank via placement group."""
|
||||
"""One FSDP2 worker per GPU; colocated with one vLLM DP rank."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -115,12 +107,6 @@ class FSDPTrainWorker:
|
||||
model_name, torch_dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
self.weight_names = [n for n, _ in model.named_parameters()]
|
||||
self.weight_dtype_names = [
|
||||
str(p.dtype).split(".")[-1] for _, p in model.named_parameters()
|
||||
]
|
||||
self.weight_shapes = [list(p.shape) for _, p in model.named_parameters()]
|
||||
|
||||
for layer in model.model.layers:
|
||||
fully_shard(layer)
|
||||
fully_shard(model)
|
||||
@@ -130,237 +116,200 @@ class FSDPTrainWorker:
|
||||
def get_rank(self):
|
||||
return self.rank
|
||||
|
||||
def get_weight_metadata(self):
|
||||
return self.weight_names, self.weight_dtype_names, self.weight_shapes
|
||||
def get_gpu_ids(self):
|
||||
"""Physical GPU id(s) Ray assigned to this worker."""
|
||||
return ray.get_gpu_ids()
|
||||
|
||||
def setup_engine(self, llm_handles):
|
||||
"""Build the trainer IPC engine on every FSDP rank.
|
||||
|
||||
Called on all ranks: rank 0 becomes the sender (drives the inference
|
||||
side and performs the no-op IPC init handshake against all DP LLM
|
||||
actors); the other ranks hold a null-client engine and only join the
|
||||
IPC handle all-gather during send_weights.
|
||||
"""
|
||||
def setup_engine(self, base_url: str):
|
||||
"""Build the trainer IPC engine. Called on every FSDP rank."""
|
||||
self.engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=IPCTrainerInitInfo(
|
||||
rank=self.rank, # FSDP rank; sender is 0
|
||||
packed=PACKED,
|
||||
packed_buffer_size_bytes=PACKED_BUFFER_SIZE_BYTES,
|
||||
),
|
||||
client=RayVLLMWeightSyncClient(llm_handles),
|
||||
client=HTTPVLLMWeightSyncClient(base_url),
|
||||
source=ModuleSource(self.model),
|
||||
)
|
||||
|
||||
def gather_and_broadcast_weights_ipc(self):
|
||||
"""All-gather full params across FSDP ranks; rank 0 sends to vLLM.
|
||||
|
||||
Called on all ranks concurrently. `send_weights` gathers each param and
|
||||
contributes to the IPC handle all-gather on every rank; only rank 0 (the
|
||||
sender) drives start/update/finish on the inference side (the other
|
||||
ranks' client RPCs no-op).
|
||||
"""
|
||||
"""Send the current weights to vLLM. Called on every FSDP rank."""
|
||||
self.engine.send_weights()
|
||||
|
||||
|
||||
@ray.remote(num_cpus=1)
|
||||
class DataParallelInferenceEngine:
|
||||
"""Manages a pool of DP-sharded vLLM LLM actors.
|
||||
def start_vllm_server(model_path: str, device_ids: str) -> subprocess.Popen:
|
||||
"""Spawn a `vllm serve` HTTP server (DP+EP) pinned to `device_ids`."""
|
||||
serve_args = [
|
||||
"vllm",
|
||||
"serve",
|
||||
model_path,
|
||||
"--served-model-name",
|
||||
SERVED_MODEL_NAME,
|
||||
"--tensor-parallel-size",
|
||||
str(INFERENCE_TP_SIZE),
|
||||
"--data-parallel-size",
|
||||
str(INFERENCE_DP_SIZE),
|
||||
"--enable-expert-parallel",
|
||||
# Pins the server to the same physical GPUs as the training ranks.
|
||||
"--device-ids",
|
||||
device_ids,
|
||||
"--enable-sleep-mode",
|
||||
"--enforce-eager",
|
||||
"--load-format",
|
||||
"dummy",
|
||||
"--gpu-memory-utilization",
|
||||
str(SERVER_GPU_MEMORY_UTILIZATION),
|
||||
"--port",
|
||||
str(SERVER_PORT),
|
||||
"--weight-transfer-config",
|
||||
'{"backend": "ipc"}',
|
||||
]
|
||||
env = os.environ.copy()
|
||||
env["VLLM_SERVER_DEV_MODE"] = "1" # exposes weight-transfer + sleep endpoints
|
||||
env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" # IPC handles over HTTP
|
||||
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
print(f"[server] Launching: {' '.join(serve_args)} (GPUs {device_ids})")
|
||||
proc = subprocess.Popen(
|
||||
serve_args,
|
||||
env=env,
|
||||
stdout=sys.stdout,
|
||||
stderr=sys.stderr,
|
||||
start_new_session=True,
|
||||
)
|
||||
|
||||
Spawns one MyLLM actor per placement group, waits for all engines to
|
||||
finish initializing, and exposes generation / weight-sync helpers.
|
||||
"""
|
||||
# Wait for the server to come up (model load can take a while).
|
||||
deadline = time.monotonic() + 1800
|
||||
while True:
|
||||
if proc.poll() is not None:
|
||||
raise RuntimeError("vLLM server exited before becoming ready.")
|
||||
try:
|
||||
if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200:
|
||||
break
|
||||
except requests.RequestException:
|
||||
pass
|
||||
if time.monotonic() > deadline:
|
||||
raise RuntimeError("vLLM server failed to start in time.")
|
||||
time.sleep(2)
|
||||
print("[server] Ready.")
|
||||
return proc
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str,
|
||||
pgs: list,
|
||||
dp_master_ip: str,
|
||||
dp_master_port: int,
|
||||
):
|
||||
dp_size = len(pgs)
|
||||
self.llm_actors = []
|
||||
for r in range(dp_size):
|
||||
sched = PlacementGroupSchedulingStrategy(
|
||||
placement_group=pgs[r],
|
||||
placement_group_capture_child_tasks=True,
|
||||
)
|
||||
actor = (
|
||||
ray.remote(num_cpus=0, num_gpus=0)(MyLLM)
|
||||
.options(scheduling_strategy=sched)
|
||||
.remote(
|
||||
model=model,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=INFERENCE_TP_SIZE,
|
||||
distributed_executor_backend="ray",
|
||||
enable_expert_parallel=True,
|
||||
gpu_memory_utilization=0.35,
|
||||
weight_transfer_config=WEIGHT_TRANSFER_CONFIG,
|
||||
enable_sleep_mode=True,
|
||||
load_format="dummy",
|
||||
dp_rank=r,
|
||||
dp_size=dp_size,
|
||||
dp_master_ip=dp_master_ip,
|
||||
dp_master_port=dp_master_port,
|
||||
)
|
||||
)
|
||||
self.llm_actors.append(actor)
|
||||
|
||||
ray.get([actor.ready.remote() for actor in self.llm_actors])
|
||||
def generate_completions(client: OpenAI, prompts: list[str]) -> list[str]:
|
||||
"""Generate completions via the OpenAI HTTP API."""
|
||||
results = []
|
||||
for prompt in prompts:
|
||||
response = client.completions.create(
|
||||
model=SERVED_MODEL_NAME,
|
||||
prompt=prompt,
|
||||
max_tokens=32,
|
||||
temperature=0,
|
||||
)
|
||||
results.append(response.choices[0].text)
|
||||
return results
|
||||
|
||||
def get_llm_actors(self):
|
||||
return self.llm_actors
|
||||
|
||||
def generate(self, prompts: list[str], sampling_params):
|
||||
"""Distribute prompts round-robin across DP ranks and collect results."""
|
||||
dp_size = len(self.llm_actors)
|
||||
per_rank: list[list[str]] = [[] for _ in range(dp_size)]
|
||||
indices: list[list[int]] = [[] for _ in range(dp_size)]
|
||||
def sleep_engine(base_url: str, level: int) -> None:
|
||||
"""Put the engine to sleep (level 1 offloads weights, drops KV cache)."""
|
||||
response = requests.post(f"{base_url}/sleep", params={"level": level}, timeout=600)
|
||||
response.raise_for_status()
|
||||
|
||||
for i, prompt in enumerate(prompts):
|
||||
rank = i % dp_size
|
||||
per_rank[rank].append(prompt)
|
||||
indices[rank].append(i)
|
||||
|
||||
refs = [
|
||||
actor.generate.remote(per_rank[r], sampling_params)
|
||||
for r, actor in enumerate(self.llm_actors)
|
||||
if per_rank[r]
|
||||
]
|
||||
all_outputs = ray.get(refs)
|
||||
def wake_up_engine(base_url: str, tags: list[str] | None = None) -> None:
|
||||
"""Wake the engine, optionally only for specific memory tags."""
|
||||
params = [("tags", tag) for tag in tags] if tags else None
|
||||
response = requests.post(f"{base_url}/wake_up", params=params, timeout=600)
|
||||
response.raise_for_status()
|
||||
|
||||
ordered = [None] * len(prompts)
|
||||
rank_idx = 0
|
||||
for r in range(dp_size):
|
||||
if per_rank[r]:
|
||||
for local_i, orig_i in enumerate(indices[r]):
|
||||
ordered[orig_i] = all_outputs[rank_idx][local_i]
|
||||
rank_idx += 1
|
||||
return ordered
|
||||
|
||||
def sleep(self, level: int = 0):
|
||||
ray.get([actor.sleep.remote(level=level) for actor in self.llm_actors])
|
||||
|
||||
def wake_up(self, tags: list[str] | None = None):
|
||||
ray.get([actor.wake_up.remote(tags=tags) for actor in self.llm_actors])
|
||||
def print_generations(label: str, prompts: list[str], outputs: list[str]) -> None:
|
||||
print("-" * 60)
|
||||
print(label)
|
||||
print("-" * 60)
|
||||
for prompt, text in zip(prompts, outputs):
|
||||
print(f"Prompt: {prompt!r}")
|
||||
print(f"Generated: {text!r}")
|
||||
print("-" * 60)
|
||||
|
||||
|
||||
def main():
|
||||
ray.init(
|
||||
runtime_env={
|
||||
"env_vars": {
|
||||
# The trainer pickles IPC handles for the HTTP client.
|
||||
"VLLM_ALLOW_INSECURE_SERIALIZATION": "1",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
assert TRAIN_GPU_FRACTION + VLLM_GPU_FRACTION <= 1.0, (
|
||||
"Train + vLLM GPU fractions must sum to at most 1.0 per bundle."
|
||||
)
|
||||
|
||||
local_model_path = snapshot_download(MODEL_NAME)
|
||||
print(f"[init] Model downloaded to {local_model_path}")
|
||||
|
||||
fsdp_master_addr = get_ip()
|
||||
fsdp_master_port = get_open_port()
|
||||
dp_master_port = get_open_port()
|
||||
dp_master_ip = get_ip()
|
||||
|
||||
# Create one placement group per DP rank (one GPU each).
|
||||
pgs = []
|
||||
for _ in range(INFERENCE_DP_SIZE):
|
||||
pg = placement_group([{"GPU": 1, "CPU": 1}])
|
||||
pgs.append(pg)
|
||||
ray.get([pg.ready() for pg in pgs])
|
||||
print(f"[init] {len(pgs)} placement groups ready.")
|
||||
|
||||
# Launch FSDP training workers, one per PG.
|
||||
scheduling = [
|
||||
PlacementGroupSchedulingStrategy(
|
||||
placement_group=pgs[r],
|
||||
placement_group_capture_child_tasks=True,
|
||||
)
|
||||
for r in range(FSDP_WORLD_SIZE)
|
||||
]
|
||||
|
||||
# Launch the training workers first so Ray reserves their GPUs; the server
|
||||
# is then pinned to those same physical GPUs.
|
||||
fsdp_workers = [
|
||||
FSDPTrainWorker.options(scheduling_strategy=scheduling[r]).remote(
|
||||
FSDPTrainWorker.remote(
|
||||
local_model_path,
|
||||
r,
|
||||
rank,
|
||||
FSDP_WORLD_SIZE,
|
||||
fsdp_master_addr,
|
||||
fsdp_master_port,
|
||||
)
|
||||
for r in range(FSDP_WORLD_SIZE)
|
||||
for rank in range(FSDP_WORLD_SIZE)
|
||||
]
|
||||
ray.get([w.get_rank.remote() for w in fsdp_workers])
|
||||
print(f"[init] {FSDP_WORLD_SIZE} FSDP workers ready.")
|
||||
print(f"[init] {FSDP_WORLD_SIZE} FSDP training workers ready.")
|
||||
|
||||
# Launch DP inference engine (spawns and initializes all LLM actors).
|
||||
inference_engine = DataParallelInferenceEngine.remote(
|
||||
model=local_model_path,
|
||||
pgs=pgs,
|
||||
dp_master_ip=dp_master_ip,
|
||||
dp_master_port=dp_master_port,
|
||||
training_gpus = sorted(
|
||||
int(g)
|
||||
for ids in ray.get([w.get_gpu_ids.remote() for w in fsdp_workers])
|
||||
for g in ids
|
||||
)
|
||||
llm_actors = ray.get(inference_engine.get_llm_actors.remote())
|
||||
print(f"[init] {INFERENCE_DP_SIZE} LLM actors ready.")
|
||||
if len(training_gpus) != INFERENCE_TP_SIZE * INFERENCE_DP_SIZE:
|
||||
raise RuntimeError(
|
||||
f"Need {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} colocated GPUs but "
|
||||
f"Ray assigned training to {training_gpus}."
|
||||
)
|
||||
device_ids = ",".join(str(g) for g in training_gpus)
|
||||
print(f"[init] Colocating training and inference on GPUs [{device_ids}].")
|
||||
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0)
|
||||
server_proc = start_vllm_server(local_model_path, device_ids)
|
||||
try:
|
||||
client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY")
|
||||
|
||||
print("[generate] Generating with dummy weights...")
|
||||
outputs = ray.get(inference_engine.generate.remote(prompts, sampling_params))
|
||||
print("-" * 60)
|
||||
print("BEFORE weight sync (dummy weights):")
|
||||
print("-" * 60)
|
||||
for output in outputs:
|
||||
print(f"Prompt: {output.prompt!r}")
|
||||
print(f"Generated: {output.outputs[0].text!r}")
|
||||
print("-" * 60)
|
||||
print("[generate] Generating with dummy weights...")
|
||||
outputs = generate_completions(client, PROMPTS)
|
||||
print_generations("BEFORE weight sync (dummy weights):", PROMPTS, outputs)
|
||||
|
||||
# --- Weight transfer ---
|
||||
# The rank-0 FSDP worker owns the trainer engine and drives the inference
|
||||
# side (init/start/update/finish) for all DP actors via the Ray client.
|
||||
print("[transfer] Initializing IPC weight transfer (all FSDP ranks)...")
|
||||
ray.get([w.setup_engine.remote(llm_actors) for w in fsdp_workers])
|
||||
# --- Weight transfer ---
|
||||
print("[transfer] Initializing IPC weight transfer (all FSDP ranks)...")
|
||||
ray.get([w.setup_engine.remote(BASE_URL) for w in fsdp_workers])
|
||||
|
||||
# Two-phase sleep/wake pattern:
|
||||
# 1. sleep(level=1) — offload weights to CPU, discard KV cache
|
||||
# 2. wake_up(tags=["weights"]) — bring weights back to GPU (KV cache still free)
|
||||
# 3. IPC weight transfer — overwrite weights, plenty of room without KV cache
|
||||
# 4. wake_up(tags=["kv_cache"]) — re-allocate KV cache for inference
|
||||
print("[sync] Sleeping engines (offload weights + free KV cache)...")
|
||||
ray.get(inference_engine.sleep.remote(level=1))
|
||||
print("[sync] Sleeping engine (offload weights + free KV cache)...")
|
||||
sleep_engine(BASE_URL, level=1)
|
||||
|
||||
print("[sync] Waking weights (KV cache stays free)...")
|
||||
ray.get(inference_engine.wake_up.remote(tags=["weights"]))
|
||||
print("[sync] Waking weights (KV cache stays free)...")
|
||||
wake_up_engine(BASE_URL, tags=["weights"])
|
||||
|
||||
# All FSDP ranks participate in the IPC handle all-gather; rank 0's engine
|
||||
# additionally drives start_weight_update / update_weights /
|
||||
# finish_weight_update on the inference side.
|
||||
print("[sync] Packed IPC transfer FSDP → vLLM...")
|
||||
ray.get([w.gather_and_broadcast_weights_ipc.remote() for w in fsdp_workers])
|
||||
print("[sync] Weight transfer complete.")
|
||||
print("[sync] Packed IPC transfer FSDP -> vLLM...")
|
||||
ray.get([w.gather_and_broadcast_weights_ipc.remote() for w in fsdp_workers])
|
||||
print("[sync] Weight transfer complete.")
|
||||
|
||||
print("[sync] Waking KV cache + scheduling...")
|
||||
ray.get(inference_engine.wake_up.remote(tags=["kv_cache", "scheduling"]))
|
||||
print("[sync] Waking KV cache + scheduling...")
|
||||
wake_up_engine(BASE_URL, tags=["kv_cache", "scheduling"])
|
||||
|
||||
print("[generate] Generating with synced weights...")
|
||||
outputs_updated = ray.get(
|
||||
inference_engine.generate.remote(prompts, sampling_params)
|
||||
)
|
||||
print("-" * 60)
|
||||
print("AFTER weight sync (real weights):")
|
||||
print("-" * 60)
|
||||
for output in outputs_updated:
|
||||
print(f"Prompt: {output.prompt!r}")
|
||||
print(f"Generated: {output.outputs[0].text!r}")
|
||||
print("-" * 60)
|
||||
print("[generate] Generating with synced weights...")
|
||||
outputs_updated = generate_completions(client, PROMPTS)
|
||||
print_generations("AFTER weight sync (real weights):", PROMPTS, outputs_updated)
|
||||
finally:
|
||||
print("[server] Shutting down...")
|
||||
server_proc.terminate()
|
||||
try:
|
||||
server_proc.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
server_proc.kill()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,196 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Demonstrates reinforcement learning using vLLM and Ray,
|
||||
with native weight syncing APIs at engine instance.
|
||||
|
||||
The script separates training and inference workloads onto distinct GPUs
|
||||
so that Ray can manage process placement and inter-process communication.
|
||||
A Hugging Face Transformer model occupies one GPU for training, whereas a
|
||||
2x tensor-parallel vLLM inference engine occupies two GPUs.
|
||||
|
||||
The example performs the following steps:
|
||||
* Load the training model on one gpu (scheduled via ray)
|
||||
* Initialize the inference model with dummy weights across
|
||||
two gpus using vLLM's tensor parallelism and Ray placement groups.
|
||||
* Generate gibberish from a list of prompts using the randomly initialized
|
||||
inference engine.
|
||||
* Update the weights of the training model and broadcast the updated weights
|
||||
to the inference engine by using a Ray collective RPC group.
|
||||
* Generating from the list of prompts after weight sync should result
|
||||
in sensible outputs.
|
||||
|
||||
This example assumes a single-node cluster with three GPUs, but Ray
|
||||
supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
|
||||
workloads. Residual GPU activity interferes with vLLM memory profiling and
|
||||
causes unexpected behavior.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import ray
|
||||
import torch
|
||||
from ray.util.placement_group import placement_group
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import WeightTransferConfig
|
||||
from vllm.distributed.weight_transfer import (
|
||||
ModuleSource,
|
||||
RayVLLMWeightSyncClient,
|
||||
WeightTransferTrainerFactory,
|
||||
)
|
||||
from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.network_utils import get_ip, get_open_port
|
||||
|
||||
MODEL_NAME = "facebook/opt-125m"
|
||||
# MODEL_NAME = "inference-optimization/Qwen3-0.6B-W4A16-G128"
|
||||
|
||||
|
||||
def get_assigned_gpu():
|
||||
"""This is a temporary workaround for a runtime bug in RCCL on ROCm."""
|
||||
if not current_platform.is_rocm():
|
||||
return 0
|
||||
assigned_gpu = int(ray.get_gpu_ids()[0])
|
||||
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
|
||||
os.environ.pop("HIP_VISIBLE_DEVICES", None)
|
||||
torch.accelerator.set_device_idx(assigned_gpu)
|
||||
return assigned_gpu
|
||||
|
||||
|
||||
class MyLLM(LLM):
|
||||
"""Configure the vLLM worker for Ray placement group execution."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1"
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class TrainModel:
|
||||
"""Ray actor that wraps the training model on a dedicated GPU."""
|
||||
|
||||
def __init__(self, model_name: str):
|
||||
assigned_gpu = get_assigned_gpu()
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
).to(f"cuda:{assigned_gpu}")
|
||||
|
||||
self.port = get_open_port()
|
||||
self.master_address = get_ip()
|
||||
|
||||
def init_weight_transfer(self, world_size, llm_handle):
|
||||
"""Build the trainer-side weight-transfer engine.
|
||||
|
||||
`trainer_init` drives the full handshake: it kicks off the inference
|
||||
side's `init_weight_transfer_engine` (via the Ray client) on a worker
|
||||
thread while opening the trainer's NCCL endpoint, so both ends
|
||||
rendezvous together. After this returns, `send_weights()` is callable.
|
||||
"""
|
||||
self.engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=NCCLTrainerInitInfo(
|
||||
master_address=self.master_address,
|
||||
master_port=self.port,
|
||||
world_size=world_size,
|
||||
rank=0, # single-GPU trainer is the sole (sender) rank
|
||||
packed=True,
|
||||
),
|
||||
client=RayVLLMWeightSyncClient(llm_handle),
|
||||
source=ModuleSource(self.model),
|
||||
)
|
||||
|
||||
def broadcast_weights(self):
|
||||
"""Push the current weights to the inference engine.
|
||||
|
||||
Drives start/update/finish on the inference side and the NCCL
|
||||
broadcast internally — one call.
|
||||
"""
|
||||
self.engine.send_weights()
|
||||
|
||||
|
||||
# Initialize Ray and set the visible devices. The vLLM engine will
|
||||
# be placed on GPUs 1 and 2.
|
||||
ray.init()
|
||||
|
||||
# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
|
||||
# Learn more about Ray placement groups:
|
||||
# https://docs.ray.io/en/latest/placement-groups.html
|
||||
# Launch the training model actor. Ray's resource scheduler will allocate
|
||||
# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
|
||||
train_model = TrainModel.remote(MODEL_NAME)
|
||||
|
||||
pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
|
||||
ray.get(pg_inference.ready())
|
||||
scheduling_inference = PlacementGroupSchedulingStrategy(
|
||||
placement_group=pg_inference,
|
||||
placement_group_capture_child_tasks=True,
|
||||
placement_group_bundle_index=0,
|
||||
)
|
||||
|
||||
# Launch the vLLM inference engine. The `enforce_eager` flag reduces
|
||||
# start-up latency.
|
||||
# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights)
|
||||
# are now native to vLLM workers.
|
||||
llm = ray.remote(
|
||||
num_cpus=0,
|
||||
num_gpus=0,
|
||||
scheduling_strategy=scheduling_inference,
|
||||
)(MyLLM).remote(
|
||||
model=MODEL_NAME,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=2,
|
||||
data_parallel_size=1,
|
||||
distributed_executor_backend="ray",
|
||||
weight_transfer_config=WeightTransferConfig(backend="nccl"),
|
||||
load_format="dummy",
|
||||
quantization="fp8",
|
||||
)
|
||||
|
||||
# Generate text from the prompts.
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
sampling_params = SamplingParams(temperature=0)
|
||||
|
||||
outputs = ray.get(llm.generate.remote(prompts, sampling_params))
|
||||
|
||||
# Generate text with the initial model. The output is expected to be nonsense
|
||||
# because the weights are randomly initialized.
|
||||
print("-" * 50)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
|
||||
print("-" * 50)
|
||||
|
||||
ray.get(llm.sleep.remote(level=0))
|
||||
|
||||
# Set up the communication channel between the training process and the
|
||||
# inference engine. The trainer engine now owns the full handshake — the
|
||||
# driver only has to hand it the vLLM actor handle and the world size.
|
||||
world_size = ray.get(llm.get_world_size.remote()) + 1 # +1 for the trainer
|
||||
ray.get(train_model.init_weight_transfer.remote(world_size, llm))
|
||||
|
||||
# Synchronize the updated weights to the inference engine. The engine drives
|
||||
# start_weight_update / update_weights / finish_weight_update on the inference
|
||||
# side internally, concurrent with the NCCL broadcast.
|
||||
ray.get(train_model.broadcast_weights.remote())
|
||||
|
||||
ray.get(llm.wake_up.remote(tags=["scheduling"]))
|
||||
|
||||
# Generate text with the updated model. The output is expected to be normal
|
||||
# because the weights are updated.
|
||||
outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
|
||||
print("-" * 50)
|
||||
for output in outputs_updated:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
|
||||
print("-" * 50)
|
||||
@@ -92,11 +92,6 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
|
||||
def shutdown(self) -> None:
|
||||
MockWeightTransferEngine.shutdown_called = True
|
||||
|
||||
@staticmethod
|
||||
def trainer_send_weights(*args, **kwargs):
|
||||
"""Mock method to simulate trainer sending weights."""
|
||||
pass
|
||||
|
||||
|
||||
def mock_create_engine(config, vllm_config, device, model):
|
||||
"""Mock factory function that returns our mock engine."""
|
||||
|
||||
@@ -77,7 +77,15 @@ class WeightSource(ABC):
|
||||
|
||||
@abstractmethod
|
||||
def metadata(self) -> list[ParamMeta]:
|
||||
"""Declare what iteration will yield, without transferring anything."""
|
||||
"""Declare what iteration will yield, without transferring anything.
|
||||
|
||||
Must agree with iteration element for element: the same parameters, in
|
||||
the same order, with the same dtypes and shapes. Backends may read both
|
||||
channels and trust that they match (dense NCCL sizes the worker's receive
|
||||
buffers and its packed chunk boundaries from this, then sends the bytes
|
||||
from iteration), so a source that disagrees between the two splits the
|
||||
stream differently on each side.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
@@ -342,32 +350,6 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def trainer_send_weights(
|
||||
iterator: Iterator[Any],
|
||||
trainer_args: dict[str, Any] | Any,
|
||||
) -> None:
|
||||
"""
|
||||
Send weights from trainer to inference workers.
|
||||
|
||||
This is a static method that can be called from the trainer process
|
||||
to send weights to all inference workers.
|
||||
|
||||
Args:
|
||||
iterator: Iterator of backend-specific items to send.
|
||||
trainer_args: Dictionary containing backend-specific arguments needed
|
||||
to send weights. The structure depends on the backend:
|
||||
- NCCL: Contains 'group', 'src', 'packed', etc.
|
||||
- IPC: Contains 'mode' ('http' or 'ray'),
|
||||
'llm_handle' (for Ray), 'url' (for HTTP), etc.
|
||||
|
||||
Example:
|
||||
>>> param_iter = ((n, p) for n, p in model.named_parameters())
|
||||
>>> engine.trainer_send_weights(param_iter, trainer_args)
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class VLLMWeightSyncClient(Protocol):
|
||||
|
||||
@@ -249,21 +249,6 @@ class IPCWeightTransferEngine(
|
||||
def shutdown(self) -> None:
|
||||
self._packed_importer.close()
|
||||
|
||||
@staticmethod
|
||||
def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
|
||||
"""Removed. Use the stateful `IPCTrainerWeightTransferEngine` instead.
|
||||
|
||||
Transitional stub kept only to satisfy the (still abstract)
|
||||
`WeightTransferEngine.trainer_send_weights`; that member is dropped from
|
||||
the worker ABC once every backend has migrated to the trainer engine.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"The static IPC trainer path has been replaced by "
|
||||
"IPCTrainerWeightTransferEngine. Build it via "
|
||||
"WeightTransferTrainerFactory.trainer_init(IPCTrainerInitInfo(...), "
|
||||
"client=..., source=...) and drive it with send_weights()."
|
||||
)
|
||||
|
||||
|
||||
class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitInfo]):
|
||||
"""Trainer-side CUDA IPC weight transfer engine.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import TYPE_CHECKING, Any, ClassVar
|
||||
from typing import TYPE_CHECKING, ClassVar
|
||||
|
||||
import torch
|
||||
from typing_extensions import Self
|
||||
@@ -226,21 +226,6 @@ class NCCLWeightTransferEngine(
|
||||
# Clean up the communicator by removing the reference
|
||||
self.model_update_group = None
|
||||
|
||||
@staticmethod
|
||||
def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
|
||||
"""Removed. Use the stateful `NCCLTrainerWeightTransferEngine` instead.
|
||||
|
||||
Transitional stub kept only to satisfy the (still abstract)
|
||||
`WeightTransferEngine.trainer_send_weights`; that member is dropped from
|
||||
the worker ABC once every backend has migrated to the trainer engine.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"The static NCCL trainer path has been replaced by "
|
||||
"NCCLTrainerWeightTransferEngine. Build it via "
|
||||
"WeightTransferTrainerFactory.trainer_init(NCCLTrainerInitInfo(...), "
|
||||
"client=..., source=...) and drive it with send_weights()."
|
||||
)
|
||||
|
||||
|
||||
class NCCLTrainerWeightTransferEngine(TrainerWeightTransferEngine[NCCLTrainerInitInfo]):
|
||||
"""Trainer-side NCCL weight transfer engine.
|
||||
|
||||
@@ -17,7 +17,7 @@ MVP limitations:
|
||||
from collections.abc import Iterable
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import TYPE_CHECKING, Any, ClassVar
|
||||
from typing import TYPE_CHECKING, ClassVar
|
||||
|
||||
import torch
|
||||
from typing_extensions import Self
|
||||
@@ -227,22 +227,6 @@ class SparseNCCLWeightTransferEngine(
|
||||
if self.model_update_group is not None:
|
||||
self.model_update_group = None
|
||||
|
||||
@staticmethod
|
||||
def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
|
||||
"""Removed. Use the stateful `SparseNCCLTrainerWeightTransferEngine`.
|
||||
|
||||
Transitional stub kept only to satisfy the (still abstract)
|
||||
`WeightTransferEngine.trainer_send_weights`; that member is dropped from
|
||||
the worker ABC once every backend has migrated to the trainer engine.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"The static sparse NCCL trainer path has been replaced by "
|
||||
"SparseNCCLTrainerWeightTransferEngine. Build it via "
|
||||
"WeightTransferTrainerFactory.trainer_init("
|
||||
"SparseNCCLTrainerInitInfo(...), client=..., source=...) and drive "
|
||||
"it with send_weights()."
|
||||
)
|
||||
|
||||
|
||||
class SparseNCCLTrainerWeightTransferEngine(
|
||||
TrainerWeightTransferEngine[SparseNCCLTrainerInitInfo]
|
||||
|
||||
Reference in New Issue
Block a user