モデル並列とデータ並列

ひとことで言うと

並列化は、計算を複数のプロセスへ割り当てて1回の学習を処理する設計です。データ並列はモデルを複製してデータを分け、モデル並列はモデルの一部を分けます。したがって、前者の主な通信対象は勾配、後者は層の境界を通る活性化です。

なぜ必要か

1台に収まるモデルなら、データ並列はスループットを上げる自然な選択です。しかし各プロセスに全モデルを置くため、モデルがメモリに収まらなければ始まりません。逆にモデル並列はパラメータを分割して収容できますが、層の境界ごとに活性化を転送する待ち時間が生じます。

「何を分けるか」は、そのまま通信量の見積もりになります。データ並列では各プロセスが計算した勾配を all-reduce などの集団通信でそろえます。勾配の総量は概ねパラメータサイズです。モデル並列では入力の活性化や中間表現を次の担当へ送り、通信量はバッチ、系列長、隠れ次元などの形状に左右されます。

収めたい制約	まず検討する方式	通信の中心
モデルは収まるが処理量を増やしたい	データ並列	勾配
モデルを1台に置けない	モデル並列	活性化

仕組み

データ並列で、モデルパラメータを $\theta$、プロセス $i$ のミニバッチ損失を $L_i$、プロセス数を $N$ とします。各プロセスは同じ $\theta$ から勾配を計算し、平均を使って更新します。

$$g = \frac{1}{N}\sum_{i=1}^{N}\nabla_{\theta}L_i, \qquad \theta \leftarrow \theta - \eta g$$

$g$ は平均勾配、$\eta$ は学習率です。all-reduce は全プロセスの値を集約し、その結果を全員に返す操作なので、更新後のモデルがそろいます。モデル並列では、ある層の出力活性化 $\mathbf{h}^{(k)}$ を次の担当へ渡し、受信側が残りの層を計算します。Megatron-LMのようなテンソル並列は、層の内部演算自体を分割するため、層境界だけでなく演算途中にも通信点が入ります。

パイプライン並列は層を連続した区間に分け、入力をマイクロバッチへ分割して各区間を流します。先頭が次のマイクロバッチを処理している間に、後段も前のマイクロバッチを処理できます。ただし開始直後と終了直前は一部の段しか動かず、全段が埋まらない時間が「気泡」です。マイクロバッチ数を増やすと気泡の割合を下げやすい一方、保持する活性化やスケジューリングの複雑さが増えます。

試験でどう問われるか

問われ方	正解に寄る条件	引っかけ
データ並列の通信	全モデルを各機に置き、勾配を集約する。規模はパラメータ側	データだけ送るので通信不要とする
モデル並列の通信	分割境界の活性化を送る。テンソル並列では途中にも通信	勾配だけを通信対象とする
パイプライン並列	マイクロバッチを流し、段の空き時間が気泡になる	分割すれば常に全段が稼働するとする
大きなバッチ	1更新あたりの勾配の統計が変わるため学習率も再調整する	バッチだけ増やせば同じ設定で必ず同じ学習になる

実装で確かめる

通信対象の違いを、単純な要素数で見積もります。実際の帯域や all-reduce のアルゴリズムまでは含めません。

import numpy as np

params = np.zeros((1024, 1024), dtype=np.float32)
activation = np.zeros((8, 128, 4096), dtype=np.float32)
print("gradient MB:", params.nbytes / 2**20)
print("activation MB:", activation.nbytes / 2**20)

この設定では勾配は4.0 MiB、活性化は16.0 MiBです。実装では、前者を集団通信し、後者を隣接する段へ送る、という見積もりの出発点になります。

取り違えやすいもの

方式	各プロセスが持つもの	主な通信	向く制約
データ並列	モデル全体と異なるデータ	勾配の集約	モデルは収まるがデータ処理を増やしたい
テンソル並列	層や演算の一部分	活性化・演算途中のテンソル	1層またはモデル全体が大きい
パイプライン並列	連続した層の区間	段間の活性化	層を順序付き区間へ分けられる
混合並列	上記の組合せ	勾配と活性化の双方	サイズとスループットを同時に調整したい

想起チェック

データ並列で各機に置くものと、主な通信対象は何か

各機にモデル全体を複製し、異なるデータで計算した勾配を集約します。通信量の基準はパラメータサイズです。

モデル並列とパイプライン並列で送るものは何か

モデルを分割した境界の活性化です。パイプライン並列ではそれをマイクロバッチ単位で流します。

パイプラインの気泡と、大バッチ化の注意点は何か

開始・終了時に段が埋まらない空き時間が気泡です。バッチを大きくすると勾配の統計が変わるため、学習率の再調整が必要になります。