论文

通过系统集成 推测解码 加速 RL 后训练 的采样轨迹

Accelerating RL Post-Training Rollouts via System-Integrated Speculative Decoding

模型推理投机采样

摘要

前沿语言模型的 RL 后训练 越来越受到自回归 rollout 生成的瓶颈,使得 rollout 加速成为核心系统挑战。许多现有的效率方法通过改变部署或优化机制来提高吞吐量,例如,通过 离策略 执行、重放或较低精度的生成。我们将 推测解码 作为 RL 采样轨迹的无损加速原语进行研究,以保留目标模型的输出分布。我们在带有 vLLM 后端的 NeMo-RL 中实现 推测解码,支持同步和异步管道,并在 RL 采样轨迹期间启用推测。这种好处可以通过各种推测机制来实现,例如预训练的 MTP 头、小型外部草案模型,甚至是 Eagle3 等传统上在 RL 阶段之后应用的技术。这为 RL 训练中最先进的 推测解码 提供了一条部署路径。在同步 RL 下 8B 规模的推理 后训练 工作负载中,推测解码 将采样轨迹吞吐量提高了 1.8 倍。使用高保真性能模拟器,我们预计将 推测解码 与异步 RL 相结合可在 235B 规模下实现高达 2.5 倍的端到端训练加速。