【大模型训练】Hybrid FLow verl
基于人类反馈的强化学习(RLHF)被广泛用于大型语言模型(LLM)的对齐。传统的强化学习可以被建模为一个数据流,其中每个节点代表一个神经网络(NN)的计算,每条边表示NN之间的数据依赖关系。RLHF通过将每个节点扩展为一个分布式的LLM训练或生成程序,并将每条边扩展为多对多的多播(multicast),从而使数据流变得更加复杂。传统的RL框架使用单个控制器来指导节点内计算和节点间通信以执行数据流,
好的,这是对您提供的论文《HybridFlow: A Flexible and Efficient RLHF Framework》的全文翻译。
HybridFlow: 一个灵活高效的RLHF框架
作者信息:
- Guangming Sheng (香港大学, gmsheng@connect.hku.hk)
- Chi Zhang (字节跳动, zhangchi.usc1992@bytedance.com)
- Zilingfeng Ye (字节跳动, yezilingfeng@bytedance.com)
- Xibin Wu (字节跳动, wuxibin@bytedance.com)
- Wang Zhang (字节跳动, zhangwang.nozomi@bytedance.com)
- Ru Zhang (字节跳动, zhangru.1994@bytedance.com)
- Yanghua Peng (字节跳动, pengyanghua.yanghua@bytedance.com)
- Haibin Lin (字节跳动, haibin.lin@bytedance.com)
- Chuan Wu (香港大学, cwu@cs.hku.hk)
摘要
基于人类反馈的强化学习(RLHF)被广泛用于大型语言模型(LLM)的对齐。传统的强化学习可以被建模为一个数据流,其中每个节点代表一个神经网络(NN)的计算,每条边表示NN之间的数据依赖关系。RLHF通过将每个节点扩展为一个分布式的LLM训练或生成程序,并将每条边扩展为多对多的多播(multicast),从而使数据流变得更加复杂。传统的RL框架使用单个控制器来指导节点内计算和节点间通信以执行数据流,这在RLHF中可能效率低下,因为分布式节点内计算的控制分发开销巨大。现有的RLHF系统采用多控制器范式,但由于其嵌套了分布式计算和数据通信,可能不够灵活。
我们提出了HybridFlow,它以一种混合的方式结合了单控制器和多控制器范式,以实现RLHF数据流的灵活表示和高效执行。我们精心设计了一套分层API,用于解耦和封装复杂RLHF数据流中的计算和数据依赖关系,从而允许高效的操作编排来实现RLHF算法,并能灵活地将计算映射到各种设备上。我们进一步设计了一个3D-HybridEngine,用于在训练和生成阶段之间高效地进行actor模型的重分片(resharding),实现了零内存冗余并显著减少了通信开销。
我们的实验结果表明,与最先进的基线系统相比,使用HybridFlow运行各种RLHF算法时,吞吐量提升了1.53倍至20.57倍。HybridFlow的源代码将在 https://github.com/volcengine/verl 上提供。
CCS概念:
- 计算方法学 → 分布式计算方法学; 机器学习。
关键词:
分布式系统,基于人类反馈的强化学习
ACM参考文献格式:
Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. 2025. HybridFlow: A Flexible and Efficient RLHF Framework. In Twentieth European Conference on Computer Systems (EuroSys ’25), March 30-April 3, 2025, Rotterdam, Netherlands. ACM, New York, NY, USA, 19 pages. https://doi.org/10.1145/3689031.3696075
1 引言
大型语言模型(LLMs),如GPT、Llama和Claude,已经彻底改变了从写作、搜索到编码等各种人工智能(AI)应用。LLMs首先通过在来自书籍、网站等的数万亿词元上进行下一个词预测来积累广泛的知识。接下来,LLMs在特定领域的数据集上通过监督微调(SFT)进行训练,以使其能够遵循人类指令。尽管经过预训练和SFT后,LLMs在自然语言任务上表现出色,但训练数据集中的有害和有偏见的内容仍可能误导LLM生成有毒和不良内容。基于人类反馈的强化学习(RLHF)被引入,以进一步将LLM与人类价值观对齐,从而构建有益且无害的AI应用。
RLHF建立在传统的RL算法之上,例如近端策略优化(PPO)和REINFORCE。广泛采用的基于PPO的RLHF系统通常由四个LLMs组成:一个actor、一个critic、一个参考策略网络和一个奖励模型。基于PPO的RLHF以迭代方式进行,每次迭代包括三个阶段:(1)使用actor模型和一批提示(prompts)生成响应;(2)通过对critic、参考策略和奖励模型进行单次前向传播来对生成的响应进行评分,从而准备训练数据;(3)通过前向和后向计算更新actor和critic,从而从人类偏好中学习。其他RLHF变体遵循类似的阶段,但涉及不同数量的模型和模型间的数据依赖关系。
传统的RL可以被建模为一个数据流,这是一个有向无环图(DAG):RL数据流中的每个节点代表一个神经网络(例如,可以是CNN或MLP的actor或critic网络)的计算;每条边表示NN计算之间的数据依赖关系(例如,critic的输出被用作actor训练的输入)。RLHF数据流更加复杂,涉及更复杂的模型(例如,用于actor/critic/参考/奖励模型的LLMs),每个模型运行不同的计算,并且它们之间的数据依赖关系也更加多样化(即,分布式模型分区之间的多播)。RLHF数据流中LLM的训练和生成需要分布式计算(例如,使用张量/流水线/数据并行)。因此,RLHF数据流中的每个节点都是一个复杂的分布式程序,对应于相应LLM的分布式计算。不同节点中的模型通常使用不同的并行策略,因为它们的工作负载各不相同。边代表数据重分片,这通常是一个多对多的多播。因此,对复杂且资源密集的RLHF进行灵活的表示和高效的执行至关重要。
传统的RL框架,如RLLib和RLLib Flow,利用分层的单控制器范式来运行RL数据流。一个中心化的控制器将数据流中的节点分配给不同的进程并协调它们的执行顺序。每个节点进程可以进一步衍生出更多的工作进程(worker)来执行计算,同样遵循单控制器范式。然而,它们只为数据并行训练提供了原语,并且受限于最多几百MB大小的神经网络。在RLHF数据流中,每个节点对应一个拥有多达数十亿操作符的LLM,使用某种复杂的并行策略进行计算。由于向分布式加速器分发操作符的巨大开销,单控制器范式效率低下。
现有的RLHF系统采用多控制器范式来管理节点内计算和节点间数据重分片。每个控制器独立管理一个设备的计算,并使用多个点对点操作来协调不同节点之间的数据依赖关系。这种多控制器范式在执行LLM计算时引入的调度开销可以忽略不计(详见§2.2)。
然而,没有中央控制,实现各种RLHF数据流是不灵活的,因为修改单个节点以适应不同的数据依赖关系需要改变所有相关节点的实现,这阻碍了代码复用。
为了解决这些限制,我们提出了HybridFlow,一个灵活高效的RLHF框架,可以轻松地表示和执行多样的RLHF数据流,实现高吞吐量。我们的关键观察是,在节点间层面利用单控制器范式可以灵活地表达各种数据依赖关系,并以最小的开销轻松协调节点间的数据重分片,而在节点内计算中集成多控制器范式可以显著提高计算效率。
我们提倡一种分层的混合编程模型来生成RLHF数据流。在节点层面,提供了多个模型类,将数据流中不同LLM的分布式计算(训练、推理和生成)封装成原语API。这些API可以无缝支持现有LLM框架的各种并行策略,包括3D并行、ZeRO和PyTorch FSDP,并在多控制器范式下执行分布式计算。在节点之间,设计了一套传输协议,以在单个控制器的协调下,向用户隐藏数据重分片的复杂性。这种编程模型抽象了分布式计算的复杂性,允许用户用几行代码实现一个RLHF数据流,并通过单个控制器的单一进程来运行RLHF。它还有效地解耦了节点内计算和节点间数据传输,允许独立优化每个模型而无需改变数据流中其他模型的代码。
actor模型的训练和生成是RLHF数据流中的主要计算。我们进一步设计了一个3D-HybridEngine,以实现actor模型训练和生成的高效执行,在训练和生成阶段之间的模型参数重分片过程中引入了零内存冗余并显著减少了通信开销。我们的混合编程模型还促进了模型在相同或不同GPU设备集上的灵活放置。这使我们能够提供一个有效的算法来优化GPU分配和模型放置,以适应任何RLHF数据流中各种模型大小和不同工作负载。我们在设计HybridFlow中的贡献总结如下:
- 我们提出了一个分层的混合编程模型,用于方便地构建RLHF数据流。该模型能够高效地执行节点内分布式计算,并灵活地进行节点间数据重分片和传输,适用于各种RLHF算法(§4)。
- 我们设计了一个3D-HybridEngine,它以高计算效率执行actor模型的训练和生成,并在训练和生成阶段之间实现了零冗余转换(§5)。
- 我们设计了一个有效的映射算法,以自动识别RLHF数据流中每个节点(模型)的优化GPU分配和放置(§6)。
- 我们进行了广泛的实验,将HybridFlow与最先进的RLHF系统在各种RLHF算法、模型大小和集群规模下进行了比较。我们的评估显示了1.53倍至20.57倍的吞吐量提升。我们已经开源了HybridFlow,并相信它将推动未来的RLHF研究和发展。
2 背景与动机
2.1 基于人类反馈的强化学习
RLHF工作流。RLHF使用一组给定提示的人类排序候选答案,将LLM的语言空间与人类价值观对齐。一个RLHF系统通常由多个模型组成,例如,一个actor、一个critic、一个参考策略和一个或多个奖励模型。actor和参考策略都是经过预训练/微调的LLM(即正在进行RLHF的LLM)。critic和奖励模型可以是不同的LLM,在人类偏好数据集上进行微调,其语言建模头被替换为标量输出头。RLHF工作流可分解为3个阶段(图1),我们以PPO为例:
- 阶段1(生成): actor使用自回归生成方式,从一批提示中产生响应。
- 阶段2(准备): 使用提示和生成的响应,critic计算它们的值,参考策略计算它们的参考对数概率,奖励模型计算它们的奖励,所有这些都通过相应模型的单次前向计算完成。
- 阶段3(学习/训练): 使用前几个阶段产生的数据批次和损失函数,通过Adam优化器更新actor和critic。
其他RLHF算法也大致遵循这3阶段工作流(图1(b)©)。Safe-RLHF引入了一个辅助的预训练损失(遵循PPO-ptx),并包含一个额外的成本模型以同时拟合人类偏好和安全标签。ReMax需要一个额外的生成过程来减少方差,并消除了数据流中的critic模型。研究人员正在积极探索新颖的RLHF算法,并将传统RL方法集成到RLHF领域。这些差异要求RLHF数据流图具有灵活的表示,以适应多样的算法需求。
(图1:3种RLHF算法的数据流图。PPO、Safe-RLHF、ReMax。阶段1、2、3分别代表生成、准备和训练。)
并行策略。LLM通过数据、流水线和张量并行进行训练和服务。在数据并行(DP)中,输入数据被分割成多个子集;每个子集由一个独立的设备(例如一个GPU)处理。ZeRO是DP训练的一种内存优化方案,它逐步将优化器状态、梯度和模型参数分片到不同的GPU上。流水线并行(PP)和张量并行(TP)将模型参数、梯度和优化器状态分布在多个GPU上。现代分布式训练框架,如Megatron-LM和MegaScale,利用3D并行或PTD并行,其中P、T、D分别代表PP、TP、DP。在3D并行中,PP大小表示模型训练中的流水线阶段数,TP大小表示一个张量被分割成的分片数,DP大小是模型副本的数量。LLM服务系统采用与训练类似的3D并行,但只分片模型参数和KVCache。
RLHF数据流中的LLM模型可能执行不同的计算,包括训练(一次前向传播、一次后向传播和模型更新)、推理(一次前向传播)和生成(多次前向传播的自回归生成)。具体来说,训练和生成在actor模型上进行,训练和推理在critic上进行,而推理在参考策略和奖励模型上进行。可以对不同模型应用不同的并行策略以适应不同的计算,从而实现最佳吞吐量。
2.2 分布式机器学习的编程模型
单控制器。它采用一个中心化的控制器来管理分布式程序的整体执行流。通过中心化的控制逻辑,用户可以将数据流的核心功能构建为单个进程(图2(b)),而控制器自动生成分布式工作进程来执行计算。凭借对硬件和数据流图的全局视图,单控制器范式允许灵活和优化的资源映射以及数据流任务间的执行顺序协调。然而,协调消息从控制器传递给所有工作进程,在大型集群上执行庞大的数据流图时会产生显著的分发开销。
多控制器。每个设备(即工作进程)都有自己的控制器。最先进的分布式LLM训练和服务系统采用多控制器范式,因其可扩展性和低分发开销(控制消息主要通过高速PCIe链路从CPU传递到GPU)。如图2(a)中采用多控制器RLHF实现的示例所示,为每个模型运行一个单独的程序,并且一个模型的所有工作进程都执行相同的程序。每个工作进程只拥有系统状态的局部视图,并需要模型之间的点对点通信(蓝色代码和箭头)来协调模型执行顺序。要在多控制器架构中实现RLHF工作流,用户必须在每个设备上运行的程序中,复杂地集成集体通信、计算和点对点数据传输的代码。这导致了计算和数据传输代码的深度嵌套,难以开发、维护和优化。在图2(a)中,每个模型执行本地计算和all_gather操作(黑色代码),而actor模型必须显式管理向critic和奖励模型的发送操作,后者也必须在其程序中的精确位置相应地实现接收操作。
(图2:RLHF系统中使用的编程模型。(a) 现有RLHF框架采用多控制器范式。(b) HybridFlow利用混合编程模型:单控制器协调模型;每个模型在分布式计算中使用多控制器范式。)
2.3 RLHF的特性
异构模型工作负载。RLHF中的actor、critic、参考和奖励模型在不同阶段可能执行训练、推理或生成,具有不同的内存占用和计算需求。对于参考策略和奖励模型,只需将其模型参数存储在GPU内存中,因为它们只执行前向传播。对于actor和critic,它们的模型参数、梯度和优化器状态都必须存储,因为它们需要进行模型训练。此外,一个小的actor模型(例如,一个7B的预训练/微调LLM)可以与更大的critic和奖励模型(例如,70B的LLMs)配对,以获得更好的对齐效果。鉴于这种异构性,在RLHF期间需要为每个模型的运行采用不同的并行策略和定制的优化。
actor训练和生成之间的计算不平衡。在RLHF数据流中,actor模型的训练和生成由两个节点表示(图1),它们通常构成了每次RLHF迭代中的大部分工作负载(例如,在HybridFlow中占总RLHF时间的58.9%)。actor训练是计算密集型的,通常需要更大的模型并行(MP)大小(即模型被分片成的数量)并将工作负载分配到更多的GPU上,例如,在8个GPU上对一个7B模型进行8路分片。对生成使用相同的并行策略(例如,相同的MP大小)可能导致GPU计算资源利用率不足,因为生成是内存密集型的。先前的研究表明,结合更大的DP大小和更小的MP大小(混合数据和模型并行),例如,将一个7B模型分成两份并在8个GPU上复制四次,可以提高生成吞吐量。尽管为actor训练和生成使用不同的并行策略可能优化两个阶段的吞吐量,但在两个阶段之间动态地重分片actor模型权重会产生显著的通信和内存开销。例如,对齐一个70B的actor模型,每次RLHF迭代需要从训练阶段向生成阶段传输140GB的模型权重,当两个阶段在不同设备上时,这可能占据一次迭代时间的36.4%。
(图3:给定模型放置方案的数据流执行。不同机器上的模型可以并发计算。同地部署的参考模型和奖励模型顺序执行。)
多样化的模型放置需求。根据模型的计算工作负载和数据依赖关系,RLHF数据流中模型的策略性设备放置是必要的。图3给出了一个模型放置方案及其对应的RLHF执行流。放置在不同设备集上的模型,如果没有数据依赖关系,可以并行执行。放置在同一组GPU上的模型,称为同地部署(colocated)模型,共享GPU内存并以分时方式顺序执行,因为如果同地部署的LLMs并发执行,很容易发生内存不足(OOM)错误。
我们观察到一个权衡:将模型放置在不同设备上允许并行处理,但考虑到RLHF中的分阶段模型执行,这可能不可避免地导致一些GPU空闲时间。在图3中,actor和critic分开放置,并行进行训练,但在其他RLHF阶段,它们有1/3的GPU时间处于空闲状态。支持各种放置策略并最大化设备利用率对于优化任何模型大小和集群规模下的RLHF性能至关重要。
2.4 现有RLHF系统的局限性
对各种RLHF数据流图的支持不灵活。现有的RLHF系统采用多控制器范式实现数据流。为了实现各种RLHF算法,用户必须浏览和管理混合了集体通信、模型计算(可能使用各种分布式训练/服务框架)和点对点数据传输的代码。这种代码结构缺乏模块化/功能封装,使得RLHF系统与特定的LLM训练和服务框架紧密耦合。因此,用户需要逐个案例地实现和优化不同的RLHF数据流,这阻碍了代码复用并增加了出错的风险。现有的RLHF框架仅支持PPO算法。此外,由于实现复杂性,支持的并行策略有限。例如,要在DeepSpeed-Chat中为LLM训练和生成集成3D并行,可能因为混合的代码结构而需要重新实现整个系统。
RLHF执行效率低下。表1总结了现有RLHF系统采用的并行策略、模型放置和执行模式。DeepSpeed-Chat和OpenRLHF采用ZeRO-3进行actor训练,采用TP进行actor生成。OpenRLHF在不同设备上为训练和生成使用不同的actor模型副本,导致了冗余的内存使用和设备间频繁的权重同步。DeepSpeed-Chat在同一组设备上为训练和生成维护相同的actor模型副本,并在训练和生成之间重分片模型权重(由于两个阶段使用了不同的并行策略),这对于大型模型仍可能产生巨大的内存和通信开销(详见§5.4)。NeMo-Aligner在actor训练和生成中使用相同的3D并行配置,导致生成吞吐量低下(§8.4)。
现有的RLHF框架仅限于一种模型放置方案,因此也只有一种RLHF执行模式,如表1所示。实现不同的放置方案很困难,需要改变模型初始化和节点间数据传输的内部逻辑(如图2中蓝色突出显示部分)。OpenRLHF和NeMo-Aligner允许在准备和学习阶段并发模型计算;在生成阶段,除actor外的模型都处于空闲状态,浪费了它们占用的GPU。DeepSpeed-Chat将所有模型同地部署在同一组设备上,每个设备根据RLHF数据流顺序执行每个模型。由于模型间工作负载不平衡,这种放置在资源利用上可能效率低下(在§8.3中评估)。
(表1:RLHF框架的比较。图示为一次PPO迭代的执行。数字1-6分别代表响应生成、奖励模型推理、参考模型推理、critic推理、actor训练和critic训练。)
2.5 设计考量
为了解决现有系统的局限性,关键问题是——如何设计一个灵活高效的编程模型来实现RLHF数据流?单控制器设计在节点间层面特别有优势,因为它在协调分布式计算中不同模型间的数据传输、执行顺序和资源虚拟化方面具有灵活性。RLHF数据流图通常只包含少数几个节点。与数据流中节点(模型)所需的分布式计算相比,从单控制器向不同节点分发控制消息产生的开销可以忽略不计。多控制器范式以其向加速器分发操作符的低延迟而闻名,可以用于每个模型的分布式计算中。基于这些洞察,我们提出了一个用于RLHF数据流实现的分层混合编程模型。我们的关键设计原则是以混合方式结合单控制器和多控制器范式。这种设计确保了RLHF数据流的灵活表达和高效执行,在节点间和节点内两个层面都保持了低控制开销。如图2(b)所示,这种范式解耦了节点内分布式计算和节点间数据传输,允许每个模型专注于本地计算而无需管理节点间通信。
3 HybridFlow 概述
图4描绘了HybridFlow的架构,它由三个主要组件组成:混合编程模型、3D-HybridEngine和自动映射算法。混合编程模型包含一套分层API,以实现RLHF数据流的灵活表达和数据流中模型的高效计算(§4)。3D-HybridEngine专为actor模型的高效训练和生成而设计,允许在两个阶段使用不同的3D并行配置,并在两个阶段之间的转换过程中实现零内存冗余和最小化的通信开销(§5)。自动映射算法确定每个模型的优化设备放置,以最大化RLHF的吞吐量(§6)。
(图4:HybridFlow的架构。用户输入->自动映射->模型放置->资源池->物理设备。LLM训练/生成引擎由3D-HybridEngine驱动,并通过传输协议与并行工作进程通信。)
我们的RLHF系统的工作流程如下。用户提供以下输入来启动RLHF系统:(i)模型规格,包括RLHF数据流中actor/critic/参考策略/奖励模型的架构和大小;(ii)数据流中模型的设备放置,这是通过在给定GPU集群配置下运行自动映射算法获得的;(iii)每个模型在每个阶段运行的并行策略,例如,对于3D并行,是一个元组(p, t, d),其中p, t, d分别代表PP大小、TP大小和DP大小。单控制器程序接收这些输入,以初始化RLHF数据流中的模型和虚拟化资源池,根据放置计划将操作/模型分派到设备,并调用由多控制器在设备上运行的函数来执行每个模型的分布式计算。
多控制器程序实现了ParallelWorker类:它根据其并行策略在分配的设备之间构建每个模型的并行组,调用3D-HybridEngine进行actor训练和生成,并且可以与现有的LLM引擎无缝集成,用于其他模型的训练、推理和生成。传输协议由单控制器程序协调,以支持在具有不同并行策略的模型之间重分片数据(包括提示、响应和RLHF中的其他模型输出)。actor在训练和生成之间的数据重分片由3D-HybridEngine处理。
4 混合编程模型
4.1 分层API
节点内:封装分布式程序。对于不同RLHF阶段中每个模型的分布式计算,我们提供了一个基类3DParallelWorker。给定分配的设备,它促进了分布式模型权重的初始化,并为每个模型建立3D并行组。一个并行组包括一组GPU来承载模型的特定并行维度,例如,TP中的不同张量分片和DP中的不同模型副本。图5(a)说明了使用我们的API初始化actor模型,而其他模型的初始化过程类似。
(图5:分层API的图示。(a) 具有3D并行配置、资源分配和3DParallelWorker初始化的模型。(b) 两个模型之间使用3D_PROTO中的collect和distribute函数进行异步数据重分片。)
继承自3DParallelWorker类,我们分别为actor、critic、参考和奖励模型提供了几个模型类。这些模型类中的每一个都封装了API来实现模型的分布式前向和后向计算、自回归生成和优化器更新,从而将分布式计算代码与和其他模型的数据依赖关系解耦。这些API可以通过重用现有LLM系统的计算脚本来轻松实现。例如,ActorWorker(actor模型的类)的update_actor函数中涉及的计算与Megatron-LM中的预训练脚本类似。一个模型类封装了实现各种RLHF算法的基本操作,例如,actor模型类中的generate_sequences用于根据提示生成响应,奖励模型类中的compute_reward用于通过前向传播评估响应。(更多API详见附录A)。
除了实现3D并行的基类3DParallelWorker,我们还为PyTorch FSDP(FSDPWorker)和ZeRO(ZeROWorker)提供了基类,以及继承自每个基类的相应模型类,以支持模型计算中的不同并行策略。图4中的ParallelWorker表示这些基类之一。
节点间:统一模型间的数据重分片实现。在采用不同并行策略、部署在不同设备上的模型之间传输数据,涉及到多对多的多播。我们通过使用@register将每个模型类中的每个操作与一个传输协议关联起来,从而统一了这种数据传输的实现。每个传输协议由一个collect函数和一个distribute函数组成,以根据每个模型的并行策略聚合输出数据和分发输入数据。
在图5(a)的示例中,update_actor操作被注册到传输协议3D_PROTO,因为actor训练使用了3D并行。在3D_PROTO中,collect函数将相应模型函数的所有输出数据(例如,从update_actor返回的损失标量)在每个DP组中收集到单个控制器,而distribute函数将输入数据(例如,update_actor的advantages)分发到每个DP组。数据重分片通过使用源模型的输出collect函数和目标模型的输入distribute函数来实现。图5(b)说明了actor(生成)和critic(推理)之间的数据重分片,其中模型的计算采用了不同的3D并行策略。单个控制器使用actor的3D_PROTO中的collect函数收集数据future(步骤1-3),并将其发送给critic(步骤4);critic使用其3D_PROTO中的distribute函数将接收到的数据future分发到每个DP组(步骤5)。然后,远程数据从actor检索到critic,每个critic的GPU只根据其DP rank获取所需的actor输出数据的本地批次(步骤6)。实际的数据传输只发生在GPU之间,避免了任何中央瓶颈。
我们提供了8种传输协议,包括3D_PROTO, DP_PROTO, ONE_TO_ALL等,涵盖了大多数数据重分片场景(详见附录B)。用户还可以通过实现自定义的collect和distribute函数来进一步扩展传输协议。
促进灵活的模型放置。我们提供了一个ResourcePool类,它虚拟化了一组GPU设备。当将一个ResourcePool实例应用于一个模型类时(图5(a)),该模型的分布式计算将被映射到这些设备上。使用相同ResourcePool实例的模型被同地部署在同一组GPU上;当在它们的模型类中应用不同的ResourcePool实例时,模型被放置在不同的GPU集上。我们假设不同的ResourcePool实例之间没有重叠。
异步数据流执行。当模型被放置在不同的设备集上时,它们的执行会在其输入可用时立即被自动触发。在图5(b)中,来自actor的数据future在控制器调用后立即返回(步骤1-3);然后控制器向critic发起新的调用,并根据传输协议分发future(步骤4-5)。当某些模型被放置在同一组设备上时,它们根据调用顺序顺序执行。通过我们的编程模型,HybridFlow能够灵活地支持多样的分布式执行模式,而无需对RLHF算法的代码进行任何更改(图6)。
4.2 不同RLHF算法的实现
我们的API使得各种RLHF算法(数据流)的开发变得流水线化。用户可以用几行代码实现一个RLHF算法,作为一个在单控制器上运行的单进程程序,该程序涉及一系列对模型分布式计算的原语API调用。图6给出了PPO、ReMax和Safe-RLHF的示例。PPO只需8行代码即可实现,通过调用模型的compute_values和generate_sequences等操作,这些操作在多控制器范式下在多个GPU上执行。为了适应集成了额外成本模型以评估安全偏好和actor预训练损失的Safe-RLHF,只需在PPO实现的基础上增加5行代码。为了适应ReMax,需要额外调用一次actor生成,并且可以删除与critic相关的代码。
(图6:PPO、ReMax和Safe-RLHF的实现。用户只需添加或删除几行代码即可适应不同的RLHF算法。)
实现灵活性。这种扩展的灵活性对于研究人员探索不同的RLHF算法至关重要:他们可以重用封装在每个模型类中的分布式计算,并根据特定算法(如GAE和compute_advantage中的KL散度以及actor和critic的损失函数)简单地调整数值计算的代码。这种流水线式的开发得益于混合编程模型。我们的模块化API设计简化了开发,促进了广泛的代码重用,并能够直接整合现有LLM训练/服务框架的代码库。它还解耦了模型计算和模型间的数据传输。分布式框架中的任何更改都不会影响RLHF算法的代码(图6),从而为每个模型的执行实现个性化优化(§5)。支持具有不同工作负载的模型的灵活放置,从而能够将RLHF数据流优化地映射到各种设备上(§6)。
5 3D-HybridEngine
我们设计了3D-HybridEngine来支持actor模型的高效训练和生成,旨在显著提高RLHF的吞吐量。
5.1 并行组
为了消除冗余的actor模型副本,我们提倡在同一组设备(分配给actor的N_a个GPU)上部署actor训练和生成阶段,并在同一份actor模型权重上顺序执行它们。尽管如此,actor训练和生成很可能采用不同的3D并行策略,即生成阶段通常需要比训练阶段更小的TP和PP大小,但需要更大的DP大小(§2.3)。3D-HybridEngine在这种情况下,能够在同一组设备上高效地进行actor训练和生成之间的模型参数重分片。
让p-t-d表示为actor训练构建的3D并行组,分别对应于承载p个流水线阶段、t个张量分片和d个模型副本的GPU集。3D-HybridEngine根据actor训练和生成的不同3D并行策略,分别为它们构建不同的并行组。我们使用p_g, t_g, d_g分别表示生成阶段的生成流水线并行组、生成张量并行组和微数据并行组的大小。d_g表示生成中模型副本数与训练中模型副本数的比率,即训练中的每个DP副本变成d_g个微DP副本,以处理d_g个微批次的提示和响应。我们有 N_a = p×t×d = p_g×t_g×d_g×d,因此 d_g = (pt) / (p_gt_g)。微DP组专门用于actor生成阶段,以提供更大的DP大小,从而充分利用设备。生成并行组表示为p_g-t_g-d_g-d。
5.2 3D-HybridEngine 工作流
在RLHF的第i次迭代的actor训练和第i+1次迭代的actor生成之间,actor模型参数需要根据两个阶段的并行组配置进行重分片,提示数据也需要分发。在RLHF的第i+1次迭代中,3D-HybridEngine在每个微DP组内收集第i次迭代中更新的actor模型参数用于生成(图7中的步骤1)。然后,将一批提示加载到每个模型副本(步骤2),生成响应(RLHF的生成阶段)。之后,3D-HybridEngine在每个微DP组内对生成结果执行all-gather操作(步骤3),并根据actor训练的3D并行策略重新划分模型参数(步骤4)。随着模型权重、提示和响应的正确重新分发,计算actor模型的损失并根据RLHF算法更新actor模型权重(步骤5)——即第i+1次迭代的actor训练阶段。
(图7:一次RLHF迭代中3D-HybridEngine的工作流。4个GPU用于actor训练和生成。训练中使用1-2-2 (p-t-d) 并行组,生成中使用1-1-2-2 (p_g-t_g-d_g-d)并行组。)
5.3 零冗余模型重分片
3D并行中的并行分组方法通常如下:PP和TP组通过为流水线阶段和张量分片分配连续的rank来形成;DP组通过以规则间隔选择rank来构建,间隔由PP大小和TP大小的乘积决定。在图8(a)中,actor训练使用1-4-2的3D并行组:所有GPU有一个PP组(为清晰起见);TP组是[G1, G2, G3, G4]和[G5, G6, G7, G8];DP组是[G1, G5], [G2, G6], [G3, G7], [G4, G8]。
假设在生成中也使用相同的并行分组方法,但并行大小不同,例如图8(a)中的1-2-2-2。在从训练到生成的转换期间,3D-HybridEngine在模型并行组之间应用all-gather操作来聚合所有参数,然后根据设备所属的生成并行组,在每个设备上仅保留一部分模型权重用于生成。在某些GPU上(例如G2, G3, G6, G7),训练和生成的模型权重之间没有重叠,需要额外的内存来维护后续训练的权重(图8(a)中的灰色框)。我们称这个系统为HybridFlow-V,即3D-HybridEngine在两个阶段使用上述常规并行分组方法。
(图8:模型权重重分片。2台机器,每台4个GPU用于actor训练和生成。(a) 训练和生成之间使用相同的分组方法(HybridFlow-V)。(b) 优化的并行分组方法(HybridFlow)。)
我们进一步为3D-HybridEngine设计了一种新的并行分组方法用于生成阶段,该方法消除了权重存储中的冗余,并导致actor模型在训练和生成之间重分片时产生最小的内存占用和通信。具体来说,我们通过以规则间隔(由 t/t_g 和 p/p_g 决定)选择rank来形成生成TP和PP组,并通过沿生成TP或PP维度顺序分配rank来构建微DP组。在图8(b)中,生成中使用1-2-2-2并行组:生成TP组是[G1, G3], [G2, G4], [G5, G7], [G6, G8];微DP组是[G1, G2], [G3, G4], [G5, G6], [G7, G8]。这种对生成并行组的策略性重排导致每个设备上的训练和生成模型权重之间产生重叠,从而能够在生成期间重用训练权重,并因模型重分片而实现零设备内存冗余。此外,3D-HybridEngine在每个微DP组内并发地执行多个all-gather操作,从而显著减少了通信开销。
5.4 转换开销
在表2中,我们比较了不同actor引擎设计在训练和生成阶段之间转换时的通信开销和内存占用。我们假设actor的模型大小为M,并使用N_a个GPU进行训练和生成。DeepSpeed-Chat中的actor引擎在转换期间在所有GPU上执行一次all-gather操作;HybridFlow-V在训练TP和PP组内执行此all-gather。这些操作的通信量对于DeepSpeed-Chat是 (tpd-1)/(tpd) * M,对于HybridFlow-V是 (tp-1)/tp * M。两个引擎在后续根据生成并行组划分模型状态之前,都会在每个GPU的内存中聚合所有模型参数,导致峰值内存使用量为M。由于它们在某些GPU上无法在生成期间重用训练权重,因此需要在这些GPU上维护训练权重,分别导致 1/(tpd) 和 1/tp 的冗余内存消耗。
(表2:训练与生成之间的转换开销比较。)
通过我们为生成阶段设计的并行分组方法,HybridFlow将all-gather操作限制在每个微DP组内。通信开销减少到 (d_g-1)/tp * M = (tp-t_gp_g)/(t_gp_g*tp) * M。每个GPU只需收集其微DP组内的远程参数,并可以在生成中重用训练权重。因此,HybridFlow中模型参数的峰值内存使用量精确匹配每个GPU在生成时的模型分片大小,消除了GPU内存使用中的任何冗余。
6 自动设备映射
我们的混合编程模型要求用户输入以下配置,这些配置被称为RLHF数据流到给定设备的映射:(a)数据流中模型的设备放置;(b)每个模型在每个阶段运行的相应并行策略。我们提供了一个高效的算法(算法1)供用户识别在给定设备集群上执行RLHF数据流的优化映射,从而最小化每次RLHF迭代的端到端延迟。
给定一个数据流D,我们首先探索给定集群中模型的所有可能放置方案P(第3行)。例如,PPO算法涉及四个模型,导致15种可能的放置方案(来自贝尔划分问题),范围从所有模型都放在不同设备上的完全独立放置(例如,OpenRLHF的放置)到将所有模型同地部署在同一组设备上(例如,DeepSpeed-Chat的放置)。我们将同地部署在同一组GPU上的模型称为同地部署集 (colocated set)。一个同地部署集中的模型可以在同一组GPU上采用不同的并行策略。我们根据同地部署模型的内存消耗,确定分配给每个同地部署模型集的最小GPU数量A_min,确保没有内存不足错误(第9行)。
接下来,从A_min中的最小GPU分配开始,我们枚举每个同地部署模型集的所有可行设备分配(第10-12行)。给定同地部署集的设备分配A和集合中模型的工作负载W,我们在auto_parallel模块中探索每个模型的优化并行策略,以最小化模型执行延迟。工作负载W包括每个模型的输入和输出形状以及计算(训练、推理或生成)。在auto_parallel中,我们利用一个模拟器模块simu来估计不同并行策略的延迟,遵循先前的研究(详见附录C)。
d_cost模块根据给定的模型放置和并行策略,通过迭代数据流图中的所有阶段并累加所有阶段的延迟,来估计RLHF数据流的端到端延迟(第17, 25行)。对于在同一同地部署集且在同一阶段进行计算的模型(例如actor和critic都在RLHF训练阶段执行模型更新),它们的执行延迟被相加(第32行)。对于在不同同地部署集的模型,它们在同一阶段内的执行可以并行化,该阶段的延迟由不同集合中的最大执行时间决定(第33行)。我们识别出模型的最佳设备放置及其相应的并行策略,实现每次RLHF迭代的最小执行时间(第18-23行)。
(算法1:RLHF数据流的设备映射算法。)
算法1的复杂度为O((N-1)!/((k-1)!(N-k)!)),其中k是数据流中的模型数量,N是运行数据流的总设备数。这是为一种放置策略(即独立放置)枚举所有可能设备分配的最坏情况复杂度,通过将N个设备分配给k个模型计算得出(称为整数划分问题)。为提高效率,我们缓存为每个模型在一定数量设备A上识别出的并行策略,以消除在不同放置策略中当模型被放置在不同设备集A上时对相同并行策略的冗余搜索。
尽管我们在运行自动映射算法时假设N个同构GPU,但算法1可以很容易地扩展,通过在simu和auto_parallel模块中考虑异构设备,来优化异构设备上的模型映射。
7 实现
HybridFlow用大约12k行Python代码(LoC)实现。
混合编程模型。分层API用1.8k LoC实现。中心化的单控制器建立在Ray之上,并使用远程过程调用(RPC)来协调不同模型的执行顺序,并根据数据流在模型之间传输数据。这些中间数据存储在TensorDict中。在我们的多控制器范式中,用于分布式计算的每个模型函数在不同设备的单独进程上运行,控制消息从每个控制器的CPU进程中继到相应的GPU。我们的实现支持Megatron-LM、PyTorch FSDP和DeepSpeed作为LLM训练和推理引擎,以及vLLM用于自回归生成。在vLLM中,我们用一个分布式管理器替换了中心化的KVCache管理器,以与多控制器范式对齐。
3D-HybridEngine。其主要逻辑在Megatron-LM和vLLM之上用2.4k LoC实现。我们将用于训练和生成阶段的actor模型权重存储在不同的内存缓冲区中,在训练期间将生成权重卸载到CPU内存,在转换期间将生成权重重新加载回GPU内存,并在生成中使用两个缓冲区。我们使用NCCL通信原语来在训练和生成之间的转换期间,在每个微DP组内收集和连接模型参数。我们在生成后将KVCache卸载到CPU内存,并在下一次迭代中将其重新加载回GPU。
自动映射算法与三个用于训练、推理和生成工作负载的模拟器一起,用1.9k LoC实现。该算法在CPU上启动RLHF数据流之前运行,以生成用于数据流初始化的设备映射和并行策略。
8 评估
8.1 实验设置
测试平台。我们在一个由16台机器(128个GPU)组成的集群上部署HybridFlow。每台机器配备8个NVIDIA A100-80GB GPU,通过600GB/s的NVLink互连。机器间带宽为200Gbps。我们的实验使用以下软件版本:CUDA 12.1, PyTorch 2.1.2, Megatron-core 0.6.0, NCCL 2.18.1, 和 vLLM 0.3.1。
模型和RLHF算法。我们运行PPO、ReMax和Safe-RLHF算法的RLHF数据流(图1)。PPO是RLHF最流行的算法之一,由actor、critic、参考策略和奖励模型组成。每个模型都是一个Llama模型,大小从7B到70B不等。Safe-RLHF有一个额外的成本模型,其架构和大小与奖励模型相同,而ReMax则消除了critic模型。我们在所有实验中使用混合精度进行actor和critic训练,即BF16用于模型参数,FP32用于梯度和优化器状态,使用Adam优化器。BF16用于模型推理和自回归生成。如未特别说明,实验结果均来自PPO。
基线。我们将HybridFlow与最先进的RLHF系统进行比较,包括DeepSpeed-Chat v0.14.0、OpenRLHF v0.2.5和NeMo-Aligner v0.2.0(详见表1)。NeMo-Aligner不支持ReMax算法。我们不与其他框架(如Trlx、HuggingFaceDDP和Collosal-Chat)进行比较,因为它们代表性较差且速度比上述基线慢。
我们使用RLHF吞吐量(tokens/秒)作为性能指标,通过将全局批次中提示和响应的总词元数除以一次RLHF迭代时间来计算。所有报告的性能数据都是在10次预热迭代后,对5次训练迭代取平均值。
数据集和超参数。我们在HuggingFace的"Dahoas/ful-hh-rlhf"数据集上进行RLHF,该数据集被广泛用于LLM对齐。由于基线系统可能在生成期间未集成连续批处理优化,为公平比较,我们强制所有生成的响应长度相同。在每个实验中,输入提示长度和输出响应长度均为1024,actor模型的输入提示全局批大小为1024。PPO epoch数为1,每个epoch的PPO更新迭代次数为8,与先前的RLHF研究一致。
8.2 端到端性能
图9、10和11分别显示了运行PPO、ReMax和Safe-RLHF时的RLHF吞吐量。这组实验中的actor、critic、参考和奖励模型大小相同,遵循以往的实践。不同模型大小实验中使用的GPU数量从运行RLHF不OOM的最小GPU数到128个GPU不等。为公平比较,实验中未启用优化器状态卸载。
(图9:PPO吞吐量。括号中的数字是HybridFlow与基线相比的加速比。)
(图10:ReMax吞吞吐量。括号中的数字是HybridFlow与基线相比的加速比。)
(图11:Safe-RLHF吞吐量。括号中的数字是HybridFlow与基线相比的加速比。)
整体性能。我们观察到HybridFlow在所有模型规模上都持续优于基线。在图9的PPO中,HybridFlow分别比DeepSpeed-Chat、OpenRLHF和NeMo-Aligner快3.67倍(最高7.84倍)、3.25倍(最高5.93倍)和12.52倍(最高20.57倍)。这主要是因为HybridFlow通过为不同计算工作负载分片模型采用不同的并行策略,有效地执行了所有RLHF阶段的生成、推理和训练。在训练70B模型时,HybridFlow实现了最高的平均9.64倍加速,因为与使用ZeRO-3训练时也产生大量机器间通信的DeepSpeed-Chat和OpenRLHF相比,HybridFlow将转换开销分别减少了高达71.2%和89.1%。由于生成引擎中缺少KVCache,NeMo-Aligner的主要性能瓶颈在于生成阶段,该阶段占其RLHF迭代时间的81.2%。在图10和11中也可以观察到类似的结果,验证了HybridFlow在运行各种RLHF算法上的效率。
可扩展性。在8个GPU上,HybridFlow至少实现了2.09倍的加速。随着GPU数量的增加,HybridFlow在各种模型规模上的强扩展效率为66.8%,通过(最大规模吞吐量/最小规模吞吐量)/(最大GPU数/最小GPU数)计算得出,对三种算法和所有模型规模取平均。在固定全局批大小的情况下扩展到大量GPU会导致每个工作进程的本地批大小变小,可能导致GPU利用率不足。在128个GPU上运行7B模型时,HybridFlow在PPO、ReMax和Safe-RLHF上仍然比最佳基线OpenRLHF快1.68倍、1.53倍和1.71倍。这归因于HybridFlow能够为不同模型和集群大小适应最佳放置策略以最小化RLHF时间。OpenRLHF在较大的GPU集群中表现更好,但在较小的集群中效率较低。
8.3 模型放置
(图12:不同放置策略下HybridFlow的吞吐量。)
(图13:13B actor和参考策略 & 70B critic和奖励模型下的放置比较。)
在这个实验中,我们在HybridFlow中实现了PPO算法的各种模型放置,模型和集群设置与§8.2相同:(i)colocate,DeepSpeed-Chat中的放置策略;(ii)standalone,OpenRLHF中的策略;(iii)split,NeMo-Aligner的同地部署策略(actor和参考策略在同一组设备上,critic和奖励模型在另一组);(iv)hybridflow,通过算法1获得的优化放置。
不同模型放置的比较。图12揭示了在不同GPU数量下,HybridFlow的优化放置是不同的。从16到64个GPU,将所有模型同地部署在同一组设备上性能最佳。对于96到128个GPU的34B模型和96个GPU的13B模型,split策略变为最优。split策略将GPU在两组模型之间平均分配,因为它们的模型大小相等。对于128个GPU上的13B模型,standalone策略实现了最高吞吐量。在这种情况下,HybridFlow为actor分配64个GPU,为critic分配32个,为参考和奖励模型各分配16个。在较小的集群中,所有模型的计算可以充分利用GPU资源;colocate策略确保了在不同RLHF阶段最大化GPU使用率。在较大的集群中,由于批大小固定且在更多GPU上DP大小更大导致计算通信比下降,colocate放置下的RLHF吞吐量无法线性扩展。Standalone和split策略在较大的集群中将模型放置在不同设备上,每个模型的DP大小更小,促进了同一阶段不同模型的并行执行。在所有情况下,我们的算法1都产生了具有最高训练吞吐量的最佳放置。
更大的critic和奖励模型。我们进一步评估了运行PPO时使用13B的actor和参考策略以及70B的critic和奖励模型(更大的critic和奖励模型有望产生更好的对齐)的模型放置。图13显示,在多达64个GPU上,colocate策略仍然平均优于其他策略44.8%。在96个GPU上,split策略实现了更高的吞吐量。当扩展到128个GPU时,算法1获得的最佳放置是将actor、参考和奖励模型同地部署在64个GPU上,同时将剩余的64个GPU分配给critic。在相同数量的GPU上,actor和参考策略的计算时间远小于critic和奖励模型,将奖励模型与actor和参考策略同地部署减少了经验准备阶段的GPU空闲时间。总的来说,在大型集群中,将actor和critic分布在不同设备上以在训练阶段并行执行,会带来更高的吞吐量。
8.4 3D-HybridEngine
(图14:actor训练和生成之间的转换时间。)
(图15:16个GPU上actor模型不同生成并行大小的时间分解。)
转换时间比较。图14显示了在各种模型规模下,actor训练和生成阶段之间的转换时间,这是在§8.2的相同设置下,重分片模型权重所需的时间。OpenRLHF的转换时间包括两个actor模型副本在不同设备间的权重同步时间。HybridFlow平均将转换时间减少了55.2%(11.7秒),在使用70B模型时,转换开销最多减少了89.1%(78.2秒),同时在不同集群规模下保持了一致的开销。这归功于我们为生成阶段设计的新并行分组方法(§5.4)。在基线方法中,所有模型参数都必须在转换期间被收集,需要多次逐层收集以防止OOM。HybridFlow在转换期间实现了零内存冗余,并且每个微DP组只需要一次all-gather操作。
转换和生成时间。我们进一步验证了在HybridFlow中为actor训练和生成使用不同并行大小的必要性。在这个实验中,所有模型都同地部署在同一组GPU上,用于生成的KVCache使用剩余的GPU内存进行分配(即尽力分配)。图15给出了在16个GPU上运行RLHF时,7B和13B模型的转换和生成时间,训练并行组为1-8-2(遵循p-t-d约定),生成TP组大小t_g从1到8变化。生成PP组大小保持在p_g=1,微DP组大小d_g计算为8/t_g。我们观察到,对于7B模型,应用较小的生成TP组大小t_g=2,对于13B模型,t_g=4,分别将生成延迟减少了60.3%和36.4%。相反,使用与训练相同的TP大小(t_g=8),遵循NeMo-Aligner的方法,由于GPU利用率不足,导致了最大的生成延迟。进一步减小t_g未能实现更高的加速,因为更小的t_g需要每个GPU维护更大的KVCache。
8.5 算法运行时间
(图16:设备映射算法的运行时间。模型大小和GPU数量同时扩展。)
图16显示了算法1的运行时间,该时间远短于实际RLHF训练所需的天数。运行时间呈线性增长,揭示了设备映射算法在模型大小和集群大小方面具有良好的可扩展性。大部分运行时间用于估计每个模型并行策略的执行延迟。对于更大的模型,有更多的并行策略可用,需要更多的模拟来为每个放置方案确定最优策略。我们缓存模型的优化并行策略,以便在不同放置方案中重复应用,这将寻找最佳放置的时间减少到最多半小时。
9 讨论
容错性。HybridFlow与现有的容错方法是正交的,并且已经集成了检查点功能。故障可以通过NCCL错误检测,静默数据损坏可以通过校验和检测。我们的编程模型使单控制器能够通过RPC协调检查点操作,从而在每个ParallelWorker组内保存模型状态。这包括保存actor/critic模型的参数、dataloader ID和随机数生成器(RNG)状态,以确保系统范围的一致性。此外,如果存在足够的健康模型副本,HybridFlow也可以采用基于冗余的容错方法,如广播参数和CPU检查点,以实现快速恢复。
放置策略洞见。我们总结了RLHF训练中模型放置和GPU分配的三个主要洞见。1)为actor模型分配更多GPU可以减少耗时的生成延迟,而这无法与其他模型并行。2)当每个模型计算都能充分利用GPU资源时,在相对小规模的集群上训练时,同地部署所有模型最有效。3)当扩展到大规模集群时(即强扩展),将actor和critic模型分布在不同设备上,以便在训练和准备阶段并行执行,将有助于实现更高的吞吐量。
资源复用。HybridFlow通过利用分时共享GPU计算,实现了在共享设备上同地部署模型。最近在DNN任务调度方面的研究开发了细粒度的资源复用技术,主要旨在实现单个任务的服务水平目标。尽管ResourcePool的实现支持同地部署模型的并行执行,但HybridFlow通常遵循顺序执行,以防止GPU资源争用或OOM问题,如第2.3节所讨论。在RLHF训练中应用GPU共享和异构资源带来了独特的挑战,因为它需要在各种任务之间平衡计算工作负载和管理复杂的数据依赖关系。研究用于RLHF训练中GPU共享的细粒度自动映射算法,结合模型卸载优化和异构设备的集成,将是未来研究的一个有前景的方向。
从对齐到推理。在用于LLM对齐的RLHF中,奖励信号由奖励模型生成。除了对齐任务,类似的算法(例如,PPO和GRPO)可以应用于其他领域,如代码生成和数学推理。对于这些任务,每个提示可能存在一个基准真相,这可以通过评估每个代码测试用例的输出值的正确性和验证数学结果的准确性来确定。因此,奖励模型可以被非神经网络的奖励模块替换,例如用于评估生成代码的沙箱环境或用于验证数学结果的奖励函数。HybridFlow可以通过将这些奖励模块包装为远程函数,并在单进程脚本中编排它们的执行,来无缝集成这些模块,为多样的强化学习应用提供一个灵活高效的框架。
10 相关工作
RL框架。已经有大量的RL框架,从为小规模DNN设计的通用RL系统到专门为LLM优化的RLHF系统。我们在§2中已经彻底审查了密切相关的工作,本节将讨论更多的RL框架。这些RL框架,与最近的RLHF系统类似,使用各种多控制器框架来实现其算法。它们建立多个长期运行的分布式程序,每个组件都通过硬编码的数据同步来协调执行顺序。Gear进一步优化了RL流水线的经验回放部分。然而,所有这些框架都无法支持LLM在RLHF中的训练、推理和生成。
LLM训练和服务系统。TorchDDP和Horovod支持数据并行训练。ByteScheduler和DeepSpeed通过通信和内存优化扩展了数据并行。许多系统通过模型并行(如张量并行和流水线并行)来优化大型模型训练,以在设备间划分模型。LLM服务系统也采用数据和模型并行来加速自回归生成,并采用了专门的优化,如连续批处理和分块预填充。请注意,所有上述框架都采用多控制器范式以实现高效计算。
数据流系统。像MapReduce、Spark、Dryad和Naiad这样的数据流系统在分析和机器学习工作负载中很流行,但它们缺乏对动态任务图的支持。Ray在一个动态任务图中统一了任务并行和actor编程模型,并实现了一个可扩展的分布式调度器和一个全局控制存储,被许多RL框架所采用。Pathways是一个用于TPU的闭源项目,旨在轻松表达复杂的并行模式和单个DNN模型内的细粒度控制流,例如流水线并行和带有稀疏计算的专家混合模型。它采用异步分布式数据流设计,即使存在数据依赖关系也能实现并行控制平面执行,从而减少了单控制器范式的分发开销。其主要关注点在于单模型训练,需要对DNN模型的每个子网络进行复杂的编译。HybridFlow可以将Pathways作为子模块集成,以实现RLHF数据流中模型的计算。
11 结论
HybridFlow是一个RLHF框架,它能够灵活地表示和高效地执行多样的RLHF算法。我们提出了一个混合编程模型,允许用户通过将不同LLM的分布式计算封装到原语API中,并隐藏节点间数据重分片的复杂性,用几行代码轻松构建RLHF数据流。我们的3D-HybridEngine确保了actor模型训练和生成的高效执行,模型参数重分片具有零内存冗余和显著减少的通信开销。此外,我们有效的映射算法优化了RLHF数据流中模型的GPU分配和放置。广泛的实验表明,在各种模型大小和集群规模下,与最先进的RLHF系统相比,HybridFlow实现了1.53倍到20.57倍的加速。
附录
A HybridFlow中的原语API
(表4:每个模型类中提供的关键函数。用户可以使用这些提供的函数用几行代码构建各种RLHF算法。)
| 模型 | API | 计算 | 解释 |
|---|---|---|---|
| Actor | generate_sequence |
自回归生成 | 基于一批提示,actor模型生成一批响应,并返回响应中每个词元的对数概率。 |
compute_log_prob |
一次前向传播 | actor模型计算提示和响应中每个词元的对数概率。此对数概率与使用相同模型精度进行生成时返回的对数概率相同。(在PPO中可选) | |
compute_loss |
一次前向传播 | actor模型根据预训练数据集计算预训练损失。 | |
update_actor |
一次前向、后向传播和模型更新 | 基于优势(advantages)、回报(returns,从compute_advantage计算得出)和预训练损失,actor模型计算训练损失并更新其权重。我们为多种RLHF算法实现了各种损失函数。 |
|
| Critic | compute_values |
一次前向传播 | critic模型为每个提示和响应计算值(values)。 |
update_critic |
一次前向、后向传播和模型更新 | 基于值和回报,critic计算平方误差损失以更新其权重。我们也为多种RLHF算法实现了critic损失。 | |
| Reference Policy | compute_ref_log_prob |
一次前向传播 | 参考模型计算提示和响应中每个词元的参考对数概率。此对数概率被用作评估actor模型散度的基准,并约束其学习过程。 |
| Reward | compute_reward |
一次前向传播 | 奖励模型进行前向计算,为给定的提示和响应集计算分数。奖励可以是词元级别或样本级别。 |
| - | compute_advantage |
数值计算 | 基于价值模型和奖励模型分别提供的值和奖励,该函数在给定的提示和当前策略模型的响应上估计优势。此计算不涉及模型前向传播。 |
B 传输协议
(表3:HybridFlow中的传输协议。)
| 传输协议 | Distribute函数 | Collect函数 | 用例 |
|---|---|---|---|
ONE_TO_ALL |
将数据广播给所有rank。 | 从所有rank收集数据。 | 所有工作进程方法有相同的输入并运行相同的代码,例如模型初始化。 |
3D_PROTO |
拆分数据,在所有DP rank间散发,并在组内广播。 | 从所有DP组的p=-1, t=0工作进程收集并连接数据。 | 模型在每个数据并行组内的多个工作进程间分片。模型输出仅存在于最后一个流水线阶段,并在数据并行组间复制。这是3D并行训练中的典型场景。 |
3D_ALL_MICRO_DP |
按微DP大小拆分数据,在所有微DP组间散发,并在组内所有rank间广播。 | 从所有微DP组的local_rank=0工作进程收集并连接数据。 | 与HybridEngine一起使用。用于处理策略模型在训练和推理之间切换时的3D并行方案。 |
3D_PP_ONLY |
将数据广播给所有rank。 | 从所有PP组的t=0, d=0工作进程收集并连接数据。 | 用于检查权重名称,因为它们在TP和DP组中是相同的。 |
DP_PROTO |
将数据拆分成批次并在所有DP rank间散发。 | 从所有DP rank收集并连接数据。 | 在数据并行模式下训练模型。 |
ALL_TO_ALL |
无操作。 | 从所有rank收集数据。 | 用于调试。用户可以手动定义每个工作进程的输入并分别检查其输出。 |
C 自动并行算法
(算法2:自动并行算法。)
算法2概述了每个模型最优并行策略的搜索过程。从每个模型的最小模型并行大小开始(以防止在与多个工作进程同地部署时OOM),我们根据GPU数量和每台机器的GPU数量U,枚举所有可行的并行配置。默认的U设置为8。我们使用simu模块根据其工作负载估计每个模型的延迟。该模块包括三个用于训练、推理和生成工作负载的模拟器,都是遵循先前研究的分析模型。训练和推理工作负载是计算密集型的,而生成工作负载是内存密集型的。对于actor模型,我们首先找到训练的并行策略并记录训练阶段的内存使用情况。在actor生成期间,根据批大小和最大序列长度计算KVCache需求。如果生成阶段的模型并行大小无法容纳参数和KVCache,我们增加它。然后,我们通过比较延迟估计,寻找具有相应KVCache分配的最优策略。开发一个考虑可变KVCache大小的综合自回归生成模拟器,可以进一步增强RLHF研究中的自动映射过程。
更多推荐
所有评论(0)