※本記事は、Stanford CS336 Language Modeling from Scratchコースの並列化に関する講義内容を基に作成されています。本講義は、Stanford大学コンピュータサイエンス学部准教授でCenter for Research on Foundation Models (CRFM)所長のPercy Liang氏と、同学部助教授のTatsunori Hashimoto氏によって行われました。
コース詳細については https://stanford-cs336.github.io/spri... で、Stanford のオンラインAIプログラムについては https://stanford.io/ai で、受講申し込みについては https://online.stanford.edu/courses/c... でご覧いただけます。
本記事では、講義の内容を要約しております。なお、本記事の内容は原講義の見解を正確に反映するよう努めていますが、要約や解釈による誤りがある可能性もありますので、正確な情報や文脈については、オリジナルの講義動画をご視聴いただくことをお勧めいたします。
Stanford Onlineは、Stanford大学工学部のグローバル・オンライン教育センター(CGOE)によって運営・管理されており、世界中の受講者に向けて学位プログラム、単位認定教育、専門資格プログラム、無料・オープンコンテンツを提供しています。
1. イントロダクション・背景
1.1 マルチマシン最適化の必要性
今日の講義は基本的なシステム講義の第2回目となり、マルチマシン最適化に焦点を当てます。今日の目標は、単一GPUのスループット最適化から、本当に大規模なモデルを訓練するために必要な複雑性と詳細を理解できるようになることです。
モデルが大きくなると、もはや単一のGPUには収まらなくなります。そのため、モデルを異なるマシンに分割する必要がありますが、同時に、これらのモデルを迅速に訓練するために、保有するすべての異なるサーバーを活用できる必要があります。私たちは計算とメモリの両方の懸念に対処しなければなりません。
先週お話しした GPU スケーリングは、GPU あたりのフロップス数が非常に上昇している超指数関数的な曲線を見ると、非常に印象的です。しかし、計算とメモリの両方を急速にスケールアウトしたい場合、単一のGPUでは十分ではありません。この曲線が上昇し続けるまでにはさらに数年待つ必要があります。今日、ここで本当に強力な言語モデルを訓練したいのであれば、マルチマシン並列処理に頼らざるを得ません。
世界最速のスーパーコンピューターを見てみると、右側に示されているように、最速のスーパーコンピューターはエクサフロップス単位の計算能力を持っています。これらが緑色の線で示されているものです。今日最大で最強の言語モデルを訓練しようとするなら、本当にこれらに頼らざるを得ません。
これがマルチマシン並列処理を考える計算面での理由です。しかし、同じことを考えるメモリの観点もあります。これら2つは、考慮すべき核となるリソースと核となる懸念事項です。
メモリの観点では、多くのモデルがかなり大きくなっています。もちろん、GPU上のメモリも成長していますが、それほど急速ではありません。単一のGPUではこれらのモデルを収容することができません。おそらく遠い将来には、これらの多くについて心配する必要がないかもしれませんが、数十億、数十億のパラメータがあり、それらは単一のGPUに非常にきれいに収まることはありません。そのため、私たちが持つメモリ制約を非常に尊重する必要があります。
1.2 単一GPU制限の課題
これらが私たちが対処しなければならない現実です。これらを処理できるようになるために私たちが持つツールは何でしょうか。GPUは、クラスクラスターで気づいているはずですが、単体では存在しません。単一のマシンには、同じ物理ラック内に複数のGPUが搭載されています。
ここに例があります。これはGPT Neo Xの論文から取ったものだと思いますが、これは古い例ですが、クラスで使用しているH100マシンにも同じ教訓が適用されます。ここには8つの異なるGPUがあります。これらは高速インターコネクトを通じてさまざまなCPUに接続されています。各GPU内では、底部にこのNVSwitchというものが見えます。これは、これら8つのGPU間の非常に高速な接続です。
しかし、これら8つのGPUが異なるマシン上のGPUと通信したい場合、ネットワーキングスイッチを通過する必要があります。HDR InfiniBandと書かれているこの紫色の線が見えます。これは、NVLink接続と比較してはるかに遅い接続です。スループットの違いを見ると、レーンあたり約8倍遅いことがわかります。
私たちが持つこのハードウェア階層は、実際にモデルを並列化する方法に大きな影響を与えることになります。これらのことを話すときに、この精神的モデルを持ち続けることができます。単一のマシン内では非常に高速な接続があります。そして、マシン間を移行すると、それは遅くなります。そして、使用しているハードウェアの種類によっては、256個のGPUがネットワーク接続された範囲を超えると、さらに別のレベルの遅さが生じる可能性があります。
1.3 計算と記憶容量の拡張要求
さて、私たちは新しい計算単位について話し始めようと思います。GPUの代わりに、新しい単位はデータセンターです。データセンター全体が、私たちが行おうとしているものになります。
そして今、私たちは2つの異なることを得るアルゴリズムとシャーディング戦略を考え出そうとします。1つ目は線形メモリスケーリングです。GPUの数をスケールアップするにつれて、訓練できる最大のモデルはそれに比例してスケールします。つまり、本当に望むなら、より大きく、より大きなモデルを訓練できるということです。
私は線形計算スケーリングも必要です。より多くのGPUを取得するにつれて、モデルを訓練するために行っている有用な計算が線形にスケールします。そして最後に、これらのアルゴリズムの多くは、これらの非常にシンプルな集合通信プリミティブをさまざまな方法で呼び出すことによって実装されることになります。
これらの並列アルゴリズムのパフォーマンス特性を考えるとき、基本的に集合通信プリミティブを数えることで推論するだけで十分です。これが、これらについて考える重要な方法の一種です。ここでは、これらのアルゴリズムの低レベル実装まで完全に降りていくことはありません。
2. ハードウェア基盤
2.1 GPU階層とネットワーキング
多くの皆さんはシステムやネットワーキングのクラスを受講してすでにこれを知っているかもしれませんが、ここでは集合通信操作に関する非常に簡単な復習をします。これを取り上げる理由は、並列化アルゴリズムのパフォーマンス特性のより細かい点を本当に理解するために知っておく必要がある、1つの特に重要な恒等式または等価性があるからです。
私たちが持つハードウェア階層は、実際にモデルを並列化する方法に大きな影響を与えることになります。単一のマシン内では非常に高速な接続があります。そして、マシン間を移行すると、それは遅くなります。そして、使用しているハードウェアの種類によっては、ネットワーク接続された約256個のGPUを超えると、さらに別のレベルの遅さが生じる可能性があります。
GPU世界では、これが一般的にどのように機能するかを示すGPT Neo Xスライドで説明したように、8個のGPUを含む単一のマシンであるノードがあり、次に互いに非常に迅速に接続するスイッチがあります。これらのマシンは、約256個のGPUまで全てが全てに接続されています。これは、マシン間で非常に高速で任意の通信を行える重要な閾値です。その上では、実際には約1つのラック分のGPUを超えると、これらのリーフスイッチとスパインスイッチという、はるかに遅い通信が必要になります。
一方、GoogleのTPU設計を見ると、実際にはマシンのネットワーキングに非常に異なるアプローチを取っています。単一のTPUチップがあり、それらはすべて隣接ノードと非常に迅速に通信します。これは、彼らがトロイダルメッシュと呼ぶ、非常に簡単に拡張可能なものです。しかし、隣接ノードとのみ通信できます。
all-reduceスライドの直後にこれについて話している理由は、all-reduceやreduce-scatterのような集合通信の種類を行うことを考えると、全て対全て接続でできるのと同じくらい効率的に、トロイダルメッシュ上でそれらを実装できるからです。純粋に集合通信のために最適化している場合、GPU ネットワーキングよりもTPUネットワーキングのようなものを考えることが理にかなっています。異なる並列化操作を経るにつれて、これの長所と短所について少し話します。
2.2 集合通信操作(Collective Communications)
まず、皆さんがおそらく聞いたことがあるall-reduceについて説明します。この場合、4つのマシン、4つのランクがあり、それぞれが独自のデータ片を持っています。そして、何らかの削減操作を実行したいとします。例えば、これらすべての入力を合計し、その出力をすべてのマシンにコピーしたいとします。これは、all-reduceしているものの総数の約2倍のコストがかかります。
broadcast操作があります。ここでは、ランク2から単一の入力を取り、それを残りのすべてのランクにコピーします。これは、通信コストの観点から、出力の総数の約1倍のオーダーになります。
そして、異なる入力を持つreductionがありますが、これは合計されて1つのマシンにのみ送信されます。そして、非常に重要な2つは、all-gatherとscatterです。これらはそれほど一般的ではないかもしれませんが。
all-gatherは、例えば、パラメータの単一のサブコンポーネントをランク0から取り、それをすべてのランクにコピーする操作です。ランク1、2、3でも同様です。つまり、これらのそれぞれが異なる部分、例えばパラメータを処理し、それらが残りのマシンにコピーされます。これは、私が持っているものを他の全員にコピーすることです。
そして、reduce-scatterは、例えば、各行を取り、それらを合計し、結果をランク0にのみ送信します。これはall-reduceの部分版です。
all-gatherとreduce-scatterは、ある意味で多くの並列化アルゴリズムが構築される原始的なものであるため、非常に重要です。
これは重要な等価性または恒等式です。この講義で重要なポイントで1回か2回参照します。all-reduceを実行したい場合、異なるGPU、A、B、C、Dがあり、各GPUが異なるデータポイントを処理しているとします。各データポイントに対して異なる勾配があり、それらの勾配を合計する必要があり、その後、すべての勾配をGPUに戻す必要があります。これは、4つのGPU間で行う可能性がある古典的なデータ並列操作です。これがall-reduceです。
しかし、重要なことは、これを2つの操作、reduce-scatterとall-gatherで置き換えることができることです。reduce-scatterは各行を合計し、行の結果をそれぞれGPU 0、1、2、3に残します。そして、all-gatherを実行して、それらを残りのGPUにコピーして戻します。各GPUは現在、パラメータの一部の完全な合計を取得し、その後、それを残りのワーカーにコピーして戻します。
帯域幅制限レジームでは、これは基本的にできる最善のことです。all-reduceで実行できる最善のことは、reduce-scatterとall-gatherから得られる帯域幅とほぼ一致します。all-reduceと右辺の両方でどれだけの通信操作が発生するかを書き出すことで、これを自分で確認できます。
2.3 GPU vs TPUアーキテクチャの比較
並列化アルゴリズムについて話す前に簡単に触れたい最後のことは、GPUとTPUについて話す唯一の場所になります。今日の議論の大部分は実際に基盤となるハードウェアを抽象化できますが、実際に1つの重要なことがあり、これを前もって言及しておくことで、これらについて話すときに後で参照できます。
GPU世界では、これが一般的にどのように機能するかを、ここのGPT Neo Xスライドで示したように、8個のGPUを含む単一のマシンであるノードがあり、その後、互いに非常に迅速に接続するスイッチがあります。これらのマシンは約256個のGPUまで全て対全てに接続されています。これは、マシン間で非常に高速で任意の通信を行える重要な閾値です。その上では、実際には約1つのラック分のGPUを超えると、これらのリーフスイッチとスパインスイッチという、はるかに遅い通信が必要になります。
一方、GoogleのTPU設計を見ると、実際にはマシンのネットワーキングに非常に異なるアプローチを取っています。単一のTPUチップがあり、それらはすべて隣接ノードと非常に迅速に通信します。これは、彼らがトロイダルメッシュと呼ぶ、非常に簡単に拡張可能なものです。しかし、隣接ノードとのみ通信できます。
all-reduceスライドの直後にこれについて話している理由は、all-reduceやreduce-scatterのような集合通信の種類を行うことを考えると、全て対全て接続でできるのと同じくらい効率的に、トロイダルメッシュ上でそれらを実装できるからです。純粋に集合通信のために最適化している場合、GPUネットワーキングよりもTPUネットワーキングのようなものを考えることが理にかなっています。異なる並列化操作を経るにつれて、これの長所と短所について少し話します。
3. データ並列化(Data Parallelism)
3.1 基本的なSGDとナイーブデータ並列化
データ並列化の出発点は単なるSGDです。非常にナイーブなバッチ確率的勾配降下を行っている場合、これを行うための公式は、ここのスライドにあるこの方程式のようになります。バッチサイズ大文字Bを取り、すべての勾配を合計し、パラメータを更新します。
ナイーブデータ並列化は、単純に、バッチサイズBを取り、それを分割し、異なるマシンに送ると言っているだけです。各マシンは合計の一部を計算します。そして、各勾配ステップの前に、すべての勾配を一緒に交換して同期し、勾配を同期してから、パラメータ更新を行います。
計算とメモリのスケーリングやこれらすべてのことについて話してきたので、これらのそれぞれがどのようなものかを話してみましょう。
計算スケーリングについて、データ並列化は非常に優れています。各マシン、各GPUはB over M個の例を取得します。バッチサイズが十分に大きければ、各GPUはかなり良いバッチサイズ、マイクロバッチサイズを取得します。そして、うまくいけば、その計算を飽和させることができます。それは良いことです。
通信オーバーヘッドは何でしょうか。毎バッチごとに、パラメータ数の2倍を送信する必要があります。all-reduceは、通信コストの観点から、all-reduceしているものの量の約2倍になることを覚えておいてください。これは、バッチサイズが大きい場合には問題ありません。バッチサイズが本当に大きい場合、時々勾配を同期する通信オーバーヘッドをマスクできます。
メモリスケーリングについては、これを全く触れていません。すべてのGPUがパラメータ数を複製する必要があります。オプティマイザー状態を複製する必要があります。メモリスケーリングにとっては非常に悪いです。メモリについて全く心配する必要がなければ、これは問題ない戦略です。しかし、実際にはメモリが問題だと思います。
ここに座っている皆さん全員が、大きなモデルをGPUに置こうとして、PyTorchが「ああ、メモリが不足しています」と言われた経験があると思います。これは、より多くのバッチサイズを適合させることができれば、データ並列がより効率的になるため、訓練でも本当に問題になります。理想的には、メモリを節約したいと思います。
3.2 メモリ使用量の問題点
ナイーブデータ並列のメモリ状況をより詳しく見てみましょう。メモリ状況は実際に見た目よりもさらに悪いです。実際に非常にひどいです。皆さんは課題1でこれを行いましたが、モデルのコピーをいくつ保存する必要があるかを考えることができます。
これは非常に大きいです。訓練を行っている精度によって、パラメータあたり約16バイトのデータを保存する必要があります。実際に、重みの約5つのコピーを保存する必要があります。これは本当にひどいです。モデルパラメータについて考えるだけなら、技術的には2バイトしか必要ないからです。では、その8倍の係数はどこから来たのでしょうか。
少なくとも、勾配が必要です。BF16で勾配を計算している場合、それはさらに2バイトです。しかし、オプティマイザー状態が現れ、それは本当に大きな問題になります。4バイトのマスター重みがあり、SGDで累積しているもの、つまり行っている中間合計があります。Adamの第1モーメント推定には4バイトまたは2バイトが必要です。Adamは履歴勾配を追跡することを覚えておいてください。そして、Adamには第2モーメント推定、過去に得た勾配の分散も必要で、それはさらに4バイトまたは2バイトが必要です。
最初は問題ないように見えたものが、実際には非常に厳しい状況になっています。
16倍を図として描くと、パラメータメモリの観点では、メモリ使用量の大部分が実際にAdamオプティマイザーのオプティマイザー状態によって支配されていることがわかります。消費されるメモリは、オプティマイザー状態に使用されるバイト数の関数になります。そして、それは一般的に、核となるパラメータと勾配のメモリ使用量よりもさらに多くなります。
64個のアクセラレータに分散された7.5Bモデルの簡単な例では、大量のメモリを使用しています。この総メモリは、少なくともGPUの数に比例して上昇します。それは全く良くありません。
しかし、この図を見ると、非常にシンプルなアイデアが得られます。明らかに、またはそうでないかもしれませんが、データ並列を行うためには、パラメータと勾配をデバイス間でコピーする必要があるようです。これは必要に思えます。しかし、すべてのオプティマイザー状態を本当にすべてのマシン上に置く必要があるでしょうか。その質問をすると、ここの2行目に到達できるかもしれません。これはオプティマイザー状態シャーディングと呼ばれます。
これができれば、少なくともこの場合、120ギガバイトの総メモリ使用量から31.4まで下げることができます。そして、勾配のシャーディングを開始できるかもしれません。すると、16.6ギガバイトのメモリ使用量まで下げることができます。そして、パラメータもシャーディングすれば、1.9ギガバイトのメモリ使用量まで下げることができます。
これは非常に良い場所になるでしょう。なぜなら、今では必要なすべてのオプティマイザー状態、パラメータ、勾配メモリを完全にシャーディングしたからです。
3.3 ZeRO最適化手法(Stage 1-3)
ZeRO Stage 1: オプティマイザー状態シャーディング
オプティマイザー状態をシャーディングすることについて、「どうやってオプティマイザー状態をシャーディングできるのか」という非常に良い質問があります。データ並列を行う場合、GPU 0はデータポイント1に責任を持つ必要があります。明らかに、すべてのパラメータについて知り、更新する必要があります。では、どうやってオプティマイザー状態をシャーディングできるのでしょうか。
ある意味で、これはZeRO、つまりゼロオーバーヘッドデータ並列オプティマイザーと呼ばれるものの非常に巧妙なアイデアです。データ並列を行っている場合でも、実際にはすべてをすべてのマシンにコピーする必要がないということを示しています。これらすべてを避けるために、通信を行う方法について非常に巧妙になることができます。
私たちが行うことは、言ったように、オプティマイザー状態を分割することです。第1と第2モーメントは現在、すべてのGPUに分割されています。しかし、全員がパラメータと勾配を持っています。なぜこれが重要なのでしょうか。
パラメータと勾配があれば、GPU 0として、すべてのパラメータと勾配を持っています。それは完全な勾配を計算するのに十分な情報です。この例では完全な勾配更新を計算できます。私にできない唯一のことは、その勾配を取ってAtom stepを取ることです。すべてのオプティマイザー状態を見ない限り、パラメータを更新できません。それが重要なアイデアです。
今、起こることは、GPU 0がすべての勾配を計算しますが、GPU 0は現在、自分が所有するシャードのパラメータの更新のみに責任を持ちます。それが重要なアイデアです。パラメータを更新する作業を分散します。そして、パラメータを同期して戻します。
これがどのように機能するか、そしてなぜゼロオーバーヘッドと呼ばれるかを、はるかに詳細に示します。ステップ1では、すべてのGPUが異なるデータポイントを取得するとしましょう。このバッチ計算をすべて簡略化します。GPU 0から4まであり、すべてのGPUが単一の例を取得し、所有する例で完全な勾配を計算します。
次に行うことは、勾配をreduce scatterすることです。各GPUが所有する勾配を送信します。ある意味で、各GPUが所有する勾配を収集します。GPU 0は、この最初の4分の1のパラメータに責任を持つとしましょう。パラメータがY軸でここにあり、X軸がGPUです。
行うことは、GPU 0が責任を持つパラメータのサブセットについて、他のすべてのGPUからのすべての勾配情報をGPU 0が持つようにreduce scatterを行うことです。今、GPU1、GPU2、GPU3からこの勾配情報を取得し、それがすべてGPU 0に削減されます。
今、GPU 0は、この最初の部分に対応するオプティマイザー状態を持っているため、自身のパラメータを更新するために必要なすべての情報を持っています。この最初の部分の完全な合計勾配を持っています。そして今、勾配と状態を使用して、パラメータの一部で勾配更新を行います。
今、このサブセットの完全に更新されたパラメータをGPU 0に持っており、行う必要があることは、更新されたパラメータをすべてのランクにall gatherで戻すことだけです。
重要なことは、reduce scatterとall gatherを行っているということです。前に言ったことを覚えているなら、reduce scatterとall gatherはall reduceと同じコストがかかります。ここで少し驚くべき魔法的なことが起こりました。以前はすべての勾配でall reduceを行って、全員の勾配が同期されるようにしていました。それはパラメータ数の2倍のコストがかかりました。
しかし、更新の方法について巧妙であれば、reduce scatterとall gatherを行うことができ、2つのステップの間で計算を行うことができます。それは同じ量の計算通信コストを与えますが、今では、少なくともオプティマイザー状態について、モデル全体でオプティマイザー状態を完全にシャーディングしました。ZeRO stage 1は、帯域幅制限レジームでは、ある意味で無料であり、メモリの勝利を与えてくれます。
ZeRO Stage 2: 勾配シャーディング
ZeRO stage 2に進みます。それはまだ比較的シンプルです。オプティマイザー状態シャーディングトリックが理にかなったことを願っています。それは非常にクールだと思います。
今、さらに多くのものをシャーディングしたいと思います。マシン間で勾配をシャーディングしたいと思います。大まかには、stage 1と同じ種類のトリックを行うことができますが、追加の複雑さが1つあります。
追加の複雑さは何でしょうか。完全な勾配ベクトルをインスタンス化することは決してできません。完全な逆伝播を行い、完全な勾配ベクトルを計算しようとすると、メモリが不足する可能性があります。最大メモリ使用量を基本的にこれで制限したいのです。これは、完全なパラメータ、シャーディングされた勾配、シャーディングされたオプティマイザー状態のようなものです。
行う必要があることは、逆伝播を行うときに、勾配ベクトルを計算しているときに、最初に完全な勾配をインスタンス化してから通信を行うことはできません。行う必要があることは、勾配を逆方向に計算するとき、レイヤー分の勾配を計算するとすぐに、それが属する対応するGPUにそれを送る必要があります。
これがどのように機能するかです。大まかには同じアイデアです。今、全員が自分のバッチコンポーネントを持っています。全員が計算グラフで段階的に逆方向に進みます。レイヤーごとに操作するとしましょう。レイヤーは異なるGPUに原子的にシャーディングされています。
計算グラフで逆方向に進むときに行うことは、レイヤーの勾配を計算した後、すぐにreduction操作を呼び出して、これを適切なワーカーに送ることです。レイヤーは何らかのワーカーに属します。この場合、GPU番号2のようなものかもしれません。その時点でそれをすぐに削減し、ワーカーに送信します。勾配はもはや必要ありません。
ランク0、1、3で勾配を保存する必要はありません。すぐにそれを解放できます。そして、このプロセスを続けます。
すべてのマシンが完全に更新された勾配を持っています。パラメータのシェアに対する完全な勾配、パラメータのシェアに対する完全なオプティマイザー状態を持っているので、各マシンはパラメータを更新でき、パラメータを一緒にall gatherして戻すことができます。
これは、すべてのレイヤーでこの種のreduction操作を行っているため、より多くの通信のように見えるかもしれません。しかし、これは少量のパラメータに対してのみです。シャーディングされています。そのため、完全な通信は同じままです。ZeRO stage 2は、レイヤーごとに同期し、勾配が適切なワーカーに適切に送信されることを確認する必要があるため、より多くのオーバーヘッドがあります。しかし、オーバーヘッドは非常に最小限です。まだ非常にシンプルで、かなり直接的です。
ZeRO Stage 3 (FSDP):
最後のZeRO stage 3は確実により複雑ですが、最大の勝利を可能にします。今では基本的にすべてが保有するGPUの数で除算されます。可能な最大の節約を得ることができます。
FSDPを聞いたことがあるなら、おそらく過去の人生のある側面でそれを使用したことがあるでしょう。FSDPはまさにZeRO stage 3です。今日、うまくいけばFSDPがどのように機能するかがわかるでしょう。
同じアイデアが適用されます。パラメータを含むすべてをシャーディングします。ZeRO stage 2と同じことを行います。これらの大きな勾配ベクトルを保持しないように、段階的に通信と計算を行います。また、計算グラフを通過しながら、フォワードパスとバックワードパスの両方について、需要に応じてパラメータを送信および要求します。
もちろん、重要なことは、可能な限り低いオーバーヘッドでこれを行うことです。FSDPについて本当に驚くべきことは、これが可能であることではなく、これが比較的低いオーバーヘッドで可能であることです。次のスライドで、なぜ低いオーバーヘッドなのかがわかります。
これは最も親しみやすいグラフィックではないかもしれませんが、これはFSDPのベビー版です。次のスライドはもう少し複雑です。しかし、概念的には、これは実際にすべてを説明します。
行うことは、モデルの重みを持ち、進むにつれてモデルの重みをall gatherすることです。各レイヤーについて、単一のGPUがすべてのパラメータを持つことはありません。GPU 0に進んでフォワードパスを実行すると言う通常のことはできません。それは不可能です。
GPU 0は、最下位レイヤーのみを所有するとしましょう。その計算を行い、停止して、すべての他のワーカーからパラメータを要求します。停止してall gatherを行います。ここにall gatherステップがあります。すべてのパラメータを収集します。
今、フォワードを行うために必要なパラメータを持っています。前進して、以前持っていなかったレイヤーを計算できます。そして今、重みを解放できます。もう重みは必要ありません。削除します。
今、次のレイヤーをall gatherし、別のフォワードを行い、重みを解放して、これを繰り返すことができます。アクティベーションは保存する必要があります。ここでアクティベーションメモリが成長しています。それは最終的な問題になります。
しかし、今のところアクティベーションを無視すれば、これは素晴らしいです。レイヤーをロードし、フォワードを行い、解放するからです。ここでのメモリオーバーヘッドは非常に低いです。
最後まで到達すると、バックワードパスでも同じことができます。バックワードを呼び出すことができ、ニューラルネットワークを通って逆方向に移動するたびに、必要なパラメータをall gatherできます。勾配が計算された後に更新するためにreduce scatterを行うことができ、今、重みを解放できます。必要のない勾配とパラメータの両方を解放できます。
最後に、完全に更新されたモデルを得ました。ここで心配する必要がある3つの異なる操作があります。all gatherがあり、別のall gatherがあり、その後、勾配更新ステップを取った後にモデルを更新するための別のreduce scatterがあります。
概念的には、これはZeRO stage 2を超える単一ステップに過ぎませんが、より多くのオーバーヘッドがあることがわかります。総通信コストは今やより高くなります。以前は、パラメータ数の2倍でした。すべてがある意味で無料でした。今はそうではありません。
パラメータ通信コストの合計3倍があり、これらの通信が完了するのを待つことに関連するコストもあります。
しかし、FSDPについて本当にクールなことは、実際に驚くほど低いオーバーヘッドであることです。いつもパラメータを送受信するというこのクレイジーなことをしているため、物事が本当に遅くなると想像するかもしれません。いつも通信をしなければなりません。
しかし、通信と計算をオーバーラップするというこの中核的なアイデアを行うことができます。GPUを動作させたいと同時に、バックグラウンドで通信も行われ、プリフェッチのようなもので、何らかの情報が必要になる時までに、それはすでにロードされ、すでに通信されており、準備ができています。
計算グラフがW1 W0 + W2 W0 × xのような感じだとしましょう。何らかの入力をyとしましょう。このような非常にシンプルな計算グラフです。そして、FSDPを実行するかもしれません。
実際に、この最後のブロック図のような計算と通信を取得します。先週nsight systemsの例を行ったので、うまくいけばこの図が明確になるでしょう。
CPUは基本的に多くのコマンドをディスパッチし、GPUの通信部分に基本的にいくつかのパラメータを取得するよう求めます。GPUに、よし、いくつかの行列乗算を行うよう言うものをディスパッチし、ある意味でGPUよりもはるかに先に実行されます。先週プロファイラーを見たときにこれを見ました。
今、デバイス上で起こる通信と計算の両方のシーケンスを見てみましょう。需要に応じて物事を収集する必要があることを覚えておいてください。
最初に、全員がレイヤー0またはここのW0の重みを持っていることを確認する必要があります。all gather 0を行い、それが完了するのを待ちます。それが完了すると、W0でフォワードステップを行うことができます。x × W0を計算できるとしましょう。
この時点で、all gather 0が終了すると同時にall gather 1が開始されます。この行列乗算を行っているときに、基本的に必要な次のパラメータをすでにロードし始めています。もちろん、通信はより遅いです。そのため、多少のギャップがありますが、初期ロードよりもはるかに速く終了します。
今、forward 1が起こることができます。そして、バックグラウンドで、再び、パラメータ番号2のロードを開始しました。ここの黄色のスライスで、forward 1に関連するパラメータを解放しています。そして今、ここの他のことは、W0が2回使用される計算を繰り返していることです。これを再び通信する必要はありません。これは非常に迅速に起こります。
必要になる前にforward 2がすでにロードされています。ここにバブルはありません。そして、番号2を解放できます。それがフォワードパス全体です。ここでギャップが比較的小さいことがわかります。
計算が起こる必要がある前に多くのロードを行うことができました。実際に必要になる前に重みの要求をキューイングするというこの非常に巧妙なことを行うことで、通信に関連するオーバーヘッドの多くを避けることができます。
そして今、forward 2の時点で、フォワードパスが完了しました。重み番号2を解放できます。そして、バックワードパスを開始します。バックワードパス用のall gather 2はすでに完了していることがわかります。そのため、backward 2、backward 0を開始できます。重み0はすでに保存されています。それは完了です。
そして、reduce scatterとall gatherなどを行う必要があるため、バックワードパスでここで高いオーバーヘッドが発生します。うまくいけば、この図を見て、このクレイジーなシャーディングを行っているにもかかわらず驚くべきことだと言うでしょう。この図に戻ると、パラメータ、勾配、オプティマイザー状態を完全にシャーディングしましたが、必要な総帯域幅は2倍ではなく3倍のみです。それはそれほど悪くないようです。
そして、実際に見られるバブルはひどくありません。通信はほぼ完全に利用されており、計算は非常に長時間停止していません。実際に持っているリソースを非常に効率的に利用しており、それはクールです。
4. モデル並列化(Model Parallelism)
4.1 パイプライン並列化(Pipeline Parallelism)
パイプライン並列化は、おそらくニューラルネットワークを分割する最も明白な方法だと思います。深いニューラルネットワークにはレイヤーがあることを知っています。レイヤーがある場合、ネットワークを分割する非常に自然な場所は、レイヤー境界で分割することです。各GPUが何らかのレイヤーのサブセットを処理し、アクティベーションを渡します。この場合、各レイヤーはGPUに属します。GPUは、アクティベーションを一方から他方に渡します。バックワードの場合、GPU 3から0に向かってバックワード勾配を渡します。
それはクールです。それは素晴らしいです。この図の何が問題でしょうか。
ほとんどのGPUがほとんどの時間アイドル状態であることがわかるはずです。これは実際に非常にひどい利用率です。
以前に説明したこの種のナイーブな並列化を行う場合、各レイヤーがフォワードを持ち、単一の例があるとしましょう。これは、この図のような図になります。
この図の異なる行は異なるレイヤーであり、異なるGPUでもあります。ここのx軸は時間で、左から右に進んでいます。何が見えますか? 最初に一番左でmy最初のレイヤーを計算します。そして、アクティベーションが第2レイヤーに渡されます。GPU 2が目を覚まし、「よし、私の番だ」と言います。仕事をして、GPU 3に渡し、次にGPU 4に渡します。
そして今、バックワードパスが始まることができ、以下同様です。人々がバブルと呼ぶこの巨大なものが見えます。これは、絶対に何もしていない大きなオーバーヘッドです。GPUがn分の1の時間アクティブであることがわかります。ある意味で、これは最悪の並列化です。4つのGPUを追加しましたが、単一GPUのスループットを得ています。
できることの1つは、行うことについて少し巧妙になることです。「よし、パイプラインを持つつもりだ。レイヤーで分割するだけでなく、各GPUで処理される必要があるもののシーケンスを持つつもりだ」と言うことです。
今、マイクロバッチがあるとしましょう。各マシンが4つの例を処理します。行うことは、最初の例、最初のデータポイントを完了でき、完了するとすぐに、そのアクティベーションを第2のGPUに送ることができます。そして、第2のデータポイントで作業を開始することができます。
そして今、通信と計算をオーバーラップしました。第2のGPUは、第1のGPUが継続して作業している間に作業を開始できます。そして今、バッチサイズが大きいことによって、バブルのサイズを潜在的に縮小できます。
バッチサイズがリソースであると前に言った理由がうまくいけばわかるでしょう。有限のバッチサイズがあり、パイプライン並列がある場合、同じバッチサイズを使用してパイプラインバブルサイズを小さくすることができ、例えば、またはデータ並列を行うために使用することができます。単一のバッチサイズを取り、2つの異なる方法で分割できるさまざまな方法があります。
今、マイクロバッチサイズがバブル時間を制御できます。実際に、持っているオーバーヘッドと有用な計算の比率は、ステージ数マイナス1をマイクロバッチ数で割ったものです。大きなバッチサイズがある場合、パイプライン並列は潜在的に効率的である可能性があります。しかし、前に言ったように、バッチサイズは有限です。望む値まで単純にそれを上げることはできません。一般的に、パイプラインは本当にひどいようです。
なぜやるのでしょうか? なぜバブルのこのコストを負担して並列化するのでしょうか?
いくつかの理由があります。パイプラインは、データ並列と比較してメモリを節約するのに役立ちます。ZeRO 3もパラメータをシャーディングしますが、これはアクティベーションもシャーディングし、それは良いことです。
パイプラインは良い通信特性も持つことができます。アクティベーションにのみ依存します。ポイントツーポイントでもあります。持っているトポロジーと持っているものによっては、パイプラインが実際にネットワークの遅い部分に対して非常に有利である可能性があります。
パイプライン並列は、遅いネットワークリンクでよく使用されます。ノード間、または時にはさまざまなラック間、またはさまざまなデータセンター間で、パイプライン並列を行うかもしれません。実際、データセンター間ではありません。さまざまなラック間で、パイプライン並列を行うかもしれません。
最近Google の人々から言われた例の1つは、実際に、TPUの大きな利点の1つは、すべての接続がはるかに大きいため、パイプライン並列をあまり行う必要がないということです。彼らはこの大きなトロイダルメッシュを持っています。256 GPUでの制限がなく、そこで突然、パイプライン並列に切り替えたくなるような遅いネットワークリンクに向かうことがありません。これは、パイプライン並列について考え始める実世界の例です。
これはNVIDIAの論文からの例です。この論文については後でより詳細に話します。異なる種類の並列化のパフォーマンス特性を示すいくつかの本当に良い仕事をしました。バッチサイズ8では、パイプライン並列サイズ、デバイス数を増やすにつれて、GPUあたりの利用率が本当に落ち始めることがわかります。
一方、128の大きなバッチサイズがある場合、かなりサイズの大きなパイプライン並列に対してかなり良い利用率を得ることができます。バッチサイズは、バブルのサイズを隠すために本当に重要です。そうでなければ、問題があります。
高度なパイプライン戦略
もちろん、異なる種類のパイプライン戦略を行うことができます。バブルの標準パターンをスケジューリングする代わりに、物事をより細かい部分に分割することができます。異なるステージ、異なるサブレイヤーを異なるデバイスに割り当て、異なる部分で異なる計算を行っています。そして、パイプラインをより良くインターリーブすることができます。
これの高度なバージョンで、少し時間をかけて話したいもの、これは非常に巧妙です、ゼロバブルパイプライニング、またはdeepseaの用語では、彼らはdual pipeと呼んでいると思います。しかし、中核的な単一のトリックは同じです。
ここで、考えてみると、勾配を計算するためにバックワードパスを行っているとしましょう。これを2つの異なるコンポーネントに分割できます。最初の部分は、アクティベーションをバックプロパゲーションすることです。残差接続を下っていくとき、基本的にアクティベーションに関する導関数を計算する必要があります。そして、パラメータに到達すると、勾配そのものも計算したいと思います。
具体的な例を与えると、左下のこの図を見てみましょう。この図では、フォワードパスが見えます。これは単一のMLPです。重みで乗算します。非線形性を行います。そして、非線形性を出力するだけです。これは、MLPのナイーブな単一部分のようなものです。
今、バックワードを見てみましょう。損失に関する導関数があります。それが入ってきます。そして、それがMLPへの入力であるxをどのように変化させるかを計算できます。これは、ある意味で、ここのアクティベーションに関する導関数です。そして、これらを計算するとき、もちろん、重みを更新するために必要な勾配を計算するためにそれらを使用できます。
しかし、重要なことは、重みの勾配を計算するこの部分です。これはいつでも行うことができます。これの依存関係はありません。計算グラフの任意の部分にこの計算のスケジューリングを再配置できます。
できることは、シリアル依存部分に対して標準のパイプライン並列を行うことです。しかし、パラメータを更新するためだけにこれらの計算を行う必要があるときはいつでも、どこでもそれらを再スケジューリングできます。
重要なアイデアは、1F-1Bパイプラインと呼ばれる、バブルサイズを削減する素晴らしい最適化されたスケジュールから始めることです。そして、これを取ることができます。このBであるバックワード部分の計算と、重みの勾配を計算するために必要な計算であるWを分離できます。そして今、元々バブルがあったところで、重みの計算であるWを行うことができます。元々これらの白いアイドル利用コンポーネントがあった部分で、今度はこれらのWで埋めることができます。
実際のシリアル依存関係が何であるかを慎重に考えることで、今度は実際にGPUから良い利用率を得ている本当に素晴らしいものを持つことができます。
明確にするために、これは恐ろしく複雑です。実際にこの方法でパイプライン並列を実装したい場合、autodiffが実際にこれらのことをどのように計算しているかに介入する必要があります。物事がどこに行くかを追跡できるキューを持つ必要があります。
最近、フロンティアラボでLLMを訓練している人との会話で面白い逸話を聞きました。彼らは、「実際に、グループには私たちのInfraWorksでパイプライン並列がどのように機能するかを理解する2人がいます。1人が去りました。そのため、私たちの訓練インフラには単一の負荷を支える人がいます」と言いました。このような話があります。
パイプライン並列は、インフラ的に非常に複雑です。ここではシンプルに見えます。興味があるなら、実装してみることをお勧めします。かなり早く、かなり複雑になります。
4.2 テンソル並列化(Tensor Parallelism)
それは、他のモデル並列化に切り替える良いメモです。これははるかにシンプルだからです。そして、これは多くのフレームワークや、本当に大きなモデルを訓練している人々でさえも、非常にきれいに利用され、この種のモデル並列化に非常に重く、または主に依存しています。
他にどのような方法でモデルを分割できるでしょうか。考えてみると、私たちが行うことの大部分は行列乗算です。大きなモデルでは、計算の大部分は行列乗算です。パラメータの大部分は行列乗算または行列です。
では、何ができるでしょうか。大きな行列乗算を並列化できれば、それは非常に良いでしょう。テンソル並列は、大きな行列を取り、それを乗算できるサブ行列のセットに分割できるというアイデアです。
上部にこの行列乗算がある場合、XとX×A = Yがあります。代わりにできることは、半分に切り上げ、Xも半分に切り上げることです。サブ行列を計算し、それらを合計すれば、最終的に答えを得ることができます。
概念的に、パイプライン並列はレイヤーのような深さ次元に沿って切断しています。テンソル並列は、これが何であるかですが、行列乗算の幅次元に沿って切断しています。
サブ行列に分解してから部分和を行います。MLPでどのようなものかの例があります。各GPUが、大きなMLP行列乗算の異なるサブ行列を処理しています。そして、必要に応じてアクティベーションを同期するために集合通信を行います。
何をするつもりでしょうか。これはMLPです。上半分と下半分、2つの異なるパスがあります。これらは行列を分割しています。この操作Y = GeLU X × Aを行いたいと思います。行列Aを A1と A2に分割します。
右側では、dropout YBを計算し、結果をZとして返したいと思います。Bも切り上げます。大きなパラメータ行列の両方を、AとBの2つの部分に切り上げました。
フォワードパスでは、入力Xを取り、それを2回コピーするだけです。各GPUは同じ入力を取得します。それらは同じ種類の、すみません、同じ行次元を持っているので、A1とA2で操作するのは問題ありません。
XA1とXA2は、いくつかのアクティベーションY1とY2を与えます。それらはB1とB2に入ります。そして、それらを合計するためにall-reduceを行います。それは前に示した図とまったく同じです。コピーしてからall-reduceし、答えZを得ます。
バックワードパスでは、今度は勾配がバックワードステップで逆方向に来るので、それは実際に逆です。このGはアイデンティティになります。両側で導関数をコピーします。そして、ずっとバックワード操作を行います。Fに到達すると、これはall-reduceです。両方のパスから2つの導関数が入ってくるからです。そして、それらを合計して戻します。
FとGは同期化バリアです。フォワードパスでは、単一のall-reduceを行います。バックワードパスでは、計算グラフの2つの異なる場所で、単一のall-reduceを行います。
行列乗算がある場所であれば、行列乗算を切り上げて、異なるデバイス間で並列化できることがうまくいけばわかるでしょう。
想像できるように、これは実際にかなり費用がかかります。レイヤーごとに存在する同期化バリアがあります。フォワードバックワードパスで2回、アクティベーション残差アクティベーション分の物を通信する必要があります。
テンソル並列は、この非常にシンプルなアイデアですが、非常に高速なインターコネクトを必要とします。覚えておくべき非常にシンプルな経験則があります。テンソル並列は単一ノード内で適用されるということです。
単一のNVIDIA GPUのボックスは8つの異なるGPUと一緒に出荷され、それらは同じボックス内に住んでいます。講義の最初に示したように、非常に高速で接続されています。これら8つのGPUは非常に迅速に互いに通信できます。
そのため、これら8つのデバイス間で非常に帯域幅を消費するテンソル並列のようなものを使用することは理にかなっています。典型的に見るのは、テンソル並列が最大8つのGPUまで適用されることです。8つのGPUが同じマシンに住んでいる場合、最小のパフォーマンス低下を与えるからです。
これは、Hugging Faceの並列化チュートリアルからの例で、テンソル並列化の異なるレベルのスループット低下を示しています。テンソル並列化を行うにつれて、10%と12%のスループットへの打撃があることがわかります。8まででは、まあ、これは管理可能かもしれません。これは、より良く並列化できることに対して支払う代価の一種です。
しかし、16デバイスに行くと、この驚異的な42%のパフォーマンス低下を得ます。32に行くと、別の65%のスループット低下があります。うまくいけば、視覚的にここで、テンソル並列化で8で止めたいことがわかるでしょう。それが本当にスウィートスポットです。手に入れることができるハードウェアインターコネクトの種類のためです。
パイプライン並列との比較
パイプライン並列と比較して、物事は今どのように比較されるでしょうか。パイプライン並列と比較すると、以前にあったこのバブルの問題を実際に扱う必要がありません。バブルを削減するためにより大きなバッチサイズを消費する必要がなく、それは良いことです。
テンソル並列を適用するのに比較的、非常にとは言いませんが、比較的低い複雑さがあります。本当に知る必要があるのは、大きな行列乗算がどこにあるか、それらを分割して異なるデバイスに住まわせることができるかだけです。フォワードとバックワード操作は同じままです。デュアルパイプライン並列のようなゼロオーバーヘッドを実装することと比較すると、これを行う方がはるかに良い状況にいることになります。
短所は、はるかに大きな通信オーバーヘッドがあることです。パイプライン並列では、マイクロバッチあたり、バッチサイズ×シーケンス長×残差次元のポイントツーポイント通信があります。テンソル並列では、レイヤーあたりその8倍があり、all-reduce通信があります。行う必要がある潜在的に非常に大量の通信です。
経験則として、前に言ったように、テンソル並列は、低遅延、高帯域幅インターコネクトがあるときはいつでも使用されます。野生では、持っているマシンの種類によって、2から16のテンソル並列を見ることになります。
4.3 各手法の特徴と制約
パイプラインまたはテンソル並列について、シーケンス並列とアクティベーションシャーディングの3番目の種類に移る前に質問はありますか。
「それらを同時に使用できるか」という質問がありました。答えは、はい、両方を使用するということです。後で例に到達すると思います。しかし、大規模な実行で見る典型的なことは、テンソル並列を非常に頻繁に見ることです。パイプライン並列はその上によく使用されます。パイプラインを行うがテンソル並列を行わない例として知っているのは、私が知る限り、DeepSeek V3だけだと思います。
単一マシン内で、例えば5つの異なるマシンがある場合、最初の20%のパラメータが最初のマシンの範囲内にあり、テンソル並列を使用します。そして、それが第2のマシンにパイプライン並列で次のステップに移ります。
質問は、マシン内でテンソル並列を行い、マシン間でパイプライン並列を行うかということでした。はい、例えば、マシン内でテンソル並列を行い、マシン間でデータとパイプライン並列の組み合わせを行うようなことをするでしょう。経験則は後で示しますが、基本的に、モデルが収まらないためにパイプライン並列を行います。モデル全体を収めることができれば、データ並列プラステンソル並列、または多分データ並列だけを行います。
これは、分散データ並列の一種です。ZeRO は、ある意味で、人々が分散データ並列を効率的に行う方法です。異なるステージがあります。ステージ1は基本的に無料です。ナイーブデータ並列と同じ通信パターンを行っていますが、オプティマイザー状態をシャーディングできます。それは素晴らしいです。いつでもそれを行うべきです。
ZeRO ステージ2は、パラメータ数の2倍です。総帯域幅消費は同じですが、逆方向に進むときに勾配の段階的解放を行う必要があることに追加のオーバーヘッドがあります。
ZeRO ステージ3はより複雑です。プログラム通信コストの3倍を行いますが、それほど悪くありません。前に見た図では、いくらかのオーバーヘッドがありました。しかし、通信パターンを非常に巧妙にマスクすれば、実際には非常に良いです。人々は、ネットワーキングパターンで比較的遅いリンクに対してもデータ並列を使用します。
これも概念的に非常にシンプルです。ここでの利点の1つは、特にデータ並列はアーキテクチャをあまり気にしないことです。これらのいずれでも、実際にトランスフォーマーをどのように実装するかについて全く話しませんでした。すべて非常に抽象化されています。これが、例えば、FSDPが非常に人気がある理由の1つです。アーキテクチャが実際に何をしているかについて深い知識や深い内省を持たずに、任意のニューラルネットワークを並列化するラッパーを書くのが非常に簡単だからです。
ここにいくつかの例があります。いつもGPUでメモリが不足しているので、例を作りました。8倍A100 80ギガバイトノードで収めることができるモデルの最大サイズを見ることができます。
ベースラインでは、かろうじて60億パラメータモデルを収めることができるかもしれません。一方、ZeRO ステージ3を使用すると、約500億パラメータモデルのようなものを収めることができます。FSDPのような巧妙なメモリ節約を行って、より大きく、より大きなモデルを収める能力には大きな節約があります。
バッチサイズという重要なリソース
データ並列には重要なリソースがあります。これは実際に覚えておいてほしい重要なアイデアです。データ並列では、バッチサイズは実際に非常に重要なリソースです。各マシンで端数例を持つことはできないため、バッチサイズより大きく並列化することはできないという意味で。
これは、バッチサイズに制限がある場合、データ並列を使用することを停止するということを意味します。バッチサイズには収穫逓減があります。課題1でバッチサイズを変えて遊んだかもしれませんが、バッチサイズを特定のポイントを超えて上げると、最適化率にかなり急速な収穫逓減が見られ始めることを知っているでしょう。
これについて書かれた論文がたくさんあります。OpenAIには、クリティカルバッチサイズと呼ばれるものについて本当に良いものがあり、特定のポイントを過ぎると、各例が最適化能力に貢献する方法に非常に急速な収穫逓減があると基本的に主張しています。
基本的に、直感は、特定のポイント以下では、多くの勾配ノイズがあり、それを減らすことは非常に価値があるということです。しかし、特定のポイントでは、分散減少ではなく、取っている勾配ステップの数によって本当に基本的に制限されるということです。
これは基本的に、データ並列だけでは任意に大きな並列化に到達することはできないということを意味します。バッチサイズは本当に重要なリソースです。基本的に、固定された最大バッチサイズがあり、それを異なる方法で使うことができます。他の種類の並列化もより大きなバッチから恩恵を受けるため、後で話すように、バッチサイズを特定の部分で使います。
データ並列には問題が残ります。ZeRO ステージ1と2はメモリをスケールさせません。ZeRO ステージ3は原理的には良いですが、遅い可能性があります。そして、おそらくより重要なことは、前の質問に関連しますが、アクティベーションメモリを削減しません。理想的には、モデルを完全に切り上げて、完全に別々に住まわせたいと思います。そうすれば、アクティベーションメモリも削減されるからです。
そこで、これらの本当に大きなモデルをこれらのGPUに収めることができるように、モデルを分割するより良い方法が欲しいのです。
これがモデル並列化をもたらします。バッチサイズを変更せずにメモリでスケールアップしたいと思います。並列化するために大きなバッチサイズを持つ必要がない、または基本的に大きなバッチサイズを持つ必要がない代替軸が欲しいのです。
行うことは、パラメータをGPU間で分割することです。ある意味で、それはZeRO 3のようなものです。しかし、もうパラメータを通信するつもりはありません。アクティベーションを渡すつもりです。それは異なります。そして、時にはアクティベーションはパラメータよりもはるかに小さくなります。それは私たちにとって非常に良いでしょう。
5. アクティベーション並列化
5.1 シーケンス並列化(Sequence Parallelism)
今話したい最後の複雑さはアクティベーションメモリです。テンソルとパイプライン並列は基本的にほとんどのものを線形に削減できます。しかし、実際にはすべてのアクティベーションメモリ使用量を削減することはできません。
これは、アクティベーションメモリを削減する方法について話しているNVIDIAの論文からの例です。本当に興味深いと思うことの1つは、モデルをより大きく、より大きくするということです。左から右に行くにつれて、積極的に並列化すれば、パラメータとオプティマイザー状態メモリは同じままでいることができることがわかります。
しかし、アクティベーションメモリは、その一部がきれいに並列化されないため、成長し続けます。持っているデバイスの数に関係なく、実際にはデバイスあたりのアクティベーションメモリの成長を本当に取り除くことはできません。
一方、再計算のようなもう少し賢いことを行えば、アクティベーションメモリを低く保つことができます。そして、それは最大のモデルの一部を並列化するのに本当に重要です。
レイヤーあたりのアクティベーションメモリは何でしょうか。以前にこのトランスフォーマーの数学と微積分のいくつかを行ったことがあります。うまくいけば、これらすべてに慣れているでしょう。必要なレイヤーあたりのアクティベーションメモリの量を計算できます。ここに便利な公式があります。これは必要なメモリの量です。sbh × 34 + 5as/hです。
これらの数字の一部は不可解ですが、実際にはそれほど不可解ではありません。左の項と右の項があることがよくわかります。左の項は、MLPと他の点ごと操作から来ています。そこからsbh × 34が来ています。これらは、残差ストリームのサイズであるhのサイズに依存します。
右側では、これを掛け合わせると、hがキャンセルされるためa²bになる項があります。これは、softmax項と注意の他の二次項に必要なメモリです。
もちろん、flash attentionを使用し、再計算を使用すれば、その第2項を大幅に削減できることを知っています。
テンソル並列を行うとしましょう。できる場所でテンソル並列を行います。MLPで行い、kq計算で行い、注意計算で行います。このようなものになります。
これはかなり良く見えますが、完全にそこまではいっていません。デバイス数であるtで割ったレイヤーあたりのアクティベーションメモリです。8で割る場合、理想的には、すべてのアクティベーションメモリを8で割りたいでしょう。
しかし、削減されていないstragglerの項、sbh × 10があることがわかります。これらが何であるかを考えると、これらは非matmulコンポーネントです。LayerNorm、Dropout、注意とMLPへの入力、これらの項すべてが残念ながらサイズと共に成長し続け、非常にきれいに並列化されません。
考える必要がある最後のことは、これまで並列化していないこれらのシンプルな点ごと操作を取り、それらを分割することです。それらを分割する非常にシンプルな方法があります。layer normを行っている場合、これらのlayer normは、シーケンス内の異なる位置間で全く相互作用しないと言うことです。
他の何にも全く気にしません。行うことは、1024の長いシーケンスがあるとしましょう。それを切り上げ、各デバイスがそのlayer normの異なる部分、またはそのdropoutの異なる部分を処理します。これらの点ごと操作は、今やシーケンス次元全体で完全に分割できます。
今シーケンス次元で物事を切り上げているため、行った並列計算が再び集約されることを確実にするために、いくらかの同期化を行う必要があります。
フォワードパスでは、これらのguysはall gatherになり、g barsはreduce scatterになります。バックワードパスでは、2つが逆になります。ある意味で、2つの間に双対性があります。
layer normでは、ものを散らばらせました。そのため、標準計算を行えるように、それらを再び一緒に集める必要があります。そして今、dropoutに到達するときはいつでも、持っている並列コンポーネントに再びそれらを散らばらせたいと思います。バックワードパスでは、それを逆に行っています。
うまくいけば、それは明確です。これは非常にシンプルなアイデアです。以前に並列化に失敗した最後のコンポーネントを並列化しているだけです。
今、これらの異なる部分すべてを一緒にまとめて、最後に到達できます。上部で始まりました。これは全く並列化がありません。テンソル並列を行いました。これにより、点ごと操作ではないすべてをtで割ることができます。
このシーケンス並列化のアイデアを適用すれば、このコンポーネントをもう一度tで割ることができます。そして、第2項を除去するフラッシュアテンションのトリックであるアクティベーション再計算のようなことを行うことができます。
簡単に取得できる最小メモリは、底部にあるもので、sbh 34 / tです。これは、異なる公式を見ているときによく使用されます。どれだけのアクティベーションメモリを使用するかについてのトランスフォーマー算術では、sbh 34のようなものをよく見ます。そして、tテンソル並列でtで割ったものがある場合、これがその種のメモリに対して得られる簡単な最小値だからです。
5.2 アクティベーションメモリの最適化
メモリについて話してきて、メモリは実際に本当に大きなモデルを訓練しているため、並列化の非常に重要な部分です。訓練をしているとき、メモリ使用量を見ると、実際にアクティベーションがメモリ使用量の本当に大きな部分であることがわかります。
標準的なフォワードバックワードパスを見ると、これはPyTorchのチュートリアルの1つからのものだと思いますが、メモリ使用量が非常に動的であることがわかります。これは一般的に興味深いプロットだと思うので、これについて話してみます。
訓練をしているとき、常にパラメータを持っています。それは静的だからです。しかし、イテレーション0では、まだオプティマイザー状態を全く持っていません。実際に、メモリ使用のその部分を持っていません。
しかし、フォワードとバックワードを行うとき、すべてのアクティベーションを蓄積するにつれて、アクティベーションが成長、成長、成長、成長するのが見えます。バックワードパスを開始すると、アクティベーションを使い上げて解放するため、アクティベーションが下がり、勾配を蓄積しています。勾配メモリ使用量が上がります。
ピークは実際にバックワードパスの途中のどこかにあり、まだすべてのアクティベーションを解放していませんが、まだ勾配を構築しています。イテレーション2では、ここで同じことが見えます。
この図のポイントは、他のすべての部分について考えたということです。パラメータについて考え、オプティマイザー状態について考え、勾配について考えました。しかし、少なくとも非常に深く、アクティベーションについて考えていません。それでは、それを行いましょう。
テンソルとパイプライン並列は、基本的にほとんどのものを線形に削減できます。しかし、実際にはすべてのアクティベーションメモリ使用量を削減することはできません。
これは、アクティベーションメモリを削減する方法について話しているNVIDIAの論文の1つからの例です。本当に興味深いと思うことの1つは、モデルをより大きく、より大きくするということです。左から右に行くにつれて、積極的に並列化すれば、パラメータとオプティマイザー状態メモリは同じままでいることができることがわかります。
しかし、アクティベーションメモリは、その一部がきれいに並列化されないため、成長し続けます。持っているデバイスの数に関係なく、実際にはデバイスあたりのアクティベーションメモリの成長を本当に取り除くことはできません。
一方、再計算のようなもう少し賢いことを行えば、アクティベーションメモリを低く保つことができます。そして、それは最大のモデルの一部を並列化するのに本当に重要です。
5.3 メモリ使用量の計算式
レイヤーあたりのアクティベーションメモリは何でしょうか。以前にこのトランスフォーマーの数学と微積分のいくつかを行ったことがあります。うまくいけば、これらすべてに慣れているでしょう。必要なレイヤーあたりのアクティベーションメモリの量を計算できます。ここに便利な公式があります。これは必要なメモリの量です。sbh × 34 + 5as/hです。
これらの数字の一部は不可解ですが、実際にはそれほど不可解ではありません。左の項と右の項があることがよくわかります。左の項は、MLPと他の点ごと操作から来ています。そこからsbh × 34が来ています。これらは、残差ストリームのサイズであるhのサイズに依存します。
右側では、これを掛け合わせると、hがキャンセルされるためa²bになる項があります。これは、softmax項と注意の他の二次項に必要なメモリです。
もちろん、flash attentionを使用し、再計算を使用すれば、その第2項を大幅に削減できることを知っています。
テンソル並列を行うとしましょう。できる場所でテンソル並列を行います。MLPで行い、kq計算で行い、注意計算で行います。このようなものになります。
これはかなり良く見えますが、完全にそこまではいっていません。デバイス数であるtで割ったレイヤーあたりのアクティベーションメモリです。8で割る場合、理想的には、すべてのアクティベーションメモリを8で割りたいでしょう。
しかし、削減されていないstragglerの項、sbh × 10があることがわかります。これらが何であるかを考えると、これらは非matmulコンポーネントです。LayerNorm、Dropout、注意とMLPへの入力、これらの項すべてが残念ながらサイズと共に成長し続け、非常にきれいに並列化されません。
今、これらの異なる部分すべてを一緒にまとめて、最後に到達できます。上部で始まりました。これは全く並列化がありません。テンソル並列を行いました。これにより、点ごと操作ではないすべてをtで割ることができます。
このシーケンス並列化のアイデアを適用すれば、このコンポーネントをもう一度tで割ることができます。そして、第2項を除去するフラッシュアテンションのトリックであるアクティベーション再計算のようなことを行うことができます。
簡単に取得できる最小メモリは、底部にあるもので、sbh 34 / tです。これは、異なる公式を見ているときによく使用されます。どれだけのアクティベーションメモリを使用するかについてのトランスフォーマー算術では、sbh 34のようなものをよく見ます。そして、tテンソル並列でtで割ったものがある場合、これがその種のメモリに対して得られる簡単な最小値だからです。
6. 実践的な並列化戦略
6.1 3D/4D並列化の組み合わせ
時間と皆さんを疲れさせることを考慮して、話さない他のいくつかの並列化戦略があります。なぜなら、すでに皆さんを並列化の実行方法に関する多くの低レベルの詳細に引きずり込んだと思うからです。
最初に話したいのは、コンテキスト並列またはリングアテンションです。以前にリングアテンションという用語を聞いたことがあるかもしれません。これは、本質的に本当に大きなアテンションの計算とアクティベーションコストの両方を分割する方法で、本質的に、キーと値を異なるマシンの周りで渡すだけです。各マシンは異なるクエリに責任を持ち、その後、キーと値がKQVの内積を計算するためにマシン間でリング状に移動します。
ここでクールなことは、flash attentionのタイリングを行ったので、すでにこれを行う方法を知っているということです。アテンションがこの種のオンラインタイル・バイ・タイルの方法で計算できることを知っています。それがリングアテンションで起こっていることです。
テンソル並列を知っている今、かなり簡単である他のことは、エキスパート並列です。エキスパート並列は、ある意味でテンソル並列に似ていると考えることができます。1つの大きなMLPをより小さなエキスパートMLP、例えば、に分割し、それらを異なるマシンに散らばらせるという意味で。
エキスパート並列とのキーの違いは、エキスパートがスパースに活性化されることです。そのため、ルーティングについて少し考える必要があります。ルーティングは、テンソル並列で以前に持っていたような全対全通信ほど予測可能ではありません。今では、おそらく1つのエキスパートが過負荷になっているかもしれません。ネットワーキングがもう少し複雑になります。
しかし、そうでなければ、概念的にはエキスパート並列についてテンソル並列と同じ世界に住んでいます。
並列化戦略の比較表
話したすべてのことを要約するために、持っている異なる種類の戦略の小さなテーブルを作りました。DDP in ZeRO 1があります。これは、行うナイーブなデータ並列の一種です。ここでは、バッチあたりいくらかのオーバーヘッドがあり、メモリスケーリングはなく、合理的な帯域幅特性がありますが、これを行うためにバッチサイズを消費します。
大きなデータ並列を持つには、大きなバッチサイズが必要です。FSDPがあります。これは、ZeRO 1の一種のより良いバージョンです。メモリスケーリングを得ることができるという意味で、しかし、異なるレイヤー間でオーバーヘッドを支払うことになります。
そのため、今度は、より高い通信コストと、利用率の悪さにつながる可能性のある同期化バリアがあります。パイプライン並列は、このバッチごとの側面にもはや依存しておらず、線形メモリスケーリングを得ることができるという点で良いです。
しかし、別の問題があります。これもバッチサイズを消費し、セットアップと使用が恐ろしく複雑です。そのため、多くの人は、可能であればパイプライン並列を避けたがります。
そして最後に、テンソル並列は、帯域幅と行う必要がある同期化の量の観点から非常に高いコストです。しかし、これはバッチサイズに影響を与えないという本当に良い特性を持っています。これは、グローバルバッチサイズに関してコストがない、使用できる1種類の並列化戦略です。それは良いことです。
いくつかの限られたリソースのバランスを取る必要があります。メモリがあります。これは1つのリソースです。帯域幅と計算があります。これは別のリソースです。そして、バッチサイズがあります。これは一種の型破りなリソースですが、効率を改善するためにこれらの異なる側面に費やすことができる限られたものとして本当に考えるべきものです。
私が先週参照したGoogleの非常に良いTPU並列化、またはTPUブック、呼んでみましょう、があります。しかし、実際に、彼らは本当に良い並列化セクションを持っています。そして、最後の例に移る前に見せたかった素晴らしい図があります。
重要な量は、前に言っていたように、バッチサイズです。バッチサイズ対持っているGPU数の比率によって、異なる種類の並列化が最適になります。彼らは、これらの各モデルに対してどれだけの通信と計算を行うことになるかについて、特定の公式を使用しています。これは、このプロットを生成するための簡略化された公式です。
バッチサイズが小さすぎる場合、多くのGPUと本当に小さなバッチサイズがある場合、効率的になる方法がないことがわかります。常に通信制限されています。これがここの下半分です。実際に、ほとんどの時間を通信に費やしています。
より多くのバッチサイズを取得するにつれて、最終的にFSDP、つまりZeRO ステージ3とMP、この場合はテンソル並列の両方を混合すれば、実際に基本的に計算制限のある場所に到達できるポイントに到達できます。そのため、今では、通信を待っているフロップスを無駄にしていません。
そして最後に、バッチサイズが大きい場所に到達すれば、純粋なデータ並列で逃げることができます。純粋なFSDPは、行う計算に費やす時間が行う通信に費やす時間よりも高いレジームに到達させます。バッチサイズが十分に大きい場合、FSDPだけで逃げることができます。
これは、なぜこれらを混合するのか、いつこれらを混合するのか、なぜバッチサイズがリソースなのかというこのアイデアの素晴らしい例です。うまくいけば、これは、これが何であるかを非常に視覚的な方法で示しています。
これらすべてを一緒にまとめると、人々が3Dまたは4D並列化と呼ぶものになります。最近5D並列化という用語を聞いたと思います。5番目の次元が何であるかはよくわかりませんでした。それについて読む必要があります。
しかし、今度は、並列化の異なる次元をすべて一緒にまとめることができます。これは本当にシンプルな経験則です。元々は昨年調べてまとめましたが、今年もまだ同じであることがわかりました。これで従うことができます。
最初に行う必要があることは、モデルとアクティベーションをメモリに収める必要があります。それを行わなければ、単純に訓練できません。これは要件です。モデルがメモリに収まるまで、モデルを分割する必要があります。テンソル並列を行います。
そして、マシンあたりのGPU数まで、非常に効率的で、非常に高速であることを知っています。そのポイントまでテンソル並列を行います。
その後、パイプライン並列を扱う意欲や帯域幅制約などによって、モデルがメモリに収まるまで、マシン間でZeRO 3またはパイプライン並列を使用することになります。
そのポイントの後、GPUが不足するまで、全体を実行することができ、唯一の目標は、手元にある総フロップス数を増やすことです。残りの方法をデータ並列でスケールします。データ並列は低帯域幅通信チャネルでよく動作し、非常にシンプルだからです。そして、それは、すべてのGPUを使用する方法を与えます。
バッチサイズが本当に小さい場合、より良い通信効率のためにバッチサイズを取引する方法があります。リソースとしてすべてのバッチサイズを消費していない場合、できることは、デバイスで勾配累積を使用することです。それは、メモリ制約があっても、基本的に効果的により大きなバッチサイズを持てるようにします。
そして、それは、マシン間で同期化する頻度が少なくなるため、より良い通信効率のためにバッチサイズを取引することができます。
シンプルな経験則。これは、何をしていても、合理的な効率でモデルを訓練することができます。
6.2 バッチサイズリソースの管理
データ並列には重要なリソースがあります。これは実際に覚えておいてほしい重要なアイデアです。データ並列では、バッチサイズは実際に非常に重要なリソースです。各マシンで端数例を持つことはできないため、バッチサイズより大きく並列化することはできないという意味で。
これは、バッチサイズに制限がある場合、データ並列を使用することを停止するということを意味します。バッチサイズには収穫逓減があります。課題1でバッチサイズを変えて遊んだかもしれませんが、バッチサイズを特定のポイントを超えて上げると、最適化率にかなり急速な収穫逓減が見られ始めることを知っているでしょう。
これについて書かれた論文がたくさんあります。OpenAIには、クリティカルバッチサイズと呼ばれるものについて本当に良いものがあり、特定のポイントを過ぎると、各例が最適化能力に貢献する方法に非常に急速な収穫逓減があると基本的に主張しています。
基本的に、直感は、特定のポイント以下では、多くの勾配ノイズがあり、それを減らすことは非常に価値があるということです。しかし、特定のポイントでは、分散減少ではなく、取っている勾配ステップの数によって本当に基本的に制限されるということです。
これは基本的に、データ並列だけでは任意に大きな並列化に到達することはできないということを意味します。バッチサイズは本当に重要なリソースです。基本的に、固定された最大バッチサイズがあり、それを異なる方法で使うことができます。他の種類の並列化もより大きなバッチから恩恵を受けるため、後で話すように、バッチサイズを特定の部分で使います。
重要な量は、前に言っていたように、バッチサイズです。バッチサイズ対持っているGPU数の比率によって、異なる種類の並列化が最適になります。彼らは、これらの各モデルに対してどれだけの通信と計算を行うことになるかについて、特定の公式を使用しています。これは、このプロットを生成するための簡略化された公式です。
バッチサイズが小さすぎる場合、多くのGPUと本当に小さなバッチサイズがある場合、効率的になる方法がないことがわかります。常に通信制限されています。これがここの下半分です。実際に、ほとんどの時間を通信に費やしています。
より多くのバッチサイズを取得するにつれて、最終的にFSDP、つまりZeRO ステージ3とMP、この場合はテンソル並列の両方を混合すれば、実際に基本的に計算制限のある場所に到達できるポイントに到達できます。そのため、今では、通信を待っているフロップスを無駄にしていません。
そして最後に、バッチサイズが大きい場所に到達すれば、純粋なデータ並列で逃げることができます。純粋なFSDPは、行う計算に費やす時間が行う通信に費やす時間よりも高いレジームに到達させます。バッチサイズが十分に大きい場合、FSDPだけで逃げることができます。
バッチサイズが本当に小さい場合、より良い通信効率のためにバッチサイズを取引する方法があります。リソースとしてすべてのバッチサイズを消費していない場合、できることは、デバイスで勾配累積を使用することです。それは、メモリ制約があっても、基本的に効果的により大きなバッチサイズを持てるようにします。そして、それは、マシン間で同期化する頻度が少なくなるため、より良い通信効率のためにバッチサイズを取引することができます。
6.3 実装上の経験則
これらすべてを一緒にまとめると、人々が3Dまたは4D並列化と呼ぶものになります。最近5D並列化という用語を聞いたと思います。5番目の次元が何であるかはよくわかりませんでした。それについて読む必要があります。
しかし、今度は、並列化の異なる次元をすべて一緒にまとめることができます。これは本当にシンプルな経験則です。元々は昨年調べてまとめましたが、今年もまだ同じであることがわかりました。これで従うことができます。
最初に行う必要があることは、モデルとアクティベーションをメモリに収める必要があります。それを行わなければ、単純に訓練できません。これは要件です。モデルがメモリに収まるまで、モデルを分割する必要があります。テンソル並列を行います。
そして、マシンあたりのGPU数まで、非常に効率的で、非常に高速であることを知っています。そのポイントまでテンソル並列を行います。
その後、パイプライン並列を扱う意欲や帯域幅制約などによって、モデルがメモリに収まるまで、マシン間でZeRO 3またはパイプライン並列を使用することになります。
そのポイントの後、GPUが不足するまで、全体を実行することができ、唯一の目標は、手元にある総フロップス数を増やすことです。残りの方法をデータ並列でスケールします。データ並列は低帯域幅通信チャネルでよく動作し、非常にシンプルだからです。そして、それは、すべてのGPUを使用する方法を与えます。
バッチサイズが本当に小さい場合、より良い通信効率のためにバッチサイズを取引する方法があります。リソースとしてすべてのバッチサイズを消費していない場合、できることは、デバイスで勾配累積を使用することです。それは、メモリ制約があっても、基本的に効果的により大きなバッチサイズを持てるようにします。そして、それは、マシン間で同期化する頻度が少なくなるため、より良い通信効率のためにバッチサイズを取引することができます。
シンプルな経験則。これは、何をしていても、合理的な効率でモデルを訓練することができます。
7. 実例・ケーススタディ
7.1 Megatron-LMの実装例
これを具体的にするために、最後にいくつかの例について話します。2021年に戻るMegatron LLMからのこの本当に素晴らしい論文と、昨年のモデルのいくつかの両方を、まさにこれらのことを図で示し、多くのアブレーションも示して、駆け抜けます。
これは、17億パラメータから1兆パラメータまでのモデルを訓練した方法の大きなテーブルです。これらすべてで素晴らしい利用率を得ています。得られる理論ピークフロップスの割合が見え、40から52%の範囲です。それはかなり良いです。
テンソル並列が1で始まり、最終的に8まで上がり、8でキャップアウトすることがわかります。テンソル並列を最初に使用しています。そして、パイプライン並列は1にとどまります。しかし、モデルが十分に大きくなると、これらの大きなモデルを収めることができません。そのため、補償するためにパイプライン並列を増やす必要があります。
そして、データ並列サイズは基本的に可能な限り大きく始まり、その後、ゆっくりと下がっていきます。パイプライン並列の量を増やすにつれて、これはある意味でバッチサイズを消費しているからです。そのため、パイプライン並列に使用されている場合、実質的に大きなバッチサイズを持つことができません。
慎重な3D並列化は、集約フロップスで線形ゲインを与えます。慎重な3D並列化を行えば、GPUあたりの非常にフラットな全体達成フロップスを見ることができ、これはより多くのGPUを追加する場合、総集約スループットでの線形スケーリングを与えています。それは素晴らしいです。
テンソル並列8はしばしば最適です。これはパイプライン並列サイズとテンソル並列サイズです。8、8でバッチサイズ128が最適であることがわかります。より小さなバッチサイズを持っている場合でも、テンソル並列サイズ8が最適のままです。
アクティベーション再計算はより大きなバッチサイズを可能にします。より大きなバッチは、パイプライン並列のオーバーヘッドをマスクするのに役立つことを覚えておいてください。そのため、アクティベーション再計算は、より多くのフロップスであっても、自分自身に対して支払うことができます。flash attentionでこの話がすでに展開されるのを見ました。
7.2 最新言語モデルの並列化戦略
最後の部分は、最近の言語モデル、彼らが何をするかです。人々の並列化戦略がどのような例であるかを見るために、いくつかの論文を調べました。
OLMoとDolma論文では、70億パラメータモデルに対してFSDPを行います。
DeepSeekの最初の論文では、テンソルシーケンスとパイプライン並列でZeRO stage 1を行います。これは私が話した標準的なものです。V3は実際に少し異なることを行います。彼らは16-wayパイプライン並列、64-wayエキスパート並列(これは一種のテンソル並列)、そして彼らのデータ並列化戦略にZeRO stage 1を行います。
Yiは、別の中国のモデルですが、再びZeRO stage 1テンソルとパイプライン並列を行います。Yi-lightningは、MoEを行っているため、テンソル並列をエキスパート並列に置き換えます。
最後に、多くの詳細を持つ最先端の分散訓練に興味がある場合、Llama3のレポートは実際に読むのが本当に興味深いです。彼らはネットワーキングをどのように行うか、何が起こるかについて多くの詳細を持っています。
そして、前に言ったような種類のことが再び見えます。テンソル並列8が見え、CPまたはこれはコンテキスト並列です。これは長いコンテキスト訓練にのみ関連し、これは最後のステップです。それを無視できます。そして、これらの最初の2つのフェーズでパイプライン並列とデータ並列が起こっています。
安定させるために行った小さなバッチサイズ訓練であるため、最初のステージも無視できます。彼らの並列化戦略に対する彼らの根拠を見ると、基本的に、必要な帯域幅の量の順序で、TP、CP、パイプライン並列、DPを行いたいと前に言ったまさにそのことが見えます。
データ並列は、シャーディングされたモデル重みの非同期フェッチを行うことができるため、これらの長いネットワーク遅延を許容できます。そのため、彼らは最大のモデルのいくつかを訓練するために私が話した戦略を使用しています。
Llama3についての面白い副注、そして、友人との単なる噂ではなく、カジュアルな会話で聞いたことがあるかもしれませんが、このような巨大なスケールでモデルを訓練するとき、多くのGPU故障があります。
彼らは故障したGPUから148回の中断があり、持っていた総中断の約30%でした。計画外のマシンメンテナンスのようなものがありました。そして、それは32の異なること、訓練の32のインスタンスの中断でした。
このような大きなモデルを訓練しているとき、アルゴリズムについて話しましたが、これらの種類のことに対処できる耐故障性アーキテクチャも必要です。
さらに怖いことは、実際に明示的なモデル故障ではなく、実際にはデータ破損であると、人々が言っているさまざまな話も聞きました。GPUは静かに故障し、ガベージデータを与え、実行を完全に台無しにする可能性があります。
7.3 大規模訓練時の障害対応
Llama3についての面白い副注、そして、友人との単なる噂ではなく、カジュアルな会話で聞いたことがあるかもしれませんが、このような巨大なスケールでモデルを訓練するとき、多くのGPU故障があります。
彼らは故障したGPUから148回の中断があり、持っていた総中断の約30%でした。計画外のマシンメンテナンスのようなものがありました。そして、それは32の異なること、訓練の32のインスタンスの中断でした。
このような大きなモデルを訓練しているとき、アルゴリズムについて話しましたが、これらの種類のことに対処できる耐故障性アーキテクチャも必要です。
さらに怖いことは、実際に明示的なモデル故障ではなく、実際にはデータ破損であると、人々が言っているさまざまな話も聞きました。GPUは静かに故障し、ガベージデータを与え、実行を完全に台無しにする可能性があります。
最後の1つの例はGemma 2です。これで終わりたかったのは、これがTPUの例だからです。彼らは、大まかにFSDPであるZeRO 3を行います。そして、彼らはモデル並列とデータ並列を行います。ここで、前に言ったように、TPUは彼らにモデル並列を少し遠くまで伸ばすことを可能にします。
すべてをまとめると、特定のポイントを超えてスケーリングするには、マルチGPU、マルチノード並列化が必要になります。単一の解決策はありません。強みを活用するために3つのアプローチすべてを組み合わせたいと思います。そして、実際にこの並列化を実行する方法についてのシンプルで解釈可能な経験則があります。