分散学習とは　1つのモデルを数万基のGPUに分ける、メタのLlama 3は4種類の分け方を併用

大きな言語モデルの学習には、膨大な計算が要ります。1基のGPUでは現実的な期間に終わりません。モデルの重みも1基のメモリには収まりません。そこで、計算とデータとモデルを多数のGPUに分け、GPU同士が途中の結果をやりとりしながら学習を進めます。これが分散学習です。

4つの分け方

メタが2024年7月に公表したLlama 3の論文は、4種類の分け方を組み合わせた「4D並列」を使ったと説明しています。

分散学習の4つの分け方（メタのLlama 3の論文から）

分け方	何を分けるか	通信の量
テンソル並列（TP）	1つの層の重みの行列を分割	最も多い。同じサーバーの中に収める
コンテキスト並列（CP）	長い入力の文章を区切って分担	多い
パイプライン並列（PP）	層の並びを前後の区間に分割	中くらい
データ並列（DP）	学習のデータを分け、同じモデルの写しで計算	遅れに強い。外側に置く

通信の多い分け方ほど、速い道でつながったGPUの近くに置きます。メタは、テンソル並列を1台のサーバーの8基（NVLinkでつながる）に収め、データ並列を一番外側に置きました。データ並列は、重みの取り寄せを前倒しするなどで、通信の遅れを吸収しやすいからです。

Llama 3 405Bの構成

論文によると、学習には最大1万6千基のH100（1基700ワット、HBM3 80GB）を使いました。1台のサーバーにGPU8基とCPU2個を積みます。8,192基の段階では、テンソル並列8、パイプライン並列16、データ並列64の組み合わせで、GPU1基あたり毎秒430兆回の計算を出しました。GPUが理論上出せる計算能力に対する実際の比率（MFU）は43%です。1万6,384基に増やした段階では41%、長い文章を学習した段階では38%でした。

ネットワークは、イーサネットの上でRDMAを使うRoCEの網です。GPU同士は毎秒400ギガビットでつながります。1つのラックに16基、192のラックで3,072基のポッドを作り、ポッドの中ではどの組み合わせでも帯域が落ちない設計です。8つのポッドを束ねて2万4千基の塊にしています。ポッドの間の帯域は7分の1に絞っているため、通信の多い仕事はポッドの中に収めるよう、仕事の割り振りを工夫しました。スイッチには、米アリスタの7800などを使っています。

規模が生む問題

分散学習は、1基の故障で全体が止まりやすい弱さを持ちます。メタは54日間で、計画外の中断が419回あったと報告しています。数万基のGPUが一斉に計算を止めたり再開したりすると、データセンターの電力の使用量が数十メガワットの単位で急に変わり、電力の網に負担をかけることも記しています。昼の気温の上昇で、処理量が1日の中で1〜2%上下したそうです。

関連する用語

• データ並列・モデル並列・パイプライン並列：分け方のそれぞれの呼び名。
• MFU：GPUの理論上の計算能力のうち、実際に学習に使えた割合。
• 集団通信：多数のGPUが一斉に途中の結果を集めたり配ったりする通信。
• チェックポイント：学習の途中の状態の保存。故障からの再開に使う。

メタは集団通信のソフトに、エヌビディアのNCCLを改良した「NCCLX」を使いました。