Skip to content

強化学習のための直接選好最適化

GitHubでこのAIアクセラレーターにアクセス

直接選好最適化(DPO)は、AIエージェントのトレーニングに使用される機械学習手法であり、従来の人間からのフィードバックによる強化学習(RLHF)に代わる、よりシンプルで安定した手段を提供します。 事前トレーニングだけでは捉えることが難しい方法でAIモデルを人間の好みに合わせるのに役立ち、モデルの有用性、無害性、および誠実さを向上させます。 指示に正確に従うこと、有害なコンテンツを回避すること、より有用な回答を生成することなどを改善できます。

一般的に、この手法は次のように機能します。

  1. モデルは、大量のテキストデータを用いて、シーケンス内で次に何が来るかを予測するようにトレーニングされます(事前トレーニング)。

  2. 人間がモデルのプロンプト出力を評価し、ランク付けします。

  3. フィードバックに基づいて、別のモデルが望ましい回答を予測するように学習します。

  4. メインモデルは、よりスコアの高い出力を生成するために、強化学習アルゴリズムを使用してファインチューニングされます。

このアクセラレーターは、直接選好最適化(DPO)を使用したLLMのファインチューニングから、そのモデルのDataRobotへのデプロイまでのプロセスを自動化します。 基本的には、ベースモデルを用い、提供されたデータセットに基づいて特定のタイプの回答を優先するように学習させ、本番環境で使用できるようにします。

アクセラレーターの具体的なアクション:

データ準備アクション

  • DataRobotレジストリから特定の選好データセットをダウンロードします。
  • Hugging Faceのデータセットライブラリを使用してCSVを読み込みます。このCSVには通常、プロンプト、選択された(望ましい)回答、および拒否された回答の3つの列が含まれています。

モデルトレーニング(DPO)

  • メモリーを節約するため、Qwen2-0.5B-Instructモデルをbfloat16精度で初期化します。

  • TRL(Transformer Reinforcement Learning)ライブラリのDPOTrainerを適用します。 これはRLHFに代わる最新の手法であり、別途報酬モデルを必要とせずにモデルを人間の好みに合わせることができます。

  • ハードウェア効率を高めるため、スクリプトはFSDP(Fully Sharded Data Parallel)用に構成されています。 これにより、モデルの重みを「シャーディング」することで複数のGPUにわたってモデルをトレーニングできるようになります。また、勾配チェックポイントを使用してVRAMの使用量をさらに削減します。

モデルの統合と保存

  • 重みの収集。専用のルーチンを適用し、FSDPのシャードを単一の統合されたモデルファイルに「集約」します。

  • データの破損や冗長な保存を避けるため、最終的なファイル書き込みは「メインプロセス」(ランク0)のみが行うようにします。

DataRobotへのデプロイ

  • モデルがローカルに保存されると、スクリプトはDataRobot APIを使用してカスタムモデルを作成し、デプロイの準備を整え、REST API経由でアクセス可能にします。