パイプライン並列処理は、ディープニューラルネットワークの分散学習戦略であり、モデルのレイヤーを連続するステージに分割し、各ステージを異なるデバイス（GPUやTPUなど）に割り当てる手法である。これは、モデル並列処理とデータ並列処理を組み合わせ、ソフトウェアパイプライン方式でマイクロバッチをステージ間で処理することで、学習スループットを向上させ、素朴なレイヤー単位の並列処理で生じるアイドル時間（バブル）を削減することを目的としている。
背景と動機
大規模なニューラルネットワーク、特に大規模言語モデルやその他の深層学習モデルの学習には、膨大な計算リソースが必要である。モデルサイズが数十億から数兆パラメータに成長するにつれ、メモリと計算の制約により単一デバイスでの学習は不可能になる。データ並列処理は各デバイスにモデル全体を複製してデータを分割するが、モデルが単一デバイスのメモリを超えると失敗する。モデル並列処理はモデル自体をデバイス間で分割するが、素朴なレイヤー単位の配置では深刻な利用率低下を招く。つまり、一度に1つのデバイスだけが計算し、他は待機するため、スループットが低くアイドル時間が長くなる。
パイプライン並列処理は、モデルをステージ（通常は連続するレイヤー群）に分割し、複数のマイクロバッチを同時に処理することでこの問題に対処する。重要な洞察は、ステージ間で計算を重ね合わせることである。ステージ1がマイクロバッチ2を計算している間に、ステージ2はマイクロバッチ1を計算できる、といった具合である。これは組立ラインに似ており、各ステージが着信データを連続的に処理し、結果を次のステージに渡す。
中核的概念
パイプライン並列処理における基本作業単位はマイクロバッチである。大きな学習バッチはより小さなマイクロバッチに分割され、パイプラインステージを順次流れる。各ステージは割り当てられたレイヤーに対して順伝播と逆伝播を実行する。パイプラインはマイクロバッチを次々と供給することで「満たされ」、初期のウォームアップ期間後は全ステージが同時に稼働し、高いスループットを達成できる。
パイプラインの効率は、ステージがアイドル状態になる時間の割合であるバブル比で測定されることが多い。p個のステージとm個のマイクロバッチを持つパイプラインでは、mがpよりはるかに大きい場合に理想的なスループットが達成され、相対的なバブルオーバーヘッドが最小化される。ただし、mを増やすと活性化のメモリ使用量も増加し、トレードオフが生じる。
主要な方式
パイプライン並列処理にはいくつかの方式が提案されており、順伝播と逆伝播のスケジューリング方法と勾配更新の処理方法が異なる。
GPipe
Googleは2019年にGPipeを導入した。これはモデルをステージに分割し、単純なスケジュールを使用する。各マイクロバッチは全ステージで順伝播を実行し、その後逆順で全ステージの逆伝播を実行する。GPipeは同期勾配更新を使用し、全ステージが全マイクロバッチの完了を待ってから重みを更新する。これにより一貫した勾配が保証されるが、ステージ数に応じてバブルが増大する。GPipeはまた、活性化の再計算または保存のいずれかを必要とし、計算とメモリをトレードオフする。
PipeDream
Microsoft Research（2019年）によるPipeDreamは、非同期スケジュールを使用し、各ステージがマイクロバッチをラウンドロビン方式で処理し、順伝播と逆伝播を可能な限り早く実行する。これによりバブルは減少するが、重みの陳腐化が生じる。つまり、ステージによって異なるバージョンの重みを使用する可能性があり、収束に悪影響を与える可能性がある。PipeDreamはまた、「重みスタッシング」と呼ばれる技術を使用して、異なるマイクロバッチに対して複数の重みバージョンを維持し、メモリオーバーヘッドを増加させる。
PipeDream-2BWとその変種
PipeDream-2BW（2020年）は、2つの重みバッファ（順伝播用と逆伝播用）を使用してメモリを削減し収束を改善することでPipeDreamを改良した。V-PipeやPipeMareなどの変種は、異なる同期戦略を探求している。
Megatron-LMとインターリーブスケジューリング
NVIDIAのMegatron-LM（2019年）は、テンソル並列処理（個々のレイヤーの分割）とパイプライン並列処理を組み合わせた。その後、Megatron-2（2020年）はインターリーブスケジュールを導入し、モデルをデバイス数よりも多くのステージに分割して、各デバイスが複数のステージを処理できるようにした。これにより、実行中のマイクロバッチ数を増やしてバブルサイズを削減するが、通信コストが増加する。
Chimeraとその他の最近の方式
Chimera（2021年）は双方向パイプラインを使用してバブルをさらに削減する。PTD-P（2021年）やZeroBubble（2023年）などの他の方式は、パイプラインバブルを排除または削減するためのより洗練されたスケジュールを探求している。
数学的および実践的考慮事項
パイプライン並列処理は、他の並列処理戦略と組み合わせられることが多い。テンソル並列処理は個々のレイヤーをデバイス間で分割し、パイプライン並列処理はモデルを垂直に分割する。データ並列処理はパイプラインをデバイスのグループ間で複製する。このハイブリッドアプローチは、Megatron-Turing NLGやOpenAIのGPT-4などのシステムに見られるように、大規模モデルの学習における標準となっている。
パイプラインの深さ（ステージ数）とマイクロバッチサイズの選択は性能に影響する。より深いパイプラインはステージあたりのメモリを削減するが、通信とバブルのオーバーヘッドを増加させる。より大きなマイクロバッチは利用率を向上させるが、より多くのメモリを消費する。最適な構成は、モデルサイズ、デバイスメモリ、および相互接続帯域幅に依存する。
応用と影響
パイプライン並列処理は、最大級のAIモデルの学習に不可欠であった。例えば、GPT-3（175Bパラメータ）はモデル並列処理とパイプライン並列処理の組み合わせで学習された。AnthropicのClaudeモデルやGoogleのPaLMも同様の技術に依存している。このアプローチは、PyTorch（torch.distributed.pipelineモジュール経由）、TensorFlow（tf.distribute経由）、JAX（pjit使用）などの主要フレームワークでサポートされている。
課題と限界
パイプライン並列処理にはいくつかの課題がある。通信オーバーヘッドは、ステージ間で活性化と勾配を送信する際に発生し、低速な相互接続ではボトルネックになる可能性がある。負荷不均衡は、ステージの計算時間が不均等な場合に発生し、効率を低下させる。逆伝播用の活性化を保存するためのメモリ圧力は、活性化の再計算によって軽減できるが、計算コストがかかる。さらに、非同期方式は古い重みによる収束問題に悩まされる可能性があり、同期方式はバブルを発生させる。
今後の方向性
パイプライン並列処理の改善に関する研究は続いており、バブルの削減、メモリの最小化、および異種ハードウェアへの適応に焦点を当てている。自動ステージ分割、動的負荷分散、混合専門家モデルとの統合などの技術は活発な研究分野である。極めて大規模なモデルの台頭とクラウドクラスターでの効率的な学習の必要性により、パイプライン並列処理は分散学習システムの主要な構成要素であり続けるだろう。