AI・機械学習
上級

Pipeline Parallelism(パイプラインパラレリズム)

ニューラルネットワークのレイヤー群を複数のデバイスに段階的に配置し、マイクロバッチを流すことでデバイス間の空き時間を最小化しながら大規模モデルを学習する分散並列化手法。

0 回閲覧
0 いいね

Pipeline Parallelism とは

Pipeline Parallelism(PP)は、ニューラルネットワークの連続するレイヤーを異なる GPU に割り当て、ミニバッチを複数のマイクロバッチに分割してパイプライン的に処理する分散学習手法である。GPipe(Google, 2019)で提案され、現在は Megatron-LM、DeepSpeed、PyTorch の主要フレームワークに実装されている。

パイプライン並列の基本原理

32 層の Transformer モデルを 4 GPU で学習する場合、各 GPU に 8 層ずつ割り当てる。GPU 0 がレイヤー 0-7 を処理し、その出力(activation)を GPU 1 へ送信、GPU 1 がレイヤー 8-15 を処理…という流れで Forward Pass が進む。Backward Pass は逆順に勾配を伝搬する。

パイプラインバブル問題

単純な PP では、1 つの GPU が計算中に残りの GPU が遊ぶ「バブル」が発生する。4 GPU のナイーブ PP ではデバイス利用率が約 25% まで低下する。この問題を解決するスケジューリング手法が複数提案されている。

スケジュールバブル率メモリ効率特徴
GPipeO(K-1)/M高(勾配蓄積)マイクロバッチ数 M を増やすとバブル減少
1F1BO(K-1)/M低(ピークメモリ削減)Forward 1 回→Backward 1 回を交互実行
Interleaved 1F1BO(K-1)/(M×V)中V 個の仮想ステージで通信増だがバブル激減
Zero Bubble PP≈0中W(Weight Update)を分離し空き時間に充当

ここで K = パイプラインステージ数、M = マイクロバッチ数。

1F1B スケジューリング

1F1B(One Forward One Backward)は PipeDream(Microsoft, 2019)で導入されたスケジュールである。ウォームアップフェーズで順次 Forward を実行した後、Forward と Backward を 1:1 で交互に実行する。GPipe と比べてピークメモリ消費を大幅に削減できるのが利点で、DeepSpeed Pipeline の標準スケジュールとなっている。

Interleaved Pipeline

Megatron-LM v3 で導入された手法で、各 GPU に連続しない複数のレイヤーグループ(仮想ステージ)を割り当てる。たとえば GPU 0 にレイヤー 0-3 と 16-19、GPU 1 にレイヤー 4-7 と 20-23 のように配置する。パイプラインの深さが実質的に倍になり、バブルが半減する代わりに、ステージ間通信回数が増加する。

Zero Bubble Pipeline Parallelism

2024 年に提案された最新手法で、Backward Pass を B(入力勾配計算)と W(重み勾配計算)に分離し、パイプラインのバブル区間に W を挿入することでバブルをほぼゼロにする。理論上のデバイス利用率が 100% に近づくが、メモリ使用量が増加するトレードオフがある。

PP と TP の使い分け

観点Pipeline ParallelismTensor Parallelism
通信量少(Activation のみ)多(AllReduce)
推奨接続ノード間 InfiniBandノード内 NVLink
スケール数十〜数百 GPU4〜8 GPU
レイテンシバブルによる増加同期通信による増加

実際の大規模学習では、ノード内は TP、ノード間は PP を適用する「TP-in-node, PP-across-node」が標準パターンである。

推論時の Pipeline Parallelism

学習時と異なり、推論時の PP は自己回帰生成のトークンごとにパイプラインを通過する必要があり、バブルが顕著になる。そのため推論では TP が好まれるが、超大規模モデルでノード数が足りない場合は PP も併用される。

FAQ

Q1: パイプライン並列のマイクロバッチ数はいくつが最適?

A1: 一般にパイプラインステージ数 K の 4 倍以上が推奨される。M が大きいほどバブル率 (K-1)/M が小さくなるが、マイクロバッチサイズが小さくなりすぎると GPU の演算効率が低下する。実測ベースで M=4K〜8K の範囲で最適値を探すのが実務的。

Q2: PP はファインチューニングでも使うべきか?

A2: フルパラメータファインチューニングでモデルが単一 GPU に収まらない場合は有効。ただし LoRA/QLoRA のようなパラメータ効率的手法を使えば単一 GPU で済むことが多く、PP の必要性は低い。

Q3: PP と FSDP は併用できるか?

A3: PyTorch 2.x 以降で PP + FSDP の組み合わせがサポートされている。各パイプラインステージ内のパラメータを FSDP でシャーディングすることで、メモリ効率をさらに改善できる。