Nvidia の堀を解体、詳細が明らかに！世界最速のスーパーコンピュータは、3072 個の AMD GPU を使用して、1 兆を超えるパラメータで LLM をトレーニングします。

AMD のソフトウェアおよびハードウェア システムは、GPT-3.5 レベルで大規模なモデルをトレーニングするためにも使用できます。

米国オークリッジ国立研究所にある世界最大のスーパーコンピュータである Frontier は、37,888 個の MI250X GPU と 9,472 個の Epyc 7A53 CPU を組み合わせています。

最近、研究者たちは、GPU の約 8% のみを使用して GPT-3.5 サイズのモデルをトレーニングしました。

研究者らは、ROCM ソフトウェア プラットフォームを使用して、AMD ハードウェア上の分散トレーニング モデルにおける多くの困難を克服し、ROCM プラットフォームを使用して AMD ハードウェア上の大規模モデル向けの最も高度な分散トレーニング アルゴリズムとフレームワークを確立しました。

NVIDIA および CUDA 以外のプラットフォームで LLM を効率的にトレーニングするための実現可能な技術フレームワークを提供することに成功しました。

トレーニング後、研究者らは Frontier で大規模モデルをトレーニングした経験を論文にまとめ、遭遇した課題と克服した困難を詳しく述べました。

論文リンク: https://arxiv.org/abs/2312.12705

研究者の見解では、1兆パラメータのLLMをトレーニングする上で最も重要な課題は、必要なメモリの量、つまり少なくとも14TBのメモリです。

単一の GPU の最大メモリは 64 GB しかないため、トレーニングを完了するには複数の AMD MI250X GPU を並行して使用する必要があります。

また、並列で GPU の数を増やすと、GPU 間の通信に対する要求が非常に高くなります。 GPU 間の帯域幅通信を効果的に利用できない場合、GPU コンピューティング リソースの大部分が無駄になります。

具体的には、研究者らはMegatron-DeepSpeed分散トレーニング フレームワークをFrontierに移植し、AMDハードウェアとROCMソフトウェア プラットフォーム上での効率的な分散トレーニングをサポートしました。

研究者らは、CUDA ベースのコードを HIP コードに変換し、ROCM プラットフォームでの JIT コンパイル エラーを回避するために DeepSpeed オペレーションを事前構築し、マスター ノードの IP アドレスを PyTorch Distributed 初期化のパラメーターとして受け入れるようにコードを変更しました。

220億パラメータのモデルでは、Frontier のトレーニング ピーク スループットは 38.38% に達し、1750億パラメータのモデルのピーク スループットの 36.14%、1兆パラメータのモデルのピーク スループットの 31.96% になりました。

研究チームは、1000B レベルのモデルをトレーニングした後、最終的に 87% のスケーリング効率を達成しました。同時に、比較のために研究者らは1750億のパラメータを持つ別のモデルもトレーニングし、スケーリング効率も89%に達した。

一方、この規模のモデルトレーニングは現在、NVIDIA ベースのハードウェアと CUDA エコシステム上で完了しているため、AMD GPU で同様のトレーニング効率とパフォーマンスを達成するには、まだやるべきことがたくさんあると研究者らは述べています。

トレーニングの詳細

GPTスタイルのモデル構造とモデルサイズ

Transformer モデルは、エンコーダー ブロックとデコーダー ブロックという 2 つの異なる部分で構成されます。

エンコーディング ブロックは、非因果的な自己注意をキャプチャするのに役立ちます。つまり、文内の各トークンは、その左側と右側のトークンに注意を払うことができます。

一方、デコード ブロックは因果的な自己注意を捉えるのに役立ちます。つまり、トークンはシーケンス内の過去のトークンにのみ注意を向けることができます。

最も単純な GPT のようなモデルは、類似のレイヤーのスタックで構成されます。

各層にはアテンション ブロックとフィードフォワード ネットワーク (FFN)2 があります。注意ブロックには 3 セットのパラメータがあり、d はモデルの隠れた次元です。 FFN モジュールには 2 つのレイヤーと重みがあり、各レイヤーには 11d^2 のパラメーターがあります。

埋め込み層はモデルの先頭にあるため、パラメータの数はおよそ 12Ld^2 です。ここで、L は層の数、d は隠れた次元です。

この式に基づいて、研究者は下の表に示すように、サイズが 22B、175B、1T の 3 つのモデルを定義できます。

メモリ要件のほとんどは、モデルの重み、オプティマイザーの状態、および勾配によって発生します。

混合精度トレーニングでは、各モデル パラメータに 6 バイトが必要です。そのうち 4 バイトは fp32 でのモデルの保存用、2 バイトは fp16 での計算用です。

オプティマイザーの状態では、fp32 で勢いを保持するためにパラメーターごとに 4 バイトが必要です。

研究者は各パラメータの fp32 勾配値を保存する必要があります。したがって、混合精度トレーニングに Adam オプティマイザーを使用する場合、最小メモリ要件は次の表のようになります。

各 Frontier ノードは 8 つの MI250X GPU で構成され、それぞれに 64 GB の HBM メモリが搭載されています。

したがって、メモリ要件表から、モデルの 1 つのコピーを適合させるには、モデルの並列化が必要であると結論付けることができます。モデルの並列処理は、潜在次元ではテンソルとフラグメント データの並列処理によって実現でき、レイヤー次元ではパイプラインの並列処理によって実現できます。

パイプラインの並列処理

パイプラインの並列処理により、モデルは p ステージに分割され、各ステージには約 L/p レイヤーが含まれます。次に、バッチはマイクロバッチに分割され、実行ステップごとに 1 つのマイクロバッチが 1 つのステージを通過します。

各ステージは GPU 上に配置されます。

最初は、最初の GPU だけが最初のマイクロバッチを処理できます。 2 番目の実行ステップでは、最初のマイクロバッチが 2 番目のステージに入り、最初のマイクロバッチが 1 番目のステージに入ることができるようになります。

このプロセスは、最後のマイクロバッチが最終段階に到達するまで繰り返されます。

その後、バックプロパゲーションが開始され、プロセス全体が逆方向に継続されます。計算の正しい順序を維持するために、各バッチの後に同期ポイントが導入され、パイプライン ステージのフラッシュが必要になります。

したがって、バッチ処理の開始時と終了時には、前段階と後段階をホストする GPU がアイドル状態になり、計算時間が無駄になったり、パイプライン バブルが発生したりします。

パイプラインの泡の割合は p-1m です。ここで、m はバッチ内のマイクロバッチの数です。

単純な GPipe スケジューリングでは、大きなパイプライン バブルが生成されます。ラインの泡立ちを抑える方法は他にもいくつかあります。

そのようなアプローチの 1 つは、PipeDream によって提案された 1F1B スケジュールです。このスケジュールでは、フォワード パス中に、最後のグループが最初のマイクロ バッチを受け取るまで、マイクロ バッチが最初に前方に流れることが許可されます。

しかし、その後、最初のバッチが後方に伝播され始め、それ以降は前方パスの後に常に後方パスが続くため、1F1B という名前が付けられています。バブルのサイズをさらに縮小するために、研究者らは、1 つのパイプライン グループを 1 つの GPU に配置するのではなく、複数の小さなパイプライン グループを 1 つの GPU に配置する段階的なスケジュールを提案しました。

1F1B プログラムのパイプラインのバブル サイズはおよそ p/m です。ここで、p はパイプライン グループの数、m はマイクロバッチの数です。

マイクロバッチの数。インターリーブを使用した 1F1B スケジュールの場合、バブル サイズは m×v p-1 です。ここで、v は単一の GPU に配置されるインターリーブ グループの数です。

シャードデータの並列処理

シャード化されたデータ並列処理では、モデル パラメータ、オプティマイザーの状態、勾配が行ごとにシャード化され、各 GPU に 1 つのパーティションが配置されます。

トレーニングは一度に 1 つのレイヤーずつ進むため、コンピューティング デバイスのメモリには 1 つの完全なレイヤーと関連する値 (オプティマイザーの状態、勾配、パラメーター) のみが必要です。

シャード化されたデータの並列処理ではこれを利用し、レイヤーを実行する前に、すべての GPU でそのレイヤーのすべての収集を実行して、すべての GPU でインスタンス化します 4b。

現在、すべての GPU に同じレイヤーのコピーがあります。次に、レイヤーはデータのバッチごとに異なる GPU で実行されます。次に、各 GPU はそのレイヤーのすべてのギャザー部分を削除し、フルギャザーによる次のレイヤーの実現を準備します。

このようにして、データの並列処理をエミュレートしますが、各 GPU がモデル全体の完全なコピーをホストするのではなく、現在アクティブなレイヤーのコピーのみをホストします。

シャード化されたデータの並列処理により、モデルが大きすぎて単一の GPU のメモリに収まらない場合でも、GPU 上での大規模モデルのデータ並列トレーニングが容易になります。

DeepSpeed の ZeRO オプティマイザーは、さまざまなレベルでシャードされたデータの並列処理をサポートします。 ZeRO-1 はオプティマイザーの状態のみをシャーディングし、ZeRO-2 は勾配とオプティマイザーの状態をシャーディングし、ZeRO-3 はオプティマイザーの状態、勾配、およびモデル パラメーターをシャーディングします。

一方、PyTorch FSDP (Fully Sharded Data Parallelism) は、3 種類のデータすべてをシャード化し、シャードされたデータ並列処理と従来のデータ並列処理を組み合わせることでハイブリッド データ並列処理をサポートします。

3D 並列処理と Megatron-DeepSpeed

モデルの並列処理を実現するために単一の並列戦略のみを使用することは、非効率的なアプローチになる可能性があります。たとえば、研究者がモデルを水平方向に分割するためにテンソルの並列処理のみを使用する場合、テンソルが薄すぎて、頻繁に完全復元通信が必要になり、トレーニングが遅くなる可能性があります。

一方、研究者がモデルをパイプラインのステージに分割しすぎると、各ステージの計算量が少なくなり、頻繁な通信が必要になります。既知の問題として、複数のノードでテンソル並列トレーニングを実行するには、低速のツリーのような allreduce が必要になることが挙げられます。

複数の並列モードをハイブリッド方式で使用すると、パフォーマンスの低下を最小限に抑えることができます。 3 次元並列処理では、テンソル、パイプライン、およびデータ (従来型およびシャード型) の並列処理技術を組み合わせて、リソースを最大限に活用します。

適切な設定を行うと、3D 並列処理によって通信と計算をオーバーラップできるため、通信の遅延が短縮されます。

人工知能分野で使われる3次元並列標準コードライブラリはMegatron-LMをベースにしています。 MegatronDeepSpeed は Megatron-LM の機能を拡張し、ZeRO-1 シャード データ並列処理や重複 1F1B パイプライン並列処理などの DeepSpeed 機能を追加します。

計画されているパイプラインは並列です。ただし、これらの標準コード ライブラリは、NVIDIA GPU および CUDA プラットフォーム用に開発されています。

研究者たちは、最も完全なフレームワークとして、ソフトウェア スタックが ROCM ソフトウェア プラットフォーム上に構築されている AMD システムである Frontier で Megatron-DeepSpeed を使用することを期待しています。

Megatron-DeepSpeed を Frontier に移植

Megatron-DeepSpeed コード ベースは Nvidia の Megatron-LM コード ベースから派生したもので、Microsoft はその後、DeepSpeed ZeRO オプティマイザー、パイプライン並列処理、MoE を追加しました。

Megatron-LM の開発は NVIDIA が担当しているため、そのコード ベースは NVIDIA GPU と CUDA 環境をターゲット プラットフォームとして開発されています。

このコード ベースを AMD プラットフォームで実行できるように移植するには、いくつかの課題があります。

1. CUDA コード: CUDA コードは AMD ハードウェアでは実行できませんが、HIP (CUDA に似た C/C++ 拡張言語) は実行できます。

研究者たちは、hipify ツールを使用して CUDA ソース コードを HIP コードに変換し、hipcc を使用して共有可能なオブジェクト (so ファイル) を構築し、pybind を使用して Python コードからこれらの共有可能なオブジェクトにアクセスしました。

2. DeepSpeed 操作: ほとんどの DeepSpeed 操作は、トレーニング パイプラインの実行中に JIT (Just in Time) コンパイルを通じて構築されます。

ただし、DeepSpeed 操作の JIT コンパイルは ROCM プラットフォームでは機能しないため、研究者は DeepSpeed のインストール時にすべての操作を事前にビルドしました。

研究者らは、実行時エラーを回避するために、Megatron-DeepSpeed コードベースのすべての JIT 機能を無効にしました。

3. PyTorch 分散環境を初期化する: Megatron-DeepSpeed は、PyTorch 分散初期化を使用して、さまざまなデータとモデルの並列グループを作成します。

初期化プロセスでは、コンピューティング ノードを「マスター」ノードとして指定する必要があり、すべての分散プロセスではその IP アドレスが必要です。

研究者らは、MASTER ADDR をパラメータとして受け入れるようにコード ベースを変更しました。

研究者らは、SLURM ノード リストから最初のノードの IP アドレスを読み取り、それを srun を使用して開始されたすべてのプロセスに引数として渡す起動スクリプトを準備しました。

初期化コードは、PyTorch 分散初期化にこの MASTER ADDR を使用します。

4. ROCM プラットフォーム ソフトウェアを通じて提供されるライブラリ/パッケージ: 研究者は AMD 開発者と協力して、APEX などのいくつかの基本的な CUDA パッケージの ROCM バージョンを入手しました。

APEX は NVIDIA の混合精度ライブラリであり、混合精度トレーニング用の Megatron-DeepSpeed コード ベースで広く使用されています。

また、Frontier のコンパイラーで使用できるように、ROCM 対応バージョンの FlashAttention および FlashAttention2 ライブラリも適応させました。 Flash-Attention 操作は、Composable Kernel ライブラリのカーネルを使用して AMDGPU に移植されました。

さまざまな配分戦略の実証分析

テンソル並列処理

テンソル並列方式では、モデルレイヤーを行ごとに分割し、各レイヤーの後に、Allreduce を通じていくつかのアクティベーション値を集計する必要があります。

各レイヤー実行後の AllReduce コストは高く、テンソル並列グループ内の GPU 間の通信帯域幅に依存します。通信量は、隠しサイズとマイクロバッチ サイズに依存します。

下の図 5 は最先端の GPU 間の通信帯域幅を示しています。ノードには 8 つの GPU があり、1 つのチップ内の GPU は 4 つの (50+50 GB/s) Infinity Fabric を介して接続されています。

チップ全体の GPU 間の帯域幅はその半分になります。ただし、ノード間の GPU 間の帯域幅は 25+25 GB/秒です。

したがって、ネットワークトポロジと構成の観点から、TP = 2 が最も高速な通信速度を持ち、TP = 4 または 8 が 2 番目に高速な通信速度を持ちます。

しかし、TP ¿ 8. 通信は低速のイーサネットを介して行われるため、通信速度が大幅に低下します。したがって、TPを[2, 4, 8]の範囲に保つことが最適な戦略であるはずです。

研究者らは 8 つの GPU を使用して、TP 値が 1 から 8 までの範囲の 14 億のモデルをトレーニングしました。結果は、TP 値が小さいほどスループットが高くなることを示しました。

観察 III.1: TP 値が大きいほど、トレーニング効果は悪くなります。

B. パイプラインの並列処理

パイプラインの並列化では、モデルをレイヤー次元に沿って分割し、連続するレイヤーをパイプライン ステージにグループ化します。マイクロバッチの実行は、あるステージから次のステージへと流れます。

パイプライン バブルは、この並列アプローチを使用した効率的なトレーニングを制限する要因となります。

研究者らは、22B パラメータと 1T パラメータを持つ 2 つのモデルについて、大きな M または大きな GBS の効果を観察し、GPU スループットへの影響を理解しました (下の図 7)。

観察 III.2: 大きなグローバル バッチ サイズまたは多数のマイクロバッチを使用してパイプライン ステージを飽和させると、パイプライン バブルのサイズが最小限に抑えられます。

パイプライン ステージ数の影響: 次に、研究者はパイプライン ステージ数がトレーニング パフォーマンスに与える影響を調査しました。直感的に言えば、パイプラインのステージが増えると、通信が発生する前の計算が少なくなります。

グローバル バッチ サイズ (マイクロ バッチの数) が固定されている場合、パイプライン ステージが増えると計算量が少なくなります。

パイプラインのステージ数が増えると、バブルのサイズも大きくなります。研究者らは、PMP を固定したままパイプライン ステージの数を増やし、グローバル バッチ サイズを比例して増やすことも試みました。

観察 III.3: グローバル バッチ サイズを一定に保ちながら、パイプライン ステージの数を増やすと、パイプライン バブルのサイズが大きくなり、トレーニング パフォーマンスが低下します。

観察 III.4: パイプライン ステージ数とミニバッチ数の比率が一定であれば、パイプライン ステージ数が増加してもトレーニング パフォーマンスは一定のままです。

最初の実験 (上の図 8a) から、パイプラインのステージ数が増えると、トレーニング パフォーマンスが低下することがわかります。ただし、グローバル バッチ サイズを調整してバブル比を固定すると、スループットを維持できます (上記の図 8b)。

研究者らは、実験、ハイパーパラメータの調整、分析を通じて、さまざまな分散戦略とソフトウェアの最適化を組み合わせた、Frontier 上で TrillionParameter モデルをトレーニングするための効率的な戦略を特定しました。

1兆パラメータモデルのトレーニング

兆パラメータモデルを訓練するための効率的な戦略

ミニバッチの数を増やすことでパイプライン ステージを飽和させる: 研究者は、DeepSpeed (Megatron バージョンではなく DeepSpeed-Megatron から) が提供するパイプライン並列処理を使用しました。このパイプライン並列アルゴリズムは、複数のステージが重なり合い、1F1B アルゴリズムを使用してバブル サイズを縮小する PipeDream のアルゴリズムです。

ただし、パイプラインのステージが飽和していない場合は、バブルのサイズが大きくなります。飽和状態を確実にするには、ミニバッチの数がパイプライン ステージの数と等しいかそれを超える必要があります。

テンソルの並列処理を単一ノード/8 つの GPU に制限する: AllReduce 操作は頻度が高すぎてレイヤーごとに実行する必要があるため、異なるノードに分散されたレイヤーにより、ノード間の GPU 間でツリーベースの AllReduce が発生し、通信の遅延が大きなボトルネックになります。

Flash-Attention v2 の使用: 通常のアテンション実装と比較して、研究者は Flash-Attention を使用することでスループットが 30% 向上することを確認しました。

ZeRO-1 オプティマイザーを使用したデータ並列処理: 研究者は ZeRO-1 を使用してデータ並列処理を実装し、メモリ オーバーヘッドを削減しました。

AWS の RCCL プラグインによる通信の安定性の向上: AWS OFI RCCL プラグインにより、EC2 開発者は AMD RCCL ベースのアプリケーションを実行するときに libfabric をネットワーク プロバイダーとして使用できるようになります。 Frontierでは、このプラグインの使用により通信の安定性が実証されています。

兆パラメータモデルのトレーニングパフォーマンス

研究者たちは、ハイパーパラメータの調整から得た教訓に基づいて、パラメータ数が 220 億から 1750 億に及ぶ一連のモデルの組み合わせを決定しました。

これら 2 つのモデルの GPU スループットに勇気づけられ、研究者は最終的に、表 V にリストされている分散戦略の組み合わせを使用して 1 兆パラメータ モデルをトレーニングし、10 回の反復を実行してそのトレーニング パフォーマンスを観察しました。

22B パラメータ モデルの場合、研究者はピーク スループット (191.5 TFLOPS) の 38.38% (73.5 TFLOPS) を抽出することができました。

175B モデルのトレーニングでは、研究者はピーク スループット (69.2 TFLOP) の 36.14% を達成しました。

最後に、1T モデルでは、ピーク スループット (61.2 TFLOP) の 31.96% が達成されます。

拡張パフォーマンス

データの並列処理とシステム内の多数の GPU の関与を通じてモデル並列トレーニングのパフォーマンスを維持することは、非常に困難な作業です。最も強力な GPU は速度の異なる通信リンクで接続されているため、ネットワークの大部分に負荷がかかるとパフォーマンスが低下する可能性があります。

そこで研究者らは、データ並列処理を通じて 175B モデルのトレーニングを 1024 GPU に、1T モデルのトレーニングを 3072 GPU に拡張し、トレーニング戦略のスケーリング効率を測定しました。

1. 弱いスケーラビリティ: 研究者らは、1T モデルで弱いスケーラビリティ実験を行うために、3200、6400、9600 のグローバル バッチ サイズを使用して、1024、2048、3072 の GPU でデータ並列トレーニングを実行しました。データ並列トレーニングにより、100% の弱いスケーラビリティ効率が達成されます (下の図 12)。

2. 強力なスケーラビリティ: 研究者は、グローバル バッチ サイズを 8000 に維持し、GPU の数を変えながら、強力なスケーラビリティ実験を実施しました。研究者らは、1024 個の GPU 上の 175B モデルで 89.93% という優れたスケーラビリティ パフォーマンスを達成しました (図 13a)。研究者らは、3072 個の GPU 上で 1 兆個のパラメータ モデルに対して 87.05% という優れたスケーラビリティ パフォーマンスを達成しました (図 13b)。

世界最速のスーパーコンピュータ

AMD 搭載の Frontier スーパーコンピュータは、1.102 ExaFlops/s の計算能力を備え、世界初の公式認定エクサスケール スーパーコンピュータとなりました。

新たに発表された世界最速スーパーコンピューターの Top500 リストで第 1 位にランクされています。

Frontier は、リストにある次の 7 つのスーパーコンピューターの合計よりも高速です。

Frontier は現在、HPL-AI ベンチマークで 6.88 エクサフロップスの混合精度パフォーマンスを実現し、世界最速の AI システムとしてもランク付けされています。

これは、脳内の 860 億個のニューロンのそれぞれが 1 秒あたり 6,800 万の命令を実行することに相当します。

Frontier スーパーコンピューターの規模の大きさは驚異的ですが、これは AMD が今年の Top500 リストで達成した数多くの成果の 1 つにすぎません。世界のトップ 10 スーパーコンピューターのうち 5 台は AMD EPYC システムを搭載しており、トップ 20 のうち 10 台は AMD EPYC システムを搭載しています。

Frontier スーパーコンピューターは HPE によって構築され、オークリッジ国立研究所 (ORNL) に設置されています。

このシステムには 9,408 個のコンピューティング ノードがあり、各ノードには 64 コアの AMD「Trento」CPU、512 GB の DDR4 メモリ、4 つの AMD Radeon Instinct MI250X GPU が搭載されています。

ノードは 74 台の HPE Cray EX キャビネットに分散されており、それぞれの重量は 8,000 ポンドです。システム全体には 602,112 個の CPU コアと 4.6 PB の DDR4 メモリが搭載されています。