GPT-175Bを例にとった大規模言語モデルの分散トレーニングの定量分析とベストプラクティス

1. Transformer 大規模言語モデルのための SOTA トレーニング技術

1. 大規模言語モデルの開発背景と課題

まず、大規模言語モデルの開発背景についてお話しします。過去数年間で、言語モデルの規模は、データ量とモデルサイズの両面で桁違いに拡大しました。データ量とモデルサイズが増加すると、いくつかの課題も発生します。

関連する課題は 2 つの部分に分けられます。まず、コンピューティングに関して、モデルのトレーニングに必要なコンピューティング FLOPS を計算するための Megatron 論文の式を以​​下に示します。GPT-3 175B モデルは、トレーニングに現在の妥当な 1.5T トークン データ量を使用し、約 128 個の DGX A100 ノード (* 技術交換のみ)、合計 1024 個の A100 カード (* 技術交換のみ)、および比較的高効率な条件下での 120 日間の継続的なトレーニングが必要であることが簡単に推測できます。これは非常に理想的な見積もりです。実際のトレーニングでは、チェックポイント、保存とロードの時間、ノードのクラッシュと再起動の時間、または異常なデバッグ損失曲線の時間など、いくつかの問題に必然的に遭遇するため、実際のトレーニング時間はこの理想的な推定時間をはるかに超えることになります。

2 番目の課題は、大規模な言語モデルを実行する際のメモリ オーバーヘッドです。たとえば、GPT-3 175B モデルのトレーニング中のメモリ オーバーヘッドは 2 つの部分に分けられます。最初の部分は、モデル パラメータ、勾配、オプティマイザなど、モデルの状態によって占有されるメモリです。このうち、メイン メモリのオーバーヘッドは、オプティマイザの状態部分、つまり Adam オプティマイザを使用するときに関係する勢いと変動です。もう 1 つの部分は、Activation によって占有されるビデオ メモリです。後ほど、この Activation が実際に占有するビデオ メモリの量と、いくつかの最適化方法を使用した後に Activation が占有するビデオ メモリの量を具体的に計算します。しかし、いずれにしても、ビデオ メモリのこれら 2 つの部分の合計は、膨大な量のビデオ メモリを占有します。モデルを実行する前に、モデルを合理的に分割し、異なる GPU に分散する必要があります。

今日は、上記の 2 つの課題に焦点を当てます。1 つはモデルを実行する方法、つまりビデオ メモリのオーバーヘッドの課題を解決することです。もう 1 つはモデルを十分に高速に実行する方法です、つまり適切な並列手段と最適化手段を選択する方法です。

次に紹介する技術は NeMO と Megatron をベースにしているため、ここでは NeMO と Megatron を紹介します。 NeMO と Megatron は、大規模言語モデル用の NVIDIA の 2 つのフレームワークです。NeMO はサポートが充実した製品指向のフレームワークで、Megatron-LM は学術研究に重点を置いたより古典的なプロジェクトです。 NVIDIA は、実際のモデル トレーニング パフォーマンスが SOTA 状態であることを保証するために、このフレームワークで多くの最適化を行いました。

2. 大規模言語モデルのためのSOTAトレーニング技術

次に、今回共有した大規模言語モデル学習技術の全体概要を紹介します。下の図で緑色でマークされているテクノロジーは、昨年または今年登場した新しいテクノロジーであり、青色でマークされているテクノロジーはより古典的なものです。並列部分の紹介には、パイプライン並列処理、テンソル並列処理、および Megatron v3 の新しいシーケンス並列処理が含まれます。エキスパート並列処理は MOE モデルの並列処理ですが、この記事では詳しく説明しません。

2 つ目の部分は、以前からよく耳にしていた ZeRO を含むビデオ メモリの最適化です。ZeRO-1、つまり Distributed Optimizer は、Megatron と NeMO に実装されています。また、チェックポイント アクティベーションでは、選択的アクティベーション チェックポイントが現在 Megatron で最も一般的に使用されているチェックポイント テクノロジであり、比較的効率的でもあります。その他には、混合精度トレーニング、カーネル フュージョン、フラッシュ アテンションなど、私たちがより一般的に使用しているものや、最新の Megatron と NeMO で更新されたパイプライン通信とテンソル通信が含まれます。

2. GPT-175Bを例に、最新のトレーニング技術の定量分析

1. GPT-3モデルの分散トレーニング技術：モデルの並列性

まず、いくつかの古典的な並列化手法を簡単に紹介しましょう。たとえば、テンソル モデル並列処理 (TP) とパイプライン モデル並列処理 (PP) です。

その中で、TP はモデルの各レイヤーをレイヤーに分割します。現在、TP は Linear レイヤーでのみ使用されます。Linear レイヤーの Shape が比較的大きい場合、TP を使用すると GPU の使用率が向上します。しかし、TP には欠点もあります。それは、各線形レイヤーを分割するため、通信の粒度が非常に細かくなることです。つまり、最適化を行わないと、TP がレイヤー分割を計算するたびに、アクティベーション マージ操作を実行するために 1 回通信する必要があります。TP の単一の通信量は少ないですが、通信頻度は非常に高く、帯域幅要件も非常に高くなります。

PP はパイプライン並列処理であり、レイヤーを分割し、異なるレイヤーを異なる GPU に配置することを意味します。 PP の粒度は TP よりも粗いため、PP の通信量は TP よりもはるかに小さくなりますが、PP にもいくつかの欠点があります。完全なパイプラインを実行するには、大きなワークロードを複数の小さなワークロードに分割する必要があります。言い換えると、パイプラインの高いスループットを並行して維持するには、比較的大きなバッチ サイズを多数の小さなバッチに分割する必要があります。さらに、パイプラインを分割したときに異なる GPU のワークロードが不均一になると、パイプライン全体の効率が非常に低くなります。

下の図は TP の概略図です。 MLP は TP に基づいて分割され、2 つの FFN で構成されます。 TP 分割を行う際に最適化が行われ、列分割は最初の Linear レイヤーで行われ、行分割は 2 番目の Linear レイヤーで行われました。そのため、2 つの Linear レイヤーの TP 分割を完了するには、前後の通信が 1 回だけ必要でした。ただし、最適化が行われない場合、TP 分割を完了するには 2 回の通信が必要になります。具体的には、順方向計算中は、活性化の平均を達成するために g 操作で All-reduce 操作が実行され、それに応じて、逆方向計算中は、勾配の平均を達成するために f 操作で All-reduce 操作が実行されます。

もう 1 つは、次の図に示すように、Self-Attention の TP 分割です。これには、Multi-head Attention と線形レイヤーの 2 つの部分が含まれます。Multi-head Attention の TP 分割は、異なるヘッドに基づいて実装され、異なるヘッドは異なる GPU に配置されます。同時に、Linear レイヤーは行に分割されます。 MLP の 2 つの線形層 TP 分割と同様に、Forward 計算中の g 操作で All-reduce が実行され、Backward 計算中の f 操作で All-reduce が実行されます。

次に、TP の最適化補足である Sequence Parallelism を紹介します。例えば、TP 分割では、MLP と Self-Attention の Linear 層のみが並列化され、Dropout や Layernorm などの他の層は並列化されません。これは、TP グループ内では、Layernorm 層の入力が同じであるため、その計算がまったく同じになり、つまり、8 つの GPU ごとに計算がまったく同じになるため、この部分の計算が冗長になるからです。ただし、シーケンス次元で入力を分割することができます。分割後、Layernorm レイヤーの入力は、Activation レイヤーを含むさまざまな GPU に分割され、これも分割されたのと同じになります。したがって、冗長な計算は行われず、ビデオ メモリの使用量を削減できます。

シーケンス並列処理の次の図に示すように、シーケンス並列処理を実装する際には、実際にはLayernorm層でシーケンス並列処理を分割し、それをLinear層でTP分割に変換し、g操作でActivationのall-gather通信を実行します。 TP を実行した後、アクティベーションは再度縮小散乱され、その後、入力がシーケンス次元で分割された状態に変換されます。通信量で見ると、Sequence Parallelism 使用後は、元の All-reduce 通信が All-gather 通信に Reduce-scatter 通信を加えたものになります。表面的には通信数が 2 倍になっていますが、実際には all-reduce の通信量は All-gather や Reduce-scatter の 2 倍近くあるため、全体の通信量はほぼ同じです。つまり、シーケンス並列処理を最適化した後、通信量は変わりませんが、計算能力とアクティベーション メモリの使用量を削減できます。

以下は 4 つの GPU パイプラインの並列処理の例です。

まず、デバイス 1 で最初の青い部分、つまり最初のミニバッチの順方向計算を計算する必要があります。計算が完了すると、デバイス 2 に送信され、最後のデバイスに送信されます。パイプラインは徐々に忙しくなります。最後のデバイスが順方向の計算を終了すると、逆方向の計算を開始し、逆方向を順方向に渡すことができます。この計算モードは、One Forward One Backward (1F1B) 状態と呼ばれます。最後に、最後のデバイスが最後のミニバッチの後方を前方に転送します。したがって、最後のデバイスだけを見ると、パイプライン全体に最初は空きスペースがいくつかあり、その後は常にビジー状態になり、最後に空きスペースがいくつかあることになります。

1F1B 計算状態では、最初の数マスはウォームアップ ステージと呼ばれ、最後の数マスはクールダウン ステージと呼ばれます。最後のデバイスだけを見ると、パイプライン全体のアイドル時間はバブル時間と呼ばれ、ウォームアップ時間とクールダウン時間の合計になります。

組立ラインのアイドル時間を削減するにはどうすればよいでしょうか?図に示されている式によれば、バブル時間を理論時間で割った比率を見る必要があります。この比率をできるだけ減らす必要があります。たとえば、ミニバッチの数をできるだけ増やす、つまり、より多くのミニバッチに分割する必要があります。ただし、最小バッチをさらに分割しても、必ずしもメリットが得られるとは限りません。ミニバッチが非常に小さい場合、GPU コンピューティングのオーバーヘッドが比較的大きくなり、単一のコンピューティング時間が長くなるため、一定のトレードオフが必要になります。そのため、PP を最適化し、段階的なパイプライン並列戦略を採用しました。

ずらしたパイプラインの並列を以下の概略図に示します。デフォルトのパイプライン並列処理とスタッガード パイプライン並列処理の違いは、デフォルトのパイプライン並列処理ではモデル レイヤー上で連続分割が実行されるのに対し、スタッガード パイプライン並列処理ではモデル レイヤー上でジャンプのような分割が実行される点です。たとえば、デバイス 1 はレイヤー 0 から 3 を計算する必要があります。上記の 1F1B コンピューティング状態によれば、デバイス 4 はコンピューティング状態に入る前にウォームアップ時間が必要です。ウォームアップ時間を短縮するにはどうすればいいですか?方法の 1 つは、小さなミニバッチに分割することであり、もう 1 つの方法は、各デバイスがより少ないレイヤーを計算して最初のアクティベーションを取得することです。たとえば、図の下部では、デバイス 1 が計算する必要があるモデル レイヤーは、レイヤー 0、1 とレイヤー 8、9 です。通信を考慮しない場合、パイプラインのバブル時間は半分に短縮されます。これも Megatron と Nemo によるパイプライン並列処理の最適化です。

次に、3 つの並列最適化を組み合わせて、通信パターンがどのようになるかを確認します。下の図は、DP、PP、TP の 3 次元での並列最適化を含む 20 層モデルを示しています。

まず、DP ディメンションでは、2 つの DP グループを設定します。 DP の各グループは異なるデータを取得するため、生成される勾配は異なります。したがって、Backward 計算が完了したら、その勾配を更新する必要があります。勾配が更新される前に、勾配の All-reduce 操作同期を実行する必要があります。各 DP グループ内で、20 層のモデルは 4 つのパイプライン並列ステージに分割され、各ステージには 5 層のモデルが含まれます。各ステージでは、関係する 5 つのレイヤーを計算し、対応する Activation を後続のステージに渡すなど、これを繰り返していく必要があります。これがパイプライン並列通信モードです。

パイプラインの並列ステージ内では、TP が分割され、5 層モデルの計算は 4 つの GPU によって維持されます。つまり、各 GPU は 1/4 のみを計算します。最初のレイヤーを計算した後、これら 4 つの GPU は、2 番目のレイヤーを計算する前に All-reduce を実行する必要があります。これが TP 通信モードです。 TP は通信量が最も大きく、通信頻度も高いため、TP の大容量通信を計算と同じ順序で実行する必要があるため、Overlap 最適化を行う方法がありません。そのため、通常は TP を最も内側の次元に設定し、NVLINK の高速帯域幅通信を使用するようにします。 DP の通信量は PP に比べてかなり大きいため、通常は DP 最適化を 2 番目の次元として設定し、NVLINK 通信も使用できます。モデルが比較的大きい場合は、TP を 8 に設定し、対応する DP 通信はノード間通信になります。

PP は通信の需要が最も小さく、通信頻度と通信量も比較的低いため、PP の最適化は一般的に最も外側の次元に設定されます。以下の図は、説明を簡単にするために若干異なっています。一般的に、クラスター ノード間の通信帯域幅が低い場合は、ノード間のパイプライン並列処理を使用できます。

2. GPT-3モデルの分散トレーニング技術：アクティベーションメモリの最適化

チェックポイントのアクティブ化は、一般的に使用されているビデオ メモリ最適化テクノロジです。簡単に言えば、チェックポイントのアクティブ化は再計算を行うことです。たとえば、順方向計算中に取得されたアクティベーションは保持されなくなります。逆方向計算を行うときは、これらのアクティベーションが再計算され、GPU メモリのオーバーヘッドを大幅に削減できます。

次の図に示すように、チェックポイントのアクティベーションを実装する方法はいくつかあります。より単純な実装の 1 つは、Transformer の各レイヤーを再計算する完全なチェックポイントです。たとえば、最後のレイヤーで Backward 計算を行う場合、この Transformer レイヤーは Backward 計算の前に Forward 計算を再実行し、Forward を再計算した後にこのレイヤーの Backward 計算を開始する必要があります。完全なチェックポイントは各 Transformer レイヤーを再計算します。

したがって、フル チェックポイントの利点は、メモリ オーバーヘッドが O(n) の複雑さに削減されることですが、欠点は、各 Transformer レイヤーを再計算する必要があるため、約 36% の追加コンピューティング オーバーヘッドが発生することです。

アクティベーションを最適化するもう 1 つの方法は、シーケンス並列処理と選択的チェックポイントを組み合わせることです。これにより、再計算のオーバーヘッドが 36% から 4% に削減されます。

前回の紹介では、シーケンス並列処理は主にドロップアウト層とレイヤーノルム層を最適化することを説明しました。単純なチェックポイントと比較して、選択的なチェックポイントでは、再計算のコスト効率が高いいくつかの OP を選択し、計算時間は比較的短いがアクティベーションのメモリ使用量が多いいくつかの OP を再計算します。たとえば、以下のサンプル画像の左側にある Self-Attention モジュールは、計算負荷が比較的小さいため Self-Attention モジュールを再計算するメリットが非常に高いことを示していますが、中間結果出力の一部によって占有されるビデオ メモリのオーバーヘッドは非常に大きいです。したがって、この部分を再計算するだけです。 Linear レイヤーや Layernorm レイヤーなどの他のレイヤーでは、他の最適化方法を使用して Activation を最適化できます。

選択的チェックポイントの中心的な考え方は、いくつかのコスト効率の高い OP を再計算し、他の並列最適化方法と組み合わせて、1+1>2 の効果を実現することです。

3. GPT-3モデルの分散トレーニング技術：分散オプティマイザ

次に、分散オプティマイザー、つまり Zero-1 を紹介します。モデルのトレーニング中、Optimizer の状態はモデルの固定メモリ使用量の主要部分となります。モデルパラメータの数が m の場合、オプティマイザのメモリオーバーヘッドは 16 * m バイトになります。 Zero-1 の考え方は、オプティマイザーの状態を異なる DP ランクに分割することです。そのため、各 DP ランクで分散オプティマイザー (Zero-1) が占有するビデオ メモリは 16 * m を DP サイズで割った値になり、ビデオ メモリのオーバーヘッドが大幅に削減され、通信モードの変更ももたらされます。

たとえば、勾配に対して直接 Reduce-scatter を実行し、その後、パラメータの部分に対して Optimizer 関連のステータス更新を実行できます。Optimizer ステータスの更新が完了したら、All-gather 操作を使用してモデルの重みを収集します。このようにして、DP 並列 All-reduce は Reduce-scatter と All-gather を組み合わせたものに変換されます。前回の紹介によると、このAll-reduceの通信量はReduce-scatterやAll-gatherの2倍なので、全体の通信量は変わらないのですが、通信回数が増えています。

次に、より一般的に使用されているのは Zero-2 と Zero-3 ですが、Megatron と Nemo ではまだサポートされていません。実際の使用では、Zero-2 と Zero-3 は毎回追加の順方向計算と逆方向計算を実行する必要があり、そのたびに通信が必要になりますが、パイプライン並列処理では大きなバッチが多数の小さな順方向計算と逆方向計算に分割されるため、大量の通信も発生します。そのため、パイプライン並列処理で Zero-2 や Zero-3 を使用する場合、通信量が大幅に増加します。そのため、パイプライン並列処理と Zero-2 や Zero-3 を同時に使用することは推奨されませんが、パイプライン並列処理は Zero-1 と同時に使用できます。

3. 大規模モデルのパフォーマンスチューニングプロセスにおけるメモリオーバーヘッド、通信オーバーヘッド、計算オーバーヘッドの分析

1. ビデオメモリのオーバーヘッド分析

次に、ビデオ メモリのオーバーヘッド、通信のオーバーヘッド、コンピューティングのオーバーヘッドを詳細に分析します。 1 つ目はビデオ メモリのオーバーヘッドです。GPT-3 を例にとると、主に 3 つの部分に分かれています。最初の部分はモデルが占有するビデオ メモリ、つまり、パラメーター、勾配、Optimizer 状態を含むモデル状態のビデオ メモリです。2 番目の部分は Activation が占有するビデオ メモリ、つまり、モデルの OP によって生成される中間値、モデルの Transformer レイヤーと Embedding レイヤーの出力です。ここでは問題を簡略化し、Transformer レイヤーのみを考慮します。3 番目の部分は、必要になる可能性のある一時的なバッファーです。ここではあまり詳しく説明しません。主に、最初の 2 つのより重要な部分に焦点を当てます。

まず、モデル パラメータのメモリ オーバーヘッド、つまり、モデルが TP と PP に分割された後、各 GPU 上の実際のパラメータがどれくらい大きくなるかを見てみましょう。まず、GPT-3 モデルの全体的な構造を見てみましょう。このモデルには 96 層のデコーダーが含まれています。デコーダーの各層には、QKV 重み、投影重み、および 2 つの FC 重みが含まれています。このような重みは 96 回繰り返されます。さらに、トークン埋め込みと位置エンコーディングの重みも含まれています。 TP 分割を実行すると、QKV 重みと投影重みが t で除算され、次に 2 つのレイヤーの FC 重みも t で除算されます。ここで、t は TP 分割の数です。次に、PP 分割が実行されます。つまり、元の n 層が p で分割されます。ここで、p は PP 分割のステージ数です。さらに、トークン埋め込みと位置エンコーディングも TP 次元で分割されます。したがって、TP が 8 で PP が 16 の場合、各 GPU に 1.37B のパラメータを保存する必要があります。

1.37B のパラメータの場合、混合精度を使用すると、モデル パラメータと勾配の両方で実際に占有されるメモリは 2.74GB になります。 Optimizer 状態が占有するビデオ メモリはパラメータ量の 16 倍であるため、Optimizer が占有するビデオ メモリは 21.9GB となります。Optimizer 状態がビデオ メモリのオーバーヘッドの大部分を占めていることがわかります。

オプティマイザー状態のメモリ オーバーヘッドは、上記の分散オプティマイザーまたは Zero テクノロジーを使用して最適化できます。その中で、Zero-1 テクノロジーはオプティマイザーの状態を DP 次元に分割し、ビデオ メモリの使用量を削減します。たとえば、DP サイズを 8 に設定すると、オプティマイザー状態でのビデオ メモリ使用量を 21.9 GB から 2.7 GB に削減できます。

Zero-1をベースにZero-2を開発すれば、グラデーションもDP次元で分割でき、グラデーションのビデオメモリ使用量を2.74GBから0.34GBに削減できる。 Zero-3 を再度作成する場合、パラメータを適切に分割して、ビデオ メモリの使用量を 0.34 GB に削減できます。

DP サイズが 8 の場合、ゼロ - 1 を使用してオプティマイザー状態を分割すると、ビデオ メモリのオーバーヘッドが大幅に削減されることがわかります。相対的に言えば、Zero-2 と Zero-3 を使用して勾配とパラメータを分割する場合のメモリ オーバーヘッドの最適化は明らかではなく、通信オーバーヘッドが大幅に増加します。したがって、ビデオ メモリが極端に不足していない限り、Zero-3 を使用しないことをお勧めします。一方では、膨大な計算および通信のオーバーヘッドが発生し、他方では、最適化のデバッグがより困難になります。

別の観点から見ると、モデルパラメータが TP と PP に分割されると、各 GPU 上のパラメータと勾配が自然に分割されるため、占有するビデオメモリが比較的小さくなり、Zero-1 を使用すると良好な効果が得られます。さらに、PP を使用した後に ZeRO-2/3 をスタックすることはお勧めしません。主な理由は、一方ではビデオ メモリのさらなる最適化があまり明確ではないこと、他方では特定の通信競合が発生することです。したがって、TP と PP を使用してから Zero-1 を使用することをお勧めします。

次に、アクティベーション メモリのオーバーヘッドについて紹介します。問題を単純化するために、単一の Transformer レイヤーのアクティベーション メモリ使用量のみを分析します。具体的な計算方法は、各 OP 計算出力のシェイプ サイズを考慮し、各 OP によって出力されたシェイプを合計して、推定ビデオ メモリ オーバーヘッド サイズを取得します。下の例の図の式に示すように、計算式から大きな問題が見つかります。式には、シーケンスの長さの2乗であるsが含まれています。sが比較的大きい場合、ビデオメモリのオーバーヘッドが非常に恐ろしくなるため、最適化する必要があります。

前述の完全チェックポイント最適化方法を使用すると、アクティベーション メモリのオーバーヘッドを効果的に削減できます。フルチェックポイントでは、各Transformerレイヤーの出力のみが保存され、中間計算結果はBackward中に直接再計算されるため、Activationメモリのオーバーヘッド計算は2 * b * s * hになります。

TP 分割を使用するということは、実際には Self-Attention と MLP を TP 次元で分割し、対応する QKV と FC の計算出力の形状を t で分割することを意味します。これに基づいてシーケンス並列処理を追加すると、2 つの Layernorm 出力の形状がさらに t で分割されます。これに基づいて、選択的チェックポイントを引き続き使用すると、自己注意の出力アクティベーションを直接ドロップするため、s の二乗を最適化できます。

したがって、全体的な再計算のオーバーヘッドは、フル チェックポイントの 36% から TP+SP+選択的チェックポイントの 4% に削減され、対応するビデオ メモリのオーバーヘッドは元の 2.86 GB から 106 MB に削減されます。

2. 通信オーバーヘッド分析

通信オーバーヘッドも DP、TP、PP の 3 つの部分に分かれています。まず、DP の通信オーバーヘッドを見てみましょう。下の図に示すように、DP の通信オーバーヘッドは、勾配通信のサイズをバス帯域幅で割り、All-reduce 係数を掛けたものに等しくなります。バス帯域幅は、実際のメッセージ サイズに基づいて見積もる必要があります。

2 つの最適化方法 (TP と TP+SP) の通信量は一定です。具体的な TP 通信オーバーヘッドは、以下の例の図のように計算されます。各 TP 通信のサイズを実際に得られたバス帯域幅で割り、All-reduce 係数を掛け、Self-attention と MLP 間の通信回数をそれぞれ 3 回掛け、Transformer レイヤー数とミニバッチ数を掛けて算出されます。

パイプライン並列処理の通信オーバーヘッドは、主に以下の例の図の緑色の部分に見られます。具体的な計算は、まず 1 回の通信にかかる時間を取得し、それを通信回数で乗算します。通信時間には、1F1B 時間に加えて、ウォームアップとクールダウンの時間が含まれます。

3. 計算オーバーヘッド分析

次に計算オーバーヘッドについて紹介します。具体的な計算式については、以下の例式を参照してください。サンプル図の左側には、Self-Attention レイヤー、FFN レイヤー、Logit レイヤーの計算の内訳が示されており、主に一部の Linear レイヤーの FLOPS を計算します。これらの FLOPS を合計すると、単一の Transformer レイヤーの FLOPS が得られ、それをバッチ数とレイヤー数で乗算すると、モデル全体の FLOPS が得られます。モデル FLOPS を使用すると、トレーニング システム全体の効率が十分かどうかを測定できます。たとえば、GPT-3 175B の場合、モデル FLOPS が A100 で 150 TFLOPS 未満 (* 技術情報のみ) であれば、トレーニングに何らかの問題がある可能性があります。120 TFLOPS 未満であれば、トレーニングに重大な異常な問題があり、速度が低下していることを意味します。

IV. 結論

最後にまとめます。 GPT-3 175B トレーニングでは、これらの最適化をどのように組み合わせればよいでしょうか?まず、いくつかの最適化項目はデフォルトで有効になっています。たとえば、Megatron では通常、混合精度、フラッシュ アテンション、および多数のデフォルトの最適化 OP を使用することをお勧めします。 2 つ目は、半精度トレーニングの選択です。たとえば、BF16 を使用するか FP16 を使用するかです。大規模なモデルのトレーニングでは、FP16 には多くの潜在的な問題と落とし穴があるため、BF16 を優先することをお勧めします。BF16 を使用すると、特に 20B を超えるモデルでは、より安定します。BF16 のみを使用することをお勧めします。

次のステップは、これらの最適化をどのように組み合わせるかです。たとえば、トレーニングの実行時にビデオ メモリのオーバーヘッドに問題がある場合は、これらの最適化を順番にオンにすることができます。まず、選択的アクティベーション チェックポイントの使用を試し、次に分散オプティマイザーをオンにして、TP を徐々にオンにすることができます。最初から TP を 8 に設定する必要はありません。ここでは、隠しサイズを TP で割った値が 1024 未満であってはならず、2048 より大きい方がより良いというルールがあり、その方がより大きなメリットが得られます。その後、PPを徐々に開きます。最初に PP を大きくしすぎず、その後、完全アクティベーション チェックポイントを使用します。

OOM に遭遇せず、GBS を十分に大きく開くことができる場合は、DP を使用して拡張してみてください。