HybridFlow: A Flexible and Efficient RLHF Framework
TL;DR 精炼摘要
HybridFlow提出了一种结合单控制器与多控制器的混合框架,以解决现有基于人类反馈的强化学习(RLHF)系统中效率低下和灵活性不足的问题。通过层次化API解耦计算与数据依赖并设计3D-HybridEngine,该框架在训练与生成阶段实现模型权重的零冗余重分区。实验结果显示,HybridFlow的吞吐量提升了1.53至20.57倍,显著优于现有系统。
摘要
Reinforcement Learning from Human Feedback (RLHF) is widely used in Large Language Model (LLM) alignment. Traditional RL can be modeled as a dataflow, where each node represents computation of a neural network (NN) and each edge denotes data dependencies between the NNs. RLHF complicates the dataflow by expanding each node into a distributed LLM training or generation program, and each edge into a many-to-many multicast. Traditional RL frameworks execute the dataflow using a single controller to instruct both intra-node computation and inter-node communication, which can be inefficient in RLHF due to large control dispatch overhead for distributed intra-node computation. Existing RLHF systems adopt a multi-controller paradigm, which can be inflexible due to nesting distributed computation and data communication. We propose HybridFlow, which combines single-controller and multi-controller paradigms in a hybrid manner to enable flexible representation and efficient execution of the RLHF dataflow. We carefully design a set of hierarchical APIs that decouple and encapsulate computation and data dependencies in the complex RLHF dataflow, allowing efficient operation orchestration to implement RLHF algorithms and flexible mapping of the computation onto various devices. We further design a 3D-HybridEngine for efficient actor model resharding between training and generation phases, with zero memory redundancy and significantly reduced communication overhead. Our experimental results demonstrate 1.53~20.57 throughput improvement when running various RLHF algorithms using HybridFlow, as compared with state-of-the-art baselines. HybridFlow source code will be available at https://github.com/volcengine/verl.
思维导图
论文精读
中文精读
1. 论文基本信息
1.1. 标题
HybridFlow: A Flexible and Efficient RLHF Framework (HybridFlow:一个灵活且高效的 RLHF 框架)
1.2. 作者
Guangming Sheng (香港大学), Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin (均为字节跳动 ByteDance), Chuan Wu (香港大学)。
1.3. 发表期刊/会议
发表于 EuroSys '25 (Twentieth European Conference on Computer Systems),这是系统架构、分布式计算和操作系统领域的顶级国际会议,具有极高的学术声誉。
1.4. 发表年份
2024年(预印本发布),正式发表于 2025年。
1.5. 摘要
基于人类反馈的强化学习 (Reinforcement Learning from Human Feedback, RLHF) 是大型语言模型 (LLM) 对齐的关键技术。然而,现有的 RLHF 系统在执行复杂的分布式数据流时,要么面临由于单控制器模式导致的调度开销过大,要么面临由于多控制器模式导致的开发灵活性不足。本文提出了 HybridFlow,一个结合了单控制器 (Single-controller) 和多控制器 (Multi-controller) 范式的混合框架。该框架通过层次化 API 解耦了计算与数据依赖,设计了 3D-HybridEngine 以实现训练与生成阶段之间模型权重的零冗余重分区 (Resharding),并提出了自动设备映射算法。实验表明,HybridFlow 的吞吐量比当前最先进的基线系统提升了 。
1.6. 原文链接
-
arXiv 链接: https://arxiv.org/abs/2409.19256
-
开源代码: https://github.com/volcengine/verl
2. 整体概括
2.1. 研究背景与动机
大语言模型 (LLM) 在预训练和微调后,虽然具备了强大的知识和指令遵循能力,但仍可能生成有毒或有害的内容。基于人类反馈的强化学习 (RLHF) 旨在将模型与人类价值观对齐。
然而,构建高效的 RLHF 系统面临以下核心挑战:
- 复杂的数据流: 典型的 RLHF 算法(如 PPO)涉及四个模型:演员 (Actor)、评论家 (Critic)、参考模型 (Reference) 和奖励模型 (Reward)。这些模型之间存在复杂的跨设备多对多通信。
- 效率与灵活性的矛盾:
- 单控制器范式 (Single-controller): 易于编程,但在分布式环境下,控制器向成千上万个加速器发送指令的开销极高。
- 多控制器范式 (Multi-controller): 运行效率高,但计算与通信逻辑深度耦合,导致开发新算法时需要修改底层通信代码,极不灵活。
- 训练与生成的工作负载不均衡: 演员模型在“生成”阶段是内存受限的,而在“训练”阶段是计算受限的。在不同阶段使用不同的并行策略可以优化性能,但频繁的模型权重重分区 (Resharding) 会带来巨大的通信和显存开销。
2.2. 核心贡献/主要发现
-
混合编程模型: 提出了在节点间使用单控制器协调数据流、在节点内使用多控制器执行分布式计算的混合模型,实现了灵活性与效率的兼顾。
-
3D-HybridEngine: 专门针对演员模型设计,支持在训练和生成阶段采用不同的 3D 并行配置,并实现了零冗余 (Zero-redundancy) 的权重切换。
-
自动设备映射 (Auto Device Mapping): 提供了一个优化算法,能够根据模型大小和硬件环境,自动决定模型在 GPU 上的放置位置和并行策略。
-
卓越的性能: HybridFlow 显著提升了 RLHF 的训练吞吐量,并支持多种 RLHF 变体算法(如 PPO、ReMax、Safe-RLHF)。
3. 预备知识与相关工作
3.1. 基础概念
- RLHF (Reinforcement Learning from Human Feedback): 利用人类的偏好数据训练一个奖励模型,再通过强化学习算法(如 PPO)优化语言模型的生成质量。
- PPO (Proximal Policy Optimization): 一种稳健的强化学习算法。在 RLHF 中,它通过对比“当前模型”和“参考模型”的输出差异,并在“奖励模型”的指导下更新“演员模型”。
- 3D 并行 (3D Parallelism): 结合了三种分布式技术:
- 数据并行 (Data Parallelism, DP): 将不同数据分给不同 GPU 处理。
- 张量并行 (Tensor Parallelism, TP): 将单个算子(如矩阵乘法)拆分到多个 GPU 上。
- 流水线并行 (Pipeline Parallelism, PP): 将模型的不同层放置在不同 GPU 上。
3.2. 前人工作
- 分布式训练框架: 如 Megatron-LM 和 DeepSpeed。它们采用多控制器范式,擅长单模型的超大规模训练,但在处理 RLHF 这种多模型交互的复杂场景时,开发者需要手动编写大量的点对点通信代码。
- 现有 RLHF 系统:
- DeepSpeed-Chat: 强制所有模型共用相同的并行策略,导致资源利用率低。
- OpenRLHF: 支持多模型分离部署,但由于模型重分区导致的通信开销巨大。
3.3. 差异化分析
HybridFlow 的核心创新在于它解耦了节点内的分布式计算和节点间的数据重分区。它允许每个模型独立选择最适合其工作负载的并行策略(TP, PP, DP 的组合),并通过统一的传输协议自动处理数据交换。
4. 方法论
4.1. 方法原理
HybridFlow 的核心思想是“层次化解耦”。它将 RLHF 的执行逻辑分为两层:
-
节点间 (Inter-node): 使用单控制器(基于 Ray 调度),通过 Python 脚本直观地描述模型间的逻辑依赖。
-
节点内 (Intra-node): 使用多控制器,每个 GPU 上的进程独立执行高性能的张量计算。
下图(原文 Figure 4)展示了 HybridFlow 的整体架构:
该图像是HybridFlow架构示意图。图中展示了RLHF数据流图、模型配置和设备配置的结构,包括用户输入、ParallelWorker、资源池及物理设备等层次。特别描述了3D-HybridEngine和自动映射的方法,以实现高效的RLHF模型训练与生成。
4.2. 核心方法详解
4.2.1. 层次化 API 与传输协议
HybridFlow 提供了一套 3DParallelWorker 类,封装了复杂的分布式初始化。对于开发者而言,只需要调用 generate_sequences 或 update_actor 等 API 即可。
传输协议 (Transfer Protocols):
为了处理模型间不一致的并行状态,作者设计了 collect (收集) 和 distribute (分发) 函数。
例如,当演员模型(使用 TP=2, DP=4)生成的响应需要传给评论家模型(使用 TP=4, DP=2)时,框架会根据注册的协议自动计算每个 GPU 需要发送和接收的数据块,而不需要用户手动编写 send/recv。
下图(原文 Figure 5)展示了这种 API 的调用逻辑和异步数据重分区过程:
该图像是示意图,展示了层次API的结构与功能。(a) 显示了Actor模型的初始化及资源池分配。(b) 说明在单控制器模式下的异步数据重分配过程,涉及多个数据传输步骤。
4.2.2. 3D-HybridEngine 与零冗余重分区
这是本文的核心技术创新。为了解决演员模型在训练和生成阶段并行策略不匹配的问题,3D-HybridEngine 提出了一种新的并行分组方法。
数学逻辑分析:
假设分配给演员模型的 GPU 总数为 。
在训练阶段,配置为 (p, t, d),即 。
在生成阶段,为了提高效率,通常需要更小的 和更大的 。
配置变为 ,其中 是微数据并行 (Micro DP) 的倍数。
零冗余重分区 (Zero-redundancy Resharding): 传统的重分区需要在每个 GPU 上全量收集模型参数,然后再按新策略切分,这会导致显存峰值达到模型全量大小 。 HybridFlow 通过重新排列 GPU 秩 (Rank),使得生成阶段的“微 DP 组”正是训练阶段“模型并行组”的一个子集。
这种方法使得在切换时:
-
通信量减少: 仅在微 DP 组内进行
all-gather。通信量从 降低到 。 -
显存零冗余: 每一块 GPU 只存储其所需的参数分片,显存峰值仅为生成所需的参数量。
下图(原文 Figure 8)展示了这种优化的权重重分区方式:
该图像是示意图,展示了在训练和生成阶段之间模型权重的重分配。图中包含两种分组方法,分别为 HybridFlow-V 和优化的并行分组方法 HybridFlow,重点展示了所有权重的收集及冗余训练权重的处理。整体过程通过不同的组进行信息交换,强调高效的计算和数据依赖关系。
4.2.3. 自动设备映射算法
为了最小化单次迭代时间 ,算法会在所有可能的放置方案 中搜索。
核心代价函数 (Algorithm 1 中的 d_cost):
其中:
-
stage表示 RLHF 的三个阶段(生成、准备、训练)。 -
simu是一个模拟器,用于评估特定并行策略下的延迟。 -
如果模型被放置在同一个设备集(
colocated),它们的延迟是累加的;如果放置在不同设备集,则取最大值(并行执行)。
5. 实验设置
5.1. 数据集
实验使用了 Dahoas/full-hh-rlhf 数据集,这是大模型对齐领域最常用的偏好数据集之一。
- 样本示例: 提示 (Prompt) 为一个问题,响应 (Response) 为模型生成的回答。
- 设置: 提示长度 1024,生成长度 1024,全局批大小 (Global Batch Size) 为 1024。
5.2. 评估指标
- RLHF 吞吐量 (Throughput):
- 概念定义: 单位时间内系统处理的词元总数。
- 数学公式:
Throughput = \frac{N_{prompts} \times (L_{prompt} + L_{response})}{T_{iteration}} - 符号解释: 是批大小, 是序列长度, 是单次迭代耗时。
- 转换时间 (Transition Time): 指模型在训练和生成状态之间切换权重所需的秒数。
- 扩展性 (Scalability): 评估随着 GPU 数量增加,吞吐量的增长比例。
5.3. 对比基线
-
DeepSpeed-Chat: 微软开发的 RLHF 框架,特点是模型强耦合。
-
OpenRLHF: 基于 Ray 的高性能框架,支持模型分离部署。
-
NeMo-Aligner: NVIDIA 开发的对齐工具包,专注于 3D 并行优化。
6. 实验结果与分析
6.1. 核心结果分析
HybridFlow 在各种算法下均表现出显著的性能优势。
以下是原文提供的 PPO 和 ReMax 算法 下的吞吐量加速比总结(基于 Figure 9):
- 在 70B 模型规模下,HybridFlow 的加速比最为突出,平均达到基线系统的 9.64倍。
- 相比于 NeMo-Aligner,加速比最高达到了 20.57倍,这主要归功于 HybridFlow 在生成阶段对 KV 缓存 (KVCache) 的高效管理和优化的并行策略。
6.2. 数据呈现 (表格)
以下是原文 Table 1 关于各系统并行策略和执行模式的对比:
| RLHF 系统 | DeepSpeed-Chat | OpenRLHF | NeMo-Aligner | HybridFlow |
|---|---|---|---|---|
| 并行策略 | 训练: ZeRO; 生成: TP | 训练与生成均支持 3D 并行 | ||
| 演员模型权重 | 从 ZeRO 到 TP 重分区 | 使用两份物理拷贝 (训练/生成各一份) | 训练与生成使用完全相同的分区 | 3D 并行切换,零冗余权重复用 |
| 模型放置方案 | 所有模型共用相同设备 (Colocate) | 每个模型放置在独立设备 (Standalone) | Actor/Ref 共用; Critic/RM 共用 (Split) | 支持各种灵活映射方案 |
| 执行模式 | 顺序执行 | 部分并行 | 部分并行 | 支持各种灵活模式 |
6.3. 权重转换时间分析
从 Figure 14 可以看出,HybridFlow 在 70B 模型上的权重转换时间比 OpenRLHF 减少了 89.1%。这直接验证了 3D-HybridEngine 减少通信冗余的有效性。
7. 总结与思考
7.1. 结论总结
HybridFlow 成功解决了 RLHF 系统在灵活性和效率之间的权衡难题。通过引入混合编程模型和 3D-HybridEngine,它不仅让研究人员能够通过几行代码实现复杂的算法(如 Safe-RLHF),还在大规模集群上实现了极高的吞吐量,为 LLM 的对齐训练提供了强有力的基础设施支持。
7.2. 局限性与未来工作
- 异构资源支持: 目前算法主要针对同构 GPU 集群。未来可以探索如何在拥有不同显存大小或计算能力的 GPU 集群中进行自动映射。
- 容错性: 虽然支持检查点 (Checkpointing),但在超大规模(万卡)级别下的自动故障恢复仍有改进空间。
7.3. 个人启发与批判
- 启发: 混合控制的思想(节点间单控制,节点内多控制)不仅适用于 RLHF,也可以推广到其他涉及多模型协作的复杂 AI 任务,如多模态模型的训练或大模型智能体 (Agent) 的编排。
- 批判: 论文虽然在吞吐量上取得了巨大提升,但并未详细讨论在“生成”阶段引入 KV 缓存重分区对长文本生成稳定性的影响。此外,自动映射算法的复杂度在模型数量极多时可能会显著增加。
相似论文推荐
基于向量语义检索推荐的相关论文。