vLLM rollout：采样为什么放在 vLLM
RLHF 训练每步都要生成 N 条 response，是整个 step 里最耗时的阶段（占 30%-70% wall time）。 OpenRLHF 把这件事单独切出来用 vLLM 做，吞吐能比"训练框架原生 generate"高 5-10 倍。 本章讲清楚 vLLM rollout 怎么集成、权重怎么同步、显存怎么省。
6.1为什么不直接用 DeepSpeed 的 generate
HuggingFace 的 model.generate() 在训练框架里跑会很慢，原因：
- 没有 KV cache 复用：每步独立分配，无 paging；
- 没有 continuous batching：每个 batch 长度都被最长那条 hold 住；
- fp32 master 占着显存：训练框架里模型权重是 fp32 master + bf16 副本，推理时其实只需要 bf16；
- 不能 chunk prefill：长 prompt 一次性 prefill 极慢。
vLLM 是专门做这件事的库，它对这 4 个问题都有针对性优化：
| 问题 | vLLM 的招 |
|---|---|
| KV cache 分配 | PagedAttention（类似 OS 的 virtual memory） |
| 批次效率 | Continuous batching（in-flight batching） |
| 显存占用 | 纯 bf16 + 用满显存的 KV cache |
| 长 prompt | Chunked prefill |
实测：同样 4B 模型、同样 prompts，vLLM 比"训练框架 .generate()" 快 5-10 倍。这一项优化是 OpenRLHF 区别于 DeepSpeed-Chat 的核心工程理由。
6.2vLLM 在 OpenRLHF 里怎么嵌入
看 openrlhf/trainer/ray/vllm_engine.py：每个 vLLM engine 就是一个 Ray Actor：
@ray.remote(num_gpus=tp_size)
class LLMRayActor:
def __init__(self, model_path, tensor_parallel_size, ...):
from vllm import LLM
self.llm = LLM(
model=model_path,
tensor_parallel_size=tensor_parallel_size,
gpu_memory_utilization=gpu_memory_utilization,
enforce_eager=enforce_eager,
enable_sleep_mode=enable_sleep,
)
def generate(self, prompts, sampling_params):
return self.llm.generate(prompts, sampling_params)
def sleep(self):
# vLLM 0.9+ API：把 KV cache + 模型权重换出显存
self.llm.sleep()
def wake_up(self):
self.llm.wake_up()
def init_weight_update_group(self, master_addr, master_port, rank, world_size):
# 和 actor 之间建一个 NCCL group，用于权重同步
...
def update_weight(self, name, dtype, shape):
# 接收 actor 广播过来的权重
...6.3vLLM 的几个关键 flag
| flag | 说明 |
|---|---|
--vllm.num_engines N |
起 N 个 vLLM engine（每个 engine 内部还有 TP） |
--vllm.tensor_parallel_size T |
每个 engine 的 TP 度。num_engines × tp_size = 总 GPU 数 |
--vllm.gpu_memory_utilization 0.5 |
vLLM 用显存的比例。colocate 时要留出空间给训练 actor，常设 0.4-0.6 |
--vllm.enforce_eager |
禁 CUDA graph（节省启动时间，但损失少量性能） |
--vllm.enable_sleep |
开 sleep mode（colocate 必须开） |
--vllm.sync_backend nccl |
权重同步用 NCCL（推荐）；可选 gloo |
--rollout.max_new_tokens 1024 |
生成最大长度 |
--rollout.temperature 1.0 |
采样温度 |
--rollout.top_p 1.0 |
nucleus 采样 |
--rollout.n_samples_per_prompt 1 |
每条 prompt 采几次（PPO=1；GRPO=4~16） |
6.4权重同步：训练 actor → vLLM rollout
这是 OpenRLHF 最关键的一段技术细节。每个 PPO step 开始前，要把训练 actor 的最新权重同步到 vLLM 内部， 但训练 actor（fp32 master + bf16）和 vLLM（bf16 + Paged KV）的布局不同，不能直接 memcpy。
同步流程
actor 和 vLLM 之间的 NCCL group A->>N: init_process_group (stateless) V->>N: init_process_group (stateless) loop 每个 PPO step Note over A: 收集自己的 bf16 权重
(从 fp32 master cast) A->>N: broadcast(param_name, param_data) N->>V: 接收到 weights V->>V: 写入 vLLM 内部的 weight tensor end
同步方式有两种
| 方式 | 说明 | flag |
|---|---|---|
| NCCL broadcast | 直接 GPU 间 broadcast，最快 | --vllm.sync_backend nccl（默认） |
| Gloo (CPU 中转) | 跨节点不通 NCCL 时备选 | --vllm.sync_backend gloo |
同步的开销 ≈ 模型大小 / NVLink 带宽。70B bf16 = 140 GB / 600 GB/s ≈ 0.23 秒。占整步时间 1% 以下，不是瓶颈。
6.5sleep mode：colocate 的命脉
colocate 模式（同一组卡既跑训练又跑 rollout）的问题：训练 actor + vLLM engine 总占用 > 100% 显存。解决：
| 阶段 | vLLM 状态 | 训练 actor 状态 |
|---|---|---|
| Rollout 阶段 | wake_up：权重 + KV cache 在显存 | offload 到 CPU |
| 训练阶段 | sleep：KV cache 释放，权重也可选 offload | reload 回 GPU 训练 |
sleep mode 是 vLLM 0.9+ 加的 API：
self.llm.sleep() # KV cache 释放，权重保留（level 1）
self.llm.sleep(level=2)# KV cache + 权重都换出（最省显存）
self.llm.wake_up() # 反向操作OpenRLHF 通过 --vllm.enable_sleep 和 --ds.enable_sleep 同时开两端的 sleep。
这是 colocate 能省 50%+ 显存的关键技术。
6.6num_engines vs tensor_parallel_size 怎么选
核心约束：
$$\text{vllm 用卡数} = \text{num\_engines} \times \text{tensor\_parallel\_size}$$
同一总卡数下两种切法：
| 方案 | num_engines | tp_size | 特点 |
|---|---|---|---|
| "多 engine 小 TP" | 4 | 2 | data parallel rollout，吞吐高；每张卡需放下完整模型 |
| "单 engine 大 TP" | 1 | 8 | 大模型才装得下；rollout 并行度低 |
经验值（8 卡 H100）：
| 模型 | 推荐配置 |
|---|---|
| ≤ 8B | num_engines=4, tp=2 或 num_engines=8, tp=1 |
| 13B-32B | num_engines=2, tp=4 |
| 70B | num_engines=1, tp=8 |
| 235B MoE | num_engines=1, tp=8（多机时 num_engines=2, tp=8） |
6.7采样多样性参数
| flag | 含义 | RLHF 推荐 |
|---|---|---|
--rollout.temperature | 采样温度 | 0.8 - 1.2（不要 0） |
--rollout.top_p | nucleus | 0.9 - 1.0 |
--rollout.top_k | top-k | -1（关）或 50 |
--rollout.max_new_tokens | 生成上限 | 1024 - 8192（看任务） |
--rollout.n_samples_per_prompt | 同 prompt 采几次 | PPO=1, GRPO=4-16 |
6.8dynamic batching 与 max_tokens_per_gpu
--train.dynamic_batch_enable
--train.max_tokens_per_gpu 16384
--rollout.max_tokens_per_gpu 32768含义：
- 静态 batch：每条样本截到
max_len，padding 浪费严重； - dynamic batch：按 token 总数打包，每个 GPU 一个 batch 不超过
max_tokens_per_gpu； - 训练和 rollout 的 max_tokens 可以不同（rollout 阶段没有反向，能装更多）。
对长 response 训练（> 4K）效果尤其显著，能省 30%+ 显存。
6.9训推 mismatch 问题
训练 actor 用 DeepSpeed kernel 算 logp，vLLM 用自家 kernel 算 logp，同样输入可能得到略不同的 logp。 这会让 PPO 的 importance ratio $\rho = \pi_\theta / \pi_\text{old}$ 偏离真值，损害稳定性。
OpenRLHF 的应对：
| 方法 | flag | 原理 |
|---|---|---|
| Importance Sampling 修正 | --algo.advantage.is_correction_enable |
在 PPO loss 里加一项 IS 修正项 |
| 训练侧重算 old_logp | 默认行为 | 用训练侧 forward 重新算一遍，丢掉 vLLM 给的 logp |
实战建议：保持默认（重算 old_logp），训练稳定性最好。
6.10多机时 vLLM 怎么部署
| 场景 | 典型放置 |
|---|---|
| 单机 8 卡 | colocate_all（actor / ref / critic / reward / vllm 同 8 卡） |
| 单机 8 卡 + 大模型 | actor/ref/critic 共置 4 卡，vLLM 独占另外 4 卡 |
| 2 节点 × 8 卡 | node 1: actor + vLLM 共置；node 2: critic + ref + reward 共置 |
| 多节点 70B 训练 | actor 跨节点 + vLLM 跨节点；ref/critic/reward 各占独立节点 |
第 8 章会详细讲三种 colocate 模式的取舍。
6.11权重同步两种后端：NCCL IPC vs Process Group
本章前面提到"NCCL broadcast 同步权重"，实际有两种实现，由 --vllm.sync_backend 和 --train.colocate_all 共同决定。来源 openrlhf/trainer/ray/ppo_actor.py:101–150：
| 后端 | 触发条件 | 原理 | 速度 |
|---|---|---|---|
NCCL IPC（use_cuda_ipc=True） | colocate_all=True 且不开 async | 同节点 GPU 通过 cudaIpcMemHandle 直接共享显存指针，零拷贝 | 最快（同节点 ~10x process group） |
Process Group（collective_rpc） | 多节点 / async / 显式 sync_backend=nccl | actor rank 0 通过 PyNCCL broadcast 把权重发到 vLLM workers | 跨节点必走这条 |
实战经验：
- 单节点 colocate → NCCL IPC，权重同步几乎不占时间；
- 多节点 → 自动退化到 process group，每次同步带宽紧时是瓶颈；
- 开
async_enable→ 强制 process group（IPC 跟异步语义不兼容）。
6.12vLLM 内部的WorkerExtension
OpenRLHF 通过 openrlhf/trainer/ray/vllm_worker_wrap.py 给 vLLM worker 注入一个 patch。注入的方法（vLLM 通过 worker_extension_cls 配置项暴露）：
| 方法 | 用途 |
|---|---|
init_weight_update_group(master_addr, master_port, rank_offset, world_size, group_name, backend) | 建立 actor rank 0 ↔ vLLM workers 的 NCCL 通信组 |
update_weight(name, dtype, shape) | 接收一个权重张量并 model.load_weights([(name, weight)]) |
update_weight_cuda_ipc(name, dtype, shape, ipc_handles) | 同上但走 IPC 句柄（零拷贝） |
reset_prefix_cache() | 同步完权重立刻清 KV cache（防止旧 prefix 复用脏数据） |
sleep / wake_up | colocate 时收 / 放权重和 KV cache |
"为什么 vLLM 默认不支持运行时权重更新？" → 因为 vLLM 设计是"加载一次模型就一直 serve"，OpenRLHF 这个 worker_extension 是给 vLLM 加 RLHF 友好的 hook，没这一层用普通 vLLM API 是发不进权重的。
6.13sleep_level 两档
本章前面提到 sleep mode 时只说了"释放显存"。实际 vLLM 的 sleep 有两档（vllm_engine.py:138–150）：
| sleep_level | 释放什么 | wake_up 时 | 用途 |
|---|---|---|---|
| 1 | 只释放 KV cache，权重留在 GPU | 只重新申请 KV | 权重不变的多轮 generate |
| 2 | 权重 + KV cache 都释放（搬 CPU 或丢） | weight broadcast + 申请 KV | colocate 时把卡完全让给训练 |
YAML 写法：--vllm.enable_sleep + --vllm.sleep_level 1/2。colocate_all 默认 sleep_level=2，省最多显存。
6.14Agent Executor：vLLM 不只是 generate
vLLM 这层不只是"输入 prompt 出 response"。2025-04 之后引入了 AgentExecutor（openrlhf/utils/agent.py），让 rollout 阶段就能跑多轮交互：
| 基类 / 实现 | 用途 |
|---|---|
AgentExecutorBase | 抽象基类；定义 execute(prompt) → trajectory |
AgentInstanceBase | 每条 prompt 一个 instance，含 reset() / step(action) |
SingleTurnAgentExecutor | 单轮（generate → reward） |
MultiTurnAgentExecutor | 多轮：generate → step（外部环境）→ 反馈 → generate ... |
用法：用户在 examples/python/agent_func.py 实现 AgentInstance.step()，CLI 通过 --reward.remote_url path/agent_func.py 加载。"agent 在 rollout 时就跑完了多轮"是 OpenRLHF 跟 TRL 区别最大的特性之一。
6.15这章带走的
- vLLM 把 rollout 时间降 5-10 倍，是 OpenRLHF 性能优势的核心；
- 每个 vLLM engine = 一个 Ray Actor，权重通过 NCCL broadcast 同步；
- sleep mode 是 colocate 能省显存的关键，必须配
--vllm.enable_sleep + --ds.enable_sleep； - num_engines × tp_size = 总 GPU；大模型用大 TP，小模型用多 engine；
- 不要把 temperature 设 0；GRPO 时 n_samples_per_prompt ≥ 4；
- dynamic batching + max_tokens_per_gpu 是长 response 训练的必开项。