分散深層学習

ひとことで言うと

分散深層学習は、1つのモデルを複数の計算資源で学習させ、処理量やメモリ容量の制約を越えるための構成です。代表的なデータ並列では、各ワーカーがモデル全体のコピーを持ち、異なるデータを処理して得た勾配を同期的に集約します。モデル並列は逆に、モデルの層やテンソルを複数のワーカーへ分割する方式です。前者は実装しやすい一方、モデル全体を各GPUへ載せる必要があります。

なぜ必要か

データ並列の狙いは、1回の更新で処理するサンプルを複数GPUへ分け、学習のスループットを上げることです。各ワーカーのローカルバッチを $b$、ワーカー数を $k$ とすると、同期SGDの1回の更新が見る実効バッチサイズは $B=kb$ になります。ここで重要なのは、GPUが $k$ 台になっても更新回数がそのまま $k$ 倍になるわけではない点です。1回の同期で、$k$ 台分のデータに対する勾配をまとめて1回だけ更新します。

同期SGDでは、各ワーカーが同じ初期パラメータから順伝播と逆伝播を行い、勾配を平均してから全ワーカーが同じ更新を適用します。したがって再現性やデバッグの基準をそろえやすい反面、最も遅いワーカーと勾配集約を全員が待つことになります。非同期方式なら待ち時間を減らせる可能性がありますが、古いパラメータに基づく勾配が混ざるため、ここで扱う同期SGDとは更新の意味が変わります。

台数を増やすと、計算は分担できますが、パラメータ勾配を全員へそろえる通信が必要です。特にモデルが大きいと、毎ステップの通信量が増え、GPUを追加しても計算時間ではなく通信時間が支配的になります。分散化は「GPU台数に比例して必ず速くなる」仕組みではなく、計算量、ネットワーク、同期の待ち時間を合わせて成立する構成です。

分散の観点	分ける対象	同期点でそろえるもの
データ並列	各ワーカーが担当する入力データ	モデルの勾配または更新
モデル並列	層・テンソルなどモデルの一部	中間テンソルや計算の順序

仕組み

ワーカー $i$ のローカルバッチを $\mathcal{B}_i$、パラメータを $\theta$、そのバッチから計算した勾配を $\mathbf{g}_i$ とします。同期データ並列の勾配集約は次の形です。

$$\bar{\mathbf{g}}=\frac{1}{k}\sum_{i=1}^{k}\mathbf{g}_i,\qquad \theta_{t+1}=\theta_t-\eta\bar{\mathbf{g}}$$

$k$ はワーカー数、$\eta$ は学習率、$\theta_t$ は時刻 $t$ のパラメータです。All-Reduce は、全ワーカーの値を足し合わせた結果を全ワーカーへ返す集団通信で、勾配の平均を実装する基本操作になります。平均を返す実装でなければ、合計を $k$ で割る処理を明示しないと、台数を増やすだけで更新量まで変わります。

小さいバッチを順番に $k$ 回処理する場合と、大きいバッチを1回処理する場合を同一視するには、各小バッチの勾配がパラメータ変化の間ほぼ変わらない、という近似が要ります。大バッチ化で更新回数が減ると、この近似を保つために線形スケーリング則を試します。基準バッチを $B_0$、基準学習率を $\eta_0$、拡大率を $r=B/B_0$ とすれば、論文の則は次です。

$$\eta=r\eta_0=\frac{B}{B_0}\eta_0$$

ただし、この式を初手から適用すると危険です。学習初期はパラメータが急速に変化するため、複数の小更新を1回の大更新へ置き換える近似が崩れやすく、論文は低い学習率から目標値へ徐々に上げる gradual warmup を使いました。実装上は、最初の数エポックまたは所定ステップで、学習率を $\eta_0$ から $r\eta_0$ へ段階的に近づけます。ウォームアップは大バッチのノイズを消す魔法ではなく、初期だけ更新を穏やかにして線形則を適用しやすくするための区間です。

大バッチが常に正解になるわけでもありません。原論文の ImageNet 実験では、一定範囲では精度を維持できましたが、バッチサイズを無制限には拡大できず、ある点を越えると精度が急速に悪化すると述べています。したがって、線形則は探索開始点であって保証ではありません。学習曲線、検証精度、1ステップの通信時間を測り、性能が折れる手前の構成を選びます。通信については、勾配計算と All-Reduce をパイプライン化して待ち時間を隠す設計もありますが、通信そのものが消えるわけではありません。

試験でどう問われるか

問われ方	正解に寄る条件	引っかけ
データ並列とモデル並列の区別	前者はモデルを複製してデータを分割、後者はモデルを分割	データを分けるだけならモデル並列とする
同期SGDの実効バッチ	ローカルバッチ $b$ と台数 $k$ から $B=kb$	台数を増やすと更新回数も同じ比率で増える
勾配の集約	各ワーカーの勾配を All-Reduce で平均し、同じ更新を適用	代表ワーカーだけが更新して他へ配らない
学習率の変更	バッチ拡大率 $r$ に合わせ $\eta=r\eta_0$ を試し、初期はウォームアップ	大バッチなら学習率を必ず固定、または必ず平方根で増やす
スケーリングの限界	通信がボトルネックになり、精度も一定点以降に悪化し得る	GPU台数に比例して無限に高速化・高精度化する

実装で確かめる

次の例では、同じパラメータに対する2ワーカーのローカル勾配を平均します。同期点の後に同じ更新を適用するため、両者のパラメータは一致します。np.stack(...).mean(axis=0) が小さな All-Reduce と平均に相当します。

import numpy as np

theta = np.array([1.0, -2.0])
local_grads = np.array([[0.4, -0.2], [0.8, 0.6]])
global_grad = local_grads.mean(axis=0)
eta = 0.1
updated = theta - eta * global_grad
workers = np.repeat(updated[None, :], repeats=2, axis=0)
print("global_grad:", global_grad)
print("same_after_sync:", np.allclose(workers[0], workers[1]))

出力は global_grad: [0.6 0.2] と same_after_sync: True です。ここで sum(axis=0) のまま割らずに更新すると、2台では勾配が2倍になり、ワーカー数を変えただけで学習率を変えたのと同じ影響が出ます。実際の分散実装では、この集約をネットワーク越しの集団通信として行うため、正しい平均だけでなく通信時間も計測対象になります。

取り違えやすいもの

用語・方式	分散深層学習との切り分け
データ並列	各ワーカーがモデル全体を持ち、データを分けて勾配を集約する
モデル並列	層やテンソルを分け、1つのモデルを複数ワーカーで実行する。1ワーカーに載らないモデルにも使う
同期SGD	各ワーカーが集約を待ってから同じ更新をする。実効バッチはローカルバッチの台数倍になる
非同期SGD	更新を待たずに進むため、ワーカー間でパラメータの世代がずれる。同期SGDの式をそのまま適用しない
勾配蓄積	1ワーカー内で複数ステップの勾配をためてから更新する。複数ワーカーの同期通信とは別の仕組み
All-Reduce	全ワーカーのテンソルを集約し、結果を全員へ返す通信操作。最適化則そのものではない

想起チェック

データ並列でワーカー数を増やすと、同期SGDの実効バッチはどうなるか

各ワーカーのローカルバッチを $b$ とすれば、1回の同期更新が見る実効バッチは $B=kb$ です。更新前に勾配を平均し、全ワーカーが同じ更新を適用します。

大バッチ化で線形スケーリング則とウォームアップを組み合わせる理由は何か

バッチを $r$ 倍にしたとき学習率を $r\eta_0$ とするのが線形則です。ただし学習初期はパラメータ変化が大きく大更新の近似が崩れやすいため、低い学習率から目標値へ徐々に上げます。

分散化しても速度向上が頭打ちになる主因は何か

各ワーカーの勾配をそろえる All-Reduce と同期待ちです。計算が速くなるほど、通信と最も遅いワーカーの待ち時間が相対的に大きくなります。

線形スケーリング則を無制限に適用してよいか

よくありません。原論文の実験でも成立する範囲があり、一定点を越えると精度が悪化します。線形則は検証すべき初期値であり、学習曲線と精度で上限を判断します。