OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework (2405.11143v1)

Abstract: As 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 LLMs 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 https://github.com/OpenLLMAI/OpenRLHF.

OpenRLHF: Making Reinforcement Learning from Human Feedback More Accessible

Introduction

Reinforcement learning from human feedback (RLHF) has been getting a lot of buzz in the AI community for its ability to align LLMs with human values and intentions. However, training these colossal models through RLHF isn't a walk in the park—especially when the model parameters exceed a whopping 70 billion. Training such large models typically means juggling multiple models and resource-intensive computations, introducing a slew of logistical challenges.

Enter OpenRLHF, an open-source framework engineered specifically for scaling RLHF. Instead of co-locating models on the same GPUs, which quickly becomes inefficient as models grow larger, OpenRLHF leverages advanced scheduling techniques using Ray, vLLM, and DeepSpeed. The framework seamlessly integrates with Hugging Face, making it a user-friendly, out-of-the-box RLHF solution.

Design of OpenRLHF

Scheduling Optimization

Traditional RLHF training frameworks like ColossalChat and DeepSpeed-Chat often fall short when dealing with models above 70 billion parameters. They rely on techniques such as Zero Redundancy Optimizer (ZeRO) to co-locate four models (actor, critic, reward, reference) on the same GPU. This setup becomes increasingly inefficient due to the limited GPU memory.

OpenRLHF addresses these challenges by leveraging Ray for scheduling, vLLM for efficient inference, and DeepSpeed for optimized training. Ray's scheduling capabilities enable OpenRLHF to distribute the models across multiple GPUs, avoiding the pitfalls of memory constraints. This approach not only improves resource utilization but also facilitates the integration of multiple reward models, a key component in several alignment strategies.

Here are some notable features:

• Flexible Model Placement: Models can be freely merged or offloaded to save GPU resources.
• Multi-Reward Models: Supports various alignment strategies like separating usefulness from harmfulness.
• Optimal Orchestration: Improves training performance by optimally coordinating GPUs.

Performance Optimization

One of the main bottlenecks in RLHF is the sample generation stage, which can consume up to 80% of the training time. OpenRLHF tackles this by employing several performance-boosting techniques:

1. Inference Optimization: Utilizes vLLM's tensor parallelism and advanced batching techniques to accelerate the sample generation process.
2. Memory Management: Offloads Adam optimizer states to CPU to free up GPU memory, allowing larger batch sizes.
3. Flash Attention 2: This method accelerates Transformer model training.
4. Padding Removal: Uses PyTorch tensor slicing to eliminate redundant padding from training samples.

The impact of these optimizations is substantial. For example, in comparative performance metrics, training a 70B model with OpenRLHF was up to 2.3x faster than a tuned version of DSChat.

Training Stability

Reinforcement learning can be notoriously unstable, especially for large models. To mitigate these risks, OpenRLHF incorporates several stability tricks in its PPO implementation. These include fine-tuning algorithmic parameters and employing techniques like gradient clipping and reward normalization. The framework also ensures weight synchronization between the ZeRO and vLLM engines using NVIDIA's NCCL, making the integration fast and reliable.

Ease of Use

Ease of use is one of OpenRLHF’s standout features. It comes with one-click trainable scripts that are fully compatible with Hugging Face models. Users simply need to specify the model and dataset paths to start RLHF training. Here is a quick example:

ray job submit -- python3 examples/train_ppo_ray.py \ --ref_num_nodes 1 \ --reward_num_nodes 1 \ --critic_num_nodes 1 \ --actor_num_nodes 1 \ --vLLM_num_engines 4 \ --vLLM_tensor_parallel_size 4 \ --pretrain meta-llama/Llama-2-70b-chat-hf \ --reward_pretrain meta-llama/Llama-2-70b-chat-hf \ --prompt_data Open-Orca/OpenOrca

This script covers the basic configuration needed to get Llama-2's 70B model running with PPO on OpenRLHF. The framework also supports a variety of algorithms out-of-the-box, including Direct Preference Optimization (DPO) and Kahneman-Tversky Optimization (KTO).

Conclusion

OpenRLHF brings significant advancements to the RLHF landscape by focusing on efficient model scheduling, robust performance optimization, and user-friendly implementation. Its design principles make it feasible to train and align LLMs beyond 70 billion parameters effectively. As AI models continue to grow in size and complexity, OpenRLHF's approach could become a crucial tool for making them more aligned with human values, scaling efficiently, and maintaining training stability.

Future iterations of OpenRLHF might push these boundaries even further, integrating more advanced algorithms and optimizations. For now, it stands as a robust framework, enabling researchers and practitioners to scale their RLHF endeavors seamlessly.