PaperlayerView exact revision on arXiv

Source-linked AI summary

HybridFlow: A Flexible and Efficient RLHF Framework

Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, Chuan Wu

arXiv:2409.19256v2cs.LGcs.DC

TL;DR

RLHF creates complex distributed dataflows whose many-to-many dependencies and intra-node computation challenge traditional controller designs. HybridFlow combines single-controller inter-node coordination with multi-controller intra-node execution through hierarchical APIs and a 3D-HybridEngine. Across various RLHF algorithms, it reports 1.53×∼20.57× throughput improvement over state-of-the-art baselines.

  • Problem

    RLHF dataflows combine distributed LLM computations with many-to-many data dependencies, while existing systems trade dispatch efficiency for inflexible dataflow implementation and limited placement support.

  • Method

    HybridFlow uses hierarchical APIs, hybrid controller coordination, flexible GPU placement, and a 3D-HybridEngine for actor training-generation resharding.

  • Results

    1.53×∼20.57× throughput improvement is reported across RLHF algorithms, model sizes, and cluster scales versus state-of-the-art baselines.

  • Takeaways & Limitations

    HybridFlow enables flexible representation and efficient execution of diverse RLHF algorithms while reducing memory redundancy and model-parameter resharding overhead.

  • Takeaways & Limitations

    Existing RLHF frameworks are limited to one model placement plan and one execution pattern, making different placements difficult to implement.

Abstract

from arXiv · show

Reinforcement Learning from Human Feedback (RLHF) is widely used in Large Language Model (LLM) alignment. Traditional RL can be modeled as a dataflow, where each node represents computation of a neural network (NN) and each edge denotes data dependencies between the NNs. RLHF complicates the dataflow by expanding each node into a distributed LLM training or generation program, and each edge into a many-to-many multicast. Traditional RL frameworks execute the dataflow using a single controller to instruct both intra-node computation and inter-node communication, which can be inefficient in RLHF due to large control dispatch overhead for distributed intra-node computation. Existing RLHF systems adopt a multi-controller paradigm, which can be inflexible due to nesting distributed computation and data communication. We propose HybridFlow, which combines single-controller and multi-controller paradigms in a hybrid manner to enable flexible representation and efficient execution of the RLHF dataflow. We carefully design a set of hierarchical APIs that decouple and encapsulate computation and data dependencies in the complex RLHF dataflow, allowing efficient operation orchestration to implement RLHF algorithms and flexible mapping of the computation onto various devices. We further design a 3D-HybridEngine for efficient actor model resharding between training and generation phases, with zero memory redundancy and significantly reduced communication overhead. Our experimental results demonstrate 1.53$\times$~20.57$\times$ throughput improvement when running various RLHF algorithms using HybridFlow, as compared with state-of-the-art baselines. HybridFlow source code will be available at https://github.com/volcengine/verl.

1 Introduction

RLHF expands traditional RL dataflows into distributed LLM computations with many-to-many data dependencies, creating challenges for both execution efficiency and flexible representation. HybridFlow addresses these challenges by combining controller paradigms, hierarchical APIs, and efficient model resharding.

  • Motivation: RLHF dataflows expand each node into distributed LLM training or generation and each edge into many-to-many data resharding.Different models can use distinct computation and parallelism strategies, making the dataflow more complex than traditional RL.
  • Limitations: Traditional single-controller RL frameworks can incur large dispatch overhead when coordinating distributed intra-node LLM computation.Existing frameworks also provide only data-parallel training primitives and target comparatively small neural networks.
  • Limitations: Existing RLHF multi-controller systems reduce LLM dispatch overhead but make dataflow changes inflexible because dependent nodes must also be modified.This coupling hinders code reuse across RLHF algorithms and data dependencies.
  • HybridFlow: HybridFlow combines single-controller inter-node coordination with multi-controller intra-node distributed computation.Hierarchical APIs encapsulate model computation and transfer protocols, supporting diverse parallelism strategies and RLHF dataflows.
  • HybridFlow: 3D-HybridEngine enables efficient actor training-generation transitions with zero memory redundancy and reduced communication overhead.The framework also supports flexible model placement across shared or separate GPU sets.
  • Results: 1.53×∼20.57× throughput improvements are reported across RLHF algorithms, model sizes, and cluster scales.The evaluation compares HybridFlow with state-of-the-art baselines.

2 Background and Motivation

RLHF combines multiple models and computation stages, while varying workloads and placement strategies make efficient execution difficult. HybridFlow’s programming model and resource-management approach target these flexibility and utilization challenges.

  • RLHF workflow: RLHF typically uses actor, critic, reference-policy, and reward models across generation, preparation, and training stages.Variants can add a cost model or generation pass, or eliminate the critic, creating different dataflow graphs.
  • Parallelism strategies: LLM training, inference, and generation use data, pipeline, tensor, or 3D parallelism according to computation and memory requirements.These strategies distribute data, model parameters, gradients, optimizer states, or KVCache across GPUs.
  • RLHF characteristics: RLHF models have heterogeneous workloads and memory footprints, so a single parallelism strategy can underutilize GPUs.Actor training is computation-bound, whereas generation is memory-bound and may benefit from smaller model parallelism with larger data parallelism.
  • Model placement: Separating models can enable concurrent computation but may leave GPUs idle during other staged RLHF operations.An example placement leaves actor and critic GPUs idle for 1/3 of their GPU time during other stages.
  • Existing-system limitations: Existing RLHF frameworks tightly couple communication and model computation, limiting supported algorithms, parallel strategies, and reusable implementations.The cited discussion states that existing frameworks support only PPO and that adding 3D parallelism may require reimplementing a system.
  • Existing-system limitations: Existing systems also constrain execution to limited placement patterns, which can waste GPUs or perform poorly with unbalanced model workloads.Some systems leave non-actor models idle during generation, while colocating all models forces sequential execution.

3 HybridFlow Overview

HybridFlow combines hierarchical APIs, a 3D-HybridEngine, and auto-mapping to represent RLHF dataflows flexibly and execute model computation efficiently. Its controller structure separates dataflow coordination from distributed model execution.

  • 3 HybridFlow Overview: HybridFlow comprises a Hybrid Programming Model, 3D-HybridEngine, and Auto-Mapping algorithm.The programming model provides hierarchical APIs for RLHF dataflows; the engine targets actor training and generation.
  • 3 HybridFlow Overview: The hybrid programming model uses hierarchical APIs to express RLHF dataflows and execute their model computations efficiently.
  • 3 HybridFlow Overview: The system accepts model specifications, device placements from auto-mapping, and per-stage parallelism strategies as inputs.
  • 3 HybridFlow Overview: A multi-controller program constructs model parallel groups and invokes the 3D-HybridEngine, while a single controller coordinates transfer protocols for data resharding.

4 Hybrid Programming Model

HybridFlow’s hierarchical APIs encapsulate distributed model computation and unify inter-model data resharding. This separation supports multiple parallelism strategies, flexible placement, and concise implementations of different RLHF algorithms.

  • 4 Hybrid Programming Model: 3DParallelWorker initializes distributed model weights and 3D parallel groups on allocated devices.Parallel groups host tensor shards and model replicas for the selected model.
  • 4 Hybrid Programming Model: Model-specific worker classes encapsulate forward, backward, generation, and optimizer operations while decoupling computation from other models’ data dependencies.
  • 4 Hybrid Programming Model: FSDPWorker and ZeROWorker base classes extend the framework beyond 3D parallelism to support alternative distributed computation strategies.
  • 4 Hybrid Programming Model: Transfer protocols pair collect and distribute functions to implement many-to-many data resharding between models using different parallelism strategies and devices.HybridFlow provides protocols including 3D_PROTO, DP_PROTO, and ONE_TO_ALL, while allowing customized protocols.
  • 4 Hybrid Programming Model: ResourcePool virtualizes GPU devices so models can be colocated or flexibly mapped without changing the RLHF algorithm code.
  • 4 Hybrid Programming Model: PPO, ReMax, and Safe-RLHF can be implemented as concise single-controller programs that invoke distributed model operations.PPO is implemented in 8 lines, and adapting to other algorithms requires adding or deleting a few lines of code.

5 3D-HybridEngine

3D-HybridEngine runs actor training and generation on the same devices while supporting different parallel configurations and efficient parameter resharding. Its grouping strategy enables weight reuse, zero redundant memory, and reduced transition communication.

  • 5 3D-HybridEngine: 3D-HybridEngine targets efficient actor training and generation by using different 3D parallel configurations for the two stages.
  • 5 3D-HybridEngine: The engine sequentially executes training and generation on the same actor-weight copy to eliminate redundant actor model copies.
  • 5 3D-HybridEngine: Training uses p-t-d groups, while generation uses p_g-t_g-d_g-d groups with micro data-parallel replicas sized to improve generation device utilization.
  • 5 3D-HybridEngine: During each RLHF iteration, the engine gathers updated parameters, distributes prompts, gathers generation results, repartitions parameters, and performs actor training.
  • 5 3D-HybridEngine: A rearranged generation grouping overlaps training and generation weights on each device, enabling weight reuse and zero redundant resharding memory.
  • 5 3D-HybridEngine: Concurrent all-gathers within micro DP groups significantly reduce communication overhead during transition between training and generation.
  • 5 3D-HybridEngine: HybridFlow’s peak parameter memory matches each GPU’s generation model partition size because generation reuses training weights.

6 Auto Device Mapping

HybridFlow’s auto-mapping algorithm searches device placements, allocations, and model-specific parallelism strategies to minimize RLHF iteration latency. It accounts for memory capacity, workload-dependent execution, and stage-level overlap, while caching repeated parallelism searches.

  • Placement search: The algorithm explores all model-placement plans, from standalone allocation to fully colocated models, and groups models sharing GPU sets.For PPO with four models, it enumerates 15 placement plans and enforces minimum allocations based on colocated-model memory consumption.
  • Parallelism search: For each feasible allocation, auto_parallel selects a latency-minimizing parallelism strategy using each model’s workload and assigned GPUs.Workloads include input and output shapes and whether computation is training, inference, or generation.
  • Cost estimation: The d_cost module estimates iteration latency by summing stage latencies, using summed colocated execution times and maximum concurrent-set latency across placements.This models stage overlap between different colocated sets while aggregating models executing together within one set.
  • Optimization objective: The algorithm returns the placement and allocation with the minimum estimated end-to-end latency per RLHF iteration.It compares the best allocation for every placement plan before selecting the global best mapping.
  • Complexity and caching: The worst-case search complexity is O((N−1)! (k−1)!(N−k)!), where k is the number of models and N the number of devices.Caching parallelism strategies for a model on A devices eliminates redundant searches across different GPU sets.
  • Scope: The mapping algorithm assumes homogeneous GPUs but can be extended to heterogeneous devices through the simulation and auto_parallel modules.The stated extension requires considering heterogeneous devices in those modules.

7 Implementation

HybridFlow implements its hierarchical hybrid programming model with a centralized controller for inter-model orchestration and multi-controller execution for distributed model computation.

  • Implementation: HybridFlow is implemented in around 12k lines of Python code, including 1.8k lines for its hierarchical APIs.The APIs provide the implementation basis for the hybrid programming model.
  • Centralized orchestration: The centralized single controller uses Ray and Remote Process Calls to coordinate model execution order and transfer data along the dataflow.Intermediate data are stored in TensorDict.
  • Distributed computation: Each model function runs in the multi-controller paradigm for distributed computation.This separates distributed model execution from centralized inter-model coordination.

8 Evaluation

HybridFlow is evaluated on PPO, ReMax, and Safe-RLHF across Llama models and cluster scales using throughput as the main metric. It consistently outperforms baselines, while optimized placement and actor resharding reduce execution overhead.

  • Experimental setup: The evaluation runs PPO, ReMax, and Safe-RLHF on a 128-GPU cluster using Llama models ranging from 7B to 70B.Baselines include DeepSpeed-Chat, OpenRLHF, and NeMo-Aligner; NeMo-Aligner does not support ReMax.
  • Metric: RLHF throughput is measured in tokens/sec by dividing global-batch prompt and response tokens by one RLHF iteration time.Reported performance averages five post-warm-up training iterations.
  • End-to-end performance: HybridFlow consistently outperforms the baselines across model scales and RLHF algorithms.The evaluation covers PPO, ReMax, and Safe-RLHF throughput figures under comparable model-scale experiments.
  • End-to-end performance: 3.67×, 3.25×, and 12.52× are HybridFlow’s average PPO speedups over DeepSpeed-Chat, OpenRLHF, and NeMo-Aligner, respectively.The corresponding maximum speedups are 7.84×, 5.93×, and 20.57×.
  • Scalability: 1.68×, 1.53×, and 1.71× are HybridFlow’s speedups over OpenRLHF for 7B models on 128 GPUs running PPO, ReMax, and Safe-RLHF, respectively.The results show benefits even when scaling a small model to a large GPU cluster.
  • 3D-HybridEngine: 55.2% (11.7s) is HybridFlow’s average reduction in actor transition time, reaching 89.1% (78.2s) for 70B models.The 3D-HybridEngine uses zero memory redundancy and one all-gather per micro DP group during transition.
  • 3D-HybridEngine: 60.3% and 36.4% are generation-latency reductions from using generation TP sizes of 2 for 7B models and 4 for 13B models, respectively.Using the training TP size of 8 produces the largest generation latency because of GPU underutilization.
  • Auto device mapping: The device-mapping algorithm’s runtime grows linearly with model and cluster size and remains much shorter than actual RLHF training.Most runtime is spent simulating execution latency for available parallelism strategies.

9 Discussions

HybridFlow incorporates checkpointing and coordinates checkpoint operations through a single controller. Its resource-multiplexing design supports collocated models, but fine-grained GPU sharing and heterogeneous-resource mapping remain future directions.

  • Fault Tolerance: HybridFlow detects failures through NCCL errors and silent data corruption through checksums.These mechanisms complement its checkpointing support.
  • Fault Tolerance: A single controller coordinates checkpoint operations via RPC across ParallWorker Groups.Checkpointed state includes actor and critic parameters, dataloader IDs, and random-number-generator states.
  • Resource multiplexing: HybridFlow supports parallel execution of collocated models but generally uses sequential execution to avoid GPU contention or out-of-memory issues.
  • Resource multiplexing: Fine-grained GPU sharing, heterogeneous-device integration, model offload optimization, and automatic GPU mapping are identified as promising future directions.

10 Related Work

Related RL frameworks span general-purpose systems for small-scale DNNs and RLHF systems optimized for LLMs. Related infrastructure also advances communication, memory, model-parallelism, and autoregressive-serving optimizations.

  • RL frameworks: Existing RL frameworks include general-purpose systems for small-scale DNNs and systems specifically optimized for LLM-based RLHF.
  • RL frameworks: Many RL and RLHF frameworks use multi-controller designs to implement their algorithms.
  • Distributed training systems: ByteScheduler and DeepSpeed extend data parallelism with communication and memory optimizations.
  • Distributed training systems: Large-model training systems use tensor and pipeline parallelism to partition models across devices.
  • LLM serving systems: LLM serving systems accelerate autoregressive generation with data and model parallelism, continuous batching, and chunked prefill.

11 Conclusion

HybridFlow targets flexible representation and efficient execution for diverse RLHF algorithms. It combines hierarchical distributed-computation APIs, efficient actor resharding, and GPU placement optimization, achieving 1.53× to 20.57× speedups over state-of-the-art RLHF systems.

  • Conclusion: HybridFlow enables flexible representation and efficient execution of diverse RLHF algorithms.
  • Conclusion: Its hybrid programming model encapsulates distributed LLM computation and hides data resharding complexity through primitive APIs.
  • Conclusion: The 3D-HybridEngine provides actor-model resharding between training and generation with zero memory redundancy and reduced communication overhead.
  • Conclusion: 1.53× to 20.57× speedup was achieved against state-of-the-art RLHF systems across various model sizes and cluster scales.

A Primitive APIs in HybridFlow

HybridFlow implements RLHF model primitives through worker classes that decouple distributed computation from user-facing operations. These primitives support common distributed inference and training operations.

  • Primitive APIs: Each RLHF model primitive is implemented by inheriting from 3DParallelWorker, FSDP Worker, or ZeROWorker.
  • Primitive APIs: The model classes decouple distributed computation code and expose fundamental RLHF operations to users.
  • Primitive APIs: The primitives support autoregressive generation, forward passes, backward passes, and model updates.

B Transfer Protocols

HybridFlow provides transfer protocols that decouple RLHF data resharding from distributed training, while its parallelism search selects model configurations using device, workload, memory, and latency constraints.

  • Transfer Protocols: Transfer protocols cover common RLHF model-resharding cases and allow users to define custom protocols with collect and distribute functions.These protocols decouple complicated data resharding from distributed training.
  • Auto Parallelism: The auto-parallelism algorithm enumerates feasible configurations from minimal model-parallel sizes using available GPUs and GPUs per machine.Minimal allocations help prevent out-of-memory errors when multiple workers are colocated.
  • Auto Parallelism: Latency estimation uses analytical simulators for training, inference, and generation workloads, distinguishing compute-bound training or inference from memory-bound generation.Actor generation additionally accounts for KVCache requirements based on batch size and maximum sequence length.
  • Auto Parallelism: HybridFlow selects the parallelism plan with the lowest estimated cost among feasible configurations.The algorithm updates the best plan when a candidate cost is lower, then returns that plan.
  • Model APIs: Model classes expose key functions that users can combine to construct diverse RLHF algorithms in a few lines of code.The provided functions are presented as programming interfaces for assembling RLHF algorithms.
Loading 2409.19256v2…