PipelineRL: Faster On-policy Reinforcement Learning for Long Sequence Generation

作者/机构: Alexandre Piché (ServiceNow AI Research), Ehsan Kamalloo (ServiceNow AI Research), Rafael Pardinas (ServiceNow AI Research), Xiaoyin Chen (Mila, Université de Montréal), Dzmitry Bahdanau (ServiceNow AI Research, Mila, McGill University, Canada CIFAR AI Chair)

A1 主要贡献

本文旨在解决扩展强化学习(RL)方法以增强大型语言模型(LLM)推理能力时遇到的核心挑战。主要问题在于,为了保持高AI加速器利用率,需要进行大规模批量生成,但这会导致产生陈旧的、离策略(off-policy)的数据,从而损害常用RL算法的性能。这些算法(如REINFORCE、PPO)被设计为使用在线策略(on-policy)或近在线策略数据进行训练,行为策略和当前策略之间差异过大会导致训练不稳定。

为了解决这一问题,本文提出了PipelineRL,一种旨在为LLM训练实现硬件效率和数据在线策略性之间更优权衡的方法。PipelineRL的核心创新点如下:

  1. 并发异步数据生成与模型训练:与传统的RL方法在生成和训练阶段之间交替进行不同,PipelineRL允许这两个过程同时进行,从而提高了硬件利用率。
  2. 新颖的“在途权重更新”(in-flight weight updates)机制:这是PipelineRL的关键技术。它允许LLM生成引擎在生成令牌序列的过程中,以最小的中断接收更新后的模型权重。这不仅最大限度地提高了加速器的利用率,也确保了训练数据的“新鲜度”,即生成的token尽可能地接近当前最新的策略。

实验证明,在128个H100 GPU上进行的长篇推理任务中,与传统的RL基线相比,PipelineRL实现了约2倍的学习速度提升,同时保持了高度在线的训练数据。此外,本文的另一项关键贡献是发布了一个可扩展、模块化的PipelineRL开源实现

图1:a) 传统RL在所有GPU上交替进行生成和训练。b) PipelineRL并发运行生成和训练,并通过在途权重更新始终使用最新的模型权重进行生成。
图1:a) 传统RL在所有GPU上交替进行生成和训练。b) PipelineRL并发运行生成和训练,并通过在途权重更新始终使用最新的模型权重进行生成。

A3 背景知识与关键观察

2.1 适用于大型语言模型的强化学习

LLM作为多步策略。强化学习(RL)常用于训练大型语言模型(LLM)以遵循人类偏好或执行长篇推理。LLM的权重可以被视为参数化一个多步策略,该策略根据提示$x$和先前生成的令牌$y_{<i}$为下一个令牌$y_i$分配概率:</p>

$$\pi(y|x) = \prod_{i=1}^n \pi(y_i|x, y_{<i}).$$ <p>近期研究表明,基础策略梯度算法(如REINFORCE 【索引27,Simple statistical gradient-following algorithms for connectionist reinforcement learning,1992】)的变体在训练LLM方面与更复杂的替代方案同样有效。给定一组提示$x_1, \dots, x_m$,REINFORCE通过遵循策略梯度$\nabla J(\pi)$的估计值$\tilde{\nabla} J(\pi)$来最大化策略$\pi$的期望回报$J(\pi)$:

公式
公式

$$\nabla J(\pi)=\frac{1}{m} \sum_{j=1}^{m}\left[\mathbb{E}_{y \sim \pi\left(\cdot \mid x_{j}\right)} \nabla \log \pi\left(y \mid x_{j}\right) R\left(x_{j}, y\right)\right]$$ $$\tilde{\nabla} J(\pi)=\frac{1}{m} \sum_{j=1}^{m} \sum_{t=1}^{T_{j}}\left(R\left(x_{j}, y_{j}\right)-v_{\phi}\left(x_{j}, y_{j, \leq t}\right)\right) \nabla \log \pi\left(y_{j, t} \mid x_{j}, y_{j,<t}\right),$$ <p>其中$R(x_j, y)$是奖励,$v_\phi(x_j, y_{j, \le t})$是通过最小化$(R(x_j, y_j) - v_\phi(x_j, y_{j, \le t}))^2$学习到的价值函数。

离策略数据处理。在大多数实际的RL设置中,由于权重滞后、量化或推理与训练软件之间的实现差异,当前策略$\pi$通常会与生成$y_k$的行为策略$\mu$不同。这种差异通常通过信任区域约束(如PPO【索引23,Proximal policy optimization algorithms,2017】)或使用重要性采样(IS)来处理。在实践中,重要性采样权重会被截断以减少估计器的方差【索引15,Safe and efficient off-policy reinforcement learning,2016;索引2,IMPALA: Scalable distributed deep-RL with importance weighted actor-learner architectures,2018】:

公式
公式

有效样本量(ESS)。有效样本量(ESS)【索引10,A note on importance sampling using standardized weights,1992】常用于量化RL中重要性采样估计器的质量。在使用离策略RL时,ESS衡量了需要多少来自当前策略$\pi$的样本才能达到与来自行为策略$\mu$的加权样本相当的性能。归一化的ESS定义为:

$$ \text{ESS} = \left(\sum_{i=1}^{N} w_{i}\right)^{2} \bigg/ N \sum_{i=1}^{N} w_{i}^{2} $$

其中$w_i$是大小为$N$的样本的重要性权重。该指标在归一化后有效范围在0到1之间,值越接近1表示采样越有效,例如,在线策略数据的ESS恰好为1。较小的ESS会导致高方差的REINFORCE梯度估计,并可能破坏学习过程的稳定性。

2.2 传统强化学习

交替生成与训练。大多数RL实现都在生成序列和在生成的数据上训练策略之间交替进行。我们将这种方法称为传统RL,并在算法1中详细描述。当训练涉及执行多于一个($G > 1$)的优化器步骤时,当前策略$\pi$会领先于用于生成数据的行为策略$\mu$。我们使用术语“滞后”(lag)来指代$\mu$和$\pi$之间的优化器步骤数。

2.3 LLM的高效序列生成

批量生成与KV缓存。Transformer模型逐个令牌地、从左到右生成序列。为了提高效率,先进的生成(推理)引擎(如vLLM和SGLang)一次处理一批序列,同时通过一种称为KV缓存的分页结构来仔细管理它们过去的键和值【索引11,Efficient memory management for large language model serving with pagedattention,2023a】。所有现代生成引擎都支持在不停止生成过程的情况下,将在进行中的请求中“在途”添加新的生成请求。根据加速器的规格,生成引擎应在数千个序列的非常大的批量大小下实现最大生成吞吐量。然而在实践中,在非常大的批量大小下,每个序列的延迟可能变得过高,KV缓存可能变得太大而无法放入加速器内存,或者请求队列管理的开销可能占主导地位。

算法1 传统RL

算法1 传统RL
算法1 传统RL
图2:生成时间和吞吐量分析。我们使用在H100 GPU上运行Qwen 2.5 7B模型的vLLM引擎进行所有测量。(a) 短提示生成吞吐量在批量大小达到256之前持续增加。(b) 随着引擎完成序列,生成批量大小逐渐减小到次优值。(c) 随着每个GPU的序列数量减少,生成时间达到平台期,吞吐量下降。我们报告了5次运行的平均值和95%置信区间。
图2:生成时间和吞吐量分析。我们使用在H100 GPU上运行Qwen 2.5 7B模型的vLLM引擎进行所有测量。(a) 短提示生成吞吐量在批量大小达到256之前持续增加。(b) 随着引擎完成序列,生成批量大小逐渐减小到次优值。(c) 随着每个GPU的序列数量减少,生成时间达到平台期,吞吐量下降。我们报告了5次运行的平均值和95%置信区间。

3. 传统RL的学习速度瓶颈

长序列生成的挑战。当训练LLM生成长序列的令牌时(例如,解决数学问题的长篇推理),强化学习可能会很慢,因为每次生成可能需要长达几分钟。在此,我们解释了为何有效扩展长序列RL(即有效利用更多的加速器$N$来使时间$t$的平均奖励$R(t)$增长得更快)是具有挑战性的。作为一个数学函数,可以将$R(t)$看作是函数$R(S)$和$S(t)$的复合,其中$S$是RL学习者在时间$t$之前处理的样本数量。一个更快的RL学习者将具有更高的学习速度$\frac{\Delta R}{\Delta t}$,我们可以将其表示为学习有效性和学习吞吐量的乘积:

$$\underbrace{\frac{\Delta R}{\Delta t}}_{\text{speed}} = \underbrace{\frac{\Delta R}{\Delta S}}_{\text{effectiveness}} \times \underbrace{\frac{\Delta S}{\Delta t}}_{\text{throughput}}$$

纯在线策略的吞吐量瓶颈。算法1中的传统RL算法在完全在线策略(即每个RL步骤只执行一个优化器步骤)时具有最高的学习有效性$\frac{\Delta R}{\Delta S}$。然而,纯在线策略情况下的吞吐量$\frac{\Delta S}{\Delta t}$可能很低,因为加速器在任何时候最多只处理批量大小为$B$的样本。增加加速器数量$N$将在提高$\frac{\Delta S}{\Delta t}$方面产生递减的回报,因为当每个加速器的样本数量$\frac{B}{N}$低于最佳范围时,每个加速器的吞吐量将下降(如图2c所示)。例如,图2a展示了在单个H100 GPU上7B Qwen模型的推理吞吐量。可以看到,吞吐量几乎呈线性增长,直到生成批量大小达到128。因此,例如使用$2N$个GPU生成32个样本并不比使用$N$个GPU生成64个样本快多少。此外,随着LLM完成较短的生成,仍在进行的较长生成会越来越少,如图2b所示。因此,为了充分利用硬件,每个加速器应生成比最佳批量大小多很多倍的序列。

离策略方法的局限性。通常,为了提高吞吐量,大多数实践者在每个RL步骤执行多个($G > 1$)优化器步骤,这需要在每个生成阶段生成$BG$个轨迹。通过这种方式,通常可以通过增加$N$来获得更高的吞吐量,直到$N$变得过小。然而,文献已知,通过使用高$G$值变得过于离策略最终会降低学习有效性$\frac{\Delta R}{\Delta S}$【索引16,Asynchronous RLHF: Faster and more efficient off-policy RL for language models,2024】。显然,在某些点上,来自旧策略的轨迹变得过于陈旧,不再是当前策略学习信号的有用来源。因此,在给定固定优化器批量大小$B$的情况下,人们通过增加$G$和$N$来扩展传统RL,直到乘积$\frac{\Delta R}{\Delta S} \frac{\Delta S}{\Delta t}$不再提高,从而达到了给定加速器数量$N$的$\frac{\Delta R}{\Delta t}$的硬性上限。

A2 方法细节

4. 使用PipelineRL突破学习速度瓶颈

PipelineRL的核心机制。PipelineRL方法与传统RL在两个方面有所不同:(1)并行异步运行训练和生成,以及(2)在每个优化器步骤后“在途”(in-flight)更新生成权重,即不停止序列生成。算法2以两个并发的Actor和Trainer进程的形式提供了PipelineRL的抽象化形式描述,这两个进程通过一个样本队列和一个高带宽的权重传输网络进行通信。

// Actor 进程:
1: function Actor
2:     S_prog ← []  // 正在进行的序列
3:     while True do
4:         S_fin, S_prog ← 从 S_prog 中弹出已完成的序列
5:         Q_train.put(S_fin)
6:         if len(S_prog) < H then
7:             向 S_prog 添加 H - len(S_prog) 个提示
8:         end if
9:         if Trainer 请求权重更新 then
10:            μ ← receive_weight_update()
11:        end if
12:        S_prog ← 使用 μ 生成下一个 token
13:    end while
14: end function

// Trainer 进程:
15: function Trainer(π, opt_state)
16:     batch ← []
17:     while True do
18:         request_actor_weight_update(π)
19:         batch ← 从 Q_train 中获取 B 个序列
20:         π, opt_state ← optimizer_step(π, opt_state, batch)
21:     end while
22: end function
图3:(a) 对于传统RL,token滞后随优化器步骤数增加。在具有N个加速器的PipelineRL中,token滞后在整个序列中变化,较早的token具有较高的滞后。每个批次中的滞后结构相同。将PipelineRL的加速器数量加倍,其他条件不变,早期token的滞后也会加倍。(b) PipelineRL的吞吐量-有效性权衡示意图,作为训练加速器T的函数;以及传统RL的权衡,作为滞后G的函数。对于相同数量的N个加速器,PipelineRL实现了更高的ΔR/ΔS × ΔS/Δt。
图3:(a) 对于传统RL,token滞后随优化器步骤数增加。在具有N个加速器的PipelineRL中,token滞后在整个序列中变化,较早的token具有较高的滞后。每个批次中的滞后结构相同。将PipelineRL的加速器数量加倍,其他条件不变,早期token的滞后也会加倍。(b) PipelineRL的吞吐量-有效性权衡示意图,作为训练加速器T的函数;以及传统RL的权衡,作为滞后G的函数。对于相同数量的N个加速器,PipelineRL实现了更高的ΔR/ΔS × ΔS/Δt。

PipelineRL的有效性-吞吐量权衡。PipelineRL的有效性-吞-吐量权衡与传统RL相反。具体来说,向PipelineRL设置中增加更多加速器会导致$\frac{\Delta S}{\Delta t}$的线性增加,但最终可能会损害$\frac{\Delta R}{\Delta S}$。在图3a中,我们展示了PipelineRL如何产生混合策略序列,其中较早的令牌比最近的令牌更偏离策略。将$N$加倍将使最早令牌的滞后以及PipelineRL批次中的平均滞后加倍。值得注意的是,PipelineRL及其传统对应方法的离策略性分布是不同的。以平均令牌滞后作为离策略性的代理,在PipelineRL中,所有批次都具有相同的离策略性,而对于传统RL,后续批次变得越来越离策略。这种差异使得难以分析地推断PipelineRL相对于基线能带来的$\frac{\Delta R}{\Delta t}$改进,因为$\frac{\Delta R}{\Delta S}$只能通过运行RL实验来凭经验估计。在补充材料中,我们展示了我们的模拟,即对于相同的最大滞后$g_{max}$,PipelineRL的学习速度可以比传统RL快1.5倍。经验上的增益可能更大,这取决于在不损害学习有效性$\frac{\Delta R}{\Delta S}$的情况下,可以多频繁地进行权重更新。

配置PipelineRL与传统RL。对于固定的批量大小$B$和加速器数量$N$,可以通过选择优化器步骤数$G$来配置传统RL,从而在学习有效性和吞吐量之间进行权衡。同样,PipelineRL的配置也主要可以简化为单个参数,即在$N$个可用加速器中用于训练的加速器数量$T$。设置更高的$T$将几乎线性地减少训练器处理$B$个序列并执行一个优化器步骤所需的时间$t_{train}$。$T$有效地决定了在所有$N-T$个加速器上使用的最佳生成批量大小$H$。使用较小的$H$会导致较低的最大生成延迟$t_{gen}$,从而减少了最大滞后$g_{max} = \lceil t_{gen}/t_{train} \rceil$。因此,使用足以产生足够训练数据的最小$H$是合理的。因此,PipelineRL的最大滞后$g_{max}$随着训练加速器数量$T$的增加而增长,因为更高的$T$需要更高的$H$,并导致更低的$t_{train}$和更高的$t_{gen}$。相反,PipelineRL的样本吞吐量随着$T$的增加而增长,直到$N-T$个加速器无法为过于强大的训练器生成足够的数据。我们建议避免$T$过高(滞后$G$非常高)和$T$过低(硬件利用率差,不如直接减少计算资源)的极端配置。图3b可视化了PipelineRL和传统RL的不同配置如何实现不同的学习有效性$\frac{\Delta R}{\Delta S}$和吞吐量$\frac{\Delta S}{\Delta t}$,其中PipelineRL的设置达到了更高的$\frac{\Delta R}{\Delta t} = \frac{\Delta S}{\Delta t} \frac{\Delta R}{\Delta S}$等值线。

架构和实现细节。我们的PipelineRL实现并发运行多个分布式vLLM生成引擎和DeepSpeed训练工作节点,形成一个三阶段流水线,如图4所示。为简单起见,算法2中省略了中间的预处理器(Preprocessor)阶段,该阶段计算强化学习中常用的人类反馈参考模型对数概率(【索引18,Training language models to follow instructions with human feedback,2022】)。PipelineRL架构是高度模块化的——任何支持PipelineRL所需的三个HTTP API端点的生成软件将来都可以轻松集成。这三个API是用于生成的常用/v1/chat/completions,用于创建权重传输进程组的/init_process_group,以及用于启动在途权重更新的/request_weight_update。PipelineRL中的关键优化包括用于快速训练的在线序列打包,以及使用环形缓冲区以在流水线早期阶段比后期阶段运行得更快时(例如,当训练器进行检查点操作时)最小化滞后。

图4:PipelineRL实现的三个流水线阶段:actor、preprocessor和trainer。早期阶段使用Redis作为流媒体代理将数据流式传输到后期阶段。
图4:PipelineRL实现的三个流水线阶段:actor、preprocessor和trainer。早期阶段使用Redis作为流媒体代理将数据流式传输到后期阶段。

A4 实验

实验环境

为了实验验证PipelineRL的高学习有效性$\frac{\Delta R}{\Delta S}$和吞吐量$\frac{\Delta S}{\Delta t}$,我们选择了一项具有挑战性的任务:训练一个基础模型(即非指令微调模型)执行长篇推理来解决数学问题。我们发现这项任务是PipelineRL的绝佳试验场,因为策略在训练过程中会发生快速变化。特别是,生成序列的长度会显著增长【索引5,DeepSeek-R1: Incentivizing reasoning capability in LLMs via reinforcement learning,2025】,这使得保持在线策略对于有效学习至关重要。

实验结果

图5:(a) PipelineRL比传统RL基线更快地达到相同的平均奖励。(b) PipelineRL实现了与G=8和G=16相同的样本效率。(c) PipelineRL比传统RL基线生成样本的速度快得多。
图5:(a) PipelineRL比传统RL基线更快地达到相同的平均奖励。(b) PipelineRL实现了与G=8和G=16相同的样本效率。(c) PipelineRL比传统RL基线生成样本的速度快得多。
图6:在图6a中,PipelineRL达到了比任何传统RL方法都高的最大滞后,但如图6b所示,有效样本量与G=8相似。这表明虽然最大滞后非常高,但根据ESS衡量,PipelineRL在很大程度上保持了在线策略。
图6:在图6a中,PipelineRL达到了比任何传统RL方法都高的最大滞后,但如图6b所示,有效样本量与G=8相似。这表明虽然最大滞后非常高,但根据ESS衡量,PipelineRL在很大程度上保持了在线策略。

PipelineRL因更高吞吐量而学习更快。我们将PipelineRL的学习速度与传统RL(G=32个优化器步骤,这是其能稳定训练的最大G值)进行比较。结果显示,PipelineRL达到相同奖励值的速度大约是基线的2倍(图5a),这是由于其样本吞吐量快了约2倍(图5c)。吞吐量提升的主要原因是,在128个GPU上,当每个GPU只需生成$32 \times 1024 / 128 = 256$个序列时,G=32实验的GPU利用率相对较低(见图2b)。将G进一步增加到64会导致训练发散(见附录图10)。

PipelineRL学习效率高。为了更好地衡量PipelineRL的学习有效性$\frac{\Delta R}{\Delta S}$,我们还运行了G=8、G=16和G=32的传统RL实验。值得注意的是,所有比较方法的奖励-样本数(R(S))曲线在G=32因偏离策略太远而变慢和不稳定之前是无法区分的。这一结果验证了PipelineRL标志性的在途权重更新不会损害序列生成过程。

PipelineRL在推理任务上与可比结果相匹配。表1比较了PipelineRL在与文献中从相同Qwen 2.5 7B模型开始训练的类似实验的测试性能。在此实验中,我们使用4096的批量大小,因为我们发现它能带来更高的性能。在数学推理基准测试MATH500【索引6,Measuring mathematical problem solving with the MATH dataset,2021】和AIME2024【索引13,NuminaMath: The largest public dataset in AI4Maths with 860k pairs of competition math problems and solutions,2024】上,PipelineRL的成功率与Open Reasoner Zero和SimpleRL Zero持平或更高。

表1:使用PipelineRL训练的模型与文献中结果的成功率比较。
表1:使用PipelineRL训练的模型与文献中结果的成功率比较。

PipelineRL更保持在线策略。为了更好地理解哪些训练方法更能保持在线策略,我们绘制了训练过程中最大滞后和ESS在线策略性度量的演变。图6a显示PipelineRL获得了比传统RL基线更高的最大滞后。值得注意的是,一些令牌的滞后超过5万个样本。然而,图6b显示,在ESS方面,PipelineRL保持了与G=8相似的在线策略性。我们进一步观察到,G=16,特别是G=32的ESS在整个训练过程中都在下降。

5.1 在途权重更新对在线策略性的影响

实验设计。在本节中,我们比较了在途权重更新的采样分布与1)具有不同最大滞后的传统RL和2)具有KV缓存重计算的在途权重更新。为此,我们保存了一组连续的检查点$C_i$,每个优化器步骤后一个。为了模拟在途权重更新,我们从一个检查点开始,每生成$L/g_{max}$个令牌就用下一个检查点更新行为策略的权重,其中$L$是最大序列长度,$g_{max}$是最大滞后。具体来说,PipelineRL的行为策略定义为:

公式
公式

其中$t_1 = 2L/g_{max}$,对于滞后$g>1$,$t_g = t_{g-1} + L/g_{max}$个令牌,因为由于训练开始时的“气泡”(bubble),第一次权重更新比后续更新花费更长时间(见图1b)。我们还使用$\hat{x}$来强调先前令牌的KV缓存是陈旧的——因为它是在先前的模型权重下计算的。然后我们计算混合行为策略$\mu_{C:C+g}$和在线行为策略$\mu_{C+g}$之间的KL散度。我们还报告了与更新了KV缓存的混合行为策略的KL散度,我们称之为PipelineRL with KV cache recomputed。为了模拟传统RL,我们从行为策略$\mu := \mu_C$中采样$N$个序列,并计算与不同滞后$k$的在线行为策略$\mu_{C+g}$的KL散度。

实验细节与结果。在此实验中,我们在OpenReasoner Zero【索引8,OpenReasoner-Zero: An open source approach to scaling up reinforcement learning on the base model,2025】数据上微调Qwen 2.5基础7B模型,共222个优化器步骤。我们考虑训练中的三个阶段来测量KL散度:从检查点0、100和190开始。最大滞后$g_{max}$设为32,最大序列长度$L$为2048。如图7所示,在所有三个训练阶段中,混合策略序列的分布与完全在线策略序列的分布非常接近。相比之下,随着滞后的增加,离策略序列表现出持续更高的散度。此外,对混合策略序列使用陈旧的KV缓存与重新计算缓存相比,仅引入了略高的散度。这支持了我们在PipelineRL中选择保留KV缓存这种更高效方法的设计决策。

图7:对于三个不同的起始检查点,通过KL散度测量,PipelineRL(无论是否重计算KV缓存)比传统RL更接近在线策略。
图7:对于三个不同的起始检查点,通过KL散度测量,PipelineRL(无论是否重计算KV缓存)比传统RL更接近在线策略。

A7 补充细节

6. 相关工作

异步与高吞吐量RL。异步和高吞吐量RL已被广泛研究。IMPALA【索引2,IMPALA: Scalable distributed deep-RL with importance weighted actor-learner architectures,2018】将行动与学习解耦,以最大化GPU利用率。与PipelineRL类似,IMPALA使用截断的重要性权重来从离策略样本中估计价值函数。此外,IMPALA在一个回合的长度内保持策略权重不变。SeedRL【索引3,SEED RL: Scalable and efficient deep-RL with accelerated central inference,2019】提出在回合中更新模型参数,导致轨迹中不同的动作由不同的策略采样。OpenAI Five【索引17,Dota 2 with large scale deep reinforcement learning,2019】使用异步PPO进行训练,在Dota 2上取得了超人表现。这些先前的工作都集中在视频游戏的RL上。与我们的工作更接近的是,【索引16,Asynchronous RLHF: Faster and more efficient off-policy RL for language models,2024】探索了用于LLM的异步RL。在他们的方法中,下一个G个优化器步骤的数据生成与前一个G个优化器步骤的训练同步,导致比传统RL更高的离策略性,这与PipelineRL不同。该研究还表明,像DPO【索引19,Direct preference optimization: Your language model is secretly a reward model,2023】这样的离线方法可以更好地容忍离策略性。

可扩展的开源RL实现。存在其他几种可扩展的开源RL实现。veRL【索引25,HybridFlow: A flexible and efficient RLHF framework,2024】通过使用一个复杂的混合生成-训练引擎,支持在同一GPU上快速切换训练和生成,从而高效地实现了传统RL。我们认为veRL的吞吐量将与我们的传统RL基线相似。如果没有混合引擎,在OpenRLHF【索引7,OpenRLHF: An Easy-touse, Scalable and High-performance RLHF Framework,2024】中,训练GPU在生成期间空闲,反之亦然。与此同时,Magistral【索引14,Magistral,2025】也引入了在途权重更新。

A5 结论与讨论

结论。我们已经展示了在途权重更新如何帮助PipelineRL突破传统两阶段RL方法的学习速度瓶颈。我们认为,特别是在长序列生成方面,这种加速很难通过其他异步RL方法实现,因为同步等待生成完成会损害吞吐量和/或学习有效性。在途更新引入的陈旧KV缓存风险可以通过每次更新后重新计算KV缓存来缓解,这可以在高GPU利用率下快速完成,但仍会降低吞吐量。

未来工作。我们相信PipelineRL可能特别适用于训练LLM在涉及多次LLM生成与环境交互的智能体行为方面表现出色。另一个有前景的未来工作方向是研究PipelineRL中最近的低滞后令牌何时有帮助,以及相反,PipelineRL在长序列中早期令牌的持续高滞后在何处会造成损害。

局限性。如果LLM被要求为相同的提示生成完全相同数量的令牌,PipelineRL相对于传统RL只会带来有限的吞吐量增长。在这种不太可能的情况下,传统RL同样能够保持恒定的生成批量大小。然而,PipelineRL稳定的平均令牌滞后和每个批次中最近令牌的低滞后可能仍会影响学习有效性。在计算资源稀缺或极其丰富的设置中,PipelineRL的吞吐量优势同样会减小。在前一种情况下,每个GPU将获得足够的生成任务以使GPU利用率很高。在后一种情况下,学习速度将不受硬件利用率的限制,而是受限于最佳可能的生成延迟和环境反馈延迟。

A6 附录

A. PipelineRL加速的分析估计

在本附录中,我们估计在最大令牌滞后$g_{max}$相同的情况下,PipelineRL相比于传统RL能快多少。我们将使用以下符号,大部分与正文相同:

A.1 单位

定义“flash”为时间单位。为了比较不同RL方法的吞吐量,采用不依赖于特定GPU型号和LLM大小的时间和吞吐量单位是很有用的。为此,我们引入一个名为flash的时间单位:

公式
公式

其中$F_{gen}$是所选LLM进行一次令牌前向传播所需的FLOPs数量,而$M$是给定GPU的最大理论FLOPs吞吐量。一个flash的含义是理论上令牌生成可以达到的最小摊销时间。因此,生成$K$个令牌至少需要$K$个flashes,尽管在更典型的约0.1的生成利用率下,它将需要10K个flashes。对于非常长的序列,$F_{gen}$可能会因注意力FLOPs成为总FLOPs的很大一部分而显著变化,但为简单起见,我们在此将忽略此细节。引入flash $f$作为单位后,我们将以tokens per flash来衡量系统吞吐量。

定义训练时间。让$\tau$为每个令牌的摊销训练时间。对于PipelineRL和传统RL,$\tau$在大规模上是相似的,因为两种方法都可以从序列打包中受益。

图8:H100在批量大小h下的利用率,作为最大理论bf16 FLOPS吞吐量的比率。我们使用(4096, h)·(h, 16384)的矩阵乘法进行测量。对于每个h,我们考虑填充到h+64以提高速度,因为根据经验,我们观察到当h是2的更高次幂(最高到128)的倍数时,利用率有很大提升。
图8:H100在批量大小h下的利用率,作为最大理论bf16 FLOPS吞吐量的比率。我们使用(4096, h)·(h, 16384)的矩阵乘法进行测量。对于每个h,我们考虑填充到h+64以提高速度,因为根据经验,我们观察到当h是2的更高次幂(最高到128)的倍数时,利用率有很大提升。

A.2 传统RL吞吐量

吞吐量公式。我们可以如下表示传统RL的吞吐量:

公式
公式

其中$t_{gen}^{conv}$和$t_{train}^{conv}$分别是生成和训练所需的时间。让我们更仔细地看一下这些项:

公式
公式

$$t_{conv}^{train}=\frac{K\tau}{N}$$

其中$h(l)$是在解码$l$步后仍在进行的序列数,而$U(h)$是批量大小为$h$时的GPU利用率。为了理解公式(11),回想一下,根据定义,在完美的GPU利用率下生成$k$个令牌需要$k$个flashes,在利用率为$U(k)$时则需要$k/U(k)$。

以tokens/flash表示。我们可以将其重写为以tokens/flash为单位的吞吐量:

公式
公式
公式
公式
公式
公式
图9:在N=128个GPU和批量大小B=N的设置下,PipelineRL和传统RL的吞吐量作为最大滞后g_max的函数。
图9:在N=128个GPU和批量大小B=N的设置下,PipelineRL和传统RL的吞吐量作为最大滞后g_max的函数。

效率下降原因。在步骤$l$中,每个GPU的批量大小较低时,即$h_N(l) = h(l)/N$,比率$h_N(l)/U(h_N(l))$作为$N$的函数只会非常缓慢地减少,因为对于现代GPU,当$x$很小时,$x/U(x)$几乎是恒定的。这是对传统RL随着$N$增长效率降低的正式解释。上述设置中的最大令牌滞后是$S-1$。

A.3 PipelineRL吞吐量

吞吐量公式。对于PipelineRL,系统吞吐量由最慢的流水线阶段决定。使用上面介绍的概念,PipelineRL的吞吐量可以估计如下:

公式
公式

$r_{pipeline}^{gen}=U(H)I$

$$r_{pipeline}^{train} = \frac{N - I}{\tau}$$

最大滞后分析。要理解PipelineRL的最大滞后,考虑这样一个事实:生成GPU在生成最长可能序列(长度为$L$)所需的时间内将产生$H\bar{L}I/L$个令牌。平均而言,这些令牌中将有$H\bar{L}I/L$个序列。因此,在最坏的情况下,即一个优化器步骤恰好在最长序列生成开始之前发生,一个长序列将在其生成开始后$g_{max} = \lceil H\bar{L}I/LB \rceil$个优化器步骤后用于训练。为了构建一个与传统RL系统具有相同滞后的等效系统,需要在保持$\lceil H\bar{L}I/LB \rceil \le S-1$的同时最大化$r_{pipeline}(H, I)$。我们发现这个问题难以解析求解,并在我们下面的研究中对所有$(H, I)$配置进行了直接搜索。

A.4 PipelineRL加速案例研究

设置与结果。为了计算PipelineRL带来的确切吞吐量提升,有必要对序列长度分布和实验使用的硬件做出假设。在下面的案例研究中,我们假设长度从1到最大长度$L$均匀分布,并使用H100作为GPU。我们在图8中可视化了GPU利用率表$U(h)$。读者可以看到$U(h)$几乎线性增长,直到$h \sim 200$,这使得可以在更少的GPU上以更高的利用率压缩生成。对于一个具有$N=128$个GPU和训练批量$B=128$的设置,我们考虑了PipelineRL所有可能的$(I, H)$配置,并将其吞吐量绘制为滞后$g_{max}$的函数。图9显示,对于$g_{max} \sim 133$,PipelineRL可以快达1.57倍。这个滞后值对于许多实际设置可能太高,但对于例如$B=2048$的更高批量大小,每个GPU生成的相同数量的序列将对应于一个实际的、低16倍的滞后$g_{max} \sim 8$。

加速机制。PipelineRL实现改进的机制如下:

显然,PipelineRL加速的根本原因是,44个生成GPU可以每flash产生16.9个令牌,这比在传统RL情况下让128个GPU每flash产生18.3个令牌(系统吞吐量为10.7)更有效。

B. 附加结果

图10:G=64发散。
图10:G=64发散。