Optimizing RLHF Training for Large Language Models with Stage Fusion

📄 arXiv: 2409.13221v3 📥 PDF

作者: Yinmin Zhong, Zili Zhang, Bingyang Wu, Shengyu Liu, Yukun Chen, Changyi Wan, Hanpeng Hu, Lei Xia, Ranchen Ming, Yibo Zhu, Xin Jin

分类: cs.LG, cs.CL, cs.DC

发布日期: 2024-09-20 (更新: 2025-04-22)

💡 一句话要点

RLHFuse：通过阶段融合优化大型语言模型的RLHF训练

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: RLHF 大型语言模型 阶段融合 GPU利用率 训练优化

📋 核心要点

  1. 现有RLHF系统受限于生成阶段的数据倾斜和训练阶段的流水线气泡，导致GPU利用率低下。
  2. RLHFuse将RLHF流程分解为细粒度子任务，通过阶段间和阶段内融合，优化GPU利用率。
  3. 实验结果表明，RLHFuse相比现有系统，训练吞吐量最高提升了3.7倍。

📝 摘要（中文）

本文提出了一种高效的RLHF训练系统RLHFuse，该系统采用阶段融合的方法。由于RLHF训练的固有特性，即生成阶段的数据倾斜和训练阶段的流水线气泡，现有的RLHF系统通常存在GPU利用率低的问题。RLHFuse打破了将RLHF工作流程视为独立任务组合的传统观念，将每个任务分解为更细粒度的子任务，并执行阶段融合以提高GPU利用率。RLHFuse包含两个关键思想。首先，对于生成和推理任务，RLHFuse将其分解为样本级子任务，从而实现高效的阶段间融合，以重叠生成和推理阶段的执行，从而缓解由长尾样本主导的原始生成瓶颈。其次，对于训练任务，RLHFuse将其分解为微批次的子任务，并执行阶段内融合，以使用融合的流水线调度并发执行训练阶段中的这些子任务，从而有效地缓解流水线气泡。实验表明，与现有系统相比，RLHFuse的训练吞吐量提高了高达3.7倍。

🔬 方法详解

问题定义：现有RLHF训练系统在处理大型语言模型时，由于生成阶段存在数据倾斜（即某些样本生成时间远长于其他样本），以及训练阶段存在流水线气泡（即GPU在某些阶段空闲），导致整体GPU利用率低下，训练效率不高。这限制了RLHF在更大规模模型上的应用。

核心思路：RLHFuse的核心思路是通过将RLHF流程中的各个阶段（生成、推理、训练）分解为更细粒度的子任务，然后通过巧妙的调度和融合这些子任务，使得GPU资源能够得到更充分的利用，从而提高整体的训练吞吐量。关键在于打破了传统RLHF流程的串行执行模式。

技术框架：RLHFuse的整体框架包含生成阶段、推理阶段和训练阶段。在生成阶段，模型生成多个候选回复；在推理阶段，奖励模型对这些回复进行打分；在训练阶段，利用这些打分来更新语言模型。RLHFuse的关键改进在于，将生成和推理阶段分解为样本级别的子任务，并进行阶段间融合，以及将训练阶段分解为微批次级别的子任务，并进行阶段内融合。

关键创新：RLHFuse的关键创新在于其阶段融合的思想。传统的RLHF系统将各个阶段视为独立的任务，串行执行。而RLHFuse通过将阶段分解为更小的子任务，并允许这些子任务在不同阶段之间重叠执行，从而隐藏了延迟，提高了GPU利用率。这种阶段融合的思想是与现有方法的本质区别。

关键设计：在生成和推理阶段的融合中，RLHFuse采用了一种基于样本的调度策略，允许奖励模型在生成阶段仍在进行时就开始对已生成的样本进行打分。在训练阶段的融合中，RLHFuse采用了一种融合的流水线调度策略，允许不同的微批次在不同的GPU上并行执行，从而减少了流水线气泡。具体的参数设置和损失函数等细节，论文中未明确说明，属于未知内容。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，RLHFuse相比于现有的RLHF系统，能够显著提高训练吞吐量，最高可达3.7倍。这一提升表明RLHFuse能够有效地解决RLHF训练中的瓶颈问题，提高GPU利用率，从而加速大型语言模型的训练过程。

🎯 应用场景

RLHFuse的潜在应用领域包括所有需要使用RLHF进行训练的大型语言模型，例如对话系统、文本生成模型等。通过提高训练效率，RLHFuse可以降低训练成本，加速模型迭代，并使得更大规模的语言模型能够通过RLHF进行优化，从而提升模型的性能和安全性。

📄 摘要（原文）

We present RLHFuse, an efficient training system with stage fusion for Reinforcement Learning from Human Feedback (RLHF). Due to the intrinsic nature of RLHF training, i.e., the data skewness in the generation stage and the pipeline bubbles in the training stage, existing RLHF systems suffer from low GPU utilization. RLHFuse breaks the traditional view of RLHF workflow as a composition of individual tasks, splitting each task into finer-grained subtasks, and performing stage fusion to improve GPU utilization. RLHFuse contains two key ideas. First, for generation and inference tasks, RLHFuse splits them into sample-level subtasks, enabling efficient inter-stage fusion to overlap the execution of generation and inference stages, thus mitigating the original generation bottleneck dominated by long-tailed samples. Second, for training tasks, RLHFuse breaks them into subtasks of micro-batches and performs intra-stage fusion to concurrently execute these subtasks in the training stage with a fused pipeline schedule, effectively mitigating the pipeline bubbles. The experiments show that RLHFuse increases the training throughput by up to $3.7\times$, compared to existing systems.