vLLM 训练-推理一致性发布说明
概述
针对大语言模型(LLM)的高性能强化学习(RL)通常会受到训练-推理不匹配的影响,即推理引擎与训练引擎之间的数值差异导致模型优化的策略与用于采样的策略不一致。这种不匹配会破坏训练的稳定性,导致奖励崩溃并增加重跑次数。vLLM 引入了一种新方法,在 rollout 和训练之间实现了位级一致的对数概率(diff = 0),在确保真正的同策略(on-policy)RL 的同时,在 DeepSeek-V4-Flash 上保持了基准 rollout 吞吐量的 91.3%。
技术挑战:训练-推理不匹配
Rollout 是 RL 训练中最耗时的部分,这促使开发者采用低精度计算、低精度 KV 缓存以及各种内核优化。这些优化通常会引入非确定性行为,或者与训练前向传播相比产生数值差异。
在 DeepSeek-V4-Flash 的实验中,基准测试的训练-推理 KL 散度在 100 个步骤内从 2e-3 上升到 29.9,导致奖励下降。vLLM 的方法通过强制训练前向传播使用与 rollout 端完全相同的 vLLM 推理内核来解决这个问题,确保两个引擎对相同的权重计算出相同的策略。
实现位级一致性
为了达到 diff 为 0,系统在四个维度上强制执行一致性:
- 内核确定性:确保相同的内核在相同的输入下始终产生相同的输出,避免非结合浮点累加问题。
- 批处理独立性:消除请求结果基于其批处理伙伴而改变的依赖关系。
- 阶段一致性:统一预填充(prefill)和解码(decode)阶段的计算路径。
- 跨 GPU 一致性:修复通信归约(communication reductions),使其不会因时序或调度而变化。
通过修改高性能内核以满足这些要求,团队在不牺牲低精度操作性能增益的情况下实现了位级一致的结果。
代理驱动的实现
为了降低将这种一致性应用于新模型的成本,vLLM 发布了一种代理技能(agent skill),可自动执行调试和对齐过程。该技能引导代理完成以下步骤:
- 定位:识别两个引擎产生分歧的第一个算子。
- 最小化复现:将分歧缩小到一个最小的、可重跑的测试用例。
- 验证:确保修复方案在保持一致性的同时不会破坏高性能基准。
结果与性能
- DeepSeek-V4-Flash:一致性版本在达到 0 diff 的同时,保持了基准 rollout 吞吐量的 91.3%。虽然基准奖励出现了回落,但一致性版本的奖励保持稳定。
- Nemotron-3.5 (Mamba 混合架构):一致性版本保持了 0.6 的 SWE-bench Pro 分数,而基准分数从 0.56 降至 0.22,这表明一致性可以防止复杂架构中的训练崩溃。
开源可用性
该流程、方案和代码现已向社区开放。这包括用于自动对齐的 train-inference-consistency 技能、DeepSeek-V4 方案以及 Nemotron-3.5-Lightning-30B-A3B 实现。目前正在进行 DeepSeek-V4.1 的开发,该版本已经实现了 0 diff,性能优化工作正在收尾中。