vLLM V1 迁移:确保强化学习中的后端正确性
ServiceNow AI 成功地将其强化学习推理引擎从 vLLM V0 迁移到 V1,重点放在后端正确性而非目标侧的校正上。团队通过解决四个关键差异,实现了 vLLM 0.8.5(V0)与 vLLM 0.18.1(V1)之间的等价:已处理的 rollout 日志概率、V1 特有的运行时默认设置、飞行中的权重更新路径,以及 lm_head 投影的精度。
迁移目标
在像 PipelineRL 这样的系统中使用 vLLM 作为 rollout 生成的推理引擎时,引擎会采样 token 并返回 logprobs,训练器使用这些 logprobs 计算关键的强化学习指标,包括策略比率、KL 散度、裁剪率和熵。logprobs 的计算方式出现任何差异都会导致训练‑推理不匹配,从而根本性地改变训练动态。
为了迁移到 vLLM V1,ServiceNow AI 确立了一个狭窄的目标:验证 V1 返回的 rollout logprobs 符合训练器的预期格式,针对 V0 基准重新运行工作负载,并仅在后端等价恢复后评估目标层面的变化。最初使用 V1 的尝试显示出相较于 V0 基准在奖励、熵和裁剪率方面的显著偏差。
后端失效模式
团队将导致训练‑推理不匹配的潜在原因划分为三层:
- 语义不匹配:后端返回的 logprobs 含义与训练器的预期不同。
- 推理路径不匹配:缓存、调度或请求处理的运行时默认设置差异导致提示走不同的执行路径。
- 目标不匹配:强化学习目标需要对残留的陈旧性或后端不匹配进行校正。
通过将前两项视为后端行为问题并首先排除,团队避免了过早使用目标侧校正来掩盖后端错误的错误做法。
vLLM V1 后端修复
日志概率语义
默认情况下,vLLM V1 在对 logits 进行后处理(如温度缩放和 top‑k/top‑p 过滤)之前,从原始模型输出返回 logprobs。PipelineRL 需要来自采样器使用的已处理分布的 logprobs。通过设置 logprobs-mode=processed_logprobs 解决了此问题,消除了 rollout logprobs 的均值偏移,使平均策略比率接近 1.0。
运行时默认设置
为确保等价,团队显式禁用了与 V0 基准路径不同的 V1 特有默认设置:
- 前缀缓存:已禁用 (
enable-prefix-caching: false),以防止在权重更新前计算的状态被复用,从而避免出现仅在 V1 中存在的缓存生命周期差异。 - 异步调度:已禁用 (
async-scheduling: false),以消除另一项仅在 V1 中存在的自由度。
飞行中的权重更新
为匹配 V0 在加载新权重时不显式失效缓存状态的行为,团队实现了以下 V1 更新路径:
await engine.pause_generation(mode="keep", clear_cache=False)
await engine_client.collective_rpc_async(
"receive_weight_update",
args=(request.model_dump_json(),),
)
await engine.resume_generation()
使用 mode="keep" 和 clear_cache=False 确保权重同步模型与 V0 基准保持一致,降低了训练后期的持续延迟。
数值等价:fp32 lm_head
即使在后端修复之后,最终的等价仍需匹配用于计算 logits 的数值路径。训练器使用 fp32 的 lm_head 进行最终投影,rollout 后端必须更新以匹配此精度。
这种精度至关重要,因为强化学习更新直接使用 token 的 logprobs;logits 的微小变化会显著影响策略比率和 KL 散度。该发现与 MiniMax-M1 技术报告以及 ScaleRL 论文相吻合,二者都指出 fp32 头部计算是大规模强化学习的必要设计选择,以防止训练/推理之间的 token 概率不匹配。
为什么后端正确性优先于目标校正
虽然截断重要性采样或重要性比重加权等工具可以校正陈旧或异步的 rollout,但在修复后端之前使用它们会混淆结果。如果使用目标侧校正来补偿破损的推理‑后端行为,训练曲线将更难解释。
团队的方法确保任何剩余的不匹配都是由于强化学习目标的设计(例如异步/离策略问题),而非推理引擎的 bug。未来的改进将聚焦于异步/离策略的清理,例如保留 rollout 时的显式行为策略 logprobs,并在优化时重新计算训练器侧的旧策略 logprobs。