HybridFlow: A Flexible and Efficient RLHF Framework

Guangming Sheng The University of Hong Kong gmsheng@connect.hku.hk
Chi Zhang ByteDance zhangchi.usc1992@bytedance.com
Zilingfeng Ye ByteDance yezilingfeng@bytedance.com
Xibin Wu ByteDance wuxibin@bytedance.com
Wang Zhang ByteDance zhangwang.nozomi@bytedance.com
Ru Zhang ByteDance zhangru.1994@bytedance.com
Yanghua Peng ByteDance pengyanghua.yanghua@bytedance.com
Haibin Lin ByteDance haibin.lin@bytedance.com
Chuan Wu The University of Hong Kong cwu@cs.hku.hk

A1 主要贡献

论文的核心问题是RLHF在LLM对齐中的复杂数据流表示和执行效率问题：传统RL框架使用单一控制器导致分布式计算开销大，现有的RLHF系统采用多控制器范式但缺乏灵活性，无法轻松适应不同RLHF算法和并行策略。研究目标是设计一个灵活高效的RLHF框架，支持多样化RLHF数据流的表示和执行，实现高吞吐量。创新点包括：提出层次化混合编程模型，结合单一控制器和多控制器范式，实现节点间灵活数据依赖协调和节点内高效计算；设计3D-HybridEngine，支持演员模型训练和生成阶段的无冗余内存和低通信开销的重分片；提供自动映射算法，优化RLHF数据流中模型的GPU分配和放置；在实验中，实现1.53×至20.57×的吞吐量提升，并开源框架。

A3 背景知识/关键Observation/设计原则

RLHF工作流程。 RLHF使用人类排名候选来对齐LLM与人类价值观，系统通常包括多个模型，如演员、评论家、参考策略和奖励模型。演员和参考是预训练/微调的LLM，评论家和奖励模型是基于人类偏好数据集微调的LLM，其语言建模头替换为标量输出头。RLHF工作流程分解为3个阶段（图1），以PPO为例：阶段1（生成）：演员使用自回归生成从一批提示产生响应；阶段2（准备）：使用提示和生成的响应，评论家计算值，参考策略计算参考对数概率，奖励模型计算奖励，所有通过各自模型的单次前向计算；阶段3（学习/训练）：使用先前阶段产生的数据和损失函数，通过Adam更新演员和评论家。其他RLHF算法（如Safe-RLHF引入辅助预训练损失和额外成本模型，ReMax需要额外生成传递以减少方差并消除评论家模型）也遵循类似3阶段工作流程，但模型数量和数据依赖不同。这些变体需要灵活的RLHF数据流图表示来适应多样化算法需求。

并行策略。 LLM使用数据、管道和张量并行进行训练和服务。在数据并行（DP）中，输入数据分割为多个子集，每个子集由单独设备处理；ZeRO逐步跨GPU分片优化器状态、梯度和模型参数。管道并行（PP）和张量并行（TP）跨多个GPU分布模型参数、梯度和优化器状态。现代框架如Megatron-LM和MegaScale使用3D并行或PTD并行，其中P、T、D分别表示PP、TP、DP大小。在3D并行中，PP大小表示管道阶段数，TP大小表示张量分片的数目，DP大小表示模型副本数。LLM服务系统使用类似3D并行，仅分片模型参数和KVCache。RLHF数据流中的模型执行不同计算，包括训练（一次前向、一次后向和模型更新）、推理（一次前向）和生成（自回归生成的多前向传递）。特别是，演员模型执行训练和生成，评论家执行训练和推理，参考策略和奖励模型执行推理。针对不同计算的不同模型应用不同的并行策略，以实现最优吞吐量。

分布式ML的编程模型。 单一控制器。 它使用集中控制器管理分布式程序的整体执行流。通过集中控制逻辑，用户可以将数据流的核心功能构建为单个进程，而控制器自动生成分布式工作者执行计算。通过硬件和数据流图的全局视图，单一控制器范式允许灵活优化的资源映射和执行顺序协调。然而，协调消息从控制器传递到所有工作者，在大型集群上执行扩展数据流图时会产生显著分发开销。多控制器。 每个设备（工作者）有自己的控制器。最先进的分布式LLM训练和服务系统采用多控制器范式，因其可扩展性和低分发开销（控制消息主要通过快速PCIe链接从CPU到GPU传递）。如图2（a）所示的多控制器RLHF实现示例，每个模型运行单独程序，一个模型的所有工作者执行相同程序。每个工作者仅拥有系统状态的本地视图，需要模型之间的点对点通信来协调执行顺序。在多控制器架构中实现RLHF工作流程，用户必须在每个设备运行的程序中精细集成集体通信、计算和点对点数据传输的代码。这导致计算和数据传输的深度嵌套代码，难以开发、维护和优化。在图2（a）中，每个模型执行本地计算和all_gather操作，而演员模型必须显式管理发送操作到评论家和奖励模型，后者必须在程序中相应实现接收操作。

RLHF特性。 异构模型工作负载。 RLHF中的演员、评论家、参考和奖励模型在不同阶段执行训练、推理或生成，具有不同的内存占用和计算需求。对于参考策略和奖励模型，仅需存储模型参数在GPU内存中，因为它们仅执行前向计算。对于演员和评论家，必须存储模型参数、梯度和优化器状态，因为它们进行模型训练。而且，小型演员模型（如7B预训练/微调LLM）可能与更大的评论家和奖励模型（如70B LLM）配对，以实现更好对齐。鉴于这种异构性，需要为每个模型运行定制并行策略和优化。演员训练和生成之间的不平衡计算。 在RLHF数据流中，演员模型的训练和生成表示每个RLHF迭代的主要工作负载（如使用HybridFlow时占总RLHF时间的58.9%）。演员训练是计算绑定的，通常需要更大的模型并行（MP）大小并分布到更多GPU，例如7B模型的8个分区在8个GPU上。使用相同并行策略（如相同MP大小）进行生成可能导致GPU计算资源利用不足，因为其内存绑定性质。先前研究显示，结合更大DP大小和更小MP大小（如将7B模型分区为2并在8个GPU上复制4次）可以改善生成吞吐量。尽管为演员训练和生成使用不同并行策略可能优化两个阶段的吞吐量，但运行时在两个阶段之间重分片演员模型权重会产生显著通信和内存开销。例如，对齐70B演员模型每个RLHF迭代需要从训练到生成传输140GB模型权重，当两个阶段在不同设备上时占迭代时间的36.4%。多样化模型放置需求。 根据模型的计算工作负载和数据依赖，需要RLHF数据流中模型的战略设备放置。图3给出了示例模型放置计划和相应RLHF执行流。放置在不同设备集上的模型如果无数据依赖可以并行执行。放置在相同GPU集上的模型（共置模型）共享GPU内存并按时间共享顺序执行，因为如果共置LLM并发执行容易发生OOM错误。我们观察到权衡：放置在不同设备上允许并行处理，但鉴于RLHF中的阶段模型执行，可能不可避免地导致一些GPU空闲时间。在图3中，演员和评论家单独放置，并行进行训练，但在其GPU时间内有1/3空闲，在其他RLHF阶段。支持各种放置策略并最大化设备利用率对于在任何模型大小和集群规模下优化RLHF性能至关重要。

现有RLHF系统的局限性。 对各种RLHF数据流图的不灵活支持。 现有RLHF系统采用多控制器范式实现数据流。为了实现各种RLHF算法，用户必须导航和管理混合集体通信、模型计算（可能使用各种分布式训练/服务框架）和点对点数据传输的代码。这种代码结构缺乏模块化/函数封装，使得RLHF系统与特定LLM训练和服务框架紧密耦合。因此，用户需要逐案实现和优化不同RLHF数据流，阻碍代码重用并增加出错风险。现有RLHF框架仅支持PPO算法。此外，由于实现复杂性，支持的并行策略有限。例如，要在DeepSpeed-Chat中纳入LLM训练和生成的3D并行，可能需要由于混合代码结构而重新实现整个系统。低效RLHF执行。 表1总结了现有RLHF系统采用的并行策略、模型放置和执行模式。DeepSpeed-Chat和OpenRLHF采用ZeRO-3用于演员训练和TP用于演员生成。OpenRLHF在不同设备上使用演员模型的不同副本进行训练和生成，导致冗余内存使用和设备间频繁权重同步。DeepSpeed-Chat在相同设备集上为训练和生成维护相同演员模型副本，并在训练和生成之间重分片模型权重（由于两个阶段使用的不同并行），这可能为大型模型产生大量内存和通信开销。NeMo-Aligner在演员训练和生成中使用相同的3D并行配置，经历低生成吞吐量。现有RLHF框架限于一种模型放置计划，因此限于一种RLHF执行模式，如表1所示。实现不同放置困难，需要改变模型初始化和节点间数据传输的内部逻辑。OpenRLHF和NeMo-Aligner允许准备和学习阶段的并发模型计算；在生成阶段，除演员外的模型空闲，浪费它们占用的GPU。DeepSpeed-Chat将所有模型共置在相同设备集上，每个设备根据RLHF数据流顺序运行每个模型。鉴于模型间不平衡工作负载，这种放置在资源利用上可能低效。

设计考虑。 为了解决现有系统的局限性，关键问题是设计一个灵活高效的编程模型来实现RLHF数据流。单一控制器设计在节点间级别特别有利，因为它在协调不同模型分布式计算之间的数据传输、执行顺序和资源虚拟化方面灵活。RLHF数据流图通常仅由几个节点组成。从单一控制器向不同节点分发控制消息的开销相对于数据流中节点（模型）所需的分布式计算而言微不足道。多控制器范式以其向加速器分发操作的低延迟而闻名，可用于每个模型的分布式计算。基于这些洞见，我们提出RLHF数据流实现的层次化混合编程模型。我们的关键设计原则是以混合方式结合单一控制器和多控制器范式。这种设计确保RLHF数据流的灵活表达和高效执行，在节点间和节点内级别保持低控制开销。如图2（b）所示，这种范式解耦节点内分布式计算和节点间数据传输，允许每个模型仅关注本地计算而不管理节点间通信。

A2 方法细节

3 HybridFlow概述

HybridFlow架构。 HybridFlow的架构包括三个主要组件：混合编程模型、3D-HybridEngine和自动映射算法。混合编程模型包括一套层次化API，以实现RLHF数据流的灵活表达和数据流中模型的高效计算。3D-HybridEngine特别设计用于演员模型的高效训练和生成，允许两个阶段的不同3D并行配置，并在两个阶段之间转换时实现零内存冗余和最小化通信开销。自动映射算法确定每个模型的优化设备放置，以最大化RLHF的吞吐量。RLHF系统的工作流程如下。用户提供以下输入启动RLHF系统：(i) 模型规格，包括RLHF数据流中演员/评论家/参考策略/奖励模型的架构和大小；(ii) 数据流中模型的设备放置，由在给定GPU集群配置下运行自动映射算法获得；(iii) 每个模型在每个阶段的并行策略，例如3D并行的(p, t, d)元组，其中p、t、d分别表示PP大小、TP大小和DP大小。单一控制器程序接收这些输入初始化RLHF数据流中的模型和虚拟化资源池，根据放置计划向设备分发操作/模型，并调用设备上的多控制器函数执行每个模型的分布式计算。多控制器程序实现ParallelWorker类：根据其并行策略在分配设备中构建每个模型的并行组，调用3D-HybridEngine进行演员训练和生成，并可与现有LLM引擎无缝集成，用于其他模型的训练、推理和生成。由单一控制器程序协调的传输协议支持具有不同并行策略的模型之间数据（包括RLHF中的提示、响应和其他模型输出）的重分片。演员在训练和生成之间的数据重分片由3D-HybridEngine处理。

4 混合编程模型

4.1 层次化API

节点内：封装分布式程序。 对于不同RLHF阶段中每个模型的分布式计算，我们提供基类3DParallelWorker。给定分配设备，它促进分布式模型权重初始化并为每个模型建立3D并行组。并行组包括一组GPU，用于托管模型的特定并行维度，例如TP中的不同张量分片和DP中的不同模型副本。图5(a)展示了使用我们的API初始化演员模型，而其他模型的初始化类似。从3DParallelWorker类继承，提供了几个模式类，分别用于演员、评论家、参考和奖励模型。这些模型类中的每个封装API实现模型的分布式前向和后向计算、自回归生成和优化器更新，将分布式计算代码与与其他模型的数据依赖解耦。这些API可以通过重用现有LLM系统的计算脚本轻松实现。例如，ActorWorker（演员模型类）的update_actor函数中涉及的计算类似于Megatron-LM中的预训练脚本。模型类封装实现各种RLHF算法的基本操作，例如演员模型类中的generate_sequences用于基于提示生成响应，以及奖励模型类中的compute_reward用于通过前向传递评估响应。除了实现3D并行的基类3DParallelWorker，我们进一步提供PyTorch FSDP (FSDPWorker)和ZeRO (ZeROWorker)的基类，以及继承每个基类的相应模型类，以支持模型计算中的不同并行策略。图4中的ParallelWorker表示这些基类之一。

节点间：统一模型之间数据重分片的实现。 涉及使用不同并行策略在不同设备上的模型之间进行多对多多播的数据传输。我们通过将每个模型类中的每个操作与传输协议关联，使用@register来统一此数据传输实现。每个传输协议包括一个collect函数和一个distribute函数，根据每个模型的并行策略聚合输出数据并分发输入数据。在图5(a)的示例中，update_actor操作注册到传输协议3D_PROTO，因为3D并行用于演员训练。在3D_PROTO中，collect函数将相应模型函数的输出数据（例如从update_actor返回的损失标量）在每个DP组中收集到单一控制器，distribute函数将输入数据分发到注册函数（例如update_actor的优势）到每个DP组。使用源模型的输出collect函数和目标模型的输入distribute函数启用数据重分片。图5(b)展示了演员（生成）和评论家（推理）之间数据重分片，其中模型计算采用不同的3D并行策略。单一控制器使用演员的3D_PROTO中的collect函数收集数据未来（步骤1-3）并发送到评论家（步骤4）；评论家使用其3D_PROTO中的distribute函数将接收的数据未来分发到每个DP组（步骤5）。然后从演员检索远程数据到评论家，每个评论家的GPU仅根据其DP排名获取所需本地批次的演员输出数据（步骤6）。实际数据传输仅发生在GPU之间，避免任何中央瓶颈。我们提供8个传输协议，包括3D_PROTO、DP_PROTO、ONE_TO_ALL等，覆盖大多数数据重分片场景。用户可以通过实现自定义collect和distribute函数进一步扩展传输协议。

促进灵活模型放置。 我们提供ResourcePool类，用于虚拟化一组GPU设备。当将ResourcePool实例应用于模型类（图5(a)）时，模型的分布式计算将映射到设备。使用相同ResourcePool实例的模型共置在相同GPU集上；当不同ResourcePool实例应用于模型类时，模型放置在不同GPU集上。我们假设不同ResourcePool实例之间无重叠。异步数据流执行。 当模型放置在单独设备集上时，一旦输入可用，它们的执行自动触发。在图5(b)中，来自演员的数据未来在控制器调用后立即返回（步骤1-3）；然后控制器启动对评论家的新调用并按照传输协议分发未来（步骤4-5）。当一些模型放置在相同设备集上时，它们基于调用顺序顺序执行。通过我们的编程模型，HybridFlow灵活支持多样化分布式执行模式，而无需更改RLHF算法的代码（图6）。

4.2 不同RLHF算法的实现

我们的API启用各种RLHF算法（数据流）的简化开发。用户可以将RLHF算法实现为在单一控制器上运行的单个进程程序，仅涉及一系列原始API调用来调用模型的分布式计算。图6给出了PPO、ReMax和Safe-RLHF的示例。PPO可以通过调用包括compute_values和generate_sequences的模型操作在仅8行中实现，这些操作在多控制器范式下在多个GPU上执行。为了适应SafeRLHF，它集成了额外成本模型来评估安全偏好和演员的预训练损失，仅需在PPO实现上添加5行代码。为了适应ReMax，需要额外调用演员生成，并可移除评论家相关代码。实现灵活性。 这种扩展灵活性对于研究人员探索不同RLHF算法至关重要：他们可以重用每个模型类中封装的分布式计算，并简单根据特定算法调整数值计算的代码，例如compute_advantage和演员及评论家的损失函数中的GAE和KL散度。这种简化开发归功于混合编程模型。我们的模块化API设计简化开发，促进广泛代码重用，并启用直接纳入现有LLM训练/服务框架的代码库。它还解耦模型计算和模型间数据传输。分布式框架中的任何变化不影响RLHF算法的代码（图6），启用每个模型执行的个性化优化。支持具有多样化工作负载的模型的灵活放置，启用RLHF数据流到各种设备的优化映射。

5 3D-HybridEngine

我们设计3D-HybridEngine支持演员模型的高效训练和生成，针对显著RLHF吞吐量改进。并行组。 为了消除冗余演员模型副本，我们主张在相同设备集上部署演员训练和生成阶段，分配给演员的N GPUs，并顺序执行它们在相同演员模型权重副本上。尽管如此，演员训练和生成可能采用不同的3D并行策略，即生成阶段通常需要比训练阶段更小的TP和PP大小但更大的DP大小。3D-HybridEngine在此上下文中启用相同设备集上演员训练和生成之间高效模型参数重分片。让p-t-d表示为演员训练构建的3D并行组，对应托管p管道阶段、t张量分片和d模型副本的GPU集。3D-HybridEngine根据各自不同的3D并行策略为演员训练和生成构建不同并行组。我们使用p_g、t_g和d_m表示生成阶段中生成管道并行组、生成张量并行组和微数据并行组的大小。r表示生成中模型副本数与训练中比率，即训练中的每个DP副本成为r微DP副本，以处理r微批次的提示和响应。我们有N = p × t × d = p_g × t_g × d_m × r，使得r = N / (p_g × t_g × d_m)。微DP组仅在演员生成阶段使用，以为完全设备利用呈现更大DP大小。生成并行组表示为p_g-t_g-d_m-r。

3D-HybridEngine工作流程。 在RLHF的迭代k的演员训练和迭代k+1的演员生成之间，需要根据两个阶段的并行组配置重分片演员模型参数并分发提示数据。在RLHF的迭代k+1中，3D-HybridEngine收集迭代k中更新的演员模型参数（图7中的步骤1），用于每个微DP组内的生成。然后，将一批提示加载到每个模型副本（步骤2），生成响应（RLHF的生成阶段）。随后，3D-HybridEngine在每个微DP组内对生成结果执行all-gather操作（步骤3），并根据演员训练的3D并行重新分区模型参数（步骤4）。随着模型权重、提示和响应正确重新分布，计算演员模型的损失并按照RLHF算法更新演员模型权重（步骤5）——迭代k+1的演员训练阶段。

零冗余模型重分片。 3D并行中的并行分组方法通常如下：PP和TP组通过为管道阶段和张量分片分配连续排名来形成；DP组通过以PP大小和TP大小乘积确定的固定间隔选择排名来构建。在图8(a)中，演员训练使用3D并行组1-4-2：所有GPU有一个PP组（为说明清晰）；TP组是[G1, G2, G3, G4]、[G5, G6, G7, G8]，DP组是[G1, G5]、[G2, G6]、[G3, G7]、[G4, G8]。假设使用相同并行分组方法但不同并行大小，例如生成中的1-2-2-2。在从训练到生成的转换中，3D-HybridEngine在模型并行组中应用all-gather操作聚合所有参数，然后仅在每个设备上保留模型权重子集用于其生成，根据设备所属的并行组。在一些GPU（如G2、G3、G6、G7）上，训练和生成模型权重无重叠，需要单独内存维护权重用于后续训练（图8(a)中的灰色框）。我们称系统为HybridFlow-V，当3D-HybridEngine在两个阶段中使用上述普通并行分组方法。我们进一步为3D-HybridEngine设计新并行分组方法，用于生成阶段，消除权重存储中的冗余，并导致由于训练和生成之间演员模型重分片的最小内存占用和通信。具体地，我们通过以d_m和p_g × t_g确定的固定间隔选择排名来形成生成TP和PP组，并通过沿生成TP或PP维度顺序分配排名来构建r微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并发执行几个all-gather操作，每个微DP组一个，导致显著减少的通信开销。

转换开销。 在表2中，我们比较不同演员引擎设计中训练和生成阶段之间通信开销和内存占用。我们假设演员的模型大小为M，并使用N GPUs进行其训练和生成。DeepSpeed-Chat中的演员引擎在转换期间跨所有GPU执行all-gather操作；HybridFlow-V在训练TP和PP组内执行此all-gather。这些操作的通信量分别为(M - 1)M / N = (N - 1)M / N用于DeepSpeed-Chat和(t - 1)M / t用于HybridFlow-V，按照[13]计算。两个引擎在随后根据生成并行组分区模型状态之前在每个GPU内存中聚合所有模型参数，导致模型参数的峰值内存使用M。因为它们无法在一些GPU上重用训练权重用于生成，训练权重需要维护，导致分别为1 M / N和1 M / t的冗余内存消耗。通过我们用于生成阶段的并行分组方法，HybridFlow将all-gather操作限制在每个微DP组内。通信开销减少到(d_m - 1) M / (d_m r) = (d_m - 1) M / (d_m N / (p_g t_g))。每个GPU仅需收集其微DP组内的远程参数，并可在生成中重用训练权重。因此，HybridFlow中模型参数的峰值内存使用精确匹配生成中每个GPU上的模型分区大小，消除GPU内存使用中的任何冗余。

6 自动设备映射

我们的混合编程模型要求用户输入以下配置，称为RLHF数据流到给定设备的映射：(a) 数据流中模型的设备放置；(b) 每个模型在每个阶段的相应并行策略。我们为用户提供高效算法（算法1）来识别在给定设备集群上执行RLHF数据流的优化映射，最小化每个RLHF迭代的端到端延迟。给定数据流G，我们首先探索给定集群中模型的所有可能放置计划P（行3）。例如，PPO算法涉及四个模型，导致15种可能放置（从Bell分区问题），从所有模型放置在不同设备上的完全独立放置（例如OpenRLHF的放置）到将所有模型共置在相同设备集上（例如DeepSpeed-Chat的放置）。我们将相同GPU集上的共置模型称为共置集。共置集中的模型可以在相同GPU集上采用不同并行策略。我们基于共置模型的内存消耗识别每个共置模型集的最小GPU分配数目min_alloc，确保无OOM错误（行9）。接下来，从min_alloc中的最小GPU分配开始，我们枚举每个共置模型集的所有可行设备分配（行10-12）。给定共置集的设备分配D和集中的模型计算工作负载W，我们在auto_parallel模块中探索每个模型的优化并行策略，最小化模型执行延迟。工作负载W包括每个模型的输入和输出形状以及计算（训练、推理或生成）。在auto_parallel中，我们利用模拟器模块simu来估计不同并行策略的延迟，按照先前研究（附录C中概述）。d_cost模块在给定模型放置和并行策略下估计RLHF数据流的端到端延迟，通过迭代数据流图中的所有阶段并求和所有阶段的延迟（行17,25）。对于相同共置集中的模型并在相同阶段涉及计算（例如演员和评论家在RLHF训练阶段都进行模型更新），它们的执行延迟求和（行32）。对于不同共置集中的模型，它们在相同阶段内的执行可以并行化，该阶段的延迟由不同集的最大执行时间确定（行33）。我们识别模型的最佳设备放置及其相应并行策略，实现每个RLHF迭代的最小执行时间（行18-23）。算法1的复杂度为O( (N -1)! (K -1)! (N - K)! )，其中K是数据流中的模型数，N是运行数据流的总设备数。这是枚举放置策略的所有可能设备分配的最坏情况复杂度（即独立放置），通过为K模型分配N设备计算（整数分区问题）。为提高效率，我们缓存每个模型在D设备上的识别并行策略，以消除在不同放置策略中模型放置在不同D GPU集上时的冗余搜索。尽管我们假设运行自动映射算法时的N同质GPU，但算法1可以通过在simu和auto_parallel模块中考虑异构设备来轻松扩展，用于优化异构设备上的模型映射。

7 实现

HybridFlow以约12k行Python代码实现。混合编程模型。 层次化API以1.8k LoC实现。集中单一控制器基于Ray构建，并使用远程进程调用（RPC）协调不同模型的执行顺序并按照数据流传输模型间数据。这些中间数据存储在TensorDict中。在我们用于分布式计算的多控制器范式中，每个模型函数在各种设备上运行单独进程，控制消息从每个控制器的CPU进程中继到相应GPU。我们的实现支持Megatron-LM、PyTorch FSDP和DeepSpeed作为LLM训练和推理引擎，以及vLLM用于自回归生成。在vLLM中，我们将集中KVCache管理器替换为分布式管理器，以与多控制器范式对齐。3D-HybridEngine。 其主要逻辑以2.4k LoC在Megatron-LM和vLLM之上实现。我们为训练和生成阶段存储演员模型权重在单独内存缓冲区，在训练期间将生成权重卸载到CPU内存，在转换期间将生成权重重新加载回GPU内存，并在生成中使用两个缓冲区。我们使用NCCL通信原语在训练和生成之间的转换期间在每个微DP组中收集和连接模型参数。我们在生成后将KVCache卸载到CPU内存，并在下一次迭代中重新加载回GPU。自动映射算法 以1.9k LoC实现，连同训练、推理和生成工作负载的三个模拟器。该算法在启动RLHF数据流之前在CPU上运行，以生成数据流初始化的设备映射和并行策略。

A4 实验环境

数据集包括"Dahoas/full-hh-rlhf"数据集（规模：广泛用于LLM对齐的HuggingFace数据集，用于RLHF训练），用途：LLM对齐；每个实验中输入提示长度和输出响应长度均为1024，全局批大小为1024，PPO epoch数为1，每个epoch的PPO更新迭代数为8。模型架构为Llama模型，关键参数：大小从7B到70B，混合精度（BF16用于模型参数，FP32用于梯度和优化器状态），Adam优化器；算法包括PPO（演员、评论家、参考策略、奖励模型）、ReMax（消除评论家模型）、Safe-RLHF（额外成本模型）。硬件配置：16台机器（128个NVIDIA A100-80GB GPU），每台机器8个GPU通过600GB/s NVLink互连，机器间带宽200Gbps。软件配置：Python代码实现，依赖CUDA12.1、PyTorch 2.1.2、Megatron-core 0.6.0、NCCL 2.18.1、vLLM 0.3.1；基线包括DeepSpeed-Chat v0.14.0、OpenRLHF v0.2.5、NeMo-Aligner v0.2.0。

A4 实验结果

端到端性能实验：内容：运行PPO、ReMax和Safe-RLHF的RLHF吞吐量（tokens/sec），模型大小从7B到70B，GPU从最小无OOM数到128；结果：HybridFlow在所有模型规模上始终优于基线，对于PPO，平均优于DeepSpeed-Chat 3.67×（最高7.84×）、OpenRLHF 3.25×（最高5.93×）、NeMo-Aligner 12.52×（最高20.57×）；对于ReMax和Safe-RLHF类似；70B模型时平均加速9.64×；可扩展性：在8 GPU上至少2.09×加速，强缩放效率66.8%；在128 GPU上7B模型仍优于最佳基线OpenRLHF 1.68×、1.53×和1.71×；分析：HybridFlow通过不同并行策略分片模型适应各种计算工作负载高效执行生成、推理和训练；减少70B模型的转换开销达71.2%和89.1%；NeMo-Aligner的生成阶段占81.2%时间为瓶颈。引用：图9 PPO吞吐量、图10 ReMax吞吐量、图11 Safe-RLHF吞吐量。

模型放置实验：内容：PPO算法下不同放置策略（colocate、standalone、split、hybridflow）的吞吐量，模型大小13B/34B，GPU16-128；结果：hybridflow放置在不同GPU数下变化，从16-64 GPU colocate最佳，96-128 GPU split或standalone最佳；对于13B演员/参考和70B评论家/奖励，colocate在64 GPU上优于其他44.8%，96 GPU split最佳，128 GPU hybridflow分配64 GPU给演员/参考/奖励和64给评论家最佳；分析：小集群中colocate最大化利用，大集群中standalone/split允许并行减少DP大小；分配更多GPU给演员减少生成延迟。引用：图12 HybridFlow在不同放置下的吞吐量、图13 13B演员/参考和70B评论家/奖励下的放置比较。

3D-HybridEngine实验：内容：演员训练和生成之间转换时间，模型7B-70B；结果：HybridFlow平均减少55.2%（11.7s），70B最高89.1%（78.2s），开销在不同规模一致；分析：新并行分组方法限制all-gather到微DP组，零冗余内存，仅一次操作避免OOM。内容：16 GPU上7B/13B模型的转换和生成时间，训练1-8-2，生成TP大小从1-8变化；结果：更小生成TP（7B为2、13B为4）减少生成延迟60.3%和36.4%，相同TP如NeMo-Aligner导致最大延迟；分析：小TP需要更大KVCache，权衡利用率。引用：图14 演员训练和生成之间的转换时间、图15 16 GPU上不同生成并行大小的时间分解。

算法运行时实验：内容：算法1运行时间，模型大小和GPU同时缩放；结果：线性增长，最大半小时；分析：大部分时间用于模拟并行策略延迟，缓存减少搜索时间。引用：图16 设备映射算法运行时。

A5 结论

论文总结HybridFlow是一个RLHF框架，支持多样化RLHF算法的灵活表示和高效执行。通过混合编程模型，用户可以通过封装不同LLM分布式计算的原始API和隐藏节点间数据重分片复杂性，在几行代码中轻松构建RLHF数据流。3D-HybridEngine确保演员模型训练和生成的零内存冗余和显著减少的通信开销。此外，有效映射算法优化RLHF数据流中模型的GPU分配和放置。实验显示HybridFlow在各种模型大小和集群规模下比最先进RLHF系统实现1.53×至20.57×加速。未来工作包括探索细粒度自动映射算法用于RLHF训练中的GPU共享，结合模型卸载优化和异构设备集成，以及将类似算法扩展到代码生成和数学推理领域，其中奖励模型可替换为非神经网络模块，如沙盒环境或奖励函数。

A6 附录

A HybridFlow中的原始API。 在HybridFlow中，我们通过继承3DParallelWorker、FSDPWorker和ZeROWorker实现了每个RLHF训练模型的原始函数。这些模型类的函数旨在解耦分布式计算代码，并为用户提供RLHF的基本操作。这种原始设计与现有分布式推理和训练框架中的自回归生成、前向传递、后向传递和模型更新操作兼容。用户可以根据算法设计自定义RLHF训练数据流（通过适应提供的函数中的数值计算），并受益于重用底层分布式计算实现。我们在表4中说明这些API的含义和实际计算。

B 传输协议。 我们实现了覆盖RLHF数据流中模型间所有常见数据重分片用例的传输协议。用户可以使用这些预定义协议生成任何RLHF数据流。而且，用户可以通过实现collect函数和distribute函数轻松定义自己的传输协议。传输协议解耦了复杂数据重分片和分布式训练。我们将p、t、d表示为工作者在管道-、张量-和数据并行组中的排名。我们在表3中说明这些预定义协议。

C 自动并行算法。 算法2概述了每个模型的最优并行策略搜索过程。从每个模型的最小模型并行大小开始（以防止与多个工作者共置时OOM），我们基于GPU数和每机器GPU数M枚举所有可行并行配置。默认M设置为8。我们使用simu模块根据工作负载估计每个模型的延迟。该模块包括训练、推理和生成工作负载的三个模拟器，所有都是按照先前研究的分析模型。训练和推理工作负载是计算绑定的，而生成工作负载是内存绑定的。对于演员模型，我们首先找到训练的并行策略并记录训练阶段的内存使用。在演员生成期间，使用批大小和最大序列长度计算KVCache需求。如果生成阶段的模型并行大小无法容纳参数和KVCache，我们增加它。然后，通过比较延迟估计寻找具有相应KVCache分配的最优策略。开发考虑可变KVCache大小的全面自回归生成模拟器可以进一步增强RLHF研究中的自动映射过程。