RollPacker: Mitigating Long-Tail Rollouts for Fast, Synchronous RL Post-Training
RollPacker: Mitigating Long-Tail Rollouts for Fast, Synchronous RL Post-Training
发表时间: 2025-09 · arXiv:2509.21009 (Alibaba)
论文标题: RollPacker:通过缓解长尾 Rollout 实现快速、同步的 RL 训练后优化
作者/机构: 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†
(†香港科技大学 ‡阿里巴巴集团 §淘宝天猫集团)
A1 主要贡献
本文针对大型语言模型(LLM)的强化学习(RL)训练后优化过程中的一个核心问题——由Rollout阶段响应长度不均衡(长尾分布)导致的严重GPU资源利用率低下(称为“气泡”)——提出了创新的解决方案。传统的同步RL训练为了保证模型精度,必须等待批次中所有响应(包括极长的响应)生成完毕才能进入下一阶段,造成了大量GPU空闲。而现有的异步或半同步方法虽然能减少空闲,但通常以牺牲训练精度为代价。
研究目标:在不牺牲模型准确性的前提下,显著提升同步RL训练后优化的端到端效率。
核心创新点:
1. 尾部批处理(Tail Batching):本文提出了一种新颖的提示(prompt)调度策略。其核心思想是,将那些会产生极长响应的“尾部提示”系统性地整合到少数几个专用的“长轮次”(long rounds)中执行。而绝大多数的“短轮次”(short rounds)则只包含响应长度均衡且较短的提示。通过这种方式,短轮次中的GPU空闲时间被大幅减少,从而加速了整体训练过程。由于该方法仅改变了训练样本的顺序,因此保留了同步训练的精度优势。
2. RollPacker系统:为了充分发挥“尾部批处理”的优势,本文设计并实现了RollPacker系统,该系统包含针对RL训练三个阶段的全面优化:
* 弹性并行规划器(Elastic Parallelism Adaptation):针对Rollout阶段,由于短轮次(使用推测执行)和长轮次的内存压力不同,该规划器能够动态调整张量并行(TP)配置,以适应变化的负载,减少Rollout延迟。
* 动态资源分配与调度的奖励器(Dynamic Resource Allocation and Scheduling for Reward):针对奖励(Reward)阶段,该调度器通过异步、逐样本的奖励计算,并根据工作负载特征(如代码执行超时、评判模型GPU共享)动态调整计算预算,有效隐藏了奖励计算的开销。
* 流式训练器(Stream-based Training):针对训练(Training)阶段,该训练器通过将Rollout阶段提前完成的样本流式传输到训练阶段,并机会性地将空闲的GPU从Rollout任务重新分配给训练任务,实现了Rollout和训练的细粒度重叠,进一步减少了GPU空闲。
通过这些设计,RollPacker在保证与标准同步RL训练相同精度的前提下,显著减少了端到端训练时间。实验结果表明,与先进的RL系统veRL和RLHFuse相比,RollPacker分别取得了高达2.56倍和2.24倍的端到端训练加速。
图1:同步和异步RL训练后优化的执行工作流。
A3 背景知识与动机
2.1 用于LLM训练后优化的RL
RL后训练工作流。强化学习(RL)已成为训练后优化LLM的关键技术。近期的进展表明,如GRPO【43,Deepseekmath: Pushing the limits of mathematical reasoning in open language models. 2024】等RL算法在提升各种领域的推理能力方面非常有效。一个RL训练后优化的工作流通常协调多个扮演不同角色的模型。Actor LLM负责生成对输入提示的响应,是主要被优化的模型。Reward LLM评估每个响应并输出一个标量奖励信号,该信号可以来自多种来源,如代码任务的沙箱执行、数学问题的基于规则的逻辑,或用于对齐任务的LLM-as-a-Judge【47,Llm-as-a-judge and reward model: What they can and cannot do, 2024】。为了进一步稳定优化过程,通常会引入一个Reference LLM作为正则化器。总的来说,该工作流包括三个阶段:(1)Rollout,Actor LLM为$P_0$个提示各生成$R_0$个响应;(2)Reward,生成的响应由相应的奖励工作器进行评估;(3)Training,Actor LLM根据计算出的奖励更新其参数,并可选择性地受参考模型的约束以减轻梯度不稳定性。
图2:Rollout阶段的特性分析:(a) 三个任务的响应呈现长尾分布;(b) 长尾Rollout在Rollout阶段造成了长时间的GPU气泡。
表1:RL训练后优化的时间分解。我们使用veRL【45】和GRPO【43】算法,以及真实世界的数据集【20, 54】,在三个任务上训练14B模型,最大长度为16k。
同步在线策略训练。为了最大化模型性能,RL训练后优化通常在同步设置下进行,即所谓的在线策略训练(on-policy training)。在此设置中,Rollout必须在训练开始前完成,并且Actor的权重只有在训练结束后才更新(见图1a)。这种同步要求确保所有响应都使用最新的模型参数生成,从而稳定训练过程并提高跨任务的可靠性【16,Interpolated policy gradient: Merging on-policy and offpolicy gradient estimation for deep reinforcement learning. 2017】。然而,正如我们的特性分析所示,这常常导致严重的流水线气泡和低利用率,尤其是在Rollout阶段。
2.2 RL训练后优化的特性分析
实验设置与表征。我们使用Qwen2.5-14B模型【38,Qwen2.5 technical report, 2025】来分析RL训练后优化的工作负载,配置最大响应长度为16k tokens,批大小为128,组大小为8,并采用GRPO算法【43】。我们使用veRL【45,verl: Volcano engine reinforcement learning for llm. 2024】在32个H800 GPU上运行数学、代码和LLM-as-a-Judge任务,数据集来自真实世界【20, 54】。对于LLM-as-a-Judge实验,我们使用了一个7B参数的评判模型。
Rollout瓶颈。表1报告了这些设置下RL训练的各阶段延迟分布。Rollout阶段在运行时间中占主导地位,在三个任务中约占每个训练步骤的70%。奖励阶段贡献了较小的一部分(5%-13%),而训练阶段占总步骤时间的21%-23%。这些结果强调了Rollout是RL训练后优化的主要瓶颈。
长尾响应与GPU气泡。RL训练后优化效率低下的一个关键来源在于响应长度的高度偏斜分布。如图2a所示,数学、代码和LLM-as-a-Judge任务的响应都呈现出明显的长尾分布:虽然大多数响应长度为中短,第75百分位数(P75)在755到1.1k tokens之间,但最长的响应可以延伸到16k tokens。在同步Rollout下,长尾响应的存在导致了GPU利用率低下。因为所有GPU都必须等待最长的响应完成,分配给较短请求的设备会进入空闲状态,产生长时间的“气泡”,浪费计算周期。图2b通过报告一个配置了双路张量并行的服务器上偶数索引GPU的SM利用率来说明了这个问题。利用率在Rollout开始时接近80%的峰值,但从未达到完全饱和,因为LLM解码本质上是内存受限的。一旦短响应完成,相应GPU的利用率迅速降至零,空闲期持续到整个批次完成。鉴于Rollout已经是训练延迟的主要贡献者,这种低效率严重拖慢了整个流水线,这个问题在文献中被广泛报道【14, 17, 19, 63】。
2.3 现有解决方案及其局限性
同步约束下的阶段重叠。先前为缓解RL训练后优化中GPU气泡所做的努力,通常分为两类:同步约束下的阶段重叠和宽松同步。该方法试图通过在同步屏障前,将长尾Rollout与其他阶段的执行进行流水线化来提高资源利用率。例如,RLHFuse【63,Rlhfuse: Efficient rlhf training for large language models with interand intra-stage fusion, 2024】将Rollout与奖励计算和参考模型推理重叠,而像ROLL【50,Reinforcement learning optimization for large-scale learning: An efficient and user-friendly scaling library, 2025】和MiMO【52,Mimo: Unlocking the reasoning potential of language model – from pretraining to posttraining, 2025】这样的框架则将每个已完成响应的奖励计算与正在进行的Rollout阶段重叠,以实现异步奖励计算。虽然这些设计在一定程度上减少了空闲气泡,但它们并未从根本上解决主导Rollout时间的长尾响应问题,性能提升有限(见§6.1)。此外,随着RL训练后优化中响应长度的持续增长【4】,奖励和参考模型推理的相对贡献减小(如表1所示,通常低于步骤运行时间的15%),即使在理想的重叠下,也无法掩盖长时间的空闲气泡。
图3:尾部批处理(Tail Batching)的图示。
宽松同步。第二类工作采取了更激进的策略,即放宽Rollout和训练之间的严格同步屏障,如图1b所示。例如,Kimi【48,Kimi k1.5: Scaling reinforcement learning with llms, 2025】通过截断长尾响应并保留已生成的tokens以便在后续步骤中继续Rollout,引入了部分Rollout。StreamRL【62,Streamrl: Scalable, heterogeneous, and elastic rl for llms with disaggregated stream generation, 2025】、AsyncFlow【17,Asyncflow: An asynchronous streaming rl framework for efficient llm post-training, 2025】和RhymeRL2【19,History rhymes: Accelerating llm reinforcement learning with rhymerl, 2025】允许一步的陈旧性(staleness),使得训练阶段可以使用稍微过时的Rollout数据,形成一种“一次性”流水线。更进一步,AReaL【14,Areal: A large-scale asynchronous reinforcement learning system for language reasoning, 2025】引入了完全异步的RL训练,其中Rollout和训练完全解耦,更新可能依赖于多步之前生成的样本。这些方法在减少GPU气泡方面很有效,但它们引入了一个根本性的权衡:通过放宽同步,它们损害了训练的在线策略性质,常常导致精度下降和稳定性降低。
总结与局限性。总而言之,现有解决方案要么通过将非瓶颈阶段与Rollout重叠来提供边际改进,要么通过放宽同步来牺牲在线策略保证。这两种方法都无法在保持同步RL训练的准确性和稳定性的同时,消除由长尾Rollout引入的低效率。
A2 方法细节
3 尾部批处理 (Tail Batching)
引言。本节我们介绍尾部批处理(tail batching),一种新颖的提示调度策略,它从根本上缓解了响应长度不平衡的问题,从而在保持在线策略RL训练语义且不损失精度的情况下减少GPU气泡。
核心问题与解决方案。GPU气泡主要由一小部分产生极长响应的提示引起。一个简单的方法是从Rollout批次中排除这些长尾响应。然而,这种方法会引入两个关键问题:(P1)Rollout阶段可能无法收集到构成有效批次所需的足够数量的提示或响应;(P2)系统性地排除长提示会扭曲训练样本分布,可能损害模型性能。尾部批处理通过两个关键技术解决了这些问题。
图4:在代码数据集【54】上训练Qwen2.5-14B模型时,尾部批处理与基线veRL的对比。(a) 五个轮次中响应长度分布的箱形图,其中带阴影的箱体是尾部批处理下的一个长轮次。须线测量1.5倍IQR。(b) 训练时间分解,总时间是连续5个步骤的累积,一个完整周期包括四个短轮次和一个长轮次。
推测执行解决P1。为了解决P1问题,尾部批处理利用了推测执行【58,Apache spark: a unified engine for big data processing. 2016】的思路,即超额提供请求,同时只选择性地保留最快完成的部分。在GRPO算法【43】下,每个Rollout步骤需要采样$P_0$个提示,每个提示产生$R_0$个响应(图3顶部)。尾部批处理并非精确启动$P_0$个提示,而是启动更多提示,并只接纳最先完成的$P_0$个。类似地,每个提示会产生超过$R_0$个响应,但只有最先完成的$R_0$个被保留。这种“竞速完成”的推测方式自然地过滤掉了长响应,产生了均衡且较短的批次,从而在保持所需批次大小的同时最小化了空闲气泡。
长提示队列解决P2。为了解决P2问题,尾部批处理保证没有提示被永久排除。如图3底部所示,在推测执行期间被中止的提示会被添加到一个长提示队列(long-prompt queue)中。一旦该队列的大小达到$P_0$,这些提示将在一个专用的长轮次(long round)中被批量调度执行,此时推测执行被禁用,以允许生成完整长度的响应。由于这类提示很少,长轮次不经常发生,并与由均衡响应组成的频繁短轮次交错进行。这种设计确保了所有提示最终都被处理,而大多数Rollout步骤保持高效。
训练精度。从统计学的角度看,尾部批处理只是将长短提示重新排序到不同的轮次中,而没有改变底层的样本分布或放宽同步。先前的研究表明,训练顺序的改变不会降低模型精度【7,Curriculum learning. 2009】,【9,Does the order of training samples matter? improving neural data-to-text generation with curriculum learning. 2021】,【13,Reverse curriculum generation for reinforcement learning. 2017】。事实上,一些近期的RL训练后优化算法明确探索了提示重排序作为一种提高效率的手段【37, 57, 59】。我们的实证结果进一步验证了这一说法:如图8所示,尾部批处理实现的精度曲线与标准同步RL训练几乎完全相同。
Rollout效率。我们通过在§2.2所述的相同设置下,使用代码数据集【54,Kodcode: A diverse, challenging, and verifiable synthetic dataset for coding. 2025】训练一个Qwen2.5-14B模型【38】,来实证验证尾部批处理在提升Rollout效率方面的有效性。推测因子设为$\eta = 1.25$,意味着在每个短轮次中,Actor LLM会推测性地启动$\eta P_0$个提示,每个提示生成$\eta R_0$个响应,同时只接受最先完成的$P_0 \times R_0$个结果。图4a比较了使用和不使用尾部批处理时,五个训练步骤中的响应长度分布。与基线方法(veRL【45】)相比,尾部批处理在前四个步骤(短轮次)中产生了更短且更均衡的响应,将最大响应长度减少了高达8.9倍。产生长响应的提示被推迟到第五步(长轮次)处理,其输出通常比基线长,但上限同样为16k tokens。这种重组显著降低了Rollout成本,并将端到端训练时间缩短了1.48倍(表2)。
图5:RollPacker的系统概览和工作流程。
新瓶颈的出现。图4b进一步分解了各阶段的训练时间。随着尾部批处理缓解了Rollout的开销,奖励和训练阶段的相对贡献变得更加突出。此外,短轮次和长轮次表现出截然不同的资源使用情况,这意味着统一的Rollout策略并非最优。这些发现促使我们接下来开发针对RL流水线每个阶段量身定制的系统级优化。
4 RollPacker系统设计
本节我们介绍RollPacker,一个高效的在线策略RL系统,其通过整体性设计来充分发挥尾部批处理的优势。我们首先概述系统,然后详细描述每个组件。
4.1 系统概览
系统组件。RollPacker集成了三个关键组件:并行规划器、奖励调度器和流式训练器,分别解决了Rollout、奖励和训练阶段的独特瓶颈。
并行规划器。采用推测执行的短轮次比长轮次产生更高的GPU内存压力。固定的张量并行(TP)配置无法适应短轮次和长轮次间变化的资源情况,导致频繁的KV缓存抢占开销和效率下降。RollPacker引入了一个并行规划器,该规划器能够动态分析工作负载并为每个步骤选择最优的TP配置,以削减Rollout开销。
奖励调度器。随着Rollout成本的降低,奖励计算变得更加显著(图4b)。为防止其成为新的瓶颈,RollPacker采用了一个奖励调度器,它将奖励计算与Rollout并行进行流水线化,同时为每个样本评估动态分配计算预算,从而有效降低其开销。
流式训练器。在禁用推测执行的长轮次中,不均衡的响应导致了明显的GPU气泡(图4a)。流式训练器在先前阶段重叠方法【50, 63】的基础上,引入了更细粒度的Rollout与训练重叠:完成的提示会立即流式传输到训练阶段,而空闲的GPU则从Rollout任务重新分配给梯度计算。为保持在线策略语义,流式训练器会仔细缩放梯度并推迟权重更新,直到整个Rollout完成,从而在保持精度的同时减少空闲时间。
工作流:离线阶段。RollPacker分两个阶段运行:离线分析阶段和在线执行阶段。在离线阶段,RollPacker在不同的TP大小、批大小和序列长度下,对Actor LLM的预填充和解码吞吐量进行基准测试。它还分析了评判LLM在不同序列长度下的GPU内存占用和运行时间成本。这些分析结果用于指导在线决策。
工作流:在线阶段。在在线阶段,RollPacker在一个同步RL作业中协调Rollout、奖励和训练(图5)。① 在Rollout期间,尾部批处理根据长提示队列的大小决定是否应用推测执行。② 接着,并行规划器结合历史作业负载和分析的性能数据,选择一个最优的TP配置。③ 与此同时,奖励调度器将评估与Rollout重叠,并为每个奖励任务动态调整预算。④ 同时,流式训练器监控Rollout进度,以决定何时将GPU从Rollout重新分配给训练。⑤ 当提示完成时,它们被流式传输到训练阶段,以便立即进行梯度计算。⑥ 一旦整个Rollout完成,流式训练器停止流式传输,累积所有计算出的梯度,并在所有可用GPU上触发同步的梯度计算和更新。⑦ 最后,更新后的Actor权重与Rollout阶段同步,然后开始下一个RL步骤。
4.2 并行规划器
短轮次的高内存压力。如§3所述,尾部批处理增加了短轮次中并发响应的数量,给GPU内存带来了更大压力。现有的LLM服务引擎【24,Efficient memory management for large language model serving with pagedattention. 2023】,【42,Sglang: Fast serving framework for large language models. 2025】通常通过抢占正在进行的请求来缓解内存压力,即将其KV缓存换出以释放GPU内存给其他请求。因此,高抢占次数表明存在严重的内存争用。图6a显示,在训练具有32k响应长度和128批大小的Qwen3-8B时,短轮次比长轮次多产生高达1.79倍的抢占,引入了巨大的计算开销。
图6:在八个H800 GPU上训练Qwen3-8B/32k时的Rollout性能。(a) 每个步骤的抢占次数。(b) 每个步骤的Rollout时间。该指标在一个由4个短轮次和1个长轮次组成的连续周期内收集。
增加TP缓解内存压力。张量并行(TP)是缓解GPU内存压力的标准技术。如图6a所示,配置更大的TP尺寸会将模型权重划分到更多GPU上,从而为KV缓存释放内存,并在短轮次和长轮次中分别将抢占次数减少21.1%和10.3%。这额外的KV缓存容量缓解了内存争用并缩短了Rollout延迟。如图6b所示,当TP=1时,短轮次的Rollout时间是长轮次的87%,抵消了尾部批处理的收益;将TP尺寸增加到2,可将短轮次持续时间减少25%。然而,更大的TP尺寸也引入了通信开销,这在Rollout受长尾响应限制的长轮次中占主导地位,从而侵蚀了性能。
自适应TP选择。大多数RL训练框架【21, 45】采用固定的并行配置,我们的分析表明这对于尾部批处理是低效的(图6)。短轮次和长轮次的最优TP尺寸不同,需要动态适应。RollPacker引入了一个并行规划器,以可忽略的开销按步骤重新配置TP尺寸。在离线阶段,规划器在没有尾部批处理的情况下分析最优TP尺寸,并将其用作训练开始时的默认配置。然后,它持续跟踪抢占次数,并使用一种轻量级启发式方法相应地调整TP尺寸:抢占次数的突然增加(例如,>1.05倍)会触发TP增加(尺寸加倍),而连续四个步骤的零抢占则触发TP减少(尺寸减半)。为了限制跨节点通信,TP组被限制在单个GPU服务器内。
4.3 奖励调度器
奖励调度器动机。随着Rollout成本的降低,奖励计算成为端到端延迟中一个不可忽视的贡献者(图4b)。为了缓解这个问题,RollPacker将奖励评估与Rollout进行流水线化,并为每个任务自适应地分配计算预算。
异步奖励计算。在这种方法中,奖励评估是异步执行的:已完成的响应被分派给奖励工作器,与正在进行的Rollout阶段并行,这与ROLL【50】和MiMo【52】类似。这种设计将奖励评估与Rollout重叠,部分隐藏了开销。然而,仅依赖这种设计不足以解决潜在的瓶颈,特别是对于代码沙箱执行和评判LLM评估。
图7:使用尾部批处理时,奖励计算引入了不可忽略的开销。(a) 代码响应的沙箱执行时间分布。(b) 分配给一个7B参数评判LLM的GPU上SM利用率随时间的变化。
代码沙箱执行。在编码任务中,实践者通常为每个测试用例设置最大执行超时。例如,在我们图4中描述的训练实验中,强制执行了30秒的超时。然而,如图7a所示,大约5%的提示达到了这个超时。这些长时间的执行最终产生零奖励,却延迟了整个奖励阶段,使其时间延长到总训练时间的近22%(图4b)。由于许多正确的响应完成得更快,固定的超时在注定失败的样本上浪费了大量计算资源。
自适应超时机制。RollPacker引入了一种自适应超时机制。对于每个测试用例,它跟踪训练期间正确响应中的最大执行时间,记为$T_{anchor}$。当一个新响应的执行时间超过此阈值时,沙箱执行会提前终止,并分配零奖励。考虑到代码执行期间可能存在的CPU争用,并为避免过于激进的截断,超时时间被放宽为 $T_{timeout} = \max(T_{min}, \min(T_{max}, \lambda \cdot T_{anchor}))$,我们根据经验设置$\lambda = 1.5$, $T_{min} = 2s$, $T_{max} = 30s$以获得良好性能。这种设计能够快速失败注定失败的案例,同时保留对有希望的响应的评估。
LLM-as-a-Judge问题。在异步奖励计算中,RL系统通常为评判LLM(judge LLM)保留固定数量的GPU(例如,总GPU的25%),以避免与其他工作器产生干扰。然而,这种策略导致了利用率低下。如图7b所示,当一个7B参数的评判LLM对响应进行评分时,其预留的GPU平均SM利用率仅为约22.6%。这种低效率源于评判模型通常处理小批量响应,使得大部分预留的计算能力处于空闲状态。
LLM-as-a-Judge 解决方案:共存与资源共享。为了提高效率,RollPacker将评判LLM与Actor LLM部署在同一GPU设备上进行并发执行。然而,这种设计引入了两个问题。首先,Rollout和奖励评估现在共享GPU资源,可能相互干扰。尽管如此,我们观察到无论是Actor LLM还是评判LLM都无法单独使GPU SM利用率饱和。为了实现高效共享,RollPacker启用了多进程服务(MPS)【33,NVIDIA Multi-Process Service (MPS) Documentation, 2025】,该服务在warp级别上划分GPU资源,允许两个模型以最小的干扰并发运行。
LLM-as-a-Judge 解决方案:内存管理。其次,在同一GPU上托管两个模型有耗尽内存的风险,因为Actor LLM已经需要大量空间用于其KV缓存。为了解决这个问题,RollPacker引入了一种分层流水线方案,以减少评判LLM的内存占用。受先前工作【6,Pipeswitch: fast pipelined context switching for deep learning applications. 2020】的启发,它将评判LLM的大部分层卸载到主机内存,并在GPU上进行激活计算时,通过PCIe同步流式传输其参数。由于Rollout很少使PCIe带宽饱和,这种流水线卸载带来的开销很小。RollPacker动态调整卸载的层数,以适应不同的输入序列长度,确保评判LLM在内存限制内运行,同时最大化利用率。
4.4 流式训练器
流式训练器动机。尽管有先前的优化,长轮次仍然因响应完成时间不均而存在空闲气泡(图4b)。流式训练器通过一种新颖的阶段重叠策略来缓解这一问题,该策略将正在进行的Rollout与梯度计算进行流水线化,从而在保持同步在线策略RL语义的同时,减少端到端延迟。
重用Rollout GPU进行训练。随着Rollout的进行,GPU利用率下降。流式训练器会缩减专用于Rollout的GPU数量,并将释放的设备重新用于训练。在重新分配的GPU上进行的训练只使用了数据并行副本的一个子集;梯度更新被推迟到Rollout完全结束后,从而保持了在线策略RL的正确性。对于工作量贡献很小的参考模型对数(reference logits),如果需要,可以通过临时交换Actor和参考模型的权重来计算。
算法1:流式训练器工作流。算法1概述了流式训练器的一般工作流程。它持续监控Rollout进度,并在已完成请求的比例超过一个阈值时触发GPU缩减(第14行)。当满足缩减标准时(第17行),一部分Rollout GPU被重新分配用于训练(第18-19行)。为了迁移正在进行的请求,RollPacker采用了一种基于重计算的策略(第20行):保留已生成的tokens,同时重计算KV缓存,以便在剩余的GPU上以最小的开销恢复Rollout【15, 19】。一旦训练实例启动,流式训练器会异步获取已完成的响应,并与正在进行的Rollout并行计算梯度(第23行)。
算法1 流式训练器
扩展标准:选择哪个GPU。训练器必须仔细选择要从Rollout重新分配到训练的GPU,因为这两个阶段通常依赖于不同的通信组拓扑。为确保正确性,紧密耦合的组,如Rollout中使用的TP组,必须保持完整,不能被分割到Rollout和训练两个任务中。在实践中,流式训练器尝试将一半的Rollout GPU重新分配给训练。在采取行动之前,它会验证这是否可能,而不会分割数据并行副本所需的通信组。如果不可能,则中止缩减尝试。
扩展标准:何时扩展。缩减Rollout GPU可能会减慢响应生成速度。通过将更多请求整合到更少的GPU上,它增大了每个设备的批处理大小,加剧了KV缓存的内存压力,最终损害了Rollout吞吐量(§4.2)。为防止这种情况,当已完成请求的比例达到20%到50%之间的里程碑(以5%为增量)时,流式训练器会结合历史响应长度分布和每个token的缓存占用,计算峰值KV缓存使用量。只有在迁移后,预计的峰值缓存需求仍在内存限制内时,才会触发GPU缩减。为简单起见,我们忽略了基于重计算的请求迁移所引入的适度开销以及缩减后解码吞吐量的轻微下降(§6.6)。
重叠流式执行。一旦满足缩减标准并重新分配了GPU,流式训练器便开始在重新分配的训练GPU上处理已完成的响应。Rollout和训练现在作为生产者-消费者对运行:Rollout生成响应,而训练通过一个与生产和消费速率对齐的流式模型来消费它们。流式训练器异步获取已完成的响应,并与正在进行的Rollout并行计算梯度,从而减少了总的步骤时间。
保持在线策略语义。流式训练器的一个关键要求是确保梯度计算在数学上等同于标准的在线策略训练流水线。这一保证分两个阶段来维持。首先,在流式执行期间,对已完成响应的梯度进行计算和缓冲,但不对模型参数或优化器状态进行任何更新。我们扩展了底层的LLM训练框架【46,Megatron-lm: Training multi-billion parameter language models using model parallelism. 2019】,在反向传播期间禁用梯度同步,以确保严格遵守在线策略约束。其次,在Rollout完成后,剩余的响应被分配到所有数据并行副本上进行梯度计算和模型更新。由于某些副本可能已经处理了部分工作负载,简单的平均会使结果产生偏差。为了纠正这一点,我们根据每个副本处理的样本数量对其上的局部梯度进行重新归一化,确保最终的更新等同于标准在线策略训练的更新。
A4 实验环境
实现细节:
- 代码:我们在ROLL【50】的基础上用约6.6千行Python代码实现了RollPacker,并将开源。
- 软件栈:Rollout阶段使用vLLM v0.8.4【24】作为服务后端;奖励阶段使用
ray.remote实现;训练阶段基于Megatron-LM v0.12.2【46】构建。
硬件配置:
* 集群:实验在一个H800集群上部署,包含16个节点,共128个GPU。
* 网络:节点间通过400 Gbps InfiniBand连接。
模型与数据集:
* 模型:使用Qwen2.5【5】系列模型,包括7B、14B和32B参数版本,最大响应长度分别配置为8k、16k和32k tokens。
* 评判模型:使用Qwen2.5-7B-Instruct作为LLM-as-a-Judge。
* 数据集:在多任务设置下进行端到端评估,数据集均匀混合了数学【20,Deepmath-103k: A largescale, challenging, decontaminated, and verifiable mathematical dataset for advancing reasoning. 2025】、代码生成【54】和多主题问答任务。
训练配置:
* RL设置:采用同步RL训练,提示批大小$P_0=128$,每提示响应数$R_0=8$。
* 资源分配:7B、14B和32B模型分别在16、32和64个GPU上训练。
* 并行策略:根据模型大小配置不同的并行策略。例如,14B模型在Rollout阶段使用TP=2,在训练阶段使用(TP, PP, CP) = (2, 2, 2)。
A4 实验结果
6.1 端到端评估
我们对比了RollPacker与当前最先进的同步RL训练后优化系统veRL【45】和RLHFuse【63】的端到端性能。
验证性能:图8显示,RollPacker与veRL的平均验证分数曲线几乎重合,证明了尾部批处理在不同模型大小和响应长度下均未损害训练精度。此外,RollPacker在训练早期表现出更快的收敛速度,我们推测这得益于更均衡的响应长度分布。
端到端延迟:图9展示了前40个训练步骤的耗时。总体而言,RollPacker在所有三个LLM上均优于veRL和RLHFuse。
* 与veRL相比,RollPacker对7B、14B和32B模型分别实现了2.03倍、2.22倍和2.56倍的加速。
* 与RLHFuse相比,加速比分别为1.14倍、1.68倍和2.24倍。
* 长轮次 vs 短轮次:在长轮次中,由于奖励调度器和流式训练器的作用,RollPacker和RLHFuse均优于veRL。但在短轮次中,得益于尾部批处理,RollPacker的优势更为显著,相比veRL和RLHFuse分别实现了2.1-3.6倍和1.2-3.2倍的加速。
图8:使用veRL和RollPacker训练Qwen2.5-7B、14B和32B模型的平均验证分数。
图9:使用veRL、RLHFuse和RollPacker训练Qwen2.5-7B、14B和32B模型的步骤时间。
表2:端到端训练加速分解。
6.2 性能分解
各优化技术贡献:表2详细分解了各项技术带来的累积加速效果。
* 尾部批处理:是性能提升的主要贡献者,有效减少了Rollout开销,尤其是在响应长度增加时效果更明显(对Qwen2.5-32B/32k设置,加速达2.21倍)。
* 奖励调度器:对短Rollout尤其有益,在8k响应长度下提供了71%的性能提升。即使在长响应下,也贡献了27%的改进。
* 并行规划器:在高内存压力场景(大模型和长序列)下效果显著,为Qwen2.5-32B/32k设置提供了额外的4%加速。
* 流式训练器:通过重叠Rollout和梯度计算,为Qwen2.5-14B/16k带来了20%的显著性能提升,并在其他设置中也有一致的正面效果。
训练步骤分解:图10分析了RollPacker的性能。在长轮次中,Rollout时间与veRL相当,但在短轮次中,RollPacker通过显著降低最大响应长度(图10a),实现了高达7.8倍的平均Rollout时间加速(图10b)。图10c显示,短轮次的节省有效降低了所有轮次的平均步骤时间。
图10:RollPacker中训练步骤时间的分解。(a) Qwen2.5-32B/32k每训练步骤的最大响应长度。(b) Qwen2.5-32B/32k每训练步骤的总Rollout时间。(c) 不同模型的步骤时间分解,比较短轮次(阴影条)和长轮次(实心条)。每个组件的时间都归一化为该模型的长轮次总时间。绝对步骤时间显示在每个条上。
6.3 尾部批处理的敏感性分析
图11展示了不同尾部批处理配置下的Rollout时间。我们分析了推测因子$\eta$对提示数(P)和响应数(R)的影响。结论是,当$\eta=1.25$并同时应用于P和R时,效果最佳,可将平均Rollout速度提高多达3.9倍,显著优于仅对P或R进行推测的配置。
图11:[尾部批处理] 不同配置的Rollout时间。我们将P和R分别固定为P0和R0,同时改变另一个参数的η。每个条形代表一个完整周期的连续短轮次和长轮次的平均迭代时间,归一化到没有尾部批处理的基线,并标注了实际时间。
6.4 并行规划器
动态TP选择:图12a显示,随着响应长度从8k增加到32k,并行规划器自适应地将TP大小从1增加到2再到4,从而显著降低了Rollout时间,与固定的TP=1相比,平均实现了1.9倍的加速。
**
💬 评论讨论
欢迎在这里分享您的想法和见解!