モデル並列性は、機械学習および深層学習において、GPUやTPUなどの単一デバイスのメモリに収まらない大規模なニューラルネットワークを訓練または推論するために使用される分散コンピューティング技術である。各デバイスがモデルの完全なコピーを保持し、異なるデータバッチを処理するデータ並列性とは異なり、モデル並列性はモデル自体を分割し、通常は異なる層または層のサブセットを異なるデバイスに割り当てる。このアプローチは、数十億または数兆のパラメータを持つ可能性があるTransformer (architecture)ベースの大規模言語モデルのようなモデルをスケールアップするために不可欠である。
概要
モデル並列性では、ニューラルネットワークはその深さ（層単位）に沿って分割されるか、より高度な形式では、幅や個々のテンソル演算などの他の次元に沿って分割される。最も単純な形式は層単位の分割であり、連続する層が異なるデバイスに配置される。順伝播中、活性化はあるデバイスから次のデバイスへ流れ、逆伝播中は勾配が逆方向に流れる。これによりパイプラインのような実行が生じ、慎重にスケジュールされない場合、アイドル時間（バブル）が発生する可能性がある。モデル並列性はしばしばデータ並列性と対比されるが、大規模モデルの訓練で使用される3D並列性のようなハイブリッドアプローチで両者を組み合わせることができる。
動機
現代のニューラルネットワーク、特に大規模言語モデルは、指数関数的にサイズが増大している。例えば、GPT-3（2020年）は1750億のパラメータを持ち、32ビット精度の重みだけで350GB以上のメモリを必要とし、A100のような高性能GPUの80GB容量をはるかに超える。混合精度や勾配チェックポイントなどのメモリ最適化を用いても、単一デバイスでの訓練は実行不可能である。モデル並列性により、研究者は単一デバイスが保持できるよりも桁違いに大きなモデルを、高速インターコネクトで接続されたデバイスクラスタ全体にモデルを分散することで訓練できる。
モデル並列性の種類
層単位（パイプライン）並列性
これは最も単純な形式であり、モデルが順次ステージに分割され、各ステージが異なるデバイスに割り当てられる。例えば、100層のネットワークを10層ずつの10ステージに分割し、各ステージを別々のGPUに配置できる。順伝播中、データはステージ1、次にステージ2というように流れる。主な課題は負荷分散とパイプラインバブルの削減である。GPipeやPipeDreamなどの技術は、マイクロバッチ処理とスケジューリングを導入して利用率を向上させる。
テンソル並列性
テンソル並列性は、個々の演算（例：行列乗算）を複数のデバイスに分割する。例えば、トランスフォーマーのアテンション機構では、クエリ、キー、バリューの射影をGPU間で分割し、結果をオールリデュース演算で結合できる。このアプローチはMegatron-LMで使用され、デバイスあたりのメモリを削減しつつ高い計算効率を維持するのに効果的であるが、高帯域幅の通信が必要である。
エキスパート並列性
混合エキスパート（MoE）モデルで使用されるエキスパート並列性は、異なるエキスパートネットワークを異なるデバイスに配置し、ルーター（ゲーティングネットワーク）がトークンを適切なエキスパートに分配する。これは、トークンごとに一部のエキスパートのみが活性化されるため、計算量を比例的に増加させずに膨大なパラメータ数を可能にするモデル並列性の一形態である。
実装上の課題
モデル並列性にはいくつかの課題が導入される：
- 通信オーバーヘッド：デバイスは活性化と勾配を交換する必要があり、特に遅いインターコネクトではボトルネックになる可能性がある。
- 負荷不均衡：層のサイズや計算要件が不均一な場合、一部のデバイスがアイドル状態になる一方で他のデバイスが動作することがある。
- パイプラインバブル：パイプライン並列性では、デバイスが前のステージからのデータを待つことがあり、利用率が低下する。
- メモリ断片化：モデルを分割するとメモリ使用量が不均一になる可能性があり、慎重な配置とスケジューリングが必要である。
- フォールトトレランス：デバイス数が多いと障害の確率が増加し、チェックポイントと回復メカニズムが必要になる。
データ並列性との比較
データ並列性では、各デバイスがモデルの完全なコピーを保持し、異なるミニバッチのデータを処理する。勾配は各ステップ後にデバイス間で平均化される。これは単純で、単一デバイスのメモリに収まるモデルに対してうまくスケールする。一方、モデル並列性は、モデル自体が大きすぎる場合に必要である。しかし、モデル並列性は、直列的な依存関係により、通信コストが高く効率が低いことが多い。実際には、大規模訓練では両方が使用される：ノード間ではデータ並列性、ノード内ではモデル並列性である。
応用
モデル並列性は、GPT-4、PaLM、LLaMAなどの大規模言語モデルの訓練と、それらの本番環境での提供に不可欠である。OpenAI、Google（Google Cloud経由）、Anthropicなどの企業は、モデルの訓練と展開にモデル並列性を依存している。また、コンピュータビジョン（例：3D CNN）や科学計算など、モデルが単一アクセラレータには大きすぎる他の領域でも使用されている。
ソフトウェアサポート
いくつかのフレームワークがモデル並列性の組み込みサポートを提供している：
- PyTorch：パイプライン並列性（例：
torch.distributed.pipeline.sync.Pipe）を持つtorch.distributedと、tensor_parallelライブラリによるテンソル並列性を提供する。 - TensorFlow：
tf.distributeがTPUStrategyやMultiWorkerMirroredStrategyなどの戦略を提供し、モデル並列性を組み込むことができる。 - Megatron-LM：NVIDIAによるテンソルおよびパイプライン並列性のための専門ライブラリ。
- DeepSpeed：Microsoft製で、オプティマイザ状態、勾配、パラメータを分割するモデル並列性の一形態であるZeRO（Zero Redundancy Optimizer）を提供する。
- JAX：
jax.shardingにより、明示的なデバイス配置と並列実行を可能にする。
ハードウェアに関する考慮事項
モデル並列性は、高速なデバイス間通信に大きく依存する。単一ノード内では、NVLinkとPCIeが高速接続を提供し、ノード間ではInfiniBandまたは高速イーサネットが使用される。NVIDIA GPUやGoogle TPUなどの専用ハードウェアには、モデル並列性を容易にする専用インターコネクト（例：NVSwitch、TPUリンク）がしばしば含まれる。CerebrasやSambaNovaなどの企業は、モデル並列性の必要性を減らすために大容量オンチップメモリを備えたシステムを設計しているが、極端なスケールでは分散アプローチが依然として必要である。
最近の進歩
最近の研究は、モデル並列性の効率向上に焦点を当てている。シーケンス並列性（シーケンス次元の分割）やコンテキスト並列性（長いシーケンス用）などの技術が開発されている。AlpaやFlexFlowなどの自動モデル分割アルゴリズムは、デバイス間の演算配置を最適化する。さらに、データ、テンソル、パイプライン並列性を組み合わせたハイブリッド並列性は、NVIDIAとMicrosoftが訓練したMegatron-Turing NLGモデル（5300億パラメータ）に例示されるように、大規模訓練の標準となっている。
限界と将来の方向性
その力にもかかわらず、モデル並列性には限界がある。通信オーバーヘッドは、小規模モデルや遅いネットワークでは支配的になる可能性がある。分割とスケジューリングの複雑さは、モデルサイズとクラスタの不均一性に応じて増加する。将来の方向性には、よりインテリジェントな自動並列化、より優れた通信圧縮、そして本質的により並列化可能な新しいアーキテクチャが含まれる。生成AIの台頭とますます大規模化するモデルにより、モデル並列性はスケーラブルな深層学習の基盤であり続けるだろう。
関連項目
参考文献
- Narayanan, D., et al. (2019). PipeDream: Generalized Pipeline Parallelism for DNN Training.
- Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.
- Rajbhandari, S., et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.
- Huang, Y., et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.