【RL】HybridFlow 1 摘要

HybridFlow是一个灵活且高效的强化学习与人类反馈 (Reinforcement Learning from Human Feedback, RLHF) 框架，旨在解决现有RLHF系统在LLM对齐中面临的效率和灵活性挑战。传统RL框架的单控制器范式在RLHF中效率低下，因为每个节点都是一个分布式LLM训练或生成程序，导致大量的控制分派开销。现有RLHF系统采用多控制器范式，虽然减少了分派开销，但由于分布式计算与数据通信的深度嵌套，使得框架不灵活，难以适应不同的RLHF算法和模型配置。

核心思想与方法 (Core Methodology)

HybridFlow的核心贡献在于提出了一种分层混合编程模型 (Hierarchical Hybrid Programming Model)，将单控制器范式用于节点间（inter-node）的协调，将多控制器范式用于节点内（intra-node）的分布式计算。这种混合范式既保证了数据流的灵活表达和协调，又实现了模型内部计算的高效性。

1. 分层API设计 (Hierarchical API Design)

   • 节点内分布式计算封装：HybridFlow为RLHF数据流中的每个LLM模型（如Actor、Critic、Reference Policy、Reward Model）提供了一组基础类（例如3DParallelWorker、FSDPWorker、ZeROWorker）。这些类封装了模型在不同RLHF阶段的分布式计算（训练、推理、生成），并支持现有的LLM并行策略，如3D并行（PTD并行，包含Pipeline Parallelism §, Tensor Parallelism (T), Data Parallelism (D)）、ZeRO和PyTorch FSDP。通过这种封装，用户可以重用现有LLM框架的计算脚本，并使模型计算逻辑与数据依赖解耦。例如，Actor模型的update_actor函数就封装了分布式训练逻辑。
   • 节点间数据重分片统一实现：针对模型间复杂的many-to-many多播数据传输（例如Actor生成响应后，需要将响应发送给Critic和Reward Model进行处理），HybridFlow引入了传输协议 (Transfer Protocols)。每个模型操作都与一个特定的传输协议关联（通过@register装饰器）。一个传输协议包含collect函数（聚合源模型的输出数据）和distribute函数（根据目标模型的并行策略分发输入数据）。例如，3D_PROTO用于3D并行模型间的数据传输，它在DP组内收集数据，并分发给目标DP组。数据传输只发生在GPU之间，避免了中央瓶颈。这种设计隐藏了用户处理数据重分片的复杂性，使得RLHF算法的实现更加简洁，用户只需调用高级API即可。
   • 异步数据流执行与灵活模型放置：通过ResourcePool类虚拟化GPU设备，用户可以灵活配置模型在设备上的放置。模型可以放置在独立的设备集合上并行执行，也可以放置在同一设备集合上（共置colocated）顺序执行以节省资源。HybridFlow的编程模型支持异步数据流执行，一旦输入可用，模型即可触发执行。

2. 3D-HybridEngine

   • 问题：Actor模型的训练和生成阶段通常需要不同的最优并行策略。例如，训练可能需要更大的模型并行度（TP/PP），而生成由于内存密集型（KVCache），可能需要更大的数据并行度或不同的TP/PP配置。在同一组设备上进行训练和生成时，actor模型参数在两种并行策略间的重分片会引入显著的通信和内存开销。
   • 解决方案：3D-HybridEngine允许Actor模型的训练和生成阶段在同一组设备上采用不同的3D并行配置。它通过重新设计并行分组方法来优化模型参数的重分片，实现**零内存冗余 (Zero Memory Redundancy)**和显著降低的通信开销。
   • 优化细节：
     • 传统的并行分组方法（如HybridFlow-V所采用的）在训练和生成阶段的TP/PP组之间可能存在不重叠的参数，导致在转换时需要all-gather全部参数，并在某些GPU上保留训练阶段的冗余权重。
     • HybridFlow的3D-HybridEngine采用了一种优化并行分组方法：在生成阶段，TP和PP组通过选择规律间隔的rank来构建，而微数据并行（micro DP）组则沿着生成TP或PP维度顺序分配rank。这种设计确保了训练和生成阶段的模型权重在每个设备上共享，从而消除了冗余内存使用，并允许在每个微DP组内进行并发的all-gather操作，大大减少了通信量。
     • 具体来说，对于模型大小为MMM，使用NaN_aNa​个GPU的Actor模型，DeepSpeed-Chat的通信量为tpd−1tpdM\frac{tpd-1}{tpd} Mtpdtpd−1​M，HybridFlow-V为tp−1tpM\frac{tp-1}{tp} Mtptp−1​M(其中t,p,dt, p, dt,p,d为训练阶段的TP, PP, DP大小)。而HybridFlow的通信量降至dg−1tpM=tp−tgpgtgpgtpM\frac{dg-1}{tp} M = \frac{tp-tgpg}{tgpgtp} Mtpdg−1​M=tgpgtptp−tgpg​M(其中tg,pg,dgtg, pg, dgtg,pg,dg为生成阶段的TP, PP, micro DP大小)。HybridFlow实现了零冗余内存开销，峰值内存占用仅为生成阶段每个GPU上模型分区的大小。

3. 自动设备映射算法 (Auto Device Mapping Algorithm)

   • 目标：在给定GPU集群上，为RLHF数据流找到最优的设备放置方案和并行策略，以最小化RLHF迭代的端到端延迟。
   • 方法：该算法（Algorithm 1）枚举所有可能的模型放置方案（共置集），并为每个共置集计算最小GPU分配 (A_min)。然后，它遍历所有可行的设备分配，并为每个模型利用auto_parallel模块（包含训练、推理、生成负载的模拟器simu）来找到最优并行策略。
   • 优化：算法会评估不同放置方案下的总延迟。对于同一共置集内且在同一阶段进行计算的模型，它们的延迟累加。对于不同共置集内的模型，如果它们在同一阶段并发执行，则该阶段的总延迟取最大值。通过缓存已识别的并行策略，减少了重复搜索。

实现与评估 (Implementation and Evaluation)

HybridFlow的实现基于Ray构建其单控制器和RPC通信机制，同时支持Megatron-LM、PyTorch FSDP、DeepSpeed作为LLM训练和推理引擎，以及vLLM进行自动回归生成。其中，vLLM的集中式KVCache管理器被替换为分布式管理器以兼容多控制器范式。

实验结果表明：

• 端到端性能：HybridFlow在PPO、ReMax和Safe-RLHF等多种RLHF算法上，对于7B到70B不同规模的Llama模型，始终优于DeepSpeed-Chat、OpenRLHF和NeMo-Aligner等现有SOTA基线。相较于基线，HybridFlow实现了1.53倍至20.57倍的吞吐量提升。尤其在70B模型上，由于其显著减少了训练到生成阶段的转换开销（最高可达89.1%），优势更为明显。
• 可扩展性：HybridFlow在8个GPU上的加速比至少为2.09倍，并在扩展到128个GPU时表现出良好的强扩展效率（平均66.8%）。
• 模型放置：通过自动映射算法，HybridFlow能动态选择最优的模型放置策略。在GPU数量较少时（如最多64个GPU），将所有模型共置在同一设备集合上通常性能最优。但在大规模集群中（如96个或更多GPU），将Actor和Critic模型分别放置在不同设备上，以实现训练阶段的并行执行，能够带来更高的吞吐量。当模型大小不一致时（如13B Actor/Reference vs. 70B Critic/Reward），共置策略在小规模集群中仍然更优，但在大规模集群中，将计算量大的Critic和Reward模型独立放置能提高整体效率。
• 3D-HybridEngine验证：实验证明，为Actor的生成阶段采用更小的TP并行度（例如7B模型为tg=2tg=2tg=2，13B模型为tg=4tg=4tg=4）能够显著降低生成延迟（分别为60.3%和36.4%），验证了3D-HybridEngine的有效性。

总结 (Conclusion)

HybridFlow通过其创新的混合编程模型，实现了RLHF数据流的灵活表达和高效执行。它提供了一套分层API，将分布式计算封装在模型内部，并通过统一的传输协议处理节点间数据重分片。特制的3D-HybridEngine在Actor训练和生成之间实现了零冗余内存和最小化通信开销的模型参数重分片。此外，其自动设备映射算法能够智能地优化GPU分配和模型放置。这些特性共同使得HybridFlow在各种RLHF算法、模型规模和集群规模下，相较于现有框架具有显著的性能优势。