HybridFlow: A Flexible and Efficient RLHF Framework

• Published in: arXiv.org
• Main Authors: Sheng, Guangming, Zhang, Chi, Ye, Zilingfeng, Wu, Xibin, Zhang, Wang, Zhang, Ru, Peng, Yanghua, Lin, Haibin, Wu, Chuan
• Format: Paper Journal Article
• Language: English
• Published: Ithaca Cornell University Library, arXiv.org 02.10.2024
• Subjects: Algorithms; Communication; Computation; Computer Science - Distributed, Parallel, and Cluster Computing; Computer Science - Learning; Control systems design; Controllers; Data communication; Large language models; Machine learning; Neural networks; Nodes; Redundancy; Source code
• ISSN: 2331-8422
• DOI: 10.48550/arxiv.2409.19256

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.