背景
豆包作为对标GPT的竞品出现后,模型的推理能力也一直在提升,强化学习(RL)对大模型复杂推理能力提升有关键作用,但其复杂的计算流程对训练和部署也带来了巨大挑战。近日,字节跳动豆包大模型团队与香港大学联合提出 HybridFlow。这是一个灵活高效的 RL/RLHF 框架,可显著提升训练吞吐量,降低开发和维护复杂度。实验结果表明,HybridFlow 在各种模型规模和 RL 算法下,训练吞吐量相比其他框架提升了 1.5 倍至 20 倍。 这篇文章也在10月入选顶会,还蛮新的,于是今天就拿来粗读一下。原文地址如下 HybridFlow: A Flexible and Efficient RLHF Framework 下面记录一下阅读笔记,如有错误欢迎指正!
摘要
这部分主要是背景问题和贡献了。 问题是RLHF被广泛用于LLM的对齐,但传统的强化学习框架无法高效处理RLHF的复杂数据流,现有的RLHF系统采用多控制器范式,但由于嵌套的分布式计算和数据通信,其灵活性受到限制,于是团队想提升灵活性。 这篇文章主要贡献是 提出HybridFlow:结合单控制器和多控制器范式,通过分层API解耦计算和数据依赖。 3D-HybridEngine:实现训练和生成阶段之间零冗余的模型参数重分配。
第一部分:Introduction(引言)
主要将来大模型的训练流程是预训练加上SFT,也就是通过海量的文本数据(如书籍、网页)进行下一个词预测,从而积累广泛的知识,再利用特定领域的数据进一步训练,使模型能够更好地理解和执行人类指令。为了避免因为训练数据的偏差,生成有害或有偏见的内容,引入RLHF技术,通过人类偏好来对齐模型行为,让AI安全可靠。 RLHF通常基于传统的强化学习算法,如近端策略优化(PPO)和REINFORCE, 它涉及多个模型:Actor(生成模型)、Critic(价值模型)、参考策略模型和奖励模型。传统强化学习可以直接建模为数据流DAG这种,但是RLHF中的数据流情况比较复杂,每个节点不仅是一个简单的神经网络,而是一个分布式训练或生成程序,每条边代表多对多的模型分片间的数据传输,这个图就会很复杂,现有RLHF框架采用多控制器管理分布式计算和数据传输,但灵活性不足,应该是修改一个节点就会会影响整个数据流,难以适配新的RLHF算法,所以整了一个新框架叫HybridFlow,为了结合单控制器和多控制器范式,人话就是单控制器管数据流全局,多控制器管分布式内部。然后3D-HybridEngine 实现训练和生成阶段之间零冗余的模型参数重分配。
第二部分:Background and Motivation(背景与动机)
又是一段RLHF介绍流程布拉布拉跳过了,这个架构主要问题是异构模型规模,例如一个小型的 Actor 模型(如7B参数)可能配合大型的 Critic 和 Reward 模型(如70B参数),大量GPU并行用自回归又会低效,而且通信开销高,现有框架(如 DeepSpeed-Chat 和 OpenRLHF)切换时内存冗余和算法适配难得问题。新提出HybridFlow,使用自动设备映射算法,根据任务负载动态调整模型的设备分配,解耦数据流描述和计算逻辑,使得用户可以灵活调整数据流结构,而无需修改底层实现,引入3D-HybridEngine,保证吞吐量的同时,降低大规模模型的训练与推理成本。
第三部分:HybridFlow Overview
三个核心混合编程模型(Hybrid Programming Model),分层 API 设计(Layered API Design),3D-HybridEngine上面说过什么用了不废话直接看流程 首先生成阶段 用户输入prompt,Actor 模型生成通过自回归生成来响应,这块整了一个优化改成并行,或者流水线并行分区减少内存压力,Actor 的生成输出被发送到 Critic 和 Reward 模型进行评估,这个输出也是批量方式,这个阶段Actor 模型在生成过程中高度依赖内存,计算负载相对较轻, Critic 和 Reward 模型处于推理状态,仅执行前向传播。 其次到了评估阶段 Critic 模型接收 Actor 的生成响应,并根据训练好的价值函数计算每个响应的价值分数,Reward打分,生成奖励信号,共同作为输入优化Actor和Critic,Critic 和 Reward 的输出被统一整合,并通过 HybridFlow 的分层 API,推理也是并行动态,确保GPU的利用率 最后训练阶段 更新 Actor 和 Critic 模型的参数这些都是正常流程,在 Critic 的梯度计算和更新中,3D-HybridEngine 提供高效的通信优化,减少梯度聚合的延迟,HybridFlow 在训练阶段切换到训练优化的并行策略,减少 Actor 和 Critic 的资源竞争
第四,五部分Hybrid Programming Model(混合编程模型)和3D-HybridEngine
提供了一种灵活、高效且易扩展的框架设计,封装复杂的分布式计算逻辑,层次化 API等用户无需关心设备映射和通信协议即可完成任务编排然后团队在减少开销方面下了不少功夫,具体实现可以见原文。 HybridFlow 支持 3D 并行策略:
- 数据并行(Data Parallelism):在多个设备上分批处理数据。
- 模型并行(Model Parallelism):将模型参数分片到多个设备。
- 流水线并行(Pipeline Parallelism):在多个设备间流水线执行任务,以显著减少内存峰值,引入微型批次(micro-batching),平衡设备间的负载。 与传统框架(如 DeepSpeed-Chat)相比,3D-HybridEngine 在参数重分配上减少了 30%-50% 的内存占用。在大规模模型训练中,通信延迟降低了 20%-40%,在运行多种 RLHF 算法时,3D-HybridEngine 提升了系统吞吐 最低提升 1.53 倍,最高提升 20.57 倍。
后面的部分
接着就是数据对比了, 在硬件扩展实验中,HybridFlow 在 8 到 128 块 GPU 的集群规模下,表现出近乎线性的扩展性 谈论的优化的设备分配策略,小型集群<64 GPUs推荐使用“colocate”策略,将模型放在同一设备上,大型集群,推荐“split”策略,分别分配 Actor 和 Critic 到不同设备,以实现更高的并行性。 讨论部分总结了 HybridFlow 的优势和局限,并提出了未来优化的方向。其高灵活性和高性能在 RLHF 中表现出色,但仍需在用户友好性、硬件适配和大规模任务中的优化上进一步探索。