OpenRLHF: 使いやすく、スケーラブルで高性能なRLHFフレームワーク

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

May 20, 2024

著者: Jian Hu, Xibin Wu, Weixun Wang, Xianyu, Dehao Zhang, Yu Cao

cs.AI

要旨

大規模言語モデル（LLM）がスケーリング則に従って拡大を続ける中、人間のフィードバックによる強化学習（RLHF）はその優れた性能から大きな注目を集めています。しかし、単一モデルの事前学習やファインチューニングとは異なり、大規模言語モデルの訓練におけるRLHFのスケーリングは、4つのモデルにわたる調整の課題を提起します。本論文では、効率的なRLHFスケーリングを可能にするオープンソースフレームワーク「OpenRLHF」を紹介します。既存のRLHFフレームワークが4つのモデルを同じGPU上に配置するのに対し、OpenRLHFはRay、vLLM、DeepSpeedを活用し、70Bパラメータを超えるモデル向けにスケジューリングを再設計し、リソース利用効率の向上と多様な訓練アプローチを実現しています。Hugging Faceとシームレスに統合し、最適化されたアルゴリズムと起動スクリプトを提供することで、ユーザーフレンドリーな即戦力ソリューションを確保しています。OpenRLHFはRLHF、DPO、リジェクトサンプリング、その他のアライメント技術を実装しています。最先端のLLM開発を支援するOpenRLHFのコードは、https://github.com/OpenLLMAI/OpenRLHF で公開されています。

English

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