通过自适应张量并行加速同步RLHF训练中的长尾生成

# 通过自适应张量并行加速同步RLHF训练中的长尾生成

来源：https://arxiv.org/html/2605.23945

Long Zhao†,§ Qinghe Wang†,§ Jiaan Zhu¶ Youhui Bai¶ Zewen Jin¶ Chaoyi Ruan‡ Shengnan Wang∗ Cheng Li¶,§

†安徽大学 ¶中国科学技术大学 ‡新加坡国立大学 ∗独立研究者
§合肥综合性国家科学中心人工智能研究院 \(2018\)

###### 摘要。
基于人类反馈的强化学习（RLHF）已成为提升模型质量的关键后训练范式。然而，同步三阶段RLHF流水线常常受到生成阶段的瓶颈制约：响应长度偏斜导致解码过程中有效批次大小快速缩小，使得GPU在少数长响应尚未完成时处于低利用率状态。主流框架采用静态张量并行（TP）配置，无法适应变化的批次特征，留下了大量未被挖掘的性能潜力。我们提出PAT，一种自适应TP方法，能够在每次RLHF迭代的生成阶段动态重新配置TP。PAT引入了两项关键技术。首先，一种基于预测器的在线重配置方法，根据离线性能分析结果决定重配置时机和目标TP配置，仅在预测的延迟收益超过重配置开销时才触发重配置。其次，一种轻量级在线重配置机制，仅更新受TP变化影响的状态和布局：通过基于成本模型的KV缓存迁移与重计算选择来适配未完成的解码状态，执行原地权重重分片，并复用缓存的通信组。我们在SGLang之上实现了PAT，并将其集成到VeRL框架中。在DeepScaleR上使用LLaMA3.1-8B和Qwen3-14B进行的评估表明，与原始的VeRL设置相比，PAT可将生成延迟降低最多34.6%，并将端到端RLHF训练迭代延迟降低最多27.2%。

RLHF训练，长尾生成，张量并行，并行重配置
††journalyear:2018
††doi:XXXXXXX.XXXXXXX
††conference:请从权限确认邮件中填写正确的会议标题；2018年6月03-05日；Woodstock, NY
††isbn:978-1-4503-XXXX-X/2018/06
††footnotetext:Long Zhao和Qing Wang对本文贡献相等。

## 1. 引言

基于人类反馈的强化学习（RLHF）(Bai et al., 2022 (https://arxiv.org/html/2605.23945#bib.bib17))已成为一种有效的后训练范式，用于使模型输出与人类偏好对齐，并显著提升大语言模型的能力。近期一些重要模型，例如DeepSeek-R1 (Guo et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib30))、Kimi K2 (Team et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib31))和OpenAI o3 (El-Kishky et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib32))，已采用或探索了RLHF训练策略以增强其模型质量。RLHF训练流水线在每次迭代中同步地经历三个阶段。首先，**生成阶段**为目标LLM根据给定提示生成响应。其次，**准备阶段**评估生成的响应，计算奖励或其他辅助信号。最后，**训练阶段**消耗准备阶段的输出，使用准备好的信号计算训练损失，并更新目标LLM。主流的RLHF框架（如VeRL (Sheng et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib7))）倾向于将所有三个阶段部署在同一组GPU服务器上。结合同步执行，这减少了跨阶段的数据移动和同步开销。然而，样本生成仍然是主要的性能瓶颈。图1 (https://arxiv.org/html/2605.23945#S1.F1) 说明了这种低效的根本原因。最初，大多数响应的长度相近且几乎同时完成——我们将这一阶段称为**对齐阶段**。随后，少数异常长的响应继续执行，进入我们所说的**尾阶段**。两个阶段的GPU资源消耗高度不平衡，在尾阶段GPU资源严重闲置。在我们的基准测试中（最大响应长度16K tokens，批次大小128），每个GPU的平均吞吐量从对齐阶段的13.71 TFLOPS下降到尾阶段仅0.11 TFLOPS。因此，加速RLHF的关键在于加速由响应长度偏斜引起的尾阶段。

![参考标题 图1.生成过程中TP/DP重配置示例。PAT在对齐阶段从(TP2, DP2)切换到尾阶段的(TP4, DP1)。](https://arxiv.org/html/2605.23945#S1.F1)

进一步的调查显示，这些框架通常在整个生成阶段使用静态的张量并行（TP）配置。尽管TP度对生成性能起着关键作用，但它始终保持不变，不随批次特征变化。在固定的GPU预算下，对齐阶段较大的有效批次大小倾向于较小的TP度以获得更高的吞吐量。相反，尾阶段少数剩余样本受益于较大的TP度，这能降低解码延迟并提高GPU利用率。这一观察结果启发我们在运行过程中动态调整TP度，特别是在有效批次大小缩小时增加TP度。然而，在生成过程中动态重配置TP并非易事，原因如下：首先，TP重配置会带来不可忽视的开销。具体来说，未完成的样本必须在新并行布局下传输或重计算，模型权重必须在新TP组内重新分片，并且通信组必须一致更新。其次，重配置的收益需要在剩余尾工作负载与重配置开销之间进行权衡。过早重配置会损害对齐阶段的吞吐量，因为较高的TP引入了额外的通信并降低了数据并行（DP）并发度¹。过晚重配置则导致剩余的尾工作不足以摊销重配置开销。因此，一个有效的系统必须判断何时重配置是有益的，并以低开销执行重配置。

¹在实践中，TP通常与DP结合使用。虽然我们主要更改TP度，但这也会导致DP配置的变化。

为了解决这些挑战，我们提出了PAT，一种自适应张量并行框架，能够在每次RLHF迭代的生成阶段内执行TP重配置。如图1 (https://arxiv.org/html/2605.23945#S1.F1) 的下半部分所示，PAT在仍有许多样本活跃的对齐阶段保持吞吐量导向的低TP/高DP配置。然后仅在预测的收益超过重配置开销时，切换到延迟导向的高TP/低DP配置。PAT基于两项关键技术：第一，PAT设计了一种基于预测器的在线切换方法，用于确定切换点和目标TP配置。预测器依靠离线性能分析的数据，估计在每种候选配置下的剩余生成延迟，并将一次性重配置开销纳入决策。PAT仅在预测的剩余尾样本延迟收益能够摊销切换成本时触发重配置。第二，PAT通过轻量级的在线运行时状态和布局适配使得重配置变得轻量。PAT并非在目标TP配置下重新初始化推理引擎，而是仅更新依赖于TP的状态和布局。对于未完成的样本，它通过基于成本模型在KV缓存迁移与重计算之间进行选择来重建解码状态。对于模型权重，它执行原地重新分片以匹配目标TP布局。它还跨迭代复用缓存的通信组，以避免重复构建组。这些机制共同将TP重配置转变为低开销操作，使得尾样本能够在新配置下高效恢复解码。

我们在SGLang (Zheng et al., 2024 (https://arxiv.org/html/2605.23945#bib.bib6))（一个广泛用于RLHF生成的推理引擎）上构建了PAT，并将其与VeRL (Sheng et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib7))（领先的RLHF框架之一）集成，便于采用。我们通过训练LLaMA3.1-8B (Dubey et al., 2024 (https://arxiv.org/html/2605.23945#bib.bib21)) 和Qwen3-14B (Yang et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib22)) 在DeepScaleR (Luo et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib24)) 上评估PAT。与原始的VeRL设置相比，PAT将生成延迟降低了最多34.6%，使得RLHF训练迭代延迟降低了最多27.2%。

## 2. 背景与动机

### 2.1. RLHF的工作流程

RLHF是将LLM与人类偏好对齐的常见方法 (Ouyang et al., 2022 (https://arxiv.org/html/2605.23945#bib.bib13); Guo et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib30))。它通常包括监督微调、奖励模型训练和强化学习。本文聚焦于强化学习步骤，包括PPO (Schulman et al., 2017 (https://arxiv.org/html/2605.23945#bib.bib26))、GRPO (Shao et al., 2024 (https://arxiv.org/html/2605.23945#bib.bib23)) 和DAPO (Yu et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib28)) 等算法，其中迭代的生成、评分和模型更新主导了总体训练成本。如图2 (https://arxiv.org/html/2605.23945#S2.F2) 所示，该步骤可以进一步分解为三个阶段。

**生成阶段。** 给定一批提示，actor模型以自回归方式生成响应。该阶段包括prefill（处理输入提示并构建KV缓存）和解码（复用KV状态生成后续token）。解码通常是内存带宽受限的，受益于较大的有效批次 (Agrawal et al., 2024 (https://arxiv.org/html/2605.23945#bib.bib8))。

**准备阶段。** 给定生成的提示-响应对，辅助模型或模块（如奖励模型和参考模型）计算奖励、对数概率或其他用于损失计算的训练信号。

**训练阶段。** actor模型在生成的样本上执行前向和反向传播，使用准备好的信号计算RL训练损失，并更新其参数。更新后的权重随后被下一个生成阶段使用。

由于这些依赖关系，主流的RLHF框架（如VeRL (Sheng et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib7)) 和RLHFuse (Ouyang et al., 2022 (https://arxiv.org/html/2605.23945#bib.bib13))）采用同地同步执行，所有阶段共享同一GPU池。这种部署减少了跨阶段的数据移动和同步开销。

![参考标题 图2.典型迭代RLHF训练的工作流程。](https://arxiv.org/html/2605.23945#S2.F2)

### 2.2. 并行策略

为了提高吞吐量并降低每设备内存压力，RLHF训练在三个阶段中采用混合的现有并行策略组合 (Sheng et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib7); Mei et al., 2024 (https://arxiv.org/html/2605.23945#bib.bib29))。在这些策略中，数据并行（DP）在多个设备上复制模型，每个副本处理批次的不同子集。训练期间，每次迭代后跨副本同步梯度 (Li et al., 2020 (https://arxiv.org/html/2605.23945#bib.bib4))。张量并行（TP）跨设备拆分模型层，使得单层内的计算可以并行执行，并降低每设备内存使用 (Shoeybi et al., 2019 (https://arxiv.org/html/2605.23945#bib.bib14))。流水线并行（PP）将模型划分为不同设备上的顺序阶段，以流水线方式传递微批次，以重叠计算和通信 (Huang et al., 2019 (https://arxiv.org/html/2605.23945#bib.bib5))。

RLHF的三个主要阶段具有不同的计算需求，因此需要不同的并行策略 (Sheng et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib7))。准备和训练阶段通常是计算密集型的，可以从混合并行（包括DP、TP和PP）中受益，具体取决于模型大小和硬件规模。相比之下，解码阶段主要是内存带宽受限。其性能通过最大化批次大小而不是计算吞吐量来优化。因此，解码阶段使用DP和TP的组合，同时倾向于尽可能大的DP度以最大化并发解码请求数量，从而提升整体吞吐量。解码时通常避免使用PP，因为自回归生成每步仅产生一个token，活跃批次已经是摊销权重加载、KV缓存访问和内核开销的主要来源。进一步的微批处理会碎片化解码工作负载，降低内存访问效率，并可能损害整体吞吐量。

### 2.3. 同步RLHF训练中的问题

![参考标题 (a) 随最大生成长度变化的端到端迭代延迟。](https://arxiv.org/html/2605.23945#S2.F3.sf1)
![参考标题 (b) LLaMA3.1-8B在DeepScaleR上生成的响应长度CDF。](https://arxiv.org/html/2605.23945#S2.F3.sf2)

图3. RLHF迭代分解与响应长度分布。

在同步RLHF训练中，生成阶段是主要瓶颈。为了量化这一点，我们为每个阶段配置了经过调优的静态并行策略，并遵循HybridFlow (Sheng et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib7)) 的设置，批次大小为128，使用LLaMA3.1-8B在DeepScaleR (Luo et al., 2025 (https://arxiv.org/html/2605.23945#bib.bib24)) 数据集上，配备八块NVIDIA A40 GPU。如图3a (https://arxiv.org/html/2605.23945#S2.F3.sf1) 所示，生成阶段消耗了总迭代时间的64.22%到81.7%，并且其占比随最大响应长度增加而增加。在所评估的最大响应长度和TP/DP配置下，解码占生成延迟的90%以上。这种解码瓶颈主要源于每个批次内的响应长度偏斜：少数序列可能比其他序列长得多 (Zhong et al., 2025b (https://arxiv.org/html/2605.23945#bib.bib9))。一旦较短的序列完成，活跃批次大小迅速下降，但生成阶段仍需等待剩余的长序列完成。因此，这些少量尾序列成为生成的关键路径，导致GPU严重利用不足 (Zhong et al., 2025a (https://arxiv.org/html/2605.23945#bib.bib12))。例如，在图3a (https://arxiv.org/html/2605.23945#S2.F3.sf1) 评估的所有最大响应长度下，仅剩一个样本的时段占生成延迟的46.29%–81.58%。在16K设置下，批次大小128时，每个GPU的平均达到吞吐量从对齐阶段的13.71 TFLOPS下降到生成接近尾声时的仅0.11 TFLOPS。这种长尾行为并非特定于某个数据集。先前的RLHF系统也已识别出长尾生成是阶段级低效的主要来源 (Zhong et al., 2025b (https://arxiv.org/html/2605.23945#bib.bib9), a (https://arxiv.org/html/2605.23945#bib.bib12))。这种行为常见于具有显著响应长度变异性的工作负载中，包括推理、代码生成和指令遵循等。