RollPacker: Mitigating Long-Tail Rollouts for Fast, Synchronous RL Post-Training

发表时间: 2025-09 · arXiv:2509.21009 (HKUST & Alibaba)

RollPacker:通过缓解长尾 Rollout 实现快速、同步的强化学习后训练

作者/机构: Wei Gao†, Yuheng Zhao†, Dakai An†, Tianyuan Wu†, Lunxi Cao†, Shaopan Xiong‡, Ju Huang‡, Weixun Wang§, Siran Yang‡, Wenbo Su§, Jiamang Wang‡, Lin Qu‡, Bo Zheng‡, Wei Wang†
(†香港科技大学, ‡阿里巴巴集团, §淘宝天猫集团)

速读

一句话结论 RollPacker 系统通过提出尾部批处理策略及配套的动态资源调度机制,在严格保持同步强化学习在策略语义和训练精度的前提下,彻底消除了长尾响应导致的 GPU 空闲气泡,使 Qwen2.5 系列大语言模型的端到端后训练速度最高提升了 2.56 倍。

要解决什么问题 大语言模型在数学推理或代码生成等任务的强化学习后训练中,为了保证训练的稳定性和跨任务的可靠性,通常必须采用严格的同步在策略(on-policy)训练。这种机制要求模型必须先完成当前批次所有提示词的响应生成,才能计算奖励并更新权重。然而,真实场景下的模型响应长度呈现极端的长尾分布,绝大多数响应长度适中,但极少数响应会达到 16K 甚至更长的 token 数量。在同步生成的约束下,所有 GPU 都必须等待批次中最长的那条响应生成完毕。这导致分配到较短请求的 GPU 会过早结束工作并陷入持久的空闲状态,产生严重的流水线气泡,使得占据整个训练时间约 70% 的生成阶段资源利用率极低。现有的优化方案陷入了进退两难的困境:一类方法尝试在同步约束下将奖励计算与生成阶段重叠,但由于奖励计算耗时占比通常不到 15%,根本无法掩盖巨大的生成气泡;另一类方法则激进地放宽同步约束,允许使用过时的数据进行异步更新,但这直接破坏了在策略训练的数学等价性,往往会导致模型训练精度下降和收敛不稳定。

怎么做的 核心思路是提出一种名为尾部批处理的调度策略,将导致气泡的长响应系统性地剥离并集中处理,再辅以三个定制化的系统组件来压榨硬件性能。首先,尾部批处理利用了推测执行的机制来构建高并发的短轮次。系统会超额下发提示词并要求生成更多响应,但只保留最先完成的目标数量。这种竞速机制自然过滤掉了耗时漫长的长响应,使得短轮次内的任务长度高度均衡,彻底消除了等待气泡。为了保证训练样本分布不被扭曲,那些在竞速中被中止的长提示词会被放入一个专属队列,当积攒到足够数量时,系统会启动一个禁用推测执行的长轮次让它们完整生成。这种仅改变样本计算顺序而不丢弃数据的做法,完美维持了在策略语义。其次,针对短轮次因并发量大导致的显存 KV 缓存频繁抢占问题,设计了弹性并行规划器。它会实时监控显存压力,当抢占次数激增时动态翻倍张量并行的规模以释放显存,在压力回落时再减半以降低通信开销。再次,为了掩盖长轮次中凸显的奖励计算开销,引入了奖励调度器。对于代码任务,它摒弃了固定的沙箱超时设置,采用自适应超时公式 $T_{timeout} = \max(\min(\lambda \cdot T_{anchor}, T_{max}), T_{min})$,其中 $T_{anchor}$ 是当前正确响应的最大耗时,以此快速掐断注定失败的死循环代码;对于使用大模型作为裁判的打分任务,它通过多进程服务将裁判模型与生成模型部署在同一张 GPU 上,并利用分层流水线将裁判模型的权重卸载到内存,按需流式加载,从而在不爆显存的情况下榨干计算力。最后,针对长轮次内部依然存在的长短不一问题,基于流的训练器实现了极细粒度的阶段重叠。它会实时监控生成进度,一旦部分请求完成,就立刻将空闲出来的 GPU 从生成任务中剥离并转交回训练任务,同时将已完成的响应异步流式传输过去计算梯度。为了坚守在策略底线,这些梯度只做累加缓冲,直到整个轮次的生成全部结束,才进行全局的梯度重归一化和权重更新。

效果如何 实验部署在包含 128 张 H800 GPU 的集群上,使用 Qwen2.5 系列的 7B、14B 和 32B 参数模型,最大上下文长度分别设为 8K、16K 和 32K,训练数据均匀混合了数学、代码和多主题问答任务。对比的基线方法有两个:一个是代表标准同步强化学习且采用固定并行策略的 veRL,另一个是代表同步约束下阶段重叠路线的 RLHFuse。量化结果显示,RollPacker 的验证集得分曲线与标准同步的 veRL 几乎完全重合,证明其在策略语义的保持上没有带来任何精度损失。在端到端训练速度上,RollPacker 取得了全面碾压的优势:相比 veRL,在 7B、14B 和 32B 模型上分别实现了 2.03 倍、2.22 倍和 2.56 倍的加速;相比 RLHFuse,最高也实现了 2.24 倍的加速。其中,仅尾部批处理一项技术,就在 32B 模型 32K 长度的极端设置下贡献了 2.21 倍的性能提升。该系统也存在一些局限性,例如在单卡共置裁判模型和生成模型时,依赖的底层多进程服务机制无法提供严格的错误隔离,一旦发生故障可能会导致双端崩溃;此外,当前的弹性并行规划器仅支持张量并行的动态调整,尚未将专家并行纳入优化空间,在更复杂的模型架构下仍有进一步挖掘的潜力。

A1 主要贡献

本文旨在解决同步强化学习(RL)后训练中因响应长度不均衡(长尾分布)导致的严重GPU资源利用率不足(称为“气泡”)问题。现有方法通常通过放宽同步来缓解此问题,但这可能损害训练精度。

研究目标:在不牺牲训练精度和在策略(on-policy)训练语义的前提下,系统性地解决同步RL后训练中由长尾rollout引起的GPU空闲问题,从而显著加速训练过程。

核心创新点
1. 提出尾部批处理(Tail Batching)策略:这是一种新颖的rollout调度策略。其核心思想是将导致长尾响应的提示(prompt)系统性地整合到少数专用的rollout步骤(称为“长轮次”)中,同时确保大多数步骤(“短轮次”)只包含长度均衡的短rollout。通过将长响应从短轮次中排除并重新调度到少数长轮次中,该策略有效减少了rollout期间的GPU空闲时间。由于仅改变了训练样本的顺序,该方法保持了训练的准确性。
2. 设计并实现RollPacker系统:该系统通过对RL三个阶段的整体优化,充分发挥了尾部批处理的优势。
* 弹性并行规划器(Parallelism Planner):针对rollout阶段,该规划器能自适应地配置并行策略。它根据短轮次和长轮次不同的内存压力,动态选择最优的张量并行(TP)配置,以适应工作负载的变化。
* 动态资源分配与调度的奖励调度器(Reward Scheduler):针对奖励(reward)阶段,该调度器将奖励计算与rollout流水线化,并根据工作负载特性(如代码执行超时、评判模型共享GPU)动态调整计算预算,以隐藏奖励计算的开销。
* 基于流的训练器(Stream Trainer):针对训练(training)阶段,该训练器实现了更细粒度的阶段重叠。它机会性地将rollout阶段空闲的GPU重新分配给训练任务,并异步地将已完成的rollout数据流式传输到训练阶段进行梯度计算,同时通过损失缩放和延迟梯度更新来保持在策略训练的语义。

通过这些优化,RollPacker在Qwen2.5系列大语言模型上,相较于veRL实现了2.03倍至2.56倍的端到端训练时间缩短,相较于RLHFuse最高实现了2.24倍的加速。

A3 背景知识与动机

2.1 用于大语言模型后训练的强化学习

图1:同步和异步RL后训练的执行工作流程。
图1:同步和异步RL后训练的执行工作流程。

2.2 RL后训练的特性分析

Table 1: RL后训练的时间分解。我们在三个任务中使用veRL [45]和GRPO [43]以及真实世界数据集[20, 54]训练14B模型,最大长度为16k。

图2:Rollout阶段的特性:(a) 三个任务的响应呈现长尾分布;(b) 长尾rollout在rollout阶段造成了持久的GPU气泡。
图2:Rollout阶段的特性:(a) 三个任务的响应呈现长尾分布;(b) 长尾rollout在rollout阶段造成了持久的GPU气泡。

2.3 现有解决方案及其局限性

A2 方法细节

3 尾部批处理 (Tail Batching)

4 RollPacker 系统设计

本节介绍RollPacker,一个为充分发挥尾部批处理优势而通过整体设计构建的高效在策略RL系统。

4.1 系统概览

4.2 并行规划器

4.3 奖励调度器

4.4 流式训练器

算法1 流式训练器

A4 实验环境

A4 实验结果

6.1 端到端评估

我们比较了RollPacker与两种SOTA同步RL后训练系统veRL【索引45,verl: Volcano engine reinforcement learning for llm, 2024】和RLHFuse【索引63,Rlhfuse: Efficient rlhf training for large language models with interand intra-stage fusion, 2024】的端到端性能。

6.2 性能分解

6.3 尾部批处理的敏感性分析

图11展示了不同尾部批处理配置下的rollout时间,通过改变推测因子$\eta$来调整每个提示的响应数(R)或提示数(P)。
* R的影响:固定P,增加R的$\eta$可以丢弃长尾响应,但只有当$\eta$增加到1.5时,rollout时间才有显著减少。
* P的影响:固定R,增加P的$\eta$会增加长轮次的频率,从而抵消短轮次带来的时间缩减。
* 结论:实验表明,对P和R都设置$\eta=1.25$是最佳选择,平均rollout速度提升高达3.9倍,优于仅固定P或R的设置。
图11:[尾部批处理] 不同配置的rollout时间。我们分别固定P和R为P0和R0,同时改变另一个参数的η。每个条形代表一个完整周期的连续短轮次和长轮次的平均迭代时间,归一化到没有尾部批处理的基线,并标注了实际时间。

6.4 并行规划器

6.5 奖励调度器

异步奖励计算本身相比同步计算在三个模型上分别带来1.48倍、1.35倍和1.18倍的加速。
* GPU共享(LLM-as-a-Judge):图13a显示,在共置评判LLM和actor LLM时,使用MPS(多进程服务)可减少干扰,带来高达1.25倍的步骤时间加速。
* 流水线化评判LLM执行:图13b显示,对于评判LLM,采用分层流水线卸载权重,相比非流水线执行,在32k token长序列下可获得1.4倍的奖励计算加速。
* 自适应超时(代码):图13c显示,自适应超时机制通过提前终止可能失败的代码执行,显著减少了短轮次中的奖励计算开销,平均加速1.6倍。
图13:[奖励调度器] LLM-as-a-Judge:(a) 有无MPS时每个步骤的归一化时间。第五步是长轮次。(b) 有无流水线执行时不同序列长度的奖励计算归一化时间。代码沙箱:(c) 有无自适应超时时,rollout和异步奖励计算的组合时间。

6.6 流式训练器性能分析

6.7 性能可扩展性

图14展示了Qwen2.5-14B/16k模型的可扩展性分析。当资源扩展到128个GPU时,RollPacker保持了强大的性能,吞吐量(样本数/秒)始终比veRL高出2.2倍。当资源翻倍时,RollPacker的吞吐量提升约1.5倍。
图14:大规模训练Qwen2.5-14B模型的端到端吞吐量(样本/秒)。

A7 补充细节

7 讨论

8 相关工作

A5 结论

本文介绍了RollPacker,一个旨在加速同步RL训练的新型RL后训练系统。我们提出了尾部批处理(tail batching)来缓解长尾rollout问题并提升资源利用率。结合尾部批处理,我们设计了分别优化rollout、奖励和训练阶段的并行规划器奖励调度器流式训练器。广泛的实验证明了RollPacker在训练效率上相比基线方法的有效性。