分散学習：データ並列・モデル並列・ZeRO

なぜ1枚のGPUでは学習できないのか

LLM の学習パイプライン で見たように、現代のモデルは数十億〜数千億パラメータに達します。問題はパラメータ本体だけではありません。学習中のGPUメモリには、次の4つが同時に乗ります。

• パラメータ本体
• 勾配（パラメータと同数）
• 最適化状態（Adam なら1次・2次モーメントでパラメータの2倍）
• アクティベーション（順伝播の中間結果。逆伝播で必要）

混合精度（FP16/BF16 で計算、FP32 でマスタ重みと状態を保持）の典型では、パラメータ1個あたり 16バイトほどかかります（FP16のパラメータ2＋勾配2、FP32のパラメータ4＋Adamモーメント2つで8）。これだけで 1B パラメータ＝約16GB、アクティベーションは別枠です。単一GPUに載らないので、複数GPUへ「何を」「どの軸で」割るかが分散学習の核心になります。

データ並列：モデルを複製しバッチを割る

最も基本的なのがデータ並列（Data Parallelism）です。全GPUに同じモデルの完全なコピーを置き、ミニバッチを GPU 台数で分割して各GPUに配ります。各GPUは自分の小バッチで順伝播・逆伝播し、ローカルな勾配を得ます。

ここで重要なのが、各GPUの勾配は別々のデータから計算されるのでそのままでは食い違う点です。パラメータの一致を保つには、更新前に全GPUの勾配を平均して揃える必要があります。この集団通信が AllReduce です。

各GPU: g_i = ∂L_i/∂w   （自分の小バッチでの勾配）
AllReduce: g = (g_0 + g_1 + ... + g_{N-1}) / N   ← 全GPUで同一値を共有
各GPU: w ← w - lr * g   （同じ更新なので重みは一致を保つ）

実装上の主役が Ring-AllReduce です。N台をリング状に並べ、勾配を N 個のチャンクに分け、Reduce-Scatter（各チャンクの合計を1台に集める）と All-Gather（結果を全台へ配る）の2段で実現します。各GPUが送受信するデータ量は 2×(N-1)/N × モデルサイズで、台数 N が増えてもほぼ一定です。だからデータ並列はスループットでよくスケールします。

モデル並列：モデル自体を割る

モデルが1台に載らないなら、モデル自体を分割します。割り方が2つあります。

テンソル並列：1つの層を割る

テンソル並列（Tensor Parallelism）は、層内の行列演算を複数GPUに横断分割します。例えば全結合 Y = X · W の重み W を列方向に [W_1, W_2] と割れば、各GPUが Y_1 = X·W_1、Y_2 = X·W_2 を計算し、結果を連結すれば Y になります。Transformer の Attention やMLPブロックは、この分割が数学的にきれいに収まるよう設計できます。

代償は通信頻度です。1つの層を割るため、順伝播・逆伝播のたびに部分結果を集約する AllReduce が層ごとに発生します。通信が極めて高頻度かつレイテンシに敏感なので、テンソル並列は NVLink で密結合した同一ノード内に閉じるのが鉄則です。

パイプライン並列：層の並びを割る

パイプライン並列（Pipeline Parallelism）は、モデルの層をステージに分けて別々のGPUに置く縦の分割です。GPU0 が第1〜8層、GPU1 が第9〜16層、というように担当を分け、データはステージ間をバケツリレーで流れます。

素朴にやると、後段のGPUは前段の出力を待つ間アイドルになります。この空白がパイプラインバブルです。対策がマイクロバッチ化で、ミニバッチを細かく刻んで次々と投入し、各ステージを常時稼働させてバブルを薄めます（GPipe / 1F1B などのスケジューリング）。通信はステージ境界のアクティベーション授受のみなので、テンソル並列より通信は軽く、ノード間にまたがらせやすいのが特徴です。

ZeRO：データ並列の「メモリ重複」を消す

ここで スケーリング則 が示す「巨大化の圧力」と、データ並列のメモリ重複が正面衝突します。ZeRO（Zero Redundancy Optimizer） は、データ並列の使いやすさ（バッチを割るだけ）を保ったまま、その重複を排除する手法です。

着想は単純です。データ並列では N 台が同一の状態を N 重に持っています。これを丸ごと持たせず、各GPUが全体の 1/N だけを保持し、必要なときだけ通信で取り寄せれば、メモリ重複が消えます。分割対象を段階的に増やしたのが ZeRO の3ステージです。

最初に効くのが ZeRO-1です。前述のとおり最適化状態は最大のメモリ食い（混合精度Adamで全16バイト中の12バイト）で、これを 1/N に割るだけで効果が大きい。ZeRO-2 は勾配も分割します。データ並列の AllReduce はもともと Reduce-Scatter と All-Gather の2段なので、各GPUが「自分の担当パラメータ分の勾配だけ」を Reduce-Scatter で受け取れば、勾配を分割保持でき通信量は据え置きです。

ZeRO-3 はパラメータ本体まで割ります。各GPUは自分の担当パラメータしか持たないので、ある層を計算する瞬間にその層の全パラメータを All-Gather で一時的に集め、使い終わったら捨てる。これによりパラメータ・勾配・状態のすべてが 1/N になり、台数を増やすほど巨大モデルが載るようになります。

通信とメモリのトレードオフを設計する

分散学習の設計は、結局メモリ・通信・計算の三者のやりくりです。代表的なメモリ節約策も、多くは通信や計算との交換で成り立ちます。

実際の超大規模学習では、これらを階層的に組みます。ノード内はテンソル並列（高速NVLinkを活かす）、ノード間はパイプライン並列、その外側をデータ並列（＋ZeRO）で包むという 3D 並列が典型です。割り当ては「通信が高頻度なものほど内側（近い）GPUへ」という原則で決まります。

まとめ：分割軸とAllReduceで全体像をつなぐ

分散学習は、一見バラバラな技術の寄せ集めに見えても、**「何をどの軸で割り、その結果どんな通信が要るか」**という1つの問いに還元できます。

この見立てを持つと、「GPUを増やしてもモデルが大きくできない」「ノードをまたぐと急に遅い」「メモリは足りるのに通信で頭打ち」といった現象が、分割軸と通信パターンの言葉で説明できます。学習対象を小さく抑える設計は パラメータ効率ファインチューニング と、最適化状態がメモリを支配する理屈は 最適化アルゴリズムの系統 と合わせて読むと、巨大モデルを学習可能にする工夫の全体像が線でつながります。