vLLM 中的原生强化学习 (RL) API
随着后训练工作负载的持续扩大，我们见证了 vLLM 作为首选推理引擎被广泛采用。然而，有两个问题反复出现
- 训练与推理之间的权重同步通常以临时方式实现，并且在不同框架之间存在重复。
- 异步强化学习系统在大规模场景下变得脆弱，尤其是在 P/D（流水线/解耦）和 DPEP（数据并行/专家并行）部署中。
在这篇文章中，我们介绍了 vLLM 的两项改进：
- 为强化学习框架提供标准接口的原生权重同步 API。
- 对异步强化学习的改进支持，包括新的暂停模式以及针对 DPEP 设置中死锁的修复。
vLLM 中的原生权重同步 API
背景
在在线强化学习场景中，vLLM 模型权重必须定期同步，以确保生成的 rollout 是基于最新或近期的模型权重版本，从而提供更有用的反馈。
传统上，这种权重加载由每个强化学习框架单独处理，通常是通过扩展 vLLM Worker 并添加自定义权重接收和加载逻辑来实现的。虽然这行得通，但会导致一些问题：
- 增加了复杂性：框架作者必须实现并维护自定义的 Worker 扩展，若能为常见的传输策略提供原生支持会更好。
- 重复劳动：大多数强化学习框架最终都有非常相似的实现（例如，打包张量传输、RPC 端点）。
- 版本锁定：框架通常使用临时的处理方式来预处理/后处理接收到的权重，以便 vLLM Worker 加载它们，这可能导致版本锁定实现。
vLLM 中的新 API
我们引入了 vLLM 原生权重同步 API 以标准化此流程。权重传输 API 由四个阶段和一个可插拔的后端组成：
- 初始化 (
init_weight_transfer_engine)：在训练器和推理 Worker 之间建立通信通道。在训练循环开始前调用一次。 - 开始权重更新 (
start_weight_update)：开始权重更新。在每个训练步骤（或批次）后调用。准备 vLLM Worker 以接收权重。 - 更新权重 (
update_weights)：从训练器向推理引擎更新全部或部分权重。可多次调用以进行分块权重传输。 - 结束权重更新 (
finish_weight_update)：完成当前的权重更新。运行任何必要的后处理（例如量化）。
相应的 API 已在 API 服务器和引擎层实现。
目前，我们支持以下后端：
- NCCL：使用 NCCL 广播操作在不同 GPU 上的训练和推理 Worker 之间传输权重。
- IPC：使用 CUDA IPC 通过共享内存句柄实现同设备权重传输。
两个后端都支持优化后的打包实现，以最大限度地减少序列化开销。
核心传输逻辑由可插拔的 WeightTransferEngine 抽象实现，将权重传输与 Worker 实现解耦，允许用户轻松引入自己的实现。核心思想是，“初始化”和“更新权重”阶段通常由强化学习框架开发者自定义并包含传输逻辑，而开始和结束阶段则属于控制消息，涉及 vLLM 中与传输无关的预/后处理。
注意：HTTP 权重传输端点要求设置
VLLM_SERVER_DEV_MODE=1。
示例
例如，以下是在 vLLM 上使用 FP8 量化通过 NCCL 进行权重传输的操作流程：
新的 API 可以按如下方式使用：
1. 配置引擎以进行权重传输。
from vllm import LLM
from vllm.config import WeightTransferConfig
llm = LLM(
model="my-model",
weight_transfer_config=WeightTransferConfig(backend="nccl"),
)2. 初始化通信状态：初始化训练器和推理引擎之间的通信状态。训练器的 Rank 0 进程和所有推理 Worker 加入一个共享的 NCCL 进程组。
from vllm.distributed.weight_transfer.base import WeightTransferInitRequest
# Initialization for inference
llm.init_weight_transfer_engine(
WeightTransferInitRequest( # <--- initialization parameters
init_info=dict(
master_address=master_address,
master_port=master_port,
rank_offset=1, # <--- offset accounts for trainer rank 0
world_size=world_size, # <--- trainer + all inference workers
)
)
)
# Initialization for training
from vllm.distributed.weight_transfer.nccl_engine import (
NCCLWeightTransferEngine,
)
group = NCCLWeightTransferEngine.trainer_init(
dict(
master_address=master_address,
master_port=master_port,
world_size=world_size,
)
)3. 从训练器发送权重：在训练器上开始权重传输。WeightTransferEngine 实现了一个 trainer_send_weights 方法，它接收可迭代的参数列表并初始化全部或部分参数的传输。用户也可以实现自己的发送功能。在这里，我们可以利用打包张量广播，通过将多个小张量合并到更大的缓冲区中来提高传输吞吐量。
from vllm.distributed.weight_transfer.nccl_engine import (
NCCLTrainerSendWeightsArgs,
NCCLWeightTransferEngine,
)
trainer_args = NCCLTrainerSendWeightsArgs(
group=group,
packed=True, # use packed broadcasting for efficiency
)
# send weights from an `AutoModelForCausalLM` instance
NCCLWeightTransferEngine.trainer_send_weights(
iterator=model.named_parameters(),
trainer_args=trainer_args,
)4. 在推理引擎中接收权重。
from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
# executed asynchronously while trainer sends weights
llm.start_weight_update()
llm.update_weights(
WeightTransferUpdateRequest(
update_info=dict(
names=names,
dtype_names=dtype_names,
shapes=shapes,
packed=True,
)
)
)
llm.finish_weight_update()自定义权重传输
权重传输 API 的主要目标之一是使强化学习框架能够实现自定义的权重传输策略。通过新 API，用户可以实现并注册自定义的 WeightTransferEngine。
from dataclasses import dataclass
from typing import Iterator, Callable, Any
from torch import Tensor
from vllm.distributed.weight_transfer.base import (
WeightTransferEngine,
WeightTransferInitInfo,
WeightTransferUpdateInfo,
)
# define custom dataclasses for initialization and update weights metadata
@dataclass
class MyInitInfo(WeightTransferInitInfo):
"""Custom initialization info."""
...
@dataclass
class MyUpdateInfo(WeightTransferUpdateInfo):
"""Custom update info."""
...
# custom weight transfer engine
class MyWeightTransferEngine(WeightTransferEngine):
init_info_cls = MyInitInfo
update_info_cls = MyUpdateInfo
def init_transfer_engine(self, init_info: MyInitInfo):
...
def receive_weights(
self,
update_info: MyUpdateInfo,
load_weights: Callable[[list[tuple[str, Tensor]]], None],
):
...
@classmethod
def trainer_send_weights(
cls,
iterator: Iterator[tuple[str, Tensor]],
trainer_args: dict[str, Any] | Any,
):
...
# finally, register the weight transfer engine
from vllm.distributed.weight_transfer import WeightTransferEngineFactory
WeightTransferEngineFactory.register_engine("my_weight_transfer", MyWeightTransferEngine)注意 trainer_send_weights 方法是可选的。它编码了训练器上使用的发送逻辑，用户不需要以这种方式构建其发送逻辑。
上述简单的 API 可以实现许多高级用例。作为一个原型，我们演示了如何实现 Etha 风格的分片权重传输，点击此处查看。
改进的异步强化学习暂停/恢复支持
在异步强化学习中，权重会在推理请求仍在进行时进行更新。通常，权重同步在异步强化学习中涉及三个操作：暂停生成、传输更新后的权重、恢复生成。用户可以选择如何处理进行中的请求（例如，中止所有正在运行的请求，或从之前生成的 token 恢复生成），以及保留还是丢弃 KV 缓存。
暂停/恢复的保持模式 (Keep Mode)
为了在推理引擎运行时安全地更新权重，vLLM 提供了 pause_generation 和 resume_generation 方法。相同功能在 HTTP 服务器中作为 POST /pause 和 POST /resume API 提供。此前，AsyncLLMEngine.pause_generation 支持两种模式：
- 中止所有请求
- 等待请求完成
我们添加了第三个选项：保持模式 (Keep Mode)。下表对比了不同模式：
| 模式 | 说明 | 客户端影响 | 异步强化学习是否可行？ |
|---|---|---|---|
中止 | 中止所有进行中的请求 | 客户端必须处理重试 | 是 |
等待 | 等待所有进行中的请求完成 | 客户端无需重试 | 否，权重更新前生成必须完成 |
保持 | 暂停进行中的请求 | 客户端无需重试 | 是 |
保持模式可以按如下方式使用：
# pause - preserve ongoing requests
await engine.pause_generation(mode="keep")
# update weights here
# resume
await engine.resume_generation()在保持模式下：
- 进行中的请求被暂停，但不会被丢弃。
- 调度器停止，但状态得到保留。
修复 DPEP 设置中的死锁
大规模异步强化学习需要在 DPEP 部署中对进行中的权重更新进行细致协调。在 vLLM 中，DPCoordinator 确保在 vLLM Rank 之间对生成进行仔细协调以防止死锁。具体来说，每个 DP Rank 在任何 DP Rank 中有活动调度请求时都会执行前向传递。
此前，vLLM 中 DP 部署的异步强化学习常导致死锁，主要是因为某些引擎收到了暂停信号，而其他引擎仍在积极处理请求并等待所有引擎加入。发生此情况的原因之一是暂停状态在 AsyncLLM 对象中跟踪，而 DP 协调消息是在 EngineCore 进程和 DPCoordinator 之间交换的。例如，对于 DP 世界大小为 2 的情况，死锁场景如下：
- API 服务器 DP Rank 0 接收到一个生成请求并将其转发给
EngineCore。API 服务器向DPCoordinator发出FIRST_REQ消息以开启新一轮。请求被转发给 DP Rank 0EngineCore，该进程开始一个新的调度步骤。 - 控制器向两个引擎发出
/pause请求。暂停状态在AsyncLLM对象中设置，新请求不再转发给EngineCore。所有 DP Rank 上的 API 服务器立即返回。 - 训练器发出权重更新请求。与此同时，DP Rank 0
EngineCore已进入前向传递，并等待其他 DP Rank 加入。（为简单起见，此处忽略start_weight_update和finish_weight_update请求）。 - 权重更新请求到达 DP Rank 1
EngineCore，副本进入 NCCL 广播集合通信，等待其他 Rank。权重更新请求在 DP Rank 0EngineCore上排队。 DPCoordinator向 DP Rank 1EngineCore发送START_DP_WAVE消息，但消息被排队。- 不同的 Rank 处于不同的集合通信中并导致死锁。
同样的场景如下所示：
我们通过两项更改来解决此问题：
1. 将暂停逻辑移入 EngineCore。不再在 AsyncLLM 入口层跟踪暂停状态，而是直接在调度器中处理。这减少了暂停和生成请求之间的竞争条件。
2. 两阶段暂停/恢复。
- 阶段 1（本地暂停）：每个引擎暂停调度，但通过尊重任何入站的
START_DP_WAVE请求继续步进，因此它仍能参与所需的前向传递。 - 阶段 2（全局暂停）：目前，所有 Rank 每 32 步执行一次全局全归约 (All-Reduce)，以检查是否有任何 DP Rank 中有挂起的请求。在同一次全归约阶段，我们也检查所有引擎是否处于“本地暂停”状态。如果所有 Rank 同意，它们将一起停止。
这确保了：
- 没有 Rank 会卡在等待中。
- 即使引擎收到暂停请求，也会尊重
START_DP_WAVE。 - 所有 Worker 表现一致。
因此，之前的场景得到了妥善处理：
- API 服务器 DP Rank 0 接收到一个生成请求并将其转发给
EngineCore。API 服务器向DPCoordinator发出FIRST_REQ消息以开启新一轮。请求被转发给 DP Rank 0EngineCore，该进程开始一个新的调度步骤。 - 控制器向两个引擎发出
/pause请求。暂停请求被转发给两个 Rank 的EngineCore。暂停请求在 DP Rank 0EngineCore中排队，直到步骤完成。注意 API 服务器此时尚未返回。 - DP Rank 0
EngineCore开始执行前向传递，Worker 在全对全集合通信中等待。 - DP Rank 1
EngineCore收到暂停请求，进入“本地暂停”状态。 DPCoordinator向 DP Rank 1EngineCore发送START_DP_WAVE消息。- DP Rank 1
EngineCore开始执行前向传递。由于两个 DP Rank 都已加入，前向传递完成。 - DP Rank 0
EngineCore处理暂停请求，进入“本地暂停”状态。 - DP Rank 0 和 DP Rank 1
EngineCore参与周期性全归约，意识到两个引擎都处于“本地暂停”状态，并进入“全局暂停”状态。 - API 服务器在
/pause调用处返回。 - 训练器发出权重更新请求。API 服务器将权重更新请求转发给
EngineCore进程。（为简单起见，此处忽略start_weight_update和finish_weight_update请求）。 - 所有 vLLM Worker 进入 NCCL 广播集合通信。训练器开始 NCCL 广播。
- 权重更新成功完成。
验证
演示新的强化学习 API
我们在 SkyRL 中演示了新强化学习 API 的使用。
在 SkyRL 中，训练器通过 HTTP 与推理引擎交互。对于权重同步，SkyRL 使用了原生权重同步 API，以及用于异步强化学习的原生 /pause 和 /resume API。与原生强化学习 API 的集成细节在文档中有所描述，并且我们演示了在原始 DAPO 配方上对 Qwen3-1.7B 进行异步训练（示例）。
大规模验证：宽 EP 设置中的完全异步强化学习
Prime-RL 团队针对 zai-org/GLM-5.1-FP8 的部署验证了这些强化学习 API，推理运行在跨 16 个 8xH200 节点的 P/D 解耦设置中——2 个 4P+4D 的副本，预填充和解码均使用 DPEP32。所有实例还配置了每节点 1TB 容量的 CPU KV 缓存卸载。使用提供缓存感知粘性路由的 vllm-router 实现了跨引擎的路由。训练器在另外 16 个 8xH200 节点上运行 BF16 模型等效版本（zai-org/GLM-5.1），在自定义数学环境上使用 IcePop 作为首选算法。此部署在训练 100 多步的过程中表现稳定，评估性能不断增长，呈现出上升的强化学习曲线，KL 不匹配稳定，权重更新进展正常。
总结
我们看到 vLLM 强化学习社区对在新的强化学习 API 之上进行构建表现出越来越浓厚的兴趣。来自 vLLM 强化学习社区的一些持续工作包括集成新的 K8s 原生权重传输引擎，以及以通用方式支持 分片感知、RDMA 原生权重传输。新开发进度在 vLLM 强化学习路线图中跟踪。
在文档中了解更多关于 vLLM 中强化学习工具的信息：
在此处尝试 vLLM 上的新 API：https://github.com/vllm-project/vllm/tree/main/examples/rl。
致谢
感谢以下团体和个人使这一切成为可能：
- Prime-RL 团队（特别是 Matej Sirovatka）和 Junjie Zhang，感谢他们在协助验证和调试大规模强化学习 API 运行方面的贡献。
- NemoRL 团队提供优化的打包张量实现。
- Robert Shaw 组织强化学习相关工作。
- Kyle Sayers 使通过逐层重新加载实现量化权重重新加载成为可能。