[論文レビュー] DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales

Zhewei Yao, Reza Yazdani Aminabadi+17 more|arXiv (Cornell University)|Aug 2, 2023

Machine Learning and Data Classification被引用数 10

ひとこと要約

DeepSpeed-Chatを導入した統一オープンソースのRLHFトレーニングパイプラインとHybrid Engineにより、数十億パラメータ級を含む非常に大規模なモデルを、単一GPUやマルチノードクラスタで、ChatGPTのようなモデルを容易かつ高速、コスト効率良く訓練できる。

ABSTRACT

ChatGPT-like models have revolutionized various applications in artificial intelligence, from summarization and coding to translation, matching or even surpassing human performance. However, the current landscape lacks an accessible, efficient, and cost-effective end-to-end RLHF (Reinforcement Learning with Human Feedback) training pipeline for these powerful models, particularly when training at the scale of billions of parameters. This paper introduces DeepSpeed-Chat, a novel system that democratizes RLHF training, making it accessible to the AI community. DeepSpeed-Chat offers three key capabilities: an easy-to-use training and inference experience for ChatGPT-like models, a DeepSpeed-RLHF pipeline that replicates the training pipeline from InstructGPT, and a robust DeepSpeed-RLHF system that combines various optimizations for training and inference in a unified way. The system delivers unparalleled efficiency and scalability, enabling training of models with hundreds of billions of parameters in record time and at a fraction of the cost. With this development, DeepSpeed-Chat paves the way for broader access to advanced RLHF training, even for data scientists with limited resources, thereby fostering innovation and further development in the field of AI.

研究の動機と目的

• ChatGPT風モデルのRLHFトレーニングを、アクセスしやすいエンドツーエンドのパイプラインを提供することで民主化する。
• InstructGPTの3段階RLHFワークフロー（SFT、Reward Model Fine-tuning、RLHF）をデータ抽象化とブレンディングとともに再現する。
• トレーニングと推論の両方を最適化する統一かつ効率的なシステム（Hybrid Engine）を提供する。
• 数十億から数百億パラメータの範囲でモデルサイズの多様性に対してスケーラビリティとコスト効果を示す。

提案手法

• 事前学習済みのHuggingFaceモデルからInstructGPT風トレーニングまで、3つのRLHFステップを実行する使いやすいスクリプトを提供する。
• データ抽象化とブレンディング機能を備えたSFT、Reward Model Fine-tuning、RLHFを順に実行するDeepSpeed-RLHFパイプラインを実装する。
• 高速推論カーネル、テンソル並列、ZeROベースのメモリ最適化、LoRAを組み合わせたDeepSpeed-Hybrid Engine (DeepSpeed-HE) を開発し、Efficient RLHFトレーニングと生成を実現する。
• EMAチェックポイントとMixture Trainingを組み込んで最終モデルの品質を向上させ、事前学習機能を保持する。
• 再利用可能なエンジンとPPOトレーナーを介してRLHFパイプラインをカスタマイズする柔軟なAPIを提供し、研究実験を促進する。
• Colossal-AIやHuggingFace DDPとのスループットとスケーラビリティをベンチマーク・比較し、単一GPUおよびマルチノード設定での改善を強調する。

実験結果

リサーチクエスチョン

• RQ1 ChatGPT風モデルのRLHFトレーニングを、規模を超えてアクセスしやすく、速く、手頃にするにはどうすればよいか？
• RQ2 大規模なアクター/報酬モデルを用いたエンドツーエンドのRLHF（SFT、RMチューニング、RLHF）を可能にするシステム設計と最適化は何か？
• RQ3 統一されたHybrid Engineは、生成と学習のフェーズで既存のフレームワークと比較してどのように性能を発揮するか？
• RQ4 数十億から数百億パラメータのモデルを訓練する際の実用的なコスト、時間、スケーラビリティの利点は何か？
• RQ5 柔軟なAPIを介してRLHFパイプラインをカスタマイズし、新しいRLHF戦略を探求できるか？

主な発見

• DeepSpeed-HEは、報告された設定で既存システムよりもRLHFトレーニングを15x超高速化している。
• 単一ノード 8x A100-40G で OPT-13B を 9時間、OPT-30B を 18時間で訓練し、Azure 上でそれぞれ $300 未満と $600 未満。
• マルチノード 64x A100-80G は OPT-13B を 1.25時間、OPT-175B を 20時間で訓練でき、コストは最大約 $5120。
• DeepSpeed-HEは、13Bを超えるモデルを含むスケーラブルなハードウェア上での訓練を可能にし、単一GPUのサポートも提供する。
• 比較では、DeepSpeed-HEが Colossal-AI より 6–19x、HuggingFace DDP より 1.4–10.5x のスピードアップをマルチGPU設定で実現し、ハードウェアごとに最大で 7.5x 大きい実用モデルサイズが可能。