分散学習のインタラクティブ・シミュレーターを公開

先日、深層学習モデルの分散学習で生じるデータ移動のボトルネックについての調査結果を発表し、GPUクラスターで並列化の方式ごとにかかる帯域幅とレイテンシのコストを詳しくモデル化した。本日、同じモデルを実装し、使いやすいインターフェースを備えたインタラクティブ・シミュレーターも公開する。

この記事では、シミュレーションツールを使って次の問いに答えながら、その機能を紹介する。2012年に AlexNet の学習に使われた GTX 580 GPU を使って、どれほど大規模なニューラルネットワークの学習を行えただろうか？言い換えれば、AlexNet を GTX 580 3GB GPU 2基で6日間学習させるのではなく、コストを度外視して GTX 580 3GB GPU だけで可能な限り大きなモデルを学習させようとしたら、AI研究所は学習をどこまでスケールさせられただろうか。

ツールの使い方

この問いに答える具体的な実験に入る前に、ツールの使い始め方を見ておこう。起動すると、次のインターフェースが表示される。

学習FLOPとモデルFLOP利用率（MFU）の関係を示すグラフが、このモデルの主な出力である。特定の学習計算量の規模で分散学習をできる限り効率よく行ったら、どれだけの利用率が得られるかを示す。ほかのタブには追加の情報がある。学習計算量を増やすにつれて並列化の各次元をどうスケールさせるのが最適か。また、特定のGPUで持続的に達成できる最大利用率の、たとえば80%を利用率が下回らないようにしたい場合、学習に費やせる学習計算量の上限はどれだけか、である。

グラフの右上の角の上にはダウンロードボタンがあり、ユーザーインターフェースに表示している要約情報だけでなく、モデルの出力の全体を入手できる。出力はJSONファイルの形式なので、標準的なライブラリで簡単に解析できるはずである。

設定タブでは、論文のさまざまなグラフに対応するプリセットを切り替えられる。各プリセットは構成のリストであり、各構成は次の要素からなる。

• 学習データセットのサイズ、臨界バッチサイズ、モデルの深さ、モデルのスパース性係数といった重要な量についてのスケーリング則。

• その構成での学習に使うGPUの仕様。

• 並列化と通信についての仮定：パイプライン並列化にどのスケジューリング方式を使うか、どの種類のネットワーク通信を計算と重ねられるか、など。

GPU間の通信は階層的にモデル化している。GPUのグループは段階的に大きくなり、大きいグループほど遅い相互接続で互いにつながっている。たとえば、1台のマシン（ノード）には最大8基の H100 GPU を載せられ、どのGPUも NVLink を介して900 GB/sで互いに通信できる。しかし、異なるマシン間の通信に使えるGPUあたりの帯域幅は、InfiniBand の帯域幅である100 GB/sまで下がる。

全体として、このツールにはネットワーク構成について次の情報が必要である。

1. ネットワークには階層がいくつあるか。ノード内とノード間の通信をモデル化する場合（既定ではこうしている）は2になり、データセンター間の通信もモデル化する場合は3に増える。

2. 各階層で、その階層の通信に使えるGPUあたりの帯域幅は、読み出しと書き込みの両方を数えてどれだけか。1秒間に1基のGPUが別のGPUに1バイトを送る場合、これを2バイト/秒の通信と数える。NVIDIA の NVLink の帯域幅の表記に合わせるため、この慣例を採用している。

3. 各階層で、その階層での1対1の通信のレイテンシはどれだけか。

ユーザーは「add new configuration」ボタンを押し、すべての設定を既定値から変えて、独自の構成を作ることもできる。右下の「run simulation」ボタンを押すと、ユーザーの新しい構成も左側に表示されるグラフに含まれる。個々のシミュレーションの実行ごとに、さまざまな構成を含めたり外したりすることもできる。

2012年に実行可能だった最大の学習はどれほどだったか？

2012年、Krizhevsky、Sutskever、Hinton は、NVIDIA GTX 580 GPU 2基で学習させた深層畳み込みニューラルネットワークを使い、ImageNet で画期的な結果を出した。このモデルはのちに AlexNet と呼ばれるようになった。120万枚の画像での学習に5〜6日かかり、それまでの手法より大幅に良い性能を示した。著者たちは、より速いGPUとより大きなデータセットを待つだけで結果を改善できるだろうと述べていた。

ここから興味深い反実仮想の問いが生まれる。2012年の AlexNet のチームが、はるかに大きな計算資源の予算を使えたとしたらどうだったか。GTX 580 を2基使うだけでなく、当時手に入ったハードウェアで可能な限り大規模な学習を組み立てようとしていたらどうなったか。このシナリオは、分散学習ツールで分析できる。GTX 580 のハードウェアでニューラルネットワークの学習をスケールさせる場合の理論上の限界をモデル化すればよい。

この独自構成の機能を使い、2012年に NVIDIA の最先端ハードウェアだった GTX 580 GPU で実験を行う。関連する仕様を簡単にまとめると次のとおりである。

表1：さまざまな情報源による GTX 580 3GB GPU の仕様。ほとんどはGPUのデータシートから直接とったものだが、いくつか（ネットワークや行列積のレイテンシなど）は、実際の値がどうだったかについての私たちの推測である。こうした推測には「推測」の印を付けた。

項目	仕様
ビット幅（32）	32
FLOP/s（32）	1.581e+12 (1.581 TFLOPS)
スレッドあたり・クロックあたりのFLOP（32）1	1/2
ビット幅（64）	64
FLOP/s（64）	197.6e+9 (197.6 GFLOPS)
スレッドあたり・クロックあたりのFLOP（64）	1/16
処理ブロックあたりのレジスタのバイト数	128 KiB
ストリーミング・マルチプロセッサの数	16
L2の帯域幅	384.8e+9 Bps
L2のバイト数	786,432バイト
グローバルメモリの帯域幅	192.4e+9 Bps
最大クロック周波数	1.544e+9 Hz (1.544 GHz)
実効利用率	85%（推測）
高帯域幅メモリ	3.0e+9バイト（3 GB）
行列積1回あたりのレイテンシ	2.25e-5秒（推測）
階層の数	2
階層ごとのネットワーク帯域幅	階層1：8e+9 Bps 階層2：1e+9 Bps
階層ごとのネットワークのレイテンシ	階層1：5e-6秒（推測） 階層2：5e-6秒（推測）
階層の大きさ	階層1：8

GTX 580 3GB GPU の仕様を私たちのツールに入力すると、AlexNet の学習に使われたGPUを使って、ピーク利用率の80%超を保ちながら3か月で組み立てられた最大の学習は1e26 FLOPだとわかる。言い換えれば、2012年の技術でも、GPT-4 ほどの大きさのモデルを3か月で学習させるのにすでに十分だったのだ。ただし、それには合計1,600万基のGPUが必要だっただろう。GPU1基あたり約300ドルとすると、GPUだけで総コストは約50億ドルになっただろう。²

この1e26 FLOPの学習をどう並列化するかも知りたいところだろう。ツールではこの問いにも答えられる。フィードフォワードブロックの最適な戦略は、1024方向のデータ並列化、32方向のパイプライン並列化、512方向のテンソル並列化である。今日のフロンティアモデルの学習では、8基のGPUからなる単一ノードの中でだけテンソル並列化を使うのが一般的な慣行であり、私たちのツールもこれが最も効率的な構成であることを裏づけている。

このインタラクティブなツールは、歴史上のシナリオに限らず、フロンティアの機械学習モデルの大規模な学習についての幅広い問いに答えるのにも役立つ。たとえば、今日の最先端のハードウェアで実行可能な最大の学習を探ったり、ネットワークのレイテンシと帯域幅が将来改善されればこうした限界がどう広がるかを調べたりできる。ツールの中でハードウェアの仕様や通信の仮定を調整すれば、さまざまな要因が学習の効率とスケーラビリティにどう影響するかをシミュレートできる。

最後に、ツールの裏側の仕組みに興味がある方は、その基になっているモデルをこのリポジトリで見られる。