GPTやLlamaなどの大規模言語モデル(LLM)は、人工知能分野に革命を起こしましたが、これらの巨大なモデルを効率的に訓練し、人間の価値観に沿わせる方法は依然として課題です。

強化学習と人間のフィードバック(RLHF)は、重要なLLM訓練方法として近年広く利用されていますが、従来のRLHFフレームワークは、柔軟性、効率性、拡張性に限界があります。

image.png

これらの問題を解決するために、バイトダンスのDoubao大規模言語モデルチームは、HybridFlowというRLHFフレームワークをオープンソース化し、LLM訓練に新たな可能性をもたらしました。

RLHFは通常、3つの段階で構成されます。

まず、actorモデルが入力されたプロンプトに基づいてテキストを生成します。次に、criticモデル、referenceモデル、rewardモデルが生成されたテキストを評価し、対応する価値、参照確率、報酬値を計算します。

image.png

最後に、これらの評価結果を使用してactorモデルを訓練し、人間の好みにより沿ったテキストを生成するようにします。従来のRLHFフレームワークは、通常、単一の制御装置を使用してデータフロー全体を管理しますが、分散計算が必要なLLMにとっては効率が低いです。

HybridFlowフレームワークは、単一制御装置と複数制御装置のモードを革新的に組み合わせ、階層的なAPI設計により複雑な計算とデータ依存関係を解きほぐすことで、RLHFデータフローの柔軟な表現と効率的な実行を実現します。

image.png

HybridFlowの主な利点は、以下の3つの点です。

様々なRLHFアルゴリズムとモデルを柔軟にサポート:HybridFlowはモジュール式のAPIを提供しており、ユーザーはPPO、ReMax、Safe-RLHFなど、様々なRLHFアルゴリズムを容易に実装および拡張できます。

効率的なモデルウェイトの再構成:3D-HybridEngineコンポーネントは、訓練と生成の段階でactorモデルのモデルウェイトを効率的に再構成し、メモリ冗長と通信オーバーヘッドを最小限に抑えます。

自動化されたモデル展開と並列化戦略の選択:Auto Mappingコンポーネントは、モデルの負荷とデータ依存関係に基づいて、モデルを異なるデバイスに自動的にマッピングし、最適な並列化戦略を選択することで、モデル展開プロセスを簡素化し、訓練効率を向上させます。

実験結果によると、HybridFlowは様々なRLHFアルゴリズムを実行する際のスループットが大幅に向上し、最大20.57倍に達しました。HybridFlowのオープンソース化は、RLHFの研究開発に強力なツールを提供し、将来のLLM技術の発展を促進します。

論文アドレス:https://arxiv.org/pdf/2409.19256