大規模モデルの学習を阻むデータ移動のボトルネック：1e28 FLOPを超えるスケーリング

はじめに

過去5年間で、大規模言語モデル（LLM）の性能は劇的に向上した。その主な原動力は、より大きなモデルと学習データセットを扱うための学習計算量の予算の急速な拡大である。私たち自身の推定によれば、フロンティアAIモデルに使われる学習計算量は、2010年から2024年にかけて毎年4〜5倍に増えてきた。このスケーリングの速さはムーアの法則をはるかに上回っており、それを維持するには三つの次元でのスケーリングが必要だった。第一に学習をより長く続けること、第二に各学習に参加するGPUの数を増やすこと、第三により高性能なGPUを使うことである。

GPUクラスタをモデルの学習に使う期間を延ばすのは比較的簡単である。¹しかし実際には、学習が6か月を超えることはめったにない。これより長い期間になると、学習に使うハードウェアとソフトウェアの両方が時代遅れになるおそれがあり、公開した途端に古くなってしまうモデルを出したい研究所はないからだ。このことが、学習をどこまで長くできるかに実際上の上限を設けている。

もう一つの方法は学習クラスタの規模を拡大することで、学習計算量のスケーリングは主にこの方法で実現されてきた。たとえば、Vaswani et al. (2017) による最初の Transformer は8基のGPUで12時間学習されたが、Llama 3.1 405B は約16,000基のGPUで約2か月間学習された。

しかし、学習クラスタの規模の拡大は、学習期間の延長ほど簡単ではない。期間が限られた学習に有効に参加できるGPUの数には、データ移動から生じる根本的な制約がある。他の条件が同じなら、一つの学習に参加するGPUが多いほど、個々のGPUの内部でも、異なるGPUの間でも、1秒あたりに移動させなければならないデータが増える。

Epoch AI の研究者 Ege Erdil と共同研究者 David Schneider-Joseph による最近の論文では、現在の技術のもとで、データ移動のボトルネックのために学習をそれ以上スケーリングできなくなりそうなのはどの時点かを分析した。この論文の推定によれば、学習期間を最大3か月と仮定すると、$2e28\text{ FLOP}$ の規模を超える学習を効率よく行うことはできない。この学習計算量の規模を超えると、データ移動にかかる時間が演算にかかる時間を上回るようになり、ハードウェアを効率よく使えなくなるからである。

効率の低下を無視し、多額の費用をかけてクラスタをさらに拡大しようとしても、いずれはネットワークのレイテンシから生じる限界に突き当たる。このレイテンシが、クラスタ内の二つのデバイスが互いにどれだけ速く通信できるかの上限を決めるからだ。インターコネクトの根本的な技術改良か、確率的勾配降下法の変種からより安価なオプティマイザへの切り替えがなければ、これらの限界は克服できない。

この記事では、分散学習中のデータ移動のパターンを調べ、大規模言語モデルのスケーリングの限界を分析する。GPU内とGPU間の両方のデータ移動を検討し、さまざまな並列化手法を論じ、より大規模な学習に向かうにつれて生じる制約を調べる。読み終えるころには、利用効率の低下やレイテンシの壁など、今後のスケーリングの取り組みが直面する課題と、これらの限界を克服するために考えられる方法を理解できるだろう。

分散学習中のデータ移動のパターン

データ移動が分散学習をどう制約するかを理解するには、まず分散学習に伴う主な種類のデータ移動を押さえておく必要がある。大まかに言えば、これらはGPU内のデータ移動とGPU間のデータ移動の二つに分けられる。順に見ていこう。

GPU内のデータ移動

現在の分散学習では、GPU内のデータ移動のコストのほぼすべてを、広帯域メモリ（HBM、VRAMと呼ばれることもある）が占めている。要するに、二つの大きな行列を掛け合わせたいなら、両方の行列をHBMから読み出し、結果をHBMに書き戻す必要があり、メモリ帯域幅は無限ではないので、これには一定の時間がかかる。

この演算に関わる行列が大きい場合（たとえば要素数が合計1,000万を超える場合）、GPUのコアで行われる演算のほうがメモリの読み出しよりかなり長くかかるので、GPU内のデータ移動のコストは無視できる。うまくカーネルを書けば、メモリの読み出しの大半を進行中の演算と重ね合わせ、演算の陰にメモリの読み出しを「隠す」ことさえできる。しかし、行列のサイズを小さくしていくと、演算時間はメモリの読み書きの時間よりも速く減っていき²、ある規模を下回るとメモリの読み書きがボトルネックになる。この規模では読み書きを隠すことができなくなり、演算の利用率が下がり始める。

これが分散学習に制約を課す理由は次のとおりである。Hoffmann et al. (2022) のChinchillaのスケーリング則によれば、計算量最適な学習では、モデルの学習計算量をパラメータ数の2乗に比例して増やす必要がある。しかし、一つのバッチについて各 Transformer ブロックに必要な演算量は、バッチサイズをモデルサイズに合わせてどう拡大できると仮定するかにもよるが、一般にパラメータ数にほぼ比例して増える。このスケーリングの食い違いのために、より大規模な学習では行列積をより小さな断片に切り分けなければならない。そしてこれらの断片が、行列積の三つの次元のいずれかに沿って十分に小さくなると、メモリ帯域幅の制約が効いてくるおそれがある。

GPU間のデータ移動

GPU間のデータ移動は、異なるGPUが計算結果を共有し、全員が同期して同じ学習に取り組めるようにするために必要である。学習を並列化する主な方法は四つあり、それぞれが参加するGPUの間に固有の通信パターンを必要とする。以下に挙げる。

• データ並列：異なるGPUが、あるバッチの中の異なる入力系列や異なるトークンを処理する。各GPUは他のGPUと独立に順伝播と逆伝播を実行できるが、重みの更新の前に勾配の情報を互いに共有する必要がある。そのため、各勾配ステップの前に1回の通信が必要になる。

• パイプライン並列：異なるGPUが、Transformer の連続する異なるブロック（層）を担当する。入力がモデルのすべてのブロックを通り抜けるまで、あるGPUの出力が次のGPUへの入力として順に渡される。そのため、パイプラインの各段の境界で通信が必要になる。加えて、計算が逐次的であるため、一つのバッチでこれを素朴に行うと、現在動いているGPU以外のすべてのGPUが遊んでしまう。これをパイプラインバブルと呼ぶ。パイプラインバブルを避けるには、バッチをマイクロバッチに切り分け、それを順にパイプラインに流し込む必要がある。

• テンソル並列：異なるGPUが、モデルの重み行列の異なる部分による行列積を担当する。モデルを横に切り分けるパイプライン並列とは違い、テンソル並列はモデルを縦に切り分ける。テンソル並列の性質上、1回か2回の行列積ごとに（正確な実装による）、GPUは互いに結果を同期させる必要がある。つまりテンソル並列は、あらゆる並列化手法の中でGPU間の通信が最も頻繁に必要になる。そのため現在は、NVIDIA の NVLink のような広帯域・低レイテンシのインターコネクト上でしか行われていない。

• エキスパート並列：この種の並列化は、混合エキスパート（MoE）アーキテクチャを持つモデルにしか当てはまらない。MoEでは、各ブロックが多数の異なる「エキスパート」で構成され、異なるエキスパートが異なるGPUに割り当てられる。安価なルーティングの仕組みが、各ブロックで各入力をどのエキスパートに送るかを決める。活性化した各エキスパートが計算を終えると、その結果がまとめられてから次の層に送られる。そのため、エキスパート並列では、エキスパートへのルーティングと、異なるエキスパートの結果をまとめるために、各ブロックで通信が必要になる。

クラスタの規模を固定してGPU間のデータ移動を最小にするには、複数の並列化手法を同時に使う必要がある。各手法のデータ移動のコストはその並列度に比例して（テンソル並列の場合はそれ以下の割合で）増える一方、クラスタの規模はすべての並列度の積に比例して増えるからだ。この条件のもとでは、すべての並列化手法を一緒に拡大することで、ある学習に必要なデータ移動の総量が最小になる。こうした方針が実際にどうなるかの例として、Falcon-180B の学習では、$64\cdot 8\cdot 8=4096\text{ GPUs}$ からなるクラスタで、64ウェイのデータ並列、8ウェイのパイプライン並列、8ウェイのテンソル並列が使われた。

スケーリングの限界

以上の要因をすべて詳細なモデルで考慮すると、現在の手法とハードウェアでは、3か月の期間で $2e28\text{ FLOP}$（計算量最適な密なモデルでは約15兆パラメータ）の規模を超える学習を良好な利用率のまま行えないことがわかる。さらに、$2e31\text{ FLOP}$ の規模を超える学習は、少なくとも確率的勾配降下法の論理的な整合性を損なわない限り、レイテンシの制約のために不可能である。

分散学習を制約する要因のうち、定量的に最も理解しやすいのはレイテンシの壁である。これは、GPUをいくつ使うかにかかわらず、決まった期間内にどの規模の学習まで完了できるかの上限を定める。レイテンシの壁が生じる基本的な理由は、モデルが大きいほど一般に学習中に多くの勾配ステップが必要になることにある。それを決まった期間内に行わなければならないなら、学習するモデルを大きくするにつれて、各勾配ステップにかけられる時間は短くならざるをえない。これが十分に進むと、いずれは勾配ステップあたりの時間が短くなりすぎ、レイテンシの制約のためにその規模のモデルは学習できなくなる。

レイテンシの壁を理解する

定量的に考えよう。モデルのパラメータ数を $N$ とし、Chinchillaのスケーリング則に従って $D=20N$ トークンで学習するとする。バッチサイズを $B$ トークンとし、1エポックを仮定すると、学習中に $n_{\text{steps}}=\frac{D}{B}=\frac{20N}{B}$ 回の勾配ステップを踏むことになる。学習にかかる時間を $t_{\text{train}}$ とすると、各勾配ステップにかけられる時間は最大で $\frac{t_{\text{train}}\cdot B}{20\cdot N}$ となる。

一方、学習のテンソル並列度が十分に高ければ、行列積のたびにGPU間の通信が必要になると予想され、ネットワークのレイテンシにもよるが、その一回一回に少なくともある時間 $t_{\text{latency}}$ がかかる。³たとえば Transformer のブロックあたり2回の行列積があり、Transformer が $n_{\text{layer}}$ 個のブロックからなり、各勾配ステップで1回の順伝播と1回の逆伝播を行う必要があるとすると、レイテンシの制約のもとで学習が実行可能であるための条件は次のようになる。

$$4\cdot n_{\text{layer}}\cdot t_{\text{latency}}\le \frac{t_{\text{train}}\cdot B}{20\cdot N}$$$$N\le \frac{t_{\text{train}}\cdot B}{80\cdot n_{\text{layer}}\cdot t_{\text{latency}}}$$

妥当な値として、$t_{\text{latency}}=9$ マイクロ秒、$n_{\text{layer}}=120$、$B=4$ 百万（＝400万）トークン（GPT-3 の学習と Falcon-180B の学習の両方で使われた値）、$t_{\text{train}}=3$ か月を代入すると、$N=400$ 兆パラメータとなる。Chinchilla最適な学習計算量に換算すると、約 $2e31\text{ FLOP}$ である。

実際には、層の数 $n_{\text{layer}}$ とバッチサイズ $B$ は一定ではなく、モデルを大きくするにつれて増えていく。しかし、この二つの量のスケーリングがおおむね打ち消し合う限り（私たちはそうなっている可能性があると考えている）、上の計算は依然として成り立つ。二つの効果が打ち消し合わない、より複雑な場合については論文で計算しているが、ここでは簡潔さを優先して省く。

現在の学習ではレイテンシは無視できる要因だが、上で論じた制約の中で、レイテンシの壁は将来克服するのが最も難しい制約になると私たちは予想している。理由は二つある。

• メモリ帯域幅やネットワーク帯域幅の限界は、学習中の利用率が悪くなるのと引き換えに超えることができる。それとは違い、レイテンシの壁は（学習期間を固定すると）モデルの学習計算量に、超えることのできない上限を課す。

• 関連するあらゆるハードウェア性能の指標（HBM帯域幅、L2キャッシュ帯域幅、ネットワーク帯域幅など）のうち、2017年の V100 の発売以来最も改善が遅いのがレイテンシである。これはデバイス内（カーネルの起動・終了のレイテンシなど）でもデバイス間（NVLink や InfiniBand のレイテンシなど）でも当てはまる。

限界を克服する方法として考えられるもの

前節のモデルは学習の規模に厳しい上限を設けているように見えるが、これらの限界を克服できる可能性のある方法は、予見できる範囲でもいくつかある。単純な方法の一つは、ハードウェアの改良でレイテンシを下げることだ。しかし先に論じたように、過去のGPUの世代におけるレイテンシの改善の傾向は期待外れだった。

もう一つの方法は、ハードウェアはそのままに、学習のやり方を変えることである。レイテンシの壁から得られるパラメータ数の上限を思い出そう。

$$N\le \frac{t_{\text{train}}\cdot B}{80\cdot n_{\text{layer}}\cdot t_{\text{latency}}}$$

この上限はChinchillaのスケーリングを仮定して導いたので、レイテンシの壁を越えるためのすぐに思いつく解決策の一つは、モデルを学習不足にすることである。つまり、$N$ パラメータのモデルを $20\cdot N$ 未満のトークンで学習する。学習に充てる決まった期間内に行う順伝播と逆伝播の回数が減るので、帯域幅とレイテンシの両方の制約が緩和され、これはうまくいくだろう。しかし、Chinchillaのスケーリング則によれば、そうした方針をとると、学習計算量をモデルの能力に変換する効率が悪くなる。そのため、これがレイテンシの壁を克服するもっともらしいシナリオだとは考えていない。

学習過程の論理的な整合性を損なわずにこの上限を引き上げたい場合、そしてモデルを学習不足にしたくない場合は、分子の変数を増やすか、分母の変数を減らすしかない。$t_{\text{latency}}$ はハードウェアで決まり、$t_{\text{train}}$ は一定と仮定しているので、使える手段はバッチサイズを増やすこととモデルの深さを減らすことの二つだけである。

バッチサイズを増やす

バッチサイズをモデルサイズとともに積極的に拡大する方法が見つかれば、レイテンシの壁を押し広げられる。またこれによって、データ並列とパイプライン並列を行う余地も大きくなり、メモリ帯域幅やネットワーク帯域幅の問題が深刻な制約になる地点も先へ押しやられる。この方法の問題は、実際には、臨界バッチサイズと呼ばれるある点を超えてバッチサイズを拡大すると、収穫逓減に陥ることである。

McCandlish et al. (2018) は、臨界バッチサイズがモデルの性能とともに大きくなることを指摘している。大きなモデルほど学習の後半で性能が高くなるので、より大きな言語モデルほど、より大きなバッチサイズに対応できるはずだと考えられる。これは有望に思える。彼らのモデルによれば、臨界バッチサイズはモデルの削減可能な損失の逆数に比例して大きくなるはずだという。これに基づき、Hoffmann et al. (2022) のスケーリング則も用いて、私たちのモデルの主要な結果ではすでに $B\sim N^{\frac{1}{3}}$ のスケーリングを認めている。レイテンシの壁を $\sim 7e31\text{ FLOP}$ より先に押しやるには、これよりさらに積極的なスケーリングが必要である。

私たちの研究が発表されるわずか数日前に発表された最近の論文 Zhang et al. (2024) は、$B=17.75\cdot D^{0.47}$（単位はトークン）というスケーリングを見いだしている。このより積極的なスケーリングを私たちのモデルに厳密に組み込むと、利用率の低下は2桁先に押しやられ、$3e30\text{ FLOP}$ 前後から始まるようになる（$2e28\text{ FLOP}$ からではなく）。もちろん、別のオプティマイザを使うなど、Zhang et al. (2024) が検討していない方法を使えば、さらに積極的なスケーリングが可能かもしれない。

全体として、もしレイテンシの壁がアルゴリズムの革新によって克服されるとすれば、その立役者として最もありそうなのは、学習のバッチサイズを私たちの仮定よりも積極的に拡大する何らかの方法だと考えている。

モデルの深さを減らす

バッチサイズをより積極的に拡大する代わりに、モデルの深さの拡大をより控えめにする方法もある。モデルの深さと幅を別々に考慮したスケーリング則についての文献は驚くほど少ない。そのため、この分野の現在のベストプラクティスを推し量るには、Hoffmann et al. (2022) のように、さまざまな規模で多数のモデルを学習している論文を参照するのが精一杯である。

Hoffmann et al. (2022) で学習されたモデルから推定したスケーリング則では、$n_{\text{layer}}\sim N^{0.27}$ 程度となる。これは私たちのモデルで臨界バッチサイズについて仮定している $N^{\frac{1}{3}}$ のスケーリングに近く、二つの効果はおおむね打ち消し合う。しかし、Hoffmann et al. (2022) はこのスケーリング方針が最適かどうかを確かめる試験を一切行っておらず、最適な方針から外れた場合のコストを定量化するのに使えるスケーリング則も示していない。そのため、これは今後モデルの深さがどう拡大されると予想できるかについての、せいぜい弱い証拠としか見なせない。

バッチサイズを予想どおりのペースで拡大し続けながら、モデルの深さをほぼ同じに保つことが可能なら、レイテンシの壁は学習計算量で何桁も先へ押しやることができる。このシナリオは、単にバッチサイズをより積極的に拡大するシナリオよりも可能性が低いと考えている。

結論

私たちのモデルからは、利用率を大きく下げずにモデルの規模とクラスタの規模を拡大できる線形スケーリングの領域が、これまでどおりのシナリオでは約 $2e28\text{ FLOP}$ の規模で終わると考えられる。これまでに知られている最大のモデルが使った計算量は約 $5e25\text{ FLOP}$ と推定されており、計算量の過去のスケーリングのペースが年約4〜5倍であることを考えると、約3年で $2e28\text{ FLOP}$ のしきい値に達することになる。

ハードウェアの改良でこのしきい値を押しやるのは、たとえ1桁だけであっても難しそうである。メモリ帯域幅⁴、ネットワーク帯域幅、ネットワークのレイテンシのすべてで、大幅な改善が同時に必要になるからだ。それに対して、このしきい値を克服するより有望な方法は、学習中のバッチサイズをより積極的に拡大することだと思われる。それがどこまで可能かは臨界バッチサイズによって制限される。大規模言語モデルの臨界バッチサイズがモデルサイズにどう依存するかは今のところ私たちにはわからないが、主要な研究所はすでに未公表の臨界バッチサイズのスケーリング則を持っているのではないかと思う。

臨界バッチサイズのスケーリング則、モデルの深さとアスペクト比のスケーリング則、そして（MoEモデルなどの）スパース性の係数のスケーリング則は、いずれも分散学習の将来を理解するうえで重要だが、どれについても確かな研究はほとんど公表されていない。これらの問いについてより信頼できる研究があれば、分散学習の将来についての私たちの理解は大きく深まるだろう。

データ移動のボトルネックと私たちの研究について、詳しくは論文を参照してほしい。