RLHFuse: Efficient RLHF Training for Large Language Models with Inter- and Intra-Stage Fusion

发表时间: 2024-09 · arXiv:2409.13221

原文: https://arxiv.org/abs/2409.13221

作者/机构: Yinmin Zhong1, Zili Zhang1, Bingyang Wu1, Shengyu Liu1, Yukun Chen2, Changyi Wan2, Hanpeng Hu2, Lei Xia2, Ranchen Ming2, Yibo Zhu2, Xin Jin1; 1北京大学 2阶跃星辰

速读

一句话结论 本文提出了一个名为 RLHFuse 的强化学习人类反馈训练框架，通过将任务拆分为更细粒度的子任务并进行阶段间与阶段内的执行融合，在保持同步训练语义的前提下，将大语言模型训练的端到端吞吐量最高提升了 3.7 倍。

要解决什么问题 现有强化学习人类反馈（RLHF）训练系统将生成、推理和训练视为不可分割的最小执行单元，导致严重的 GPU 资源浪费，具体卡在两个机制上。第一是生成阶段的数据倾斜。由于模型生成的响应长度呈现长尾分布，且推理任务必须等待整个批次的生成任务全部完成才能开始，导致在处理少数长尾样本时，大量 GPU 处于空闲等待状态；同时，大语言模型生成的解码阶段受限于显存带宽，需要大批量才能维持高利用率，少数长尾样本无法喂饱 GPU。第二是训练阶段的流水线气泡。随着模型规模增大，必须采用更高的流水线并行度，但在常用的 1F1B（一前向一反向）调度策略下，气泡比例会随着并行阶段数增加而急剧上升，当阶段数接近微批次数时，近一半的 GPU 算力会在等待中被白白消耗，而现有的交错调度方案又会引入额外的通信开销并加剧阶段间的不平衡。

怎么做的 核心思路是将传统的任务级工作流打破，下钻到样本级和微批次级的子任务视角，通过数据感知的阶段间融合与模型感知的阶段内融合来榨干硬件算力。针对生成阶段的长尾卡点，框架将生成和推理拆分为样本级子任务，设计了阶段间融合机制。系统会监控生成进度，当剩余未完成的样本数低于特定迁移阈值 $R_t$ 时，触发轻量级迁移，将这些长尾样本的键值缓存（KV Cache）通过网络集中转移到少数几个指定的 GPU 实例上继续解码。这样不仅让接收端保持较高的批大小以维持计算效率，更重要的是将大部分 GPU 资源提前释放，用于启动后续的推理任务，实现了生成与推理的执行重叠。目标实例的数量 $m$ 受到性能和显存的双重约束，满足 $m \ge \max(R_t / BS_{max}, R_t \times M / C)$，其中 $BS_{max}$ 是使 GPU 饱和的最大批大小，$M$ 是最大输出长度样本的显存消耗，$C$ 是可用显存。针对训练阶段的气泡卡点，框架利用了 RLHF 训练中 Actor（生成动作的主模型）和 Critic（评估动作的价值模型）两个模型相互独立的特性，将训练拆分为微批次级子任务，设计了阶段内融合机制。由于这两个模型参数量和最优并行策略通常不同，属于异构模型，传统的对称双向流水线调度无法直接使用。系统首先通过合并流水线阶段对齐两者的 GPU 数量，然后将寻找最优融合流水线调度矩阵 $S$ 的过程转化为一个流作业调度问题。为了在庞大的搜索空间中避开局部最优，框架采用模拟退火算法，以最小化总延迟为首要目标，以降低峰值激活显存为次要目标，自动为任意规模和并行配置的异构模型对生成双向流水线调度方案，让 Actor 和 Critic 的前向与反向微批次相互填补对方的气泡。

效果如何 实验在 32 个节点共 256 张 H100 GPU 的集群上搭建，采用 LLaMA 系列 13B 到 65B 规模的模型（组合成 13B/33B 到 65B/33B 等不同 Actor/Critic 异构对），在 HH-RLHF 数据集上进行测试。对比基线包括代表仅支持 ZeRO-3 数据并行路线的 DeepSpeed-Chat，以及代表支持定制化 3D 并行但无子任务级优化路线的 ReaLHF。在全局批大小为 512 的设置下，RLHFuse 的端到端样本吞吐量比 DeepSpeed-Chat 高出 2.5 倍至 3.7 倍，比 ReaLHF 高出 1.4 倍至 2.4 倍。消融实验证明，仅靠阶段融合技术本身就能带来 1.2 倍至 1.4 倍的相对吞吐量提升，其中生成与推理阶段的重叠带来了最高 1.6 倍的局部加速，而阶段内融合将训练阶段耗时缩短了最高 1.3 倍，且峰值激活显存逼近串行执行的理论下界。该方法的代价在于引入了长尾样本迁移、数据传输和模型权重重新分配的开销，但得益于高速网络和 all-to-all 通信优化，这些任务切换开销在总迭代时间中的占比不到 3%，基本可以忽略不计。

A1 主要贡献

本文旨在解决现有强化学习人类反馈（RLHF）训练系统在生产部署中GPU利用率低下的问题。核心问题源于RLHF训练的内在特性：生成阶段的数据倾斜（长尾效应）和训练阶段的流水线气泡。现有系统将RLHF流程中的每个任务视为最小执行单元，忽略了子任务级别的优化机会。

研究目标：提出一个名为RLHFuse的全新RLHF训练框架，通过将任务分解为更细粒度的子任务并进行阶段融合，来显著提升GPU利用率和训练吞吐量。

核心创新点：

  1. 子任务级优化视角：首次提出将RLHF工作流从传统的任务级视角转变为更细粒度的子任务级视角，从而开辟了新的优化空间。
  2. 数据感知的阶段间融合（Inter-Stage Fusion）：
    • 针对生成阶段由长尾样本导致的数据倾斜问题，RLHFuse将生成和推理任务拆分为样本级的子任务。
    • 设计了一种轻量级迁移算法，在生成阶段进入低效的长尾处理阶段时，自动将剩余的长尾样本迁移到少数GPU上处理。
    • 释放出的GPU资源被提前用于启动推理任务，实现了生成阶段和推理阶段的执行重叠，有效缓解了生成瓶颈。
  3. 模型感知的阶段内融合（Intra-Stage Fusion）：
    • 针对训练阶段因流水线并行（PP）导致的大量气泡问题，RLHFuse将训练任务分解为微批次（micro-batch）级别的子任务。
    • 利用RLHF训练阶段包含两个独立模型（Actor和Critic）的特性，设计了一种融合流水线调度策略，在同一组GPU上并发执行这两个模型的训练子任务，让它们相互填充对方的流水线气泡。
    • 提出了一种轻量级的流水线调度生成算法，能够为任意规模和并行配置的Actor和Critic模型生成接近最优的调度方案。
  4. 一系列系统优化：RLHFuse集成了针对RLHF各阶段的系统优化，使其成为一个高效且可扩展的生产级框架。

通过这些创新，RLHFuse与现有最先进系统相比，训练吞吐量最高可提升3.7倍。

图1：上图：RLHF训练工作流。下图：现有RLHF训练框架的训练迭代分解。
图1：上图：RLHF训练工作流。下图：现有RLHF训练框架的训练迭代分解。

A3 背景知识与问题分析

2.1 背景知识

RLHF模型构成：一个完整的RLHF过程通常涉及四个大语言模型（LLM）：Actor、Ref、Critic和RW。

RLHF工作流：图1展示了RLHF单次训练迭代的工作流，包含三个主要阶段：生成、推理和训练。

LLM并行化：为了扩展训练，通常会结合使用三种正交的并行策略：

实践中，通常将这三种策略组合使用，记为(dp, pp, tp)。

2.2 RLHF训练中的问题

生成阶段的数据倾斜：生成样本的响应长度呈现长尾分布，即每批中总有少数样本比其他样本长得多。由于数据依赖性，推理任务必须等待生成任务全部完成，导致即使只有少数长尾样本，也会迫使后续任务等待。这造成了严重的GPU资源浪费，因为LLM生成的解码阶段是内存密集型（memory-bound），需要大批量才能保持高GPU利用率。如图2（左）所示，在LMSYS-Chat-1M数据集中，不同模型的输出长度均呈现长尾分布，P99.9的长度是中位数的十倍以上。图2（右）展示了内部模型训练的迭代时间分解，生成长尾样本的时间（深蓝色条）占了总生成时间的一半以上，并且随着最大生成长度的增加，问题愈发严重。

图2：左图：LMSYS-Chat-1M数据集中模型的输出长度CDF。垂直虚线表示P99.9的输出长度。右图：在不同最大输出长度下，内部模型和数据集的RLHF训练迭代分解。
图2：左图：LMSYS-Chat-1M数据集中模型的输出长度CDF。垂直虚线表示P99.9的输出长度。右图：在不同最大输出长度下，内部模型和数据集的RLHF训练迭代分解。

训练阶段的流水线气泡：随着模型规模增大，需要更高的流水线并行度（PP）来扩展训练。然而，流水线气泡的比例随PP规模增加而增加，显著降低训练效率。在常用的1F1B流水线调度（如图3上）中，气泡比例为 (N-1)/(N-1+M)，其中N是PP阶段数，M是微批次数。当N接近M时（这在训练数百亿参数模型时很常见），气泡比例接近50%，意味着一半的GPU处于空闲状态。尽管Interleaved 1F1B调度【索引16，Megatron-lm: Training multi-billion parameter language models using model parallelism，2020】能缓解此问题（如图3下），但它引入了额外的通信开销且加剧了PP阶段间的不平衡，因此实践中通常将交错因子K设为2，优化空间有限。其他旨在完全消除气泡的工作要么破坏了同步训练语义【索引25，Pipedream: Generalized pipeline parallelism for dnn training，2019，ACM SOSP】，要么依赖不切实际的假设【索引21，Zero bubble pipeline parallelism，2023，arXiv】。

1F1B:

图3：4个流水线阶段和4个微批次的1F1B和Interleaved 1F1B流水线调度时间线。在此示例中，Interleaved 1F1B在每个流水线阶段部署两个模型块。
图3：4个流水线阶段和4个微批次的1F1B和Interleaved 1F1B流水线调度时间线。在此示例中，Interleaved 1F1B在每个流水线阶段部署两个模型块。

Interleaved 1F1B:

图3：4个流水线阶段和4个微批次的1F1B和Interleaved 1F1B流水线调度时间线。在此示例中，Interleaved 1F1B在每个流水线阶段部署两个模型块。
图3：4个流水线阶段和4个微批次的1F1B和Interleaved 1F1B流水线调度时间线。在此示例中，Interleaved 1F1B在每个流水线阶段部署两个模型块。

总结：这些问题源于RLHF的固有特性，难以通过任务级优化解决。现有系统将任务视为最小执行单元，忽略了其内部结构，从而错失了巨大的优化机会。

A2 方法细节

3. RLHFuse 概览

RLHFuse通过将任务分解为更细粒度的子任务来解决上述问题，其架构如图4所示，主要包括两个核心组件：执行计划生成器和运行时执行器。

图4：RLHFuse架构。
图4：RLHFuse架构。

4. 数据感知的阶段间融合

4.1 机遇与挑战：

图5：生成和推理阶段的串行（上）和融合（下）执行时间线。
图5：生成和推理阶段的串行（上）和融合（下）执行时间线。

4.2 融合执行计划：
该计划的目标是在不影响原始生成任务执行时间的前提下，最大化两个阶段的重叠，从而最小化总执行时间。

5. 模型感知的阶段内融合

5.1 机遇与挑战：

图6：(a) Chimera中使用的对称双向流水线调度示例，有4个流水线阶段和4个微批次。(b) 我们的方法可以融合两个分别有4个和2个流水线阶段的不同模型的示例。
图6：(a) Chimera中使用的对称双向流水线调度示例，有4个流水线阶段和4个微批次。(b) 我们的方法可以融合两个分别有4个和2个流水线阶段的不同模型的示例。

5.2 融合流水线调度：
我们提出了一种轻量级高效的算法来自动生成异构模型下的融合流水线调度。

表1：问题定义中的关键符号。
表1：问题定义中的关键符号。

6. 系统实现与优化

系统实现：

执行计划生成器：

运行时执行器：

系统优化：

A4 实验评估

实验环境

模型 参数量 层数 注意力头数 隐藏层维度
LLaMA-13B 13.0B 40 40 5120
LLaMA-33B 32.5B 60 52 6656
LLaMA-65B 65.2B 80 64 8192
表2：LLM规格。
表2：LLM规格。
表2：LLM规格。

实验结果

7.1 端到端性能比较

我们将RLHFuse与以下RLHF训练框架进行比较：

图7：不同生成和模型大小设置下，RLHF训练系统的端到端吞吐量。
图7：不同生成和模型大小设置下，RLHF训练系统的端到端吞吐量。

实验结论 (图7)：

7.2 性能分解分析

图8：不同生成和模型大小设置下的RLHF迭代分解。
图8：不同生成和模型大小设置下的RLHF迭代分解。

实验结论 (图8)：

7.3 消融研究

图9：在不同迁移率和模型设置下，RLHFuse的融合生成和推理阶段执行时间。最大生成长度设置为1024。
图9：在不同迁移率和模型设置下，RLHFuse的融合生成和推理阶段执行时间。最大生成长度设置为1024。

表3：在不同模型、流水线阶段和微批次设置下，由不同算法找到的流水线调度的延迟加速和峰值激活内存成本。
表3：在不同模型、流水线阶段和微批次设置下，由不同算法找到的流水线调度的延迟加速和峰值激活内存成本。

7.4 案例研究

我们展示了一个由RLHFuse为65B/33B模型设置生成的融合流水线调度方案。

这个复杂的调度方案在时间和内存上都达到了最优，手动设计几乎是不可能的，进一步证明了我们提出的融合流水线调度生成算法的有效性。

图10：由RLHFuse生成的65B Actor和33B Critic模型的融合流水线调度。两个模型分别有16和8个PP阶段，微批次数等于PP。上图：GPU执行时间线。下图：峰值激活内存使用时间线。水平虚线表示串行1F1B执行的内存成本。
图10：由RLHFuse生成的65B Actor和33B Critic模型的融合流水线调度。两个模型分别有16和8个PP阶段，微批次数等于PP。上图：GPU执行时间线。下图：峰值激活内存使用时间线。水平虚线表示串行1F1B执行的内存成本。

A5 结论

本文提出了RLHFuse，一个用于LLM的高效RLHF训练框架。通过重新审视RLHF的内在特性和任务依赖，RLHFuse在保留原始同步训练语义的同时，从更细粒度的子任务级别视角来审视工作流。这为高效的阶段间和阶段内融合执行创造了机会，极大地缓解了现有系统中数据倾斜和流水线气泡的问题。评估结果表明，RLHFuse可以显著提高GPU利用率，与当前最先进的系统相比，训练吞吐量最高可提升3.7倍。

方法细节中引用的参考文献汇总

以下是在“方法细节”章节（§3, §4, §5, §6）中引用的参考文献及其在原文中的描述：

  1. 【13】Realhf: Optimized rlhf training for large language models through parameter reallocation (Z. Mei et al., 2024, arXiv)

    • 引用段落：§3 RLHFuse Overview
    • 原文描述："We formalize the training iteration of RLHF workflow as a dataflow graph G(V , E) following previous work [13]."

  2. 【16】Megatron-lm: Training multi-billion parameter language models using model parallelism (M. Shoeybi et al., 2020)

    • 引用段落：§3 RLHFuse Overview, §6 System Implementation and Optimizations
    • 原文描述 (1)："For each vi, RLHFuse will assign it a tailored parallel strategy Pi = (d p, pp,t p) based on the model and cluster configurations to maximize the GPU utilization." (隐式引用，该工作是3D并行的基础)
    • 原文描述 (2)："Optimizing the parallel strategy for LLM training is a well-studied problem [13, 16, 33, 34]."
    • 原文描述 (3)："We implement RLHFuse based on Megatron-LM [16] with 7K lines of code in Python, C++, and CUDA."
    • 原文描述 (4)："Then we follow the guidelines in the Megatron-LM [16] paper to prune the design space and brute-force search the optimal strategy."

  3. 【20】Chimera: efficiently training largescale neural networks with bidirectional pipelines (S. Li and T. Hoefler, 2021, SC)

    • 引用段落：§5.1 Opportunities and Challenges
    • 原文描述："Inspired by the bi-directional pipeline schedule [20] from Chimera, we can break the two training tasks into subtasks of micro-batches and co-locate these subtasks to mutually fill the pipeline bubbles of each other."

  4. 【25】Pipedream: Generalized pipeline parallelism for dnn training (D. Narayanan et al., 2019, ACM SOSP)

    • 引用段落：§6 System Implementation and Optimizations
    • 原文描述："An optimal solution can then be found through optimizations [25, 33, 34]."

  5. 【26】Distserve: Disaggregating prefill and decoding for goodput-optimized large language model serving (Y. Zhong et al., 2024, arXiv)

    • 引用段落：§4.2 Fused Execution Plan
    • 原文描述："During the memorybound decoding phase of generation, as long as the batch size does not exceed a certain threshold that saturates the GPU (denoted as BSmax), the latency of each decoding iteration remains almost constant [26,27]."

  6. 【27】Fast distributed inference serving for large language models (B. Wu et al., 2023, arXiv)

    • 引用段落：§4.2 Fused Execution Plan
    • 原文描述："During the memorybound decoding phase of generation, as long as the batch size does not exceed a certain threshold that saturates the GPU (denoted as BSmax), the latency of each decoding iteration remains almost constant [26,27]."

  7. 【28】Efficient memory management for large language model serving with pagedattention (W. Kwon et al., 2023, SOSP)

    • 引用段落：§4.2 Fused Execution Plan
    • 原文描述："During LLM generation, the key-value cache is maintained for each token position for autoregressive generation. When the sequence length is long, this memory consumption becomes non-negligible [28]."

  8. 【29】Flow shop scheduling: theoretical results, algorithms, and applications (H. Emmons and G. Vairaktarakis, 2012)

    • 引用段落：§5.2 Fused Pipeline Schedule
    • 原文描述："It can be reduced to a special case of flow job scheduling problem which proves to be NP-hard [29]."

  9. 【30】Optimization by simulated annealing (S. Kirkpatrick et al., 1983, science)

    • 引用段落：§5.2 Fused Pipeline Schedule
    • 原文描述："Instead, we use simulated annealing [30] to search for a better solution."

  10. 【33】Alpa: Automating inter- and Intra-Operator parallelism for distributed deep learning (L. Zheng et al., 2022, USENIX OSDI)

    • 引用段落：§6 System Implementation and Optimizations
    • 原文描述 (1)："Optimizing the parallel strategy for LLM training is a well-studied problem [13, 16, 33, 34]."
    • 原文描述 (2)："An optimal solution can then be found through optimizations [25, 33, 34]."

  11. 【34】Galvatron: Efficient transformer training over multiple gpus using automatic parallelism (X. Miao et al., 2022, VLDB Endowment)

    • 引用段落：§6 System Implementation and Optimizations
    • 原文描述 (1)："Optimizing the parallel strategy for LLM training is a well-studied problem [13, 16, 33, 34]."
    • 原文描述 (2)："An optimal solution can then be found through optimizations [25, 33, 34]."

  12. 【35】Open mpi: Goals, concept, and design of a next generation mpi implementation (E. Gabriel et al., 2004, Recent Advances in Parallel Virtual Machine and Message Passing Interface)

    • 引用段落：§6 System Implementation and Optimizations
    • 原文描述："For intra-stage fusion, we use MPI [35] to parallelize the simulated annealing with different random seeds."

  13. 【36】Optimized primitives for collective multi-gpu communication (NVIDIA, 2024)

    • 引用段落：§6 System Implementation and Optimizations
    • 原文描述："After finding the optimal solution, it will generate the sequences of NCCL [36] operations on each device following the schedule."

  14. 【37】Orca: A distributed serving system for {Transformer-Based} generative models (G.-I. Yu et al., 2022, USENIX OSDI)

    • 引用段落：§6 System Implementation and Optimizations
    • 原文描述："It supports continuous batching [37] to allow the short samples to leave the batch early."

  15. 【38】Sglang: Efficient execution of structured language model programs (L. Zheng et al., 2024)

    • 引用段落：§6 System Implementation and Optimizations
    • 原文描述："For samples that share the same prompt, it supports prefix sharing [38] to avoid key-value cache duplication."

  16. 【39】High-dimensional continuous control using generalized advantage estimation (J. Schulman et al., 2018)

    • 引用段落：§6 System Implementation and Optimizations
    • 原文描述："During the post-processing of the inference stage, one metric called Generalized Advantage Estimation [39] is needed which requires recursive computation along the output length dimension for each sample."

  17. 【40】On optimizing the communication of model parallelism (Y. Zhuang et al., 2024)

    • 引用段落：§6 System Implementation and Optimizations
    • 原文描述："This communication pattern is well-studied [40] and we adopt a similar allto-all approach to minimize cross-node communication."