ReaL: Efficient RLHF Training of Large Language Models with Parameter Reallocation

📄 arXiv: 2406.14088v2 📥 PDF

作者: Zhiyu Mei, Wei Fu, Kaiwei Li, Guangju Wang, Huanchen Zhang, Yi Wu

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

发布日期: 2024-06-20 (更新: 2025-04-24)

备注: 11 pages (20 pages with references and the appendix), 17 figures. Accepted by MLSys 25

🔗 代码/项目: GITHUB

💡 一句话要点

提出ReaL，通过参数重分配高效训练大型语言模型的RLHF。

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

关键词: 强化学习 人类反馈 大型语言模型 参数重分配 并行化训练

📋 核心要点

  1. 现有RLHF训练方法直接沿用LLM监督训练的并行策略，无法有效应对RLHF训练中复杂的工作负载和依赖关系，导致效率低下。
  2. 论文提出参数重分配技术，通过动态调整训练集群中LLM参数的分配方式，优化不同工作负载的并行化策略，提升训练效率。
  3. 实验结果表明，ReaL在LLaMA模型上实现了高达3.58倍的加速，并且在长上下文场景中，性能优于基于Megatron-LM的启发式方法。

📝 摘要（中文）

本文提出了一种名为参数重分配（parameter ReaLlocation）的新技术，旨在提升大型语言模型（LLM）的RLHF训练效率。与LLM的监督训练相比，RLHF训练更为复杂，涉及多种计算工作负载以及LLM实例间的复杂依赖关系。简单地沿用监督训练的并行化策略不足以满足RLHF的需求，导致训练效率低下。ReaL通过在训练集群中动态地重新分配LLM参数，从而为不同的工作负载调整并行化策略。ReaL引入了执行计划的概念，该计划定义了专门为RLHF训练设计的细粒度资源分配和并行化策略。ReaL采用定制的搜索算法和轻量级运行时估计器，自动为RLHF实验实例发现高效的执行计划。运行时引擎通过有效地并行计算和重新分配参数来部署所选计划。在高达700亿参数和128个GPU的LLaMA模型上的实验结果表明，与基线方法相比，ReaL实现了高达3.58倍的加速。此外，在长上下文场景中，ReaL生成的执行计划比基于Megatron-LM的启发式方法平均提高了81%的性能。ReaL的源代码已公开。

🔬 方法详解

问题定义：现有RLHF训练方法，特别是针对大型语言模型（LLM）的训练，直接沿用监督学习的并行化策略。然而，RLHF训练涉及多个LLM实例之间的复杂依赖关系和多样化的计算负载，固定的并行化策略无法充分利用计算资源，导致训练效率低下。现有方法的痛点在于无法根据RLHF训练的动态特性进行自适应的资源分配和并行化。

核心思路：论文的核心思路是参数重分配（Parameter ReaLlocation），即在训练过程中动态地调整LLM参数在训练集群中的分布。通过这种方式，可以根据当前的工作负载和计算需求，优化并行化策略，从而提高整体训练效率。这种动态调整使得系统能够更好地适应RLHF训练的复杂性和变化性。

技术框架：ReaL系统的整体框架包括以下几个主要模块：1) 执行计划生成器：该模块负责搜索并生成高效的执行计划，该计划定义了细粒度的资源分配和并行化策略。它使用定制的搜索算法和轻量级运行时估计器来评估不同计划的性能。2) 运行时引擎：该模块负责部署选定的执行计划，包括并行化计算和重新分配参数。它有效地管理计算资源，并确保按照计划执行训练过程。3) 参数重分配机制：该机制负责在训练集群中动态地重新分配LLM参数，以适应不同的工作负载。

关键创新：ReaL的关键创新在于其动态参数重分配机制和执行计划的概念。与传统的静态并行化策略不同，ReaL能够根据RLHF训练的实际需求，自适应地调整资源分配和并行化策略。执行计划的引入使得系统能够以细粒度的方式控制资源的使用，从而实现更高的效率。此外，轻量级的运行时估计器能够快速评估不同执行计划的性能，加速搜索过程。

关键设计：ReaL的关键设计包括：1) 执行计划的表示：执行计划定义了每个计算任务的资源分配和并行化策略，例如使用多少个GPU、如何划分数据等。2) 搜索算法：该算法用于在可能的执行计划空间中搜索最优解。论文可能采用了诸如遗传算法、强化学习等方法来优化搜索过程。3) 运行时估计器：该估计器用于预测不同执行计划的性能，以便搜索算法能够快速评估候选解。4) 参数重分配策略：该策略定义了如何将LLM参数在训练集群中重新分配，以适应不同的计算需求。具体的参数设置、损失函数和网络结构等细节取决于所使用的LLM模型和RLHF算法。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ReaL在LLaMA模型（高达700亿参数）和128个GPU上的实验结果显示，与基线方法相比，ReaL实现了高达3.58倍的加速。在长上下文场景中，ReaL生成的执行计划比基于Megatron-LM的启发式方法平均提高了81%的性能。这些结果表明ReaL能够显著提升RLHF训练的效率，尤其是在处理大规模模型和复杂任务时。

🎯 应用场景

ReaL技术可广泛应用于各种需要通过人类反馈进行强化学习的大型语言模型训练场景，例如对话系统、文本生成、代码生成等。通过提高训练效率，ReaL能够降低训练成本，加速模型迭代，并促进更复杂、更智能的LLM应用的发展。该技术对于需要处理长上下文或复杂交互的LLM应用具有重要意义。

📄 摘要（原文）

Reinforcement Learning from Human Feedback (RLHF) is a pivotal technique for empowering large language model (LLM) applications. Compared with the supervised training process of LLMs, the RLHF training process is much more sophisticated, requiring a diverse range of computation workloads with intricate dependencies between multiple LLM instances. Therefore, simply adopting the fixed parallelization strategies from supervised training for LLMs can be insufficient for RLHF and result in low training efficiency. To overcome this limitation, we propose a novel technique named parameter ReaLlocation, which dynamically adapts the parallelization strategies for different workloads during training by redistributing LLM parameters across the training cluster. Building upon this idea, we introduce ReaL, a pioneering system for efficient RLHF training. ReaL introduces the concept of an execution plan, which defines a fine-grained resource allocation and parallelization strategy particularly designed for RLHF training. Based on this concept, ReaL employs a tailored search algorithm with a lightweight run-time estimator to automatically discover an efficient execution plan for an instance of RLHF experiment. Subsequently, the runtime engine deploys the selected plan by effectively parallelizing computations and redistributing parameters. We evaluate ReaL on the LLaMA models with up to 70 billion parameters and 128 GPUs. The experimental results demonstrate that ReaL achieves speedups of up to $3.58\times$ compared to baseline methods. Furthermore, the execution plans generated by ReaL exhibit an average of $81\%$ performance improvement over heuristic approaches based on Megatron-LM in the long-context scenario. The source code of ReaL is publicly available at https://github.com/openpsi-project/ReaLHF .