PAI-ChatLearn: 柔軟で使いやすく、大規模な RLHF 効率トレーニングフレームワーク (Alibaba Cloud の最新の実践)

PAI-ChatLearn は、Alibaba Cloud Machine Learning Platform の PAI チームが独自に開発した、柔軟で使いやすい大規模モデル RLHF 効率トレーニングフレームワークです。大規模モデルに対する SFT (教師あり指導の微調整)、RM (報酬モデル)、RLHF (人間のフィードバックに基づく強化学習) の完全なトレーニングプロセスをサポートします。 PAI-ChatLearn は、トレーニングと推論に異なるバックエンドの使用をサポートし、モデルごとに異なる並列戦略と柔軟なリソース割り当てを構成でき、大規模 (175B + 175B) モデルの効率的な RLHF トレーニングをサポートし、業界のフレームワークと比較してパフォーマンスが大幅に向上しているため、ユーザーはモデル効果のチューニングに集中できます。

1. 大規模モデルトレーニング方法の進化

大規模モデルの急速な発展により、モデルのトレーニング方法（特にディープラーニングと人工知能の分野）が継続的に進化しました。モデルのサイズが大きくなるにつれて、単一のデバイス (GPU や CPU など) のメモリと計算能力ではモデル全体を処理するのに十分ではなくなるため、タスクと計算負荷を分散するための並列化戦略が必要になります。以下では、モデルトレーニングのデータ並列性、モデル並列性、タスク並列性についてそれぞれ紹介します。

1. データの並列処理

初期のディープラーニングモデルは、古典的な畳み込みニューラルネットワークモデルの AlexNet や ResNet など比較的小規模で、単一の GPU でスムーズに実行できました。しかし、画像データの量が急増するにつれ、単一の GPU のトレーニングパフォーマンスでは徐々に需要を満たせなくなってきました。この課題に対処するために、モデルトレーニングではデータ並列戦略を採用してトレーニングを加速し始めました。

データ並列処理とは、トレーニングデータセットを複数の小さなバッチに分割し、モデルのコピーを複数のプロセッサまたはデバイスで同時にトレーニングするプロセスです。各デバイスにはモデルの完全なコピーがあり、各トレーニングステップの最後に、モデルパラメーターは何らかの形式の通信 (勾配平均化やパラメーターサーバーなど) を通じて同期されます。

データ並列化戦略自体も一連の進化を遂げてきました。たとえば、Zero Redundancy Optimizer (ZeRO) や Fully Sharded Data Parallel (FSDP) などの概念が登場しています。これらの概念の核となるのは、計算プロセスは依然としてデータ並列性の原則に従いながら、トレーニング中に状態情報、勾配、またはモデルパラメータを分散的に保存するというものです。トレーニングデータセットを複数のミニバッチに分割し、モデルのコピーを複数のプロセッサまたはデバイスで同時にトレーニングします。各デバイスにはモデルの完全なコピーがあり、各トレーニングステップの最後に、モデルパラメーターは何らかの形式の通信 (勾配平均化やパラメーターサーバーなど) を通じて同期されます。

このアプローチにより、各 GPU のメモリ負荷が大幅に軽減され、より大きなモデルをトレーニングできるようになります。しかし、モデルのサイズが大きくなるにつれて、データ並列戦略では、増加する同期負荷を処理するために、より効率的な通信メカニズムとより大きな帯域幅が必要になります。

2. モデルの並列処理

モデルのサイズが大きくなり続けると、純粋なデータ並列処理だけではトレーニングのニーズを満たすことができなくなります。私たちが直面している課題は、ビデオメモリ容量の不足だけでなく、計算能力の限界でもあります。そのため、モデルの並列処理は、大規模なディープラーニングモデルのトレーニングにおける一般的なソリューションとなっています。

モデル並列処理は、モデル分割方法に応じて、主にパイプライン並列処理とテンソル並列処理の 2 つの形式に分けられます。

テンソル並列処理
モデルのテンソルは複数のフラグメントに分割され、各コンピューティングユニット (GPU など) は、その一部の保存と計算を担当します。このアプローチにより、各コンピューティングユニットのメモリフットプリントを効果的に削減できるため、より大きなモデルのトレーニングをサポートできます。
デメリット:デバイス間の通信が多くなる可能性があるため、NVLink を使用するなど、高帯域幅のマシン内通信に適しています。
パイプラインの並列処理
モデルはレイヤーに分割されており、各コンピューティングユニットはモデル内のレイヤーの一部の計算を担当します。このアプローチにより、単一のコンピューティングユニットのメモリ負荷が軽減され、より大規模なモデルトレーニングのサポートにも役立ちます。通信要件は比較的小さく、インフラストラクチャが弱く帯域幅が低い環境に適しています。
デメリット: アイドル時間 (バブル) が大量に発生する可能性があり、待機時間が長くならないようにパイプラインを慎重に設計する必要があります。
その他のモデル並列バリアント
Mixture of Experts (MoE) におけるエキスパート並列処理は、エキスパートモジュールを異なるコンピューティングユニットに分散します。
シーケンス並列処理は、トレーニング中にアクティベーション関数の出力を分散して、ビデオメモリの使用量を削減します。

単一モデルの並列処理では、モデルのトレーニングを効率的にサポートすることが難しい場合がよくあります。 GPT3-175B などの大型モデルでは、並列戦略を組み合わせる必要があります。ハイブリッド並列処理は、データ並列処理、モデル並列処理、タスク並列処理の利点を組み合わせて、より効率的なトレーニングを実現します。たとえば、データ並列処理を使用して複数のデバイス上でモデルの異なるコピーをトレーニングし、各デバイス内でモデル並列処理を使用してモデルの異なる部分を処理することができます。

3. タスクの並列処理

モデルが進化し続けるにつれて、モデルのトレーニング方法にいくつかの新しいトレンドが生まれています。タスクの並列処理は、モデルまたはデータの並列処理を直接目的としたものではなく、異なるタスクを異なるプロセッサまたはデバイスに割り当てることを目的としています。大規模モデルのトレーニングでは、タスクの並列処理を使用して、ハイパーパラメータの検索、モデルの評価、異なるモデルの同時トレーニングなど、複数の異なるトレーニングプロセスを同時に実行できます。タスクの並列処理により、リソースの使用率とトレーニングの効率が向上しますが、効果的なタスクのスケジュールと管理のメカニズムが必要です。

マルチタスク/マルチモーダルモデルトレーニング
Google が提案した Pathways アーキテクチャなど、異なるタスク間でパラメータが共有される可能性があるタスク次元では、並列トレーニングが必要です。もう 1 つの例は、音声やテキストなどのさまざまな種類のデータを処理できるマルチモーダルモデルである GPT-4v です。
大規模モデルは、モデルのパフォーマンスを向上させるために RLHF と組み合わされます。たとえば、ChatGPT と GPT4 はすべて RLHF を実装しています。
シングルプログラムマルチデータ (SPMD) モードからマルチプログラムマルチデータ (MPMD) モードに移行します。

（１）RLHF

RLHF (Reinforcement Learning from Human Feedback) は、初期言語モデルの微調整タスクを RL 強化学習問題としてモデル化し、人間のフィードバック信号を使用して最適化することで、生成を人間の期待にさらに沿わせ、モデルのパフォーマンスと信頼性を向上させます。

トレーニング手順:

ステップ1: 言語モデルの事前トレーニング
言語モデルを事前トレーニングし、手動で作成されたコーパスに対して SFT (教師あり微調整) を実行して微調整します。
オープンソースの事前トレーニング済みモデル: Qwen、LLaMA、Baichuan など。
ステップ2: 報酬モデル(RM)のトレーニング
トレーニング報酬モデルは、モデルによって生成されたテキストが人間の好みに合っているかどうかを判断します。
オープンソースの報酬モデル: Ziya-LLaMA-7B-Reward など
ステップ3: 強化学習を使用して微調整する
人間のフィードバック信号と組み合わせて、強化学習を使用してモデルを微調整します。
コミュニティモデル: ChatGPT、GPT4、GPT4-Turbo など。

（２）ChatLearnフレームワーク構築の必要性

RLHF トレーニングでは、ポリシー、報酬、価値、その他のモデルを同時にトレーニングする必要があります。 Megatron-LMやDeepSpeedなどの既存のフレームワークは、単一モデルのトレーニング/推論に適しています。既存の RLHF フレームワークは、単一マシンのトレーニングまたは同じデバイスへの複数のモデルの配置のみをサポートしています。サポートされるモデルの最大スケールは制限されており、70B + 70B 以上のスケールはサポートされておらず、ハイブリッドフレームワークのトレーニング/推論をサポートすることはできません。

2. PAI-ChatLearnフレームワークの紹介

このセクションでは、主に PAI-ChatLearn フレームワークの基礎となる実装の詳細、インターフェースの定義方法、大規模で効率的な RLHF トレーニングの実装方法について紹介します。

1. PAI-ChatLearnフレームワーク

PAI_ChatLearn は、次の機能を備えた大規模な RHF トレーニングのための柔軟で効率的なフレームワークです。

SFT（教師あり微調整）、RM（報酬モデル）、RLHFフルプロセストレーニングをサポート
ポリシーモデルと値モデルの推論とトレーニングを分離し、効率的な重み同期をサポートします。
トレーニングと推論の特性が異なるため、必要なリソースと並列戦略も異なります。
HF トレーニングプロセスは比較的複雑であり (たとえば、ポリシーモデルは生成とトレーニングの両方に関係します)、推論モデルとトレーニングモデルを分離し、異なる並列戦略を使用して効率的なトレーニングと推論を実現するなど、モデルはセグメント化されます。
モデルパラメータの最終的な一貫性を確保し、より効率的なトレーニングを実現するために、高速チャネルベースのパラメータ同期メカニズムが実装されています。
トレーニングと推論は、環境の分離を実現するためにさまざまなバックエンドをサポートします。たとえば、トレーニングではバックエンドとして Megatron を使用し、推論では VIM フレームワークを使用できます。この組み合わせにより、より良い結果が得られる可能性があります。
各モデルは、並列戦略とリソースの独立した構成、および分散アクターの設計をサポートします。
あらゆるサイズのモデル（175B+175Bなど）のトレーニングをサポート
単一マシンのパースペクティブプログラミングと分散実行により、ユーザーはコードを作成するときに、基盤となるモデルの割り当てと実行について心配する必要がなくなります。

次の図は、PAI_ChatLearn に実装されている RLHF トレーニングプロセスの図です。

2. PAI-ChatLearnアーキテクチャの紹介

（１）API層

APIレイヤーはモデル構築のためのインターフェースを提供する

RLHF モジュールは、汎用抽象 RLHF モジュールインターフェイスです。このインターフェースを通じて、ユーザーはいくつかのメソッドのカプセル化を実装するだけで、さまざまなモデルの構築を完了できます。
RLHF Configはe2eトレーニングパラメータを設定します
モデル構成は各モデルのパラメータを設定します

（２）エンジン層

エンジン層は、特定のリソースの割り当てとスケジュール、およびトレーニングと推論の実際の実行を担当します。

DistActor は、単一の分散タスクとリソースを抽象化するために使用され、各モデルの実行環境とリソース割り当てを担当し、柔軟な拡張を実現できます。
トレーニングと推論では、異なるバックエンドを使用できます。たとえば、推論には PyTorch または vLLM フレームワークを使用し、トレーニングバックエンドには Megatron、DeepSpeed、または独自開発のフレームワークを選択できます。
トレーニングのパフォーマンスを確保するために、コンピューティング、通信、ビデオメモリの最適化など、さまざまな最適化も実行されます。

3. PAI-ChatLearnトレーニングコアプロセス

PAI-ChatLearn トレーニングのコアプロセスは次のとおりです。まず chatlearn を初期化し、次にモデルを定義します (RLHF トレーニングプロセスには 2 つのトレーニングモデルと 4 つの推論モデルが含まれます)。次に、エンジンとデータセットを定義し、engine.learn() を呼び出してトレーニングを開始します。

4. PAI-ChatLearnトレーニング構成

実行環境の構成

PAI-ChatLearn トレーニングフレームワークでは、ユーザーは YAML ファイルを通じて動作環境全体を構成します。PAI 実行環境 DLC とローカル環境の両方を柔軟に構成できます。

タスク配布プロセスでは、内部最適化戦略が採用されます。つまり、「除外」機能を通じて、配布する必要のないコンポーネントを明示的に指定できます。これにより、不要な通信オーバーヘッドが効果的に削減され、全体的な効率が向上します。

モデル構成

各モデルには独立した構成オプションがあり、ユーザーは必要に応じてリソースの割り当てと戦略をカスタマイズできます。 RLHF では、6 つの異なるモデル構成がサポートされています。

RLHF トレーニング構成

ユーザーは、実際の実行回数、トレーニングバッチサイズ、生成されたタスクのベースラインサイズなどのパラメータを定義できます。さらに、チェックポイントを定期的に保存する時間間隔、データの保存パス、モデルのパフォーマンスを評価する時点も設定できます。

PAI-ChatLearn フレームワークは明確な構成プロセスを提供します。ユーザーはモデルを定義し、対応する構成を追加するだけで、簡単にトレーニングプロセスを開始できます。

5. PAI-ChatLearnインターフェースモデルの定義

モデル構築プロセスを加速するために、開発プロセスを簡素化および高速化するように設計された一連のサブクラスを提供します。これらのサブクラスは、PyTorch、Meta、または VIM に基づくフレームワークに強力なサポートを提供します。同時に、ユーザーはこれらのサブクラスに基づいて HF モデルを構築するか、親クラスに直接基づいて HF モデルを構築するかを選択できます。

PAI-ChatLearn フレームワークでは、モデルのカプセル化が非常に簡単になります。ユーザーは、初期化 (init) ステップなどのモデルのセットアップと構築プロセスにのみ集中する必要があります。カプセル化レイヤーは、異なるフレームワーク間のインターフェースの違いを隠すように設計されているため、ユーザーはフレームワーク内で 1 つのカプセル化レイヤーのみを実行する必要があります。たとえば、ユーザーが推論のためだけに、または完全なトレーニングサイクルのために前進ステップを実装する必要がある場合、ラッパーによってこのプロセスが簡素化されます。標準モデルトレーニングプロセスを実装するには、ユーザーは定義されたトレーニングモジュールをモデルに統合するだけで済みます。

さらに、このフレームワークでは、さまざまな実行要件に対応するために、基盤となるインターフェースを柔軟に選択できます。ユーザーは、実行エンジンとして Megatron、バックエンドとして VIM を選択することも、当チームが開発した高性能コンピュータ HC をバックエンドリソースとして使用することもできます。この自由に切り替える機能により、ユーザーは特定のニーズと好みに基づいて最も適切なツールとリソースを選択できます。

3. PAI-ChatLearnの使用例

このセクションでは、オープンソースのモデル例に基づいて、RLHF モデルの具体的なプロセスとトレーニング結果について説明します。

1. オープンソースのトランスフォーマーモデルの例

（１）画像とモデルの準備

鏡：

PAI-DLCは、ChatLearnが提供するDockerfileを参照して、既製のイメージを提供するか、イメージをカスタマイズします。

モデル：

Huggingface や ModeScope などのオープンソースモデルをダウンロードします。

HuggingFace トランスフォーマーモデルの場合は、ChatLearn が提供するツールを呼び出して、トランスフォーマーモデルを、モデルを変更するための並列戦略をサポートする Megatron-LM 形式に変換します。

（２）トレーニング手順

ステップ1: SFTモデルをトレーニングする

SFT トレーニング用のデータを準備します。

形式: {'query': 質問、'response': 返信}。

トレーニング方法:

変換されたモデルが構築された後、ChatLearn が呼び出され、SFT トレーニングが実行されます。

SFTモデルをトレーニングして保存する

トレーニング済みの SFT モデルがある場合は、この手順をスキップできます。

ステップ 2: 報酬モデル (RM) をトレーニングします。

RM トレーニング用のデータを準備します。

形式: {'query': 質問、'response': [response1、response2、.....]、'score': [score1、score2、.....]}。

トレーニング方法:

RM を定義します。一般的なメイン構造は SFT を再利用し、ヘッドレイヤーを追加します。

トレーニングされた SFT モデルパラメータを使用して RM を初期化します。

RM トレーニングを実行するには、ChatLearn に電話します。

トレーニング済みのモデルを保存します。

ステップ 3: 強化学習を使用して微調整します。

RL トレーニング用のデータを準備します。

形式: {"prompt": 質問}。

トレーニング方法:

SFT モデルと構造が一致するポリシー/参照モデルを定義します。

SFT モデルパラメータを使用してポリシー/参照モデルを初期化します。

報酬/価値モデルを定義します。その構造は報酬モデルと一致します。

RM モデルパラメータを使用して報酬/価値モデルを初期化します。

RLHF トレーニングを実行するには、ChatLearn に電話します。

トレーニング済みのモデルを保存します。

ステップ 4: 推論。

オフラインバッチ予測。

予測する必要があるデータセットを準備します。

ChatLearn を呼び出して、オフラインバッチ予測用にトレーニング済みの RLHF モデルを読み込みます。

オンライン予測。

ChatLearn ツールを呼び出して、トレーニング済みの RLHF モデルを Transformer 形式に変換します。

PAI-EAS または vLLM を呼び出して、オンライン予測用のモデルを展開します。

2. ビクーニャ13B+13Bトレーニング効果

ベースモデルは Vicuna 13B を使用し、RM も 13B スケールを使用します。

SFT、RM、RLHF は、HH (有用かつ無害) データセットでトレーニングされます。

RLHF トレーニングでは 16 個の GPU が使用されました。

RLHF トレーニングにおける報酬モデルのスコアを図に示します。

GPT-4 API は MT-Bench で使用され、SFT モデルと RLHF モデルのパフォーマンスを評価し、平均スコアが 11% 向上しました。

ChatLearnのパフォーマンス:
さまざまなスケールの RLHF でトレーニングされた ChatLearn のパフォーマンスが A800 でテストされます。
ポリシーと報酬モデルは同じスケール構成を使用します。
DeepSpeed-chatと比較すると、7B+7B～30B+30BのスケールはDeepSpeed-chatより48%～82%高速です。
サイズが 66B + 66B になると、DeepSpeed-chat は 32 GPU で OOM になりますが、ChatLearn は正常にトレーニングできます。
ChatLearn は、175B+175B などの大規模なモデルトレーニングをサポートできます。

3. PAI-ChatLearn効果

Qwen-14Bへの影響:

ポリシーと報酬モデルは両方とも 14B です。

基本評価の結果は上図に示されており、同じ規模の他のオープンソースモデルよりも優れています。

幅広いトピックを網羅したテストデータセットが手動評価用に構築され、RLHF でトレーニングされたモデルは SFT モデルよりも大幅に優れたパフォーマンスを発揮しました。テスト問題には、知識、言語理解、創造性、数学、コーディング、その他の能力が含まれます。

4. 質疑応答

Q1: 報酬モデルには、科学的正確性、コンプライアンス、人間的な表現など、複数のスコアリングメカニズムがあります。多次元報酬モデルの評価に関する実践的なエンジニアリング経験はありますか?

A1: Baichuan2 の技術レポートでは、モデルの HF トレーニングでは複数の報酬モデルが使用されていると述べられています。たとえば、数学の分野では、既存の報酬モデルは数学的なスコアリングにおいて十分な精度がなく、効果が理想的でない可能性があります。報酬モデルのトレーニングを支援するために、より強力な数学的機能を備えた別のモジュールを具体的に使用します。

一部のオープンソースフレームワークは、タイル化された設計のため、ビデオメモリの不足などのさまざまな課題に直面し、拡張が困難です。

PAI-ChatLearn フレームワークは分離されたアーキテクチャ設計を採用しており、さまざまな段階でのさまざまな拡張を容易にします。報酬モデル部分に一連の処理ロジックを追加するだけで、複数の報酬モデルの組み合わせを実現できます。各モデルは独立した分散アクターとしてカプセル化されるため、新しい分散アクターを追加し、その結果を連結し、最後に KL ダイバージェンスまたはその他の方法を使用して計算するだけで済みます。