OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework

2024-05-20

arXiv.org

AI Summary

Existing RLHF frameworks suffer from inference bottlenecks, deployment complexity, challenges in multi-model orchestration, and low resource utilization—hindering the accessibility and scalable deployment of large language model (LLM) alignment. To address these limitations, this paper introduces the first open-source RLHF training framework specifically designed for LLM alignment. It features a cross-GPU heterogeneous scheduling architecture that decouples the reward model, policy model, reference model, and value model for independent deployment. The framework natively supports multiple alignment paradigms—including RLHF, DPO, and rejection sampling—within a unified interface. Leveraging Ray for elastic task orchestration, it tightly integrates vLLM (for high-throughput inference) and DeepSpeed (for efficient training), while maintaining native compatibility with the Hugging Face ecosystem. Experiments demonstrate substantial improvements in training throughput and GPU memory efficiency for models ≥70B parameters. The framework delivers an out-of-the-box, end-to-end alignment solution and has been open-sourced, gaining broad adoption in the research and engineering community.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

As large language models (LLMs) continue to grow by scaling laws, reinforcement learning from human feedback (RLHF) has gained significant attention due to its outstanding performance. However, unlike pretraining or fine-tuning a single model, scaling reinforcement learning from human feedback (RLHF) for training large language models poses coordination challenges across four models. We present OpenRLHF, an open-source framework enabling efficient RLHF scaling. Unlike existing RLHF frameworks that co-locate four models on the same GPUs, OpenRLHF re-designs scheduling for the models beyond 70B parameters using Ray, vLLM, and DeepSpeed, leveraging improved resource utilization and diverse training approaches. Integrating seamlessly with Hugging Face, OpenRLHF provides an out-of-the-box solution with optimized algorithms and launch scripts, which ensures user-friendliness. OpenRLHF implements RLHF, DPO, rejection sampling, and other alignment techniques. Empowering state-of-the-art LLM development, OpenRLHF's code is available at url{https://github.com/OpenRLHF/OpenRLHF}.

Problem

• Addresses inference bottlenecks in RLHF frameworks
• Reduces complexity barriers for newcomers to RLHF
• Improves training efficiency and scalability in RLHF

Innovation

• Built on Ray, vLLM, DeepSpeed, Transformers
• Simplified design with clear code structure
• Superior training efficiency with speedups