分散学習を1から — データ並列・モデル並列・通信がボトルネックになるとき

なぜ1台に載らないのかをメモリの内訳から出発し、データ並列とall-reduce、ZeRO/FSDPが何を分割するのか、テンソル並列とパイプライン並列を順に。最後は通信量と計算量の比で「どこで頭打ちになるか」を自分で見積もり、勾配蓄積・NCCLの設定・ハングの切り分けまで降ります。

比喩: 百人で一冊の辞書を書き換える

百人で1冊の辞書を改訂する仕事を考えます。分担して書くのは簡単です。難しいのは突き合わせのほうで、各自の修正を全員に反映しないと矛盾した辞書ができてしまう。1項目直すたびに全員で会議を開けば、会議の時間が執筆の時間を超えます。人数を増やすほど会議が長くなり、ある人数から先は増員が逆効果になる。

分散学習で起きるのも同じことです。計算を分けるのは難しくありません。難しいのは、分けた結果を全員のモデルに一致させて戻すところ。この記事の主役は演算器ではなく通信で、見積もるべき比は「計算時間 ÷ 通信時間」です。

なぜ1台に載らないのか

まずメモリの内訳を数えます。パラメータ数を $N$ とし、混合精度でAdamを使う典型的な構成を置きます。パラメータをbf16で持つと $2N$ バイト、勾配も同じ精度なら $2N$。ここに、更新の精度を保つためのfp32マスターパラメータ $4N$、Adamの一次モーメント $4N$、二次モーメント $4N$ が乗ります。

つまりこの式は、重み1個を学習させるには、重みそのもののほかに「予備の高精度版」と「過去を覚えておくメモ2枚」がくっついてくると言っています。$M_{\text{state}}$ は学習が終わるまで居座り続けるメモリ、$N$ は重みの個数です。

「パラメータ1個あたり16バイト」。$N = 7 \times 10^9$ なら約112 GB——まだ1つも計算していないのに、状態だけで加速器1枚のメモリを超えます。しかも $N$ に比例するこの項とは別に、逆伝播のために取っておく中間出力（活性値）が要ります。

$$M_{\text{act}} \sim L \cdot B \cdot s \cdot d$$

つまり「活性値は、層を深くしても、まとめて流す量を増やしても、文を長くしても、同じ勢いで膨らむ」ということ。足し算ではなく掛け算で効くのがこの項の性格です。

$L$ は層数、$B$ はマイクロバッチ、$s$ は系列長、$d$ は隠れ次元。層数とバッチと系列長の積で効くので、モデルが同じでもバッチを倍にすれば倍になります。ここを削る定石が活性値チェックポイントで、層の入力だけ残して中間を捨て、逆伝播時に再計算する。メモリは大きく減り、計算は再計算ぶん（おおむね順伝播1回ぶん）増えます。

要するに、1台に載らない理由は2つあり、打ち手も別です。$16N$ の状態は分割で減らし、活性値は再計算とバッチ調整で減らす。 この区別を最初に付けておくと、あとの選択がぶれません。

データ並列とall-reduce

いちばん素直な分割はデータのほうです。$K$ 台がそれぞれモデルの完全なコピーを持ち、別々のミニバッチで順伝播と逆伝播をして、勾配だけを平均して揃える。これがデータ並列で、モデルが1台に載るあいだは常に第一候補です。

揃える操作をall-reduceと呼びます。全員の値を足して（reduce）、その結果を全員が持つ（broadcast）。素朴に1台へ集めて配り直すとその1台の帯域が詰まるので、実装はリングall-reduceを使います。$K$ 台を輪にして、データを $K$ 個の塊に切り、隣へ送りながら足し込む段（reduce-scatter）と、完成した塊を隣へ回す段（all-gather）の2周。1台が送るバイト数は

つまり「どの台も、モデル全体の勾配をおよそ2回ぶん送り出すだけ」ということ。$V_{\text{comm}}$ は1台が1ステップで回線に流すバイト数です。

$b$ は勾配1要素のバイト数です。$K$ が大きくなっても $2Nb$ に漸近するだけで、台数にほぼ依存しません。 32台でも1024台でも1台が送る量は変わらない——リングall-reduceが分散学習の標準になった理由がこれです（ただし段数は $K$ に比例するので、レイテンシは効きます）。

通信量と計算量の比

頭打ちの位置は、1ステップの計算時間と通信時間を比べれば出ます。行列積のコストで見たとおり、学習の計算量はトークン1個あたり約 $6N$ FLOPs。1台が1ステップで処理するトークン数を $D$、1台の実効演算性能を $P$ とすると計算時間は $6ND/P$、通信時間は上の $2Nb/\mathrm{BW}$ です。比を取ると

つまりこの分数は「1ステップの計算が、通信の何倍あるか」を表しています。1を大きく上回っていれば通信は計算の裏に隠れ、1を下回れば台数を足しても待ち時間が伸びるだけ、ということ。

$N$ が消えました。 モデルを大きくしても計算と通信は同じ割合で増えるので、比は変わらない。効くのは3つだけです——1台あたりのトークン数 $D$、相互接続の帯域 $\mathrm{BW}$、そして勾配の精度 $b$。