总目录›TRL›第 10 章
Chapter 10

vLLM 集成：online sampling 加速

📌 commit 8221a9f把 rollout 速度提 3-4×

PPO 和 GRPO 训练里，rollout（generate）通常占 50% 以上时间。TRL 2024 后期增加了对 vLLM 的原生支持，把 rollout 替换成 vLLM 的 PagedAttention + continuous batching 引擎。

10.1为什么 generate 慢

HuggingFace 的 generate 在训练时跑：

vLLM 全部解决：continuous batching、paged KV cache、专用 generator 进程。

10.2两种集成模式

模式vllm_modevLLM 在哪适用
Colocate"colocate"跟训练同一张卡（不同进程）单机
Server"server"独立 GPU 跑 vLLM server多机大模型

10.3Colocate 模式

cfg = GRPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="colocate",
    vllm_gpu_memory_utilization=0.4,    # ★ vLLM 用多少显存
    vllm_max_model_len=4096,
    vllm_dtype="bfloat16",
)

训练时 GPU 一部分给 vLLM、一部分给训练 op。每个 step 流程：

  1. 训练侧把当前 policy weight 同步给 vLLM；
  2. vLLM 跑 rollout 拿 responses；
  3. 训练侧拿 responses 算 reward + update；
  4. 循环。

10.4Server 模式

vLLM 跑在独立机器 / GPU 上当 HTTP 服务，训练机通过网络请求 rollout。配 vllm_mode="server" 时：

# 1) 先起 vLLM server（独立终端 / 节点）
trl vllm-serve --model meta-llama/Llama-3-8B --tensor-parallel-size 2

# 2) 训练侧
cfg = GRPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
    vllm_server_host="10.0.0.5",
    vllm_server_port=8000,
)

优势：训练 GPU 全用来训，vLLM 专心 rollout，吞吐最大化。 劣势：多了网络一跳；要管两套进程。

10.5weight 同步：怎么把训练侧的最新 policy 推给 vLLM

这是集成的关键技术问题。TRL 实现：

对应代码：trl/extras/vllm_client.py 和 trl/scripts/vllm_serve.py。

10.6性能对比

Qwen2.5-7B、4 GPU、GRPO、G=8、max_completion=2K：

配置每 step rollout 时间tokens/s
HF generate（无 vLLM）~42 s~3.9k
vLLM colocate~13 s~12k
vLLM server（独立 2 GPU）~7 s~22k

10.7调优技巧

问题处理
Colocate 模式 OOM调小 vllm_gpu_memory_utilization（默认 0.3-0.4）
rollout 还是慢开 vllm_tensor_parallel_size=2/4 用 TP
weight 同步慢用 colocate 而非 server
server 模式连不上开防火墙、检查 host/port；vLLM server 启动有 30s 准备时间
训练吞吐没上去瓶颈可能在 update 而非 rollout，开 profile 验证

10.8什么时候不用 vLLM

10.9VLLMClient：trainer 怎么和 vLLM 说话

server / colocate 都通过 trl/generation/vllm_client.py:58 VLLMClient 这一层。它的方法约定如下：

方法作用
VLLMClient(base_url="http://localhost:8000") 构造 HTTP 客户端，server 模式用
client.generate(prompts) → dict 批量生成，返回 {"prompt_ids", "completion_ids", "logprobs", "logprob_token_ids"}
client.init_communicator(device="cuda") 建 NCCL / XCCL 通信子，准备权重广播
client.update_model_params(model) 把 trainer 当前 policy 权重广播到 vLLM workers

"为什么需要 init_communicator？" —— vLLM workers 不会自己监听 trainer，需要 trainer 主动建一个 collective group 把权重 push 过去。

10.10权重同步真实路径：PyNccl + StatelessProcessGroup

trl/scripts/vllm_serve.py:36–150 里的 WeightSyncWorkerExtension 描述了 server 端怎么收权重：

  1. 启动时 vLLM workers 通过 StatelessProcessGroup 建一个临时 NCCL 通信组（不挂在 PyTorch DDP 上）；
  2. Trainer 是这个 group 的 rank world_size - 1（最高 rank），专门负责广播；
  3. 每次需要同步，trainer 调 broadcast(param)，vLLM 各 worker recv 然后调 model.load_weights([(name, weight)]) 注入；
  4. 整个过程零磁盘 IO + 零 HTTP body 数据（HTTP 只发"开始同步"信号）。

这比"trainer save → vLLM load checkpoint"快几十倍——GRPO 每 step 一次同步，磁盘路径根本跑不动。

10.11colocate 模式的"两段 CUDA_VISIBLE_DEVICES"

colocate 模式下训练和 vLLM 在同一进程内，但必须用不同的 GPU（vLLM 内部 worker fork 时会独占 device）。设法：

# 8 卡机：前 4 卡给训练，后 4 卡给 vLLM colocate
# 注意是同一个 CUDA_VISIBLE_DEVICES，但 vLLM 内部按 vllm_tensor_parallel_size 抢后面几张
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
trl grpo \
    --use_vllm true \
    --vllm_mode colocate \
    --vllm_gpu_memory_utilization 0.3 \
    --vllm_tensor_parallel_size 4 \
    --num_processes 4   # accelerate 只给训练用 0-3 卡

实战经验：vllm_gpu_memory_utilization 0.3 留 70% 显存给 vLLM 自己的 KV cache + weights；更小（0.2）容易 OOM，更大（0.5）训练 batch 撑不开。

10.12哪些 trainer 支持 vLLM？

不是所有 trainer 都接了 vLLM。截至本书基线：

支持 vLLM 的 trainer位置
GRPOTrainer ⚡ trl/trainer/grpo_trainer.py（stable）
RLOOTrainer ⚡ trl/trainer/rloo_trainer.py（stable）
OnlineDPOTrainer ⚡ trl/experimental/online_dpo/
NashMDTrainer ⚡ trl/experimental/nash_md/
XPOTrainer ⚡ trl/experimental/xpo/

"⚡" 是 TRL 官方文档自己的标记。所有 online 系（需要在训练中实时采样的）都接了 vLLM；DPO / KTO / ORPO / SFT / Reward 不需要。

10.13这章你需要带走的