翻译：HybridFlow: A Flexible and Efficient RLHF Framework（未完成）
原文连接：HybridFlow: A Flexible and Efficient RLHF Framework
AI 翻译的内容，还没有仔细看和贴上图片。
还缺乏足够的背景知识，补充完成后再来看。
Abstract
人类反馈强化学习（RLHF）被广泛用于大语言模型（LLM）对齐。传统强化学习可以被建模为一种数据流，其中每个节点表示一个神经网络（NN）的计算，每条边表示神经网络之间的数据依赖。RLHF 使数据流变得更加复杂：它将每个节点扩展为一个分布式 LLM 训练或生成程序，并将每条边扩展为多对多组播。传统 RL 框架使用单个控制器来执行数据流，同时指挥节点内计算和节点间通信；但在 RLHF 中，由于分布式节点内计算会产生大量控制分发开销，这种方式可能效率较低。现有 RLHF 系统采用多控制器范式，但由于分布式计算和数据通信相互嵌套，这种方式可能不够灵活。
我们提出 HybridFlow，它以混合方式结合单控制器和多控制器范式，从而实现 RLHF 数据流的灵活表示与高效执行。我们精心设计了一组分层 API，用于在复杂 RLHF 数据流中解耦并封装计算和数据依赖，使得系统能够高效编排操作以实现 RLHF 算法，并能灵活地将计算映射到各种设备上。我们进一步设计了 3D-HybridEngine，用于在训练阶段和生成阶段之间高效地对 actor 模型进行重新分片，它具有零内存冗余，并显著降低通信开销。实验结果表明，与最先进的基线系统相比，在使用 HybridFlow 运行多种 RLHF 算法时，吞吐量提升了 1.53× 到 20.57×。HybridFlow 源代码将在 https://github.com/volcengine/verl 开放。
1 Introduction
GPT、Llama 和 Claude 等大语言模型已经革新了各种人工智能应用，从写作、搜索到编码。LLM 首先在来自书籍、网站等来源的数万亿 token 上进行预训练，通过下一词预测来积累广泛知识。随后，LLM 会在特定领域数据集上通过监督微调进行训练，从而能够遵循人类指令。尽管在经过预训练和监督微调之后，LLM 在自然语言任务上表现出卓越能力，但训练数据集中的有害和偏置内容仍可能误导 LLM 生成有毒或不受欢迎的内容。因此，引入了人类反馈强化学习，用于进一步使 LLM 与人类价值观对齐，从而构建有帮助且无害的 AI 应用。
RLHF 建立在传统强化学习算法之上，例如近端策略优化和 REINFORCE。广泛采用的基于 PPO 的 RLHF 系统通常由四个 LLM 组成：actor、critic、reference policy network 和 reward model。基于 PPO 的 RLHF 按迭代进行，每次迭代包含三个阶段：第一，使用 actor 模型对一批 prompt 生成 response；第二，通过 critic、reference policy 和 reward model 的一次前向传播，对生成的 response 进行打分，从而准备训练数据；第三，通过前向和反向计算更新 actor 和 critic，从人类偏好中学习。其他 RLHF 变体通常遵循类似阶段，但涉及不同数量的模型以及模型之间不同的数据依赖关系。
传统 RL 可以被建模为一个数据流，也就是一个有向无环图：RL 数据流中的每个节点表示一个神经网络的计算，例如 actor 或 critic 网络，这些网络可以是 CNN 或 MLP；每条边表示神经网络计算之间的数据依赖，例如 critic 的输出被用作 actor 训练的输入。RLHF 数据流更加复杂，因为其中涉及更复杂的模型，例如作为 actor、critic、reference 和 reward model 的 LLM；每个模型运行不同的计算，并且模型之间存在更加多样的数据依赖，也就是分布式模型分区之间的组播。RLHF 数据流中 LLM 的训练和生成需要分布式计算，例如使用张量并行、流水线并行和数据并行。因此，RLHF 数据流中的每个节点都是一个复杂的分布式程序，对应于相应 LLM 的分布式计算。不同节点中的模型通常使用不同的并行策略，因为它们的工作负载不同。边表示数据重新分片，这通常是多对多组播。因此，灵活表示并高效执行复杂且资源密集的 RLHF 是非常必要的。
传统 RL 框架如 RLLib 和 RLLib Flow 使用层次化的单控制器范式来运行 RL 数据流。一个集中式控制器将数据流中的节点分配到不同进程，并协调它们的执行顺序。每个节点进程还可以进一步派生更多 worker 来执行计算，同样遵循单控制器范式。然而，这些框架只提供用于数据并行训练的原语，并且仅限于最多数百 MB 大小的神经网络。在 RLHF 数据流中，每个节点对应一个 LLM，其包含多达数十亿个算子，并使用某种复杂并行方式进行计算。由于需要向分布式加速器分发大量算子，单控制器范式效率较低。
现有 RLHF 系统采用多控制器范式来管理节点内计算和节点间数据重新分片。每个控制器独立管理一个设备上的计算，并使用多个点对点操作来协调不同节点之间的数据依赖。在执行 LLM 计算时，这种多控制器范式引入的分发开销可以忽略不计。
然而，在缺少中央控制的情况下，实现各种 RLHF 数据流并不灵活，因为为了适配不同的数据依赖而修改单个节点时，需要同时修改所有依赖节点的实现，从而阻碍代码复用。
为了解决这些限制，我们提出 HybridFlow，这是一个灵活且高效的 RLHF 框架，可以方便地表示并执行多样化的 RLHF 数据流，同时获得较高吞吐量。我们的关键观察是：在节点间层面使用单控制器范式，可以灵活表达各种数据依赖，并以极小开销方便地协调节点间数据重新分片；而在节点内计算中集成多控制器范式，则可以显著提升计算效率。
我们主张使用一种层次化混合编程模型来生成 RLHF 数据流。在节点层面，我们提供多个模型类，将数据流中不同 LLM 的分布式计算，包括训练、推理和生成，封装成原语 API。这些 API 可以无缝支持现有 LLM 框架中的多种并行策略，包括 3D 并行、ZeRO 和 PyTorch FSDP，并在多控制器范式下执行分布式计算。在节点之间，我们设计了一组传输协议，用于隐藏数据重新分片的复杂性，并由单控制器进行协调。这种编程模型抽象掉了分布式计算的复杂性，使用户只需几行代码即可实现一个 RLHF 数据流，并通过单控制器的单个进程运行 RLHF。它还有效地解耦了节点内计算和节点间数据传输，使得每个模型可以独立优化，而无需修改数据流中其他模型的代码。
actor 模型的训练和生成构成了 RLHF 数据流中的主要计算。我们进一步设计了 3D-HybridEngine，用于高效执行 actor 模型的训练和生成，并在训练阶段与生成阶段之间进行模型参数重新分片时实现零内存冗余和显著降低通信开销。我们的混合编程模型还支持将模型灵活放置到相同或不同的 GPU 设备集合上。这使我们能够提供一种有效算法，为任意 RLHF 数据流优化 GPU 分配和模型放置，适配不同模型规模和不同工作负载。
我们在设计 HybridFlow 时的贡献总结如下：
我们提出了一种层次化混合编程模型，用于方便地构建 RLHF 数据流。该编程模型能够对各种 RLHF 算法实现高效的节点内分布式执行，以及灵活的节点间数据重新分片和传输。
我们设计了 3D-HybridEngine，用于以高计算效率执行 actor 模型的训练和生成，并在训练阶段和生成阶段之间实现零冗余切换。
我们设计了一种有效的映射算法，可以自动识别 RLHF 数据流中每个节点，即每个模型，优化后的 GPU 分配和放置方案。
我们进行了广泛实验，在各种 RLHF 算法、模型大小和集群规模下，将 HybridFlow 与最先进的 RLHF 系统进行比较。评估结果显示，HybridFlow 实现了 1.53× 到 20.57× 的吞吐提升。
我们已经开源 HybridFlow，并相信 HybridFlow 能够促进未来 RLHF 研究与开发。
2 Background and Motivation
2.1 Reinforcement Learning from Human Feedback
RLHF 将 LLM 的语言空间与人类价值观对齐，它使用给定 prompt 下由人类排序的候选结果集合。一个 RLHF 系统通常包含多个模型，例如 actor、critic、reference policy，以及一个或多个 reward model。actor 和 reference 都是预训练或微调后的 LLM，也就是正在接受 RLHF 的 LLM。critic 和 reward model 可以是不同的 LLM，它们在人工偏好数据集上微调，并将语言建模头替换为标量输出头。RLHF 工作流可以分解为三个阶段，这里以 PPO 为例。
第一阶段是生成阶段。actor 使用自回归生成，从一批 prompt 中生成 response。
第二阶段是准备阶段。使用 prompt 和生成的 response，critic 计算它们的 value，reference policy 计算它们的参考 log probability，reward model 计算它们的 reward。这些都通过相应模型的一次前向计算完成。
第三阶段是学习或训练阶段。actor 和 critic 使用前面阶段产生的一批数据以及损失函数，通过 Adam 进行更新。
其他 RLHF 算法大体上也遵循三阶段工作流。Safe-RLHF 在 PPO-ptx 的基础上引入一个辅助预训练损失，并额外包含一个 cost model，用于同时拟合人类偏好和安全标签。ReMax 需要额外一次生成过程用于方差降低，并从数据流中移除了 critic 模型。研究者正在积极探索新的 RLHF 算法，并将传统 RL 方法集成到 RLHF 领域。这些变化要求系统能够灵活表示 RLHF 数据流图，以适应多样化的算法需求。
LLM 使用数据并行、流水线并行和张量并行进行训练和服务。在数据并行中，输入数据被划分为多个子集，每个子集由一个独立设备处理，例如 GPU。ZeRO 是一种针对数据并行训练的内存优化方案，它逐步将优化器状态、梯度和模型参数分片到多个 GPU 上。流水线并行和张量并行则将模型参数、梯度和优化器状态分布到多个 GPU 上。现代分布式训练框架如 Megatron-LM 和 MegaScale 使用 3D 并行或 PTD 并行，其中 P、T、D 分别代表流水线并行、张量并行和数据并行。在 3D 并行中，PP size 表示模型训练中的流水线阶段数量，TP size 表示一个张量被划分成的分片数量，DP size 表示模型副本数量。LLM 服务系统采用与训练类似的 3D 并行方式，但只对模型参数和 KVCache 进行分片。
RLHF 数据流中的 LLM 模型可能执行不同计算，包括训练、推理和生成。训练包括一次前向传播、一次反向传播和模型更新；推理包括一次前向传播；生成则是包含多次前向传播的自回归生成。具体来说，actor 模型执行训练和生成，critic 执行训练和推理，reference policy 和 reward model 执行推理。不同模型针对不同计算可以采用不同并行策略，以实现最优吞吐量。
2.2 Programming Model for Distributed ML
单控制器范式使用一个集中式控制器来管理分布式程序的整体执行流程。借助集中式控制逻辑，用户可以将数据流的核心功能构建为一个单进程程序，而控制器自动生成分布式 worker 来执行计算。由于拥有硬件和数据流图的全局视图，单控制器范式可以灵活地进行资源映射优化和数据流任务之间的执行顺序协调。然而，当在大规模集群上执行庞大的数据流图时，协调消息需要从控制器传递到所有 worker，从而产生显著的分发开销。
多控制器范式中，每个设备，也就是每个 worker，都有自己的控制器。最先进的分布式 LLM 训练和服务系统采用多控制器范式，因为它具有可扩展性且分发开销低；控制消息主要通过快速 PCIe 链路从 CPU 传递到 GPU。在多控制器 RLHF 实现中，每个模型运行一个独立程序，同一模型的所有 worker 执行相同程序。每个 worker 只拥有系统状态的局部视图，并需要在两个模型之间进行点对点通信，以协调模型执行顺序。为了在多控制器架构中实现 RLHF 工作流，用户必须在每个设备运行的程序中细致地整合集体通信、计算和点对点数据传输代码。这会导致计算和数据传输代码深度嵌套，使开发、维护和优化都变得困难。
2.3 RLHF Characteristics
RLHF 中的 actor、critic、reference 和 reward model 可能在不同阶段执行训练、推理或生成，因此具有不同的内存占用和计算需求。对于 reference policy 和 reward model，由于它们只执行前向传播，因此 GPU 内存中只需要存储模型参数。对于 actor 和 critic，由于它们需要进行模型训练，因此必须存储模型参数、梯度和优化器状态。此外，在 RLHF 中，一个较小的 actor 模型，例如 7B 的预训练或微调 LLM，可以与更大的 critic 和 reward model，例如 70B LLM，配对使用，以获得更好的对齐效果。由于这种异构性，在 RLHF 运行过程中，每个模型都需要不同的并行策略和定制优化。
在 RLHF 数据流中，actor 模型的训练和生成由两个节点表示，它们通常占据每次 RLHF 迭代中的大部分工作负载。actor 训练是计算受限的，通常需要更大的模型并行规模，也就是模型被划分的分区数量，并将工作负载分布到更多 GPU 上。例如，可以将 7B 模型划分为 8 个分区并放在 8 张 GPU 上。对于生成阶段，如果使用相同的并行策略，例如相同的模型并行规模，由于生成是内存受限的，可能会导致 GPU 计算资源利用不足。已有研究表明，将更大的数据并行规模与更小的模型并行规模结合使用，例如将 7B 模型划分为两份并在 8 张 GPU 上复制四次，可以提升生成吞吐量。虽然为 actor 训练和生成使用不同并行策略可能优化两个阶段的吞吐量，但在运行时在两个阶段之间重新分片 actor 模型权重，会带来显著的通信和内存开销。例如，对齐一个 70B actor 模型需要在每次 RLHF 迭代中从训练阶段向生成阶段传输 140GB 模型权重；当两个阶段位于不同设备上时，这最多会占用一次迭代时间的 36.4%。
根据模型的计算负载和数据依赖，在 RLHF 数据流中策略性地放置模型是必要的。放在不同设备集合上的模型，如果不存在数据依赖，就可以并行执行。放在同一组 GPU 上的模型被称为共置模型，它们共享 GPU 内存，并以时间共享方式顺序执行，因为如果共置的 LLM 并发执行，很容易发生显存不足错误。
我们观察到一种折中：将模型放置在不同设备上允许并行处理，但由于 RLHF 的分阶段模型执行方式，不可避免地会造成一些 GPU 空闲时间。支持多种放置策略并最大化设备利用率，对于在任意模型规模和集群规模下优化 RLHF 性能至关重要。
2.4 Limitations of existing RLHF systems
现有 RLHF 系统采用多控制器范式来实现数据流。为了实现各种 RLHF 算法，用户必须处理和管理混合在一起的代码，这些代码包含集体通信、模型计算，以及点对点数据传输，其中模型计算还可能使用不同的分布式训练或服务框架。这种代码结构缺乏模块化和函数封装，使 RLHF 系统与特定 LLM 训练和服务框架紧密耦合。因此，用户需要针对不同 RLHF 数据流逐一实现和优化，阻碍代码复用，并增加出错风险。现有 RLHF 框架只支持 PPO 算法。此外，由于实现复杂度高，它们支持的并行策略也有限。例如，如果想在 DeepSpeed-Chat 中为 LLM 训练和生成加入 3D 并行，可能必须重新实现整个系统，因为原有代码结构混杂在一起。
现有 RLHF 系统采用的并行策略、模型放置和执行模式存在效率问题。DeepSpeed-Chat 和 OpenRLHF 在 actor 训练中采用 ZeRO-3，在 actor 生成中采用 TP。OpenRLHF 在不同设备上为训练和生成使用两份 actor 模型副本，这会产生冗余内存使用，并导致设备之间频繁进行权重同步。DeepSpeed-Chat 在同一组设备上为训练和生成维护同一份 actor 模型，但由于两个阶段使用不同并行方式，因此需要在训练和生成之间重新分片模型权重，这对于大模型仍然可能产生巨大的内存和通信开销。NeMo-Aligner 在 actor 训练和生成中使用相同的 3D 并行配置，因此生成吞吐量较低。
现有 RLHF 框架通常只能支持一种模型放置方案，因此也只能支持一种 RLHF 执行模式。实现一种不同的放置方式很困难，因为需要修改模型初始化和节点间数据传输的内部逻辑。OpenRLHF 和 NeMo-Aligner 允许在准备阶段和学习阶段并发进行模型计算；但在生成阶段，除 actor 以外的模型都处于空闲状态，从而浪费它们占用的 GPU。DeepSpeed-Chat 将所有模型共置在同一组设备上，每个设备按照 RLHF 数据流顺序运行每个模型。当模型之间工作负载不平衡时，这种放置方式在资源利用上可能效率较低。
2.5 Design Considerations
为了解决现有系统的限制，关键问题是：如何设计一种灵活且高效的编程模型来实现 RLHF 数据流？单控制器设计在节点间层面尤其有优势，因为它能够灵活协调不同模型分布式计算之间的数据传输、执行顺序和资源虚拟化。RLHF 数据流图通常只包含少量节点。与数据流中各节点，也就是模型，所需的分布式计算相比，从单控制器向不同节点分发控制消息所带来的开销可以忽略不计。而多控制器范式以低延迟将算子分发给加速器而著称，可以用于每个模型的分布式计算。基于这些观察，我们提出一种层次化混合编程模型，用于实现 RLHF 数据流。我们的关键设计原则是以混合方式结合单控制器和多控制器范式。这种设计可以保证 RLHF 数据流既能灵活表达，又能高效执行，并在节点间和节点内两个层面保持较低的控制开销。该范式解耦了节点内分布式计算和节点间数据传输，使每个模型只需专注于本地计算，而无需管理节点间通信。
3 HybridFlow Overview
HybridFlow 的架构由三个主要组件组成：Hybrid Programming Model、3D-HybridEngine 和 Auto-Mapping algorithm。混合编程模型包含一组分层 API，用于灵活表达 RLHF 数据流，并高效执行数据流中的模型计算。3D-HybridEngine 专门用于高效执行 actor 模型的训练和生成，允许两个阶段采用不同的 3D 并行配置，并在两个阶段之间切换时实现零内存冗余和最小化通信开销。自动映射算法用于确定每个模型的优化设备放置方案，以最大化 RLHF 吞吐量。
RLHF 系统的工作流如下。用户需要提供以下输入来启动 RLHF 系统：第一，模型规格，包括 RLHF 数据流中 actor、critic、reference policy 和 reward model 的架构和大小；第二，在给定 GPU 集群配置下运行自动映射算法得到的数据流中各模型的设备放置；第三，在每个阶段运行每个模型所使用的并行策略，例如 3D 并行中的三元组 $(p,t,d)$，其中 $p$、$t$、$d$ 分别表示 PP size、TP size 和 DP size。单控制器程序接收这些输入，用于初始化 RLHF 数据流中的模型和虚拟化资源池，根据放置计划将操作或模型分发到设备，并调用设备上由多个控制器运行的函数，以执行每个模型的分布式计算。
多控制器程序实现 ParallelWorker 类：它根据每个模型的并行策略，在分配的设备之间构建该模型的并行组；调用 3D-HybridEngine 进行 actor 的训练和生成；并且可以与现有 LLM 引擎无缝集成，用于其他模型的训练、推理和生成。传输协议由单控制器程序协调，用于支持在采用不同并行策略的模型之间重新分片数据，包括 RLHF 中的 prompt、response 和其他模型输出。actor 在训练和生成之间的数据重新分片由 3D-HybridEngine 处理。
4 Hybrid Programming Model
4.1 Hierarchical APIs
对于 RLHF 不同阶段中每个模型的分布式计算，我们提供了一个基类 3DParallelWorker。给定分配好的设备后，它能够进行分布式模型权重初始化，并为每个模型建立 3D 并行组。一个并行组包含一组 GPU，用于承载模型的某个特定并行维度，例如 TP 中的不同张量分片和 DP 中的不同模型副本。
继承自 3DParallelWorker 类，我们分别提供了 actor、critic、reference 和 reward model 等模型类。每个模型类都封装了 API，用于实现模型的分布式前向和反向计算、自回归生成以及优化器更新，从而将分布式计算代码与其他模型之间的数据依赖解耦。这些 API 可以通过复用现有 LLM 系统中的计算脚本轻松实现。例如，ActorWorker 中 update_actor 函数涉及的计算与 Megatron-LM 中的预训练脚本类似。一个模型类封装了用于实现各种 RLHF 算法的基础操作，例如 actor 模型类中的 generate_sequences 用于基于 prompt 生成 response，reward model 类中的 compute_reward 用于通过一次前向传播评估 response。
除了实现 3D 并行的基类 3DParallelWorker，我们还提供了用于 PyTorch FSDP 的 FSDPWorker 和用于 ZeRO 的 ZeROWorker，以及继承这些基类的对应模型类，以支持模型计算中的不同并行策略。
在采用不同并行策略并放置在不同设备上的模型之间进行数据传输时，会涉及多对多组播。我们通过将每个模型类中的每个操作与一个传输协议关联起来，统一了这种数据传输实现。每个传输协议都由一个 collect 函数和一个 distribute 函数组成，分别根据每个模型的并行策略聚合输出数据和分发输入数据。数据重新分片通过源模型输出的 collect 函数和目标模型输入的 distribute 函数实现。实际数据传输只发生在 GPU 之间，从而避免任何中心瓶颈。
我们提供了 8 种传输协议，包括 3D_PROTO、DP_PROTO、ONE_TO_ALL 等，它们覆盖了大部分数据重新分片场景。用户还可以通过实现自定义的 collect 和 distribute 函数来扩展传输协议。
我们提供了一个 ResourcePool 类，用于虚拟化一组 GPU 设备。当将一个 ResourcePool 实例应用于某个模型类时，该模型的分布式计算会被映射到这些设备上。使用同一个 ResourcePool 实例的模型会共置在同一组 GPU 上；如果模型类使用不同的 ResourcePool 实例，则模型会放置在不同的 GPU 集合上。我们假设不同 ResourcePool 实例之间没有重叠。
当模型放置在不同设备集合上时，只要它们的输入可用，就会自动触发执行。当一些模型放置在同一组设备上时，它们会根据调用顺序顺序执行。借助我们的编程模型，HybridFlow 可以灵活支持多样化的分布式执行模式，而无需修改 RLHF 算法代码。
4.2 Implementation of different RLHF algorithms
我们的 API 使各种 RLHF 算法，也就是各种数据流，能够被流畅地开发出来。用户可以用几行代码将一个 RLHF 算法实现为运行在单控制器上的单进程程序，该程序由一系列原语 API 调用组成，用于调用各模型的分布式计算。
PPO 可以只通过调用若干模型操作来实现，例如 compute_values 和 generate_sequences，这些操作会在多个 GPU 上以多控制器范式执行。为了适配 Safe-RLHF，后者集成了一个额外的 cost model 来评估安全偏好，并为 actor 引入预训练损失，只需要在 PPO 实现之上增加少量代码。为了适配 ReMax，只需要额外调用一次 actor 生成，并移除 critic 相关代码。
这种扩展灵活性对于研究者探索不同 RLHF 算法至关重要：他们可以复用每个模型类中封装的分布式计算，并只需根据具体算法调整数值计算代码，例如 compute_advantage 中的 GAE 和 KL divergence，以及 actor 和 critic 的损失函数。这种简化开发流程得益于混合编程模型。我们的模块化 API 设计简化了开发，促进了广泛的代码复用，并允许直接整合现有 LLM 训练和服务框架的代码库。它还解耦了模型计算和模型之间的数据传输。分布式框架中的任何变化都不会影响 RLHF 算法代码，从而允许对每个模型的执行进行个性化优化。系统还支持具有不同工作负载的模型灵活放置，从而能够将 RLHF 数据流优化映射到各种设备上。
5 3D-HybridEngine
我们设计 3D-HybridEngine，以支持 actor 模型的高效训练和生成，从而显著提升 RLHF 吞吐量。
5.1 Parallel Groups
为了消除冗余的 actor 模型副本，我们主张将 actor 训练和生成阶段部署在分配给 actor 的同一组设备上，并在同一份 actor 模型权重上顺序执行。然而，actor 的训练和生成很可能采用不同的 3D 并行策略；也就是说，与训练阶段相比，生成阶段通常需要更小的 TP 和 PP size，但需要更大的 DP size。3D-HybridEngine 能够在这种情况下，在同一组设备上高效地完成 actor 训练和生成之间的模型参数重新分片。
令 $p-t-d$ 表示为 actor 训练构建的 3D 并行组，分别对应承载流水线阶段、张量分片和模型副本的 GPU 集合。3D-HybridEngine 根据训练和生成阶段不同的 3D 并行策略，分别为 actor 训练和生成构建不同的并行组。我们使用 $p_g$、$t_g$ 和 $d_g$ 分别表示生成阶段中的生成流水线并行组大小、生成张量并行组大小和微数据并行组大小。$d_g$ 表示生成阶段中的模型副本数相对于训练阶段的比例，也就是说，训练中的每个 DP 副本在生成阶段会变成 $d_g$ 个 micro-DP 副本，用于处理 $d_g$ 个 prompt 和 response 的微批次。micro-DP 组只在 actor 生成阶段使用，用于获得更大的 DP size，从而充分利用设备。生成并行组表示为 $p_g-t_g-d_g-d$。
5.2 3D-HybridEngine Workflow
在 RLHF 第 $i$ 次迭代的 actor 训练与第 $i+1$ 次迭代的 actor 生成之间，actor 模型参数需要根据两个阶段的并行组配置进行重新分片，并且 prompt 数据也需要被分发。在第 $i+1$ 次 RLHF 迭代中，3D-HybridEngine 会收集第 $i$ 次迭代中更新后的 actor 模型参数，以便在每个 micro-DP 组内进行生成。然后，一批 prompt 会被加载到每个模型副本中，并生成 response。随后，3D-HybridEngine 会在每个 micro-DP 组内对生成结果执行 all-gather 操作，并根据 actor 训练所需的 3D 并行方式重新划分模型参数。当模型权重、prompt 和 response 都被正确重新分布后，系统会根据 RLHF 算法计算 actor 模型损失并更新 actor 模型权重，也就是进行第 $i+1$ 次迭代的 actor 训练阶段。
5.3 Zero redundancy model resharding
3D 并行中的传统并行分组方式通常如下：PP 组和 TP 组通过为流水线阶段和张量分片分配连续 rank 来形成；DP 组则通过按固定间隔选择 rank 构建，该间隔由 PP size 和 TP size 的乘积决定。假设训练和生成使用相同的并行分组方式，但并行大小不同。在从训练切换到生成时，3D-HybridEngine 会在模型并行组之间执行 all-gather 操作，以聚合所有参数，然后根据每个设备所属的生成并行组，仅保留该设备生成所需的模型权重子集。在某些 GPU 上，训练权重和生成权重之间没有重叠，因此需要额外内存保存后续训练所需的权重。当 3D-HybridEngine 在两个阶段使用上述普通并行分组方式时，我们称该系统为 HybridFlow-V。
我们进一步为 3D-HybridEngine 设计了一种新的生成阶段并行分组方法，它消除了权重存储冗余，并在训练和生成之间重新分片 actor 模型时实现最小的内存占用和通信开销。具体来说，我们通过按照固定间隔选择 rank 来形成生成 TP 和 PP 组，该间隔分别由 $t/t_g$ 和 $p/p_g$ 决定；并通过沿生成 TP 或 PP 维度顺序分配 rank 来构建 micro-DP 组。这种对生成并行组的策略性重新排列，使每个设备上的训练权重和生成权重发生重叠，从而能够在生成阶段复用训练权重，并在模型重新分片时实现设备内存使用的零冗余。此外，3D-HybridEngine 会并发执行多个 all-gather 操作，每个 micro-DP 组内执行一个，从而显著降低通信开销。
5.4 Transition overhead
我们比较了不同 actor 引擎设计在训练和生成阶段之间切换时的通信开销和内存占用。假设 actor 模型大小为 $M$，并使用 $N_a$ 张 GPU 进行训练和生成。DeepSpeed-Chat 中的 actor 引擎在切换期间会跨所有 GPU 执行一次 all-gather 操作；HybridFlow-V 则在训练 TP 和 PP 组内执行该 all-gather。两种引擎都会先在每张 GPU 内存中聚合所有模型参数，然后再根据生成并行组划分模型状态，因此模型参数的峰值内存使用为 $M$。由于在某些 GPU 上它们无法在生成期间复用训练权重，因此还必须在这些 GPU 上保留训练权重，从而产生冗余内存消耗。
使用我们为生成阶段设计的并行分组方法后，HybridFlow 将 all-gather 操作限制在每个 micro-DP 组内。每张 GPU 只需要在其 micro-DP 组内收集远端参数，并且可以在生成中复用训练权重。因此，HybridFlow 中模型参数的峰值内存使用恰好等于生成阶段每张 GPU 上的模型分片大小，消除了 GPU 内存使用中的任何冗余。
6 Auto Device Mapping
我们的混合编程模型要求用户输入以下配置，这些配置被称为将 RLHF 数据流映射到给定设备上的映射方案：第一，数据流中模型的设备放置；第二，在每个阶段运行每个模型所对应的并行策略。我们提供一种高效算法，帮助用户识别在给定设备集群上执行 RLHF 数据流的优化映射方案，从而最小化每次 RLHF 迭代的端到端延迟。
给定一个数据流，算法首先探索给定集群中模型的所有可能放置方案。例如，PPO 算法涉及四个模型，因此会产生 15 种可能的放置方式，这来自 Bell 分区问题。这些放置方式从完全独立放置，即所有模型都放在不同设备上，到将所有模型共置在同一组设备上。我们将共置在同一组 GPU 上的模型称为一个共置集合。共置集合中的模型可以在同一组 GPU 上采用不同并行策略。算法会根据共置模型的内存消耗，为每个共置模型集合识别需要分配的最小 GPU 数量，以确保不会发生显存不足错误。
接下来，算法从最小 GPU 分配开始，枚举每个共置模型集合的所有可行设备分配。给定一个共置集合的设备分配和其中模型的计算工作负载，算法会在 auto_parallel 模块中为每个模型探索优化并行策略，以最小化模型执行延迟。工作负载包括每个模型的输入输出形状以及计算类型，例如训练、推理或生成。在 auto_parallel 中，我们使用模拟器模块估计不同并行策略的延迟。
d_cost 模块在给定模型放置和并行策略的情况下，通过遍历数据流图中的所有阶段并累计所有阶段延迟，来估计 RLHF 数据流的端到端延迟。对于处于同一共置集合且在同一阶段执行计算的模型，它们的执行延迟会相加。对于处于不同共置集合的模型，它们在同一阶段内的执行可以并行，因此该阶段的延迟由不同集合中的最大执行时间决定。最终，算法识别出具有对应并行策略的最佳模型设备放置，从而实现每次 RLHF 迭代的最小执行时间。
算法复杂度取决于模型数量和用于运行数据流的设备总数。这是枚举某个放置策略下所有可能设备分配时的最坏情况复杂度。为了提高效率，我们缓存每个模型在某个设备数量下识别出的并行策略，从而在模型于不同放置策略中被放置到相同数量 GPU 上时，避免重复搜索相同的并行策略。
虽然我们在运行自动映射算法时假设 GPU 是同构的，但该算法可以很容易扩展到异构设备上的模型映射优化，只需在模拟器和自动并行模块中考虑异构设备即可。
7 Implementation
HybridFlow 使用大约 1.2 万行 Python 代码实现。
混合编程模型中的分层 API 使用 1.8 千行代码实现。集中式单控制器构建在 Ray 之上，并使用远程过程调用来协调不同模型的执行顺序，以及按照数据流在模型之间传输数据。这些中间数据存储在 TensorDict 中。在我们的多控制器分布式计算范式中，每个模型函数都作为一个独立进程在多个设备上运行，控制消息从每个控制器的 CPU 进程传递到对应 GPU。我们的实现支持 Megatron-LM、PyTorch FSDP 和 DeepSpeed 作为 LLM 训练和推理引擎，并支持 vLLM 用于自回归生成。在 vLLM 中，我们将集中式 KVCache 管理器替换为分布式管理器，以适配多控制器范式。
3D-HybridEngine 的主要逻辑使用 2.4 千行代码实现，建立在 Megatron-LM 和 vLLM 之上。我们将 actor 模型在训练阶段和生成阶段的权重存储在不同内存缓冲区中，在训练期间将生成权重卸载到 CPU 内存，在切换期间再将生成权重重新加载回 GPU 内存，并在生成中使用两个缓冲区。我们使用 NCCL 通信原语，在训练和生成之间切换时，在每个 micro-DP 组内收集和拼接模型参数。我们在生成之后将 KVCache 卸载到 CPU 内存，并在下一次迭代中重新加载回 GPU。
自动映射算法使用 1.9 千行代码实现，并包含用于训练、推理和生成工作负载的三个模拟器。该算法在 RLHF 数据流启动前运行在 CPU 上，用于生成数据流初始化所需的设备映射和并行策略。
8 Evaluation
8.1 Experimental Setup
我们将 HybridFlow 部署在一个包含 16 台机器的集群上，共 128 张 GPU。每台机器配备 8 张 NVIDIA A100-80GB GPU，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。
我们运行 PPO、ReMax 和 Safe-RLHF 算法的 RLHF 数据流。PPO 是 RLHF 中最流行的算法之一，由 actor、critic、reference policy 和 reward model 组成。每个模型都是 Llama 模型，大小从 7B 到 70B 不等。Safe-RLHF 额外包含一个 cost model，其结构和大小与 reward model 相同；ReMax 则移除了 critic model。所有实验中，actor 和 critic 训练使用混合精度，即模型参数使用 BF16，梯度和优化器状态使用 FP32，并使用 Adam 优化器。模型推理和自回归生成使用 BF16。如果未特别说明，实验结果来自 PPO。
我们将 HybridFlow 与最先进的 RLHF 系统进行比较，包括 DeepSpeed-Chat、OpenRLHF 和 NeMo-Aligner。NeMo-Aligner 不支持 ReMax 算法。我们没有将 HybridFlow 与 Trlx、HuggingFaceDDP 和 Collosal-Chat 等其他框架比较，因为它们代表性较低，并且已有研究报告显示它们慢于上述基线。
我们使用 RLHF 吞吐量作为性能指标，其计算方法是将全局 batch 中 prompt 和 response 的总 token 数除以一次 RLHF 迭代时间。所有报告的性能数值都是在预热 10 次迭代之后，对 5 次训练迭代取平均得到的。
我们在 HuggingFace 的 “Dahoas/full-hh-rlhf” 数据集上执行 RLHF，该数据集被广泛用于 LLM 对齐。由于基线系统在生成期间可能没有集成 continuous batching 优化，为了公平比较，我们强制所有生成的 response 具有相同长度。在每个实验中，输入 prompt 长度和输出 response 长度均为 1024，actor 模型输入 prompt 的全局 batch size 为 1024。PPO epoch 数为 1，每个 epoch 中 PPO update iteration 数为 8，与以往 RLHF 研究保持一致。
8.2 End-to-End performance
实验展示了在运行 PPO、ReMax 和 Safe-RLHF 时的 RLHF 吞吐量。在这组实验中，actor、critic、reference 和 reward model 具有相同大小，这遵循了以往实践。不同模型大小的实验中，使用的 GPU 数量从能够在不发生显存不足情况下运行 RLHF 的最小 GPU 数，到 128 张 GPU 不等。为了公平比较，实验中没有启用优化器状态卸载。
总体来看，HybridFlow 在所有模型规模上都始终优于基线。对于 PPO，HybridFlow 相比 DeepSpeed-Chat、OpenRLHF 和 NeMo-Aligner 分别平均提升 3.67×、3.25× 和 12.52×。这主要是因为 HybridFlow 能够通过对模型使用不同并行策略来适配各种计算工作负载，从而在 RLHF 的所有阶段中高效执行生成、推理和训练。训练 70B 模型时，HybridFlow 获得最高平均加速比 9.64×，因为 HybridFlow 与 DeepSpeed-Chat 和 OpenRLHF 相比，分别将切换开销最多降低了 71.2% 和 89.1%；而后两者在使用 ZeRO-3 训练时还会产生大量跨机器通信。由于生成引擎中缺少 KVCache，NeMo-Aligner 的主要性能瓶颈位于生成阶段，该阶段最多占其 RLHF 迭代时间的 81.2%。类似结果也出现在 ReMax 和 Safe-RLHF 中，验证了 HybridFlow 在运行各种 RLHF 算法时的效率。
HybridFlow 在 8 张 GPU 上至少实现了 2.09× 加速。随着 GPU 数量增加，HybridFlow 在不同模型规模上的强扩展效率平均为 66.8%。在固定全局 batch size 的情况下扩展到大量 GPU，会导致每个 worker 的本地 batch size 变小，从而可能造成 GPU 利用不足。在 128 张 GPU 上运行 7B 模型时，HybridFlow 在 PPO、ReMax 和 Safe-RLHF 上仍分别比最佳基线 OpenRLHF 快 1.68×、1.53× 和 1.71×。这可以归因于 HybridFlow 能够针对不同模型和集群规模自适应选择最佳放置策略，以最小化 RLHF 时间。OpenRLHF 在较大 GPU 集群中表现更好，但在较小集群上效率较低。
8.3 Model Placement
在这个实验中，我们在 HybridFlow 中实现了 PPO 算法的多种模型放置方式，并使用与前面相同的模型和集群设置。第一种是 colocate，即 DeepSpeed-Chat 的放置策略；第二种是 standalone，即 OpenRLHF 的放置策略；第三种是 split，即 NeMo-Aligner 的共置放置策略，其中 actor 和 reference policy 放在同一组设备上，critic 和 reward model 放在另一组设备上；第四种是 hybridflow，即由算法得到的优化放置。
不同模型放置方式的比较显示，在不同 GPU 数量下，HybridFlow 的优化放置方式会发生变化。从 16 到 64 张 GPU，将所有模型共置在同一组设备上能获得最佳性能。对于 34B 模型在 96 到 128 张 GPU 上，以及 13B 模型在 96 张 GPU 上，split 策略成为最优。由于模型大小相同，split 策略在两组模型之间均分 GPU。对于 13B 模型在 128 张 GPU 上，standalone 策略获得最高吞吐量。在这种情况下，HybridFlow 为 actor 分配 64 张 GPU，为 critic 分配 32 张 GPU，并分别为 reference 和 reward model 分配 16 张 GPU。在较小集群中，所有模型的计算都可以充分利用 GPU 资源；colocate 策略确保了不同 RLHF 阶段中的最大 GPU 使用率。在较大集群中，由于 batch size 固定，随着更多 GPU 上 DP size 变大，计算与通信比下降，colocate 放置下的 RLHF 吞吐量无法线性扩展。standalone 和 split 策略在较大集群中将模型放在不同设备上，并为每个模型使用较小 DP size，从而促进同一阶段中不同模型的并行执行。在所有情况下，算法都能够产生具有最高训练吞吐量的最佳放置方案。
我们进一步评估了使用 13B actor 和 reference policy，以及 70B critic 和 reward model 运行 PPO 时的模型放置情况。更大的 critic 和 reward model 预期能产生更好的对齐效果。结果显示，在最多 64 张 GPU 时，colocate 策略仍平均比其他策略高 44.8%。在 96 张 GPU 时，split 策略获得更高吞吐量。当扩展到 128 张 GPU 时，算法得到的最佳放置方案是将 actor、reference 和 reward model 共置在 64 张 GPU 上，同时将剩余 64 张 GPU 分配给 critic。在相同 GPU 数量下，actor 和 reference policy 的计算时间远小于 critic 和 reward model，因此将 reward model 与 actor 和 reference policy 共置，可以减少经验准备阶段的 GPU 空闲时间。总体而言，在大规模集群中，将 actor 和 critic 分布到不同设备上，以便在训练阶段并行执行，可以带来更高吞吐量。
8.4 3D-HybridEngine
实验比较了不同模型规模下 actor 训练和生成阶段之间的切换时间，也就是在相同设置下将模型权重从训练重新分片到生成的时间。OpenRLHF 的切换时间包括位于不同设备上的两份 actor 模型副本之间的权重同步时间。HybridFlow 平均将切换时间降低了 55.2%，也就是 11.7 秒，并且在 70B 模型上最多将切换开销降低 89.1%，也就是 78.2 秒，同时在不同集群规模上保持一致的开销。这归因于我们为生成阶段设计的新并行分组方法。在基线方法中，所有模型参数都必须在切换期间被收集，为了防止显存不足，需要逐层多次收集。HybridFlow 在切换期间实现零内存冗余，并且每个 micro-DP 组只需要一次 all-gather 操作。
我们进一步验证了在 HybridFlow 中，actor 训练和生成需要使用不同并行大小。在该实验中，所有模型都共置在同一组 GPU 上，生成所需的 KVCache 使用剩余 GPU 内存进行尽力分配。实验展示了在 16 张 GPU 上使用 7B 和 13B 模型运行 RLHF 时，在训练并行组固定的情况下，改变生成 TP 组大小所得到的切换时间和生成时间。我们观察到，对 7B 模型使用较小的生成 TP 组大小 $t_g=2$，对 13B 模型使用 $t_g=4$，分别使生成延迟降低了 60.3% 和 36.4%。相反，若像 NeMo-Aligner 一样在生成阶段使用与训练阶段相同的 TP size，则由于 GPU 利用不足，会得到最大的生成延迟。进一步减小 $t_g$ 无法获得更高加速，因为更小的 $t_g$ 需要每张 GPU 维护更大的 KVCache。
8.5 Algorithm Runtime
设备映射算法的运行时间显著短于实际 RLHF 训练所需的数天时间。运行时间呈线性增长，显示出该设备映射算法在模型大小和集群大小增长时具有良好可扩展性。大部分运行时间花在估计每个模型的不同并行策略执行延迟上。更大的模型会有更多可用并行策略，因此需要更多模拟来为每种放置方案识别最优策略。通过缓存模型的最优并行策略，并在不同放置中复用，最佳放置搜索时间最多可减少到半小时以内。
9 Discussions
Fault Tolerance
HybridFlow 与现有容错方法正交，并且已经集成了 checkpointing。故障可以通过 NCCL 错误检测，静默数据损坏可以通过校验和检测。我们的编程模型使单控制器能够通过 RPC 协调 checkpoint 操作，从而在每个 ParallelWorker 组内保存模型状态。这包括保存 actor 和 critic 模型参数、dataloader ID 以及随机数生成器状态，以确保系统级一致性。此外，如果有足够健康的模型副本可用，HybridFlow 也可以采用基于冗余的容错方法，例如参数广播和 CPU checkpoint，以实现快速恢复。
Placement Insights
我们总结出 RLHF 训练中模型放置和 GPU 分配的三个主要经验。第一，为 actor 模型分配更多 GPU 可以降低耗时较长且无法与其他模型并行的生成延迟。第二，当每个模型计算都可以充分利用 GPU 资源时，在相对较小规模集群上将所有模型共置是最有效的。第三，当扩展到大规模集群时，也就是进行强扩展时，将 actor 和 critic 模型分布在不同设备上，以便在训练和准备阶段并行执行，有助于获得更高吞吐量。
Resource multiplexing
HybridFlow 通过对 GPU 计算采用时间共享，支持在共享设备上共置模型。近期 DNN 任务调度研究已经发展出细粒度资源复用技术，主要目标是实现单个任务的服务级目标。虽然 ResourcePool 实现支持共置模型的并行执行，但 HybridFlow 通常仍采用顺序执行，以避免 GPU 资源争用或显存不足问题。在 RLHF 训练中应用 GPU 共享和异构资源面临独特挑战，因为它需要平衡各种任务之间的计算工作负载，并管理它们复杂的数据依赖。在 RLHF 训练中研究用于 GPU 共享的细粒度自动映射算法，并结合模型卸载优化和异构设备集成，将是一个有前景的未来方向。
From alignment to reasoning
在用于 LLM 对齐的 RLHF 中，奖励信号由 reward model 生成。除了对齐任务之外，类似算法，例如 PPO 和 GRPO，也可以应用于其他领域，例如代码生成和数学推理。对于这些任务，每个 prompt 可能存在一个 ground truth，可以通过评估每个代码测试用例的输出值是否正确，或者验证数学结果是否准确来确定。因此，reward model 可以被非神经网络奖励模块替代，例如用于评估生成代码的沙箱环境，或者用于验证数学结果的奖励函数。HybridFlow 可以通过将这些奖励模块包装为远程函数，并在单进程脚本中编排它们的执行，来无缝集成这些奖励模块，从而为多样化强化学习应用提供灵活且高效的框架。
10 Related Work
RL frameworks
已有大量 RL 框架，从为小规模 DNN 设计的通用 RL 系统，到专门为 LLM 优化的 RLHF 系统。我们已经在前文中详细讨论了紧密相关工作，并在本节讨论更多 RL 框架。这些 RL 框架与近期 RLHF 系统类似，使用混杂的多控制器框架来实现算法。它们建立多个长期运行的分布式程序，每个组件通过硬编码的数据同步来协调执行顺序。Gear 进一步优化了 RL 流水线中的经验回放部分。然而，所有这些框架都不能支持 RLHF 中的 LLM 训练、推理和生成。
LLM training and serving systems
TorchDDP 和 Horovod 支持数据并行训练。ByteScheduler 和 DeepSpeed 通过通信和内存优化扩展了数据并行。大量系统通过张量并行和流水线并行等模型并行方式，将模型划分到不同设备上，从而优化大模型训练。LLM 服务系统也采用数据并行和模型并行来加速自回归生成，并使用 continuous batching 和 chunked prefill 等专门优化。需要注意的是，上述所有框架都采用多控制器范式来实现高效计算。
Dataflow systems
MapReduce、Spark、Dryad 和 Naiad 等数据流系统广泛用于分析和机器学习工作负载，但它们缺乏对动态任务图的支持。Ray 在单个动态任务图中统一了任务并行和 actor 编程模型，并实现了可扩展的分布式调度器和全局控制存储，因此被许多 RL 框架采用。Pathways 是一个面向 TPU 的闭源项目，旨在方便表达单个 DNN 模型内部的复杂并行模式和细粒度控制流，例如流水线并行和具有稀疏计算的混合专家模型。它采用异步分布式数据流设计，即使存在数据依赖，也能支持控制平面的并行执行，从而降低单控制器范式的分发开销。它的主要关注点是单模型训练，需要对 DNN 模型的每个子网络进行复杂编译。HybridFlow 可以将 Pathways 集成为子模块，用于实现 RLHF 数据流中各模型的计算。
11 Conclusion
HybridFlow 是一个 RLHF 框架，能够灵活表示并高效执行多样化的 RLHF 算法。我们提出了一种混合编程模型，通过将不同 LLM 的分布式计算封装为原语 API，并隐藏节点之间数据重新分片的复杂性，使用户能够用几行代码轻松构建 RLHF 数据流。我们的 3D-HybridEngine 保证了 actor 模型训练和生成的高效执行，并在模型参数重新分片时实现零内存冗余和显著降低通信开销。此外，我们的有效映射算法优化了 RLHF 数据流中的 GPU 分配和模型放置。大量实验表明，在不同模型大小和集群规模下，与最先进的 RLHF 系统相比，HybridFlow 实现了 1.53× 到 20.57× 的加速。
Acknowledgments
我们感谢 shepherd Y. Charlie Hu 以及匿名审稿人的建设性反馈。我们感谢 Xin Liu、Yangrui Chen 和 Ningxin Zheng 对本项目提出的深刻反馈。本工作部分得到了 ByteDance Research Collaboration Project，以及香港 RGC 项目 HKU 17204423 和 C7004-22G（CRF）的支持。
评论