Text generation and understanding based on GPT model

概要

GPT モデルはテキスト生成分野のさまざまなタスクに対応できます。たとえば、テキスト補完、自由な質疑応答、クローゼテスト、論文執筆、要約作成、小説執筆、詩の創作などです。インターネット上で大きな注目を集めている AI プロダクトである ChatGPT も、GPT テキスト生成モデルをベースにしています。GPT 大規模モデルはこれらのアプリケーション分野で高い性能を示しますが、トレーニングコストが非常に高いという課題があります。

GPT ベースモデルはパラメータ数が多く、トレーニングと推論に必要なハードウェアリソースの消費が過大であるため、MoE に基づくスパーストレーニングが現在最も競争力のあるコスト削減・効率化の手法となっています。MoE の正式名称は Mixture of Experts で、Expert は Transformer モデルの MLP 層に対応します。トレーニング時に複数の MLP から 1 つを選択して活性化します（下図参照）。これにより、計算量（FLOPS/Bytes）を増やさずに MLP モジュールの数を増やしてモデルパラメータを増加させ、下流タスクでのモデルの汎化性能を向上させることができます。同じ品質（検証セット損失およびゼロショット NLU 下流タスク性能）の Dense モデルと比較して、MoE を適用したスパース Transformer モデルは、トレーニングスループットで約 1.2 倍、推論スループットで約 1.3 倍の性能向上を実現します。スパースアーキテクチャの全体設計では、純粋な Transformer Decoder アーキテクチャの GPT と MoE を有機的に組み合わせることにしました。これは、MoE と Decoder の組み合わせが Encoder との組み合わせよりも通常優れているためです。具体的には、Encoder はランダムマスキングにより言語モデルを学習しますが、このランダムマスクトークンがエキスパートのルーティング選択を不均衡にします。一方で、Decoder 系の GPT モデルは Encoder 系の Bert モデルよりも幅広い使用シーンがあるため、GPT+MoE 技術アーキテクチャを採用し、スタンドアロン環境で最もエネルギー効率に優れた大規模 GPT モデルのトレーニングと推論を実現するソフトウェアとハードウェアの統合適応技術の可能性を探ります。

現在比較的成熟した分散 MoE エキスパートルーティング技術に基づき、Switch Transformer [2] の top-1 ルーティング機構を採用します。各 Expert に以下の softmax 関数に基づいて確率値を与え、最も高い確率（top-1）の Expert をネットワークの FFN 層として選択します。ここで W_r はルーティング選択時に学習する必要があるパラメータです。

GPT-MoE トレーニングと推論のエネルギー効率分析

ベース事前トレーニングモデルのトレーニングと推論のパフォーマンス分析

任意の Dense GPT モデルに対して、より高速なトレーニングと推論を実現する対応するスパース MoE GPT モデルが存在します。私たちの目標は、スタンドアロン環境のような制約のあるハードウェア条件下でこの GPT-MoE モデル設定を見つけ、さらに MoE アルゴリズムの改善によりトレーニングのエネルギー効率を向上させることです。Dense モデルとスパースモデルのトレーニングと推論のパフォーマンスを比較し、Dense モデルと同等のエネルギー効率を持つスパースモデルを見つけます。

8 つの GPT モデルのパラメータ数、モデル構造、トレーニングハイパーパラメータを以下の表に示します。

GPT モデル
パラメータ数
レイヤー数
ヘッド数
Hidden Size
学習率
トークンバッチ
1.3B Dense
1.3B
24
32
2048
2e-4
1M
2.7B Dense
2.7B
32
32
2560
1.6e-4
1M
3.6B Dense
3.6B
30
32
3072
1.6e-4
1M
0.35B+MoE-64
6.7B
24
16
1024
3e-4
0.5M
1.3B+MoE-32
13B
24
32
2048
2e-4
1M
1.3B+MoE-64
27B
24
32
2048
1.6e-4
1M
2.7B+MoE-64
56B
32
32
2560
1.6e-4
1M
3.6B+MoE-64
75B
30
32
3072
1.6e-4
1M

下図に示すように、1.3B+MoE-32/64 モデルは同じステップで 1.3B Dense モデルよりも低い検証セット損失を示し、1.3B+MoE-64 モデルの損失は 2.7B Dense モデルよりもさらに低くなっています。

5 つのモデルの中で、0.35B+MoE-64 のトレーニングスループットが最も速く、他のモデルの約 2 倍です。残りの 4 つのモデルの中では、1.3B Dense と 1.3B+MoE-32 のスループット速度が高く、1.3B+MoE-64 と 2.7B Dense の速度は同等です。以下の通りです。

推論スループット速度の面では、1.3B Dense の GPU メモリ消費量が最も少なく、0.35B+MoE-64 のレイテンシが最も低くなっています。

input_len=20
output_len = 128
batch_size = 1
モデル
レイテンシ (ms)
メモリ (MB)
GPU 数
1.3B Dense
399.66
9476
1
2.7B Dense
753.37
17340
1
3.6B Dense
777.54
22558
1
0.35B+MoE-64
356.22
15772
1
1.3B+MoE-32
581.34
33294
1
1.3B+MoE-64
586.18
57880
1
2.7B+MoE-64
742.43
61054
2
3.6B+MoE-64
662.57
42724
4

上記のチャート分析から、2.7B Dense モデルに対応するエネルギー効率の良いスパースモデルは、1.3B Dense モデルをベースに 32 または 64 エキスパートの MoE を搭載したモデルであると大まかに判断できます。以下では、1.3B+MoE-32/64 と 2.7B Dense モデルのコストパフォーマンスに焦点を当てて分析します。シングルマシンで 200 時間の事前トレーニング後、Tensorboard を使用して事前トレーニングの検証セット損失曲線を描画します。検証セット損失が 2.16 に達した時点で、1.3B+MoE-64 モデルの収束速度は 2.7B Dense モデルの 1.17 倍であり、1.3B+MoE-32 の収束速度は 2.7B Dense モデルより約 15% 遅いことが分かりました。以下の図の通りです。

下図の Faster Transformer ベースのシングル GPU 推論パフォーマンスを見ると、1.3B+MoE-32 と 1.3B+MoE-64 のスループット速度は同等で、2.7B Dense モデルを上回っています。ベースサイズが 1.3B に過ぎないため、これは期待通りの結果です。

中国語ゼロショット NLU 効果評価

中国語テキスト生成効果評価

テキスト補完

詩生成

オンライン体験アドレス: https://www.modelscope.cn/models/PAI/nlp_gpt3_text-generation_0.35B_MoE-64/summary

広告コピー生成

オンライン体験アドレス: https://www.modelscope.cn/models/PAI/nlp_gpt3_text-generation_1.3B_MoE-32/summary

作文生成

オンライン体験アドレス: https://www.modelscope.cn/models/PAI/nlp_gpt3_text-generation_1.3B_MoE-64/summary

独自開発 GPT-MoE アルゴリズム革新と実験分析

背景

Top-1 Gating は現在最も主流で効果的なルーティングアルゴリズムですが、明らかな欠点もあります。たとえば、Top-1 Gating では各トークンが 1 つのエキスパートにしか処理されないため、一部のエキスパートは多くのトークンを処理する一方で、ごく少数のトークンしか処理しないエキスパートも出てきます。その結果、処理量の少ないエキスパートは十分な情報を得られず、十分にトレーニングされないという問題が生じます。

エネルギー効率の高いエキスパートルーティング

そこで、私たちは新しいルーティングアルゴリズムを独自に開発しました（下図参照）。このアルゴリズムでは、エキスパートが固定数のトークン（capacity）を能動的に選択でき、同じトークンを異なるエキスパートが同時に処理できるため、各エキスパートを十分にトレーニングできます。

最終的なトークン表現は加重和方式を採用し、異なるエキスパートが生成した表現を重み付けして合計し、Expert Residual モジュールを通じて最終的なトークン表現を得ます。この表現は、複数のエキスパートの共同作用により、より頑健なものとなります。

1. エキスパートのトークンに対する選好を計算する:

2. L-Softmax アルゴリズムによるエキスパート重みのトレーニング
トレーニング中に L-Softmax 損失を使用してエキスパートの重み W_e を最適化し、各エキスパートがトークンに対して識別可能な選好を持つようにします。

3. 各エキスパートが固定数のトークンを選択する:
ここで、各エキスパートが処理できるトークンの最大数を事前に決定し、各エキスパートが処理するトークンのインデックスを記録します。

4. 最終出力を計算する:
各エキスパートが対応するトークンの表現をインデックスに基づいて計算し、異なるエキスパートが同じトークンに対して生成した表現を重み付けして合計し、最終的に Expert Residual モジュールを通じて最終出力を生成します。

実験分析

下図は、独自開発アルゴリズムと Top-1 Gating および s-BASE アルゴリズムの検証セット損失がトレーニングステップとともに変化するグラフです。独自開発アルゴリズムの検証セット損失は、Top-1 Gating および s-BASE アルゴリズムの検証セット損失よりも常に低いことが確認でき、独自開発アルゴリズムの有効性が証明されました。

同時に、検証セット損失が初めて 2.7 を下回った時点で、独自開発アルゴリズムの速度は s-BASE の 1.48 倍であり、モデルトレーニングのオーバーヘッドを大幅に削減しました。

さらに、独自開発アルゴリズム、Top-1 Gating、および s-BASE のトレーニングスループットも分析しました。下図に示すように、独自開発アルゴリズムのトレーニングスループットは s-BASE の 1.17 倍高くなっています。

PAI DLC ベースの GPT-MoE 事前トレーニング

Rapidformer は、EasyNLP のさまざまな Transformer モデルにトレーニング高速化機能を提供しており、これは DeepSpeed と Megatron を有機的に統合することで実現されています（下図参照）。

GPT-MoE 大規模モデルの事前トレーニングで使用する主なトレーニング高速化コア技術は以下の通りです。

混合精度トレーニング（Mixed Precision Training） 混合精度トレーニングのメリットは主に以下の 2 点です。1. GPU メモリ使用量の削減。FP16 のメモリ使用量は FP32 の半分であるため、トレーニング中の GPU メモリを半分節約できます。2. トレーニングと推論の計算高速化。FP16 はメモリ節約だけでなく、モデルのトレーニング時間も短縮できます。具体的な原理は下図の通りで、パラメータ更新時の誤差逆伝播での丸め誤差を避けるために FP32 バックアップを維持することが核心です。さらに、Loss Scaling を使用してオーバーフローエラーを軽減します。

選択的活性化再計算（Selective Activation Recomputation） ニューラルネットワークの中間にいくつかのチェックポイントを設定します。チェックポイント以外のすべての中間結果を破棄し、誤差逆伝播で微分を取得する際は最寄りのチェックポイントから中間結果の再計算を開始します。これにより GPU メモリを節約しつつ、最初から計算する手間を回避できます。実際の使用では、活性化値は大きいものの計算量が少ないレイヤーもあるため、このような活性化値を選択的に除外し、重要な活性化値を保持して再計算コストを削減する必要があります。

ゼロ冗長オプティマイザー（ZeRO） ZeRO は大規模分散ディープラーニング向けの新しいメモリ最適化技術です。ZeRO には 3 つの主要な最適化ステージがあり、それぞれオプティマイザー状態、勾配、パラメータの分割に対応します。ここでは ZeRO-1 のオプティマイザー状態分割を使用しています。

シーケンス並列処理（Sequence Parallelism） 長いシーケンスを分割してトレーニングを高速化する技術です。テンソル並列処理をベースに、Transformer コアの LayerNorm 層と Dropout 層の入力をSequence Length 次元で分割し、各デバイスがそれぞれの計算だけを処理すればよいようにします。これには 2 つのメリットがあります。1. LayerNorm と Dropout の計算を各デバイスに分散し、計算リソースの無駄を削減する。2. LayerNorm と Dropout が生成する活性化値も各デバイスに分散され、GPU メモリのオーバーヘッドをさらに削減する。