分散トレーニングは、機械学習および深層学習において、GPUやサーバー全体などの複数のコンピューティングデバイスが連携してモデルをトレーニングする手法である。このアプローチは、数十億から数兆のパラメータを持ち、膨大なデータを必要とする大規模言語モデルなどの現代の人工知能システムにとって不可欠である。計算負荷を分散することで、トレーニング時間を数ヶ月から数日、あるいは数時間に短縮でき、単一のデバイスには収まらない規模のモデルを扱うことが可能になる。この手法は、数十年にわたり高性能コンピューティングの基盤となってきた並列コンピューティングの原理に基づいており、プロセッサ速度の物理的限界とモデルサイズの増大により、AI分野では主流の必須技術となっている。
分散トレーニングの必要性は、メモリ容量と計算速度という2つの主要な制約から生じる。単一のGPUやTPUには有限のメモリしかなく、保存・トレーニングできるモデルのサイズが制限される。同様に、単一のデバイスで数百万のトレーニング例を処理するのに要する時間は、実用的でないほど長くなる可能性がある。分散トレーニングは、モデルとデータを複数のデバイスに分割し、並列処理を可能にすることでこれらの問題に対処する。しかし、これにより通信、同期、フォールトトレランスに関連する複雑さが生じ、これらがこの分野の中心的な課題となっている。理論上の高速化はアムダールの法則によって制限され、最大の改善は、通信オーバーヘッドや逐次依存関係など、並列化できないワークロードの部分によって制限される。
データ並列性
データ並列性は、分散トレーニングで最も広く使用される形態である。このアプローチでは、モデルが各デバイスに複製され、トレーニングデータセットがより小さなバッチに分割され、各デバイスが異なるデータサブセットを同時に処理する。各デバイスがローカルな勾配を計算した後、これらの勾配は全デバイスで集約され、モデルパラメータを更新する。これには、通常、all-reduceなどの集合通信操作を使用した同期ステップが必要であり、全デバイスからの勾配を合計して結果をブロードキャストする。
データ並列性の主な利点は、その単純さとスケーラビリティである。トランスフォーマーベースのモデルを含むあらゆるモデルアーキテクチャに、大幅な変更なしで適用できる。しかし、デバイス数が増えるにつれて、通信オーバーヘッドがボトルネックになる可能性があり、特にパラメータ数が多いモデルでは顕著である。この問題を軽減するために、勾配圧縮、非同期更新、ローカル勾配蓄積などの技術が開発されている。PyTorchやTensorFlowなどのフレームワークは、データ並列性の組み込みサポートを提供しており、実務者が容易に利用できる。
モデル並列性
モデル並列性は、モデルが単一デバイスのメモリに収まらない場合に使用される。このアプローチでは、モデルの異なる部分が異なるデバイスに配置され、データはモデルを順次通過し、各デバイスが前方および後方パスの一部を計算する。これは、NVIDIAやAMDの最先端GPUのメモリ容量を超える数百億のパラメータを持つ大規模言語モデルに特に関連する。
モデル並列性は、各デバイスがレイヤーのサブセットを処理するレイヤー単位の分割や、単一レイヤーの計算をデバイス間で分割するレイヤー内分割など、いくつかの方法で実装できる。後者は、アテンションメカニズムとフィードフォワードネットワークを並列化できるトランスフォーマーモデルで一般的である。しかし、モデル並列性は、後のレイヤーが計算している間、前のレイヤーのデバイスがアイドル状態になるため、利用率の不均衡を引き起こすことが多い。これは、デバイス間で計算を重ね合わせるパイプライン並列性によって軽減できる。
パイプライン並列性
パイプライン並列性は、データ並列性とモデル並列性の要素を組み合わせたハイブリッドアプローチである。この方法では、モデルがステージに分割され、各ステージが異なるデバイスに割り当てられる。トレーニングデータはマイクロバッチで処理され、パイプラインをずらして流れるため、複数のデバイスが異なるマイクロバッチを同時に処理できる。これにより、純粋なモデル並列性に伴うアイドル時間が削減され、ハードウェア利用率が向上する。
パイプライン並列性の顕著な例として、2019年にGoogleの研究者が導入したGPipeフレームワークがあり、この技術を使用して大規模モデルを効率的にトレーニングできることを実証した。もう1つの変種は、Microsoftが開発したPipeDreamで、非同期更新を使用してスループットをさらに向上させる。パイプライン並列性は非常に深いモデルに特に効果的であるが、パイプラインスケジュールの管理とステージ間の通信処理に課題を導入する。ステージ数とマイクロバッチサイズの選択は、パフォーマンスに大きな影響を与える可能性がある。
集合通信
集合通信は分散トレーニングの基盤であり、デバイスがデータを交換し、計算を同期できるようにする。最も一般的な操作には、all-reduce、all-gather、broadcast、reduce-scatterがある。これらの操作は、NVIDIAのNCCL（NVIDIA Collective Communications Library）やMessage Passing Interface（MPI）などのライブラリに実装されており、InfiniBandやEthernetなどの高速インターコネクト向けに最適化されている。
データ並列性では、all-reduce操作が勾配の集約に使用される。例えば、N個のデバイスがある場合、各デバイスが勾配ベクトルを計算し、all-reduceは全デバイスにわたって要素ごとの合計を計算し、結果を各デバイスに配布する。この操作は、ツリーベースまたはリングベースのアルゴリズムを使用して最適化でき、通信時間を削減する。通信トポロジーの選択とネットワークの帯域幅は、分散トレーニングのスケーラビリティにおける重要な要素である。2024年時点で、CoreWeaveやAmazon Web Servicesが提供するような数千GPUを備えたクラスターは、通信オーバーヘッドを最小限に抑えるために高帯域幅インターコネクトに依存している。
同期および非同期トレーニング
分散トレーニングは、同期アプローチと非同期アプローチに分類できる。同期トレーニングでは、すべてのデバイスがローカルデータで勾配を計算し、モデルを更新する前に他のすべてのデバイスが完了するのを待つ。これにより、モデルが全デバイスで一貫していることが保証されるが、ハードウェアのばらつきやネットワーク輻輳による遅いデバイス（ストラグラー）によって遅延する可能性がある。同期トレーニングは、収束特性を保証するため、ほとんどの深層学習フレームワークの標準である。
一方、非同期トレーニングでは、デバイスが他のデバイスを待たずに独立してモデルを更新できる。これによりスループットが向上する可能性があるが、デバイスが古いモデルパラメータを使用するため、勾配の陳腐化のリスクが生じ、収束が遅くなったり不安定になったりする可能性がある。これらの問題に対処するために、勾配の陳腐化バウンディングや弾性平均化などの技術が提案されている。実際には、all-reduceを使用した同期トレーニングが、ほとんどの大規模トレーニングジョブで好まれており、単純さとパフォーマンスのバランスが良い。
ハードウェアとインフラストラクチャ
分散トレーニングには、高性能を達成するための専用ハードウェアとインフラストラクチャが必要である。最も一般的なセットアップは、各サーバーに複数のGPUを搭載し、高速ネットワークで接続されたサーバークラスターである。NVIDIAはGPU市場を支配しており、A100およびH100シリーズがAIトレーニングで広く使用されている。AMDもMI250Xなどの競争力のあるGPUを提供し、IntelはGaudiアクセラレータでこの分野に参入している。Amazon Web Services、Azure、Google Cloudなどのクラウドプロバイダーは、ユーザーがオンデマンドで分散トレーニングクラスターをレンタルできるマネージドサービスを提供しており、社内インフラストラクチャの必要性を減らしている。
GPUに加えて、CerebrasのウェハースケールエンジンやGraphcoreのIPU（Intelligence Processing Unit）などの専用ハードウェアが、分散トレーニングを加速するために開発されている。これらのシステムは、外部通信の必要性を減らすオンチップ通信ネットワークを備えていることが多い。さらに、インターコネクト技術の選択は重要であり、InfiniBandは低レイテンシと高帯域幅を提供し、Ethernetはよりコスト効果が高いが低速である。2023年時点で、GPT-4などの最大のAIトレーニング実行は、数万のGPUを使用したと推定されており、現代の分散トレーニングの規模を浮き彫りにしている。
ソフトウェアフレームワークと技術
分散トレーニングの実装を簡素化するために、いくつかのソフトウェアフレームワークが開発されている。PyTorchのDistributedDataParallel（DDP）は人気のある選択肢であり、データ並列性のためのシンプルなAPIを提供する。TensorFlowはtf.distributeモジュールを提供し、MirroredStrategyやMultiWorkerMirroredStrategyなどのさまざまな戦略をサポートする。Uberが開発したHorovodや、Microsoftが開発したDeepSpeedなどのより高度なライブラリは、勾配圧縮、混合精度トレーニング、ZeRO（Zero Redundancy Optimizer）などの追加の最適化を提供し、オプティマイザ状態、勾配、パラメータをデバイス間で分割することでメモリ使用量を削減する。
これらのフレームワークは、分散通信の複雑さの多くを抽象化し、研究者がモデル開発に集中できるようにする。しかし、デバッグとパフォーマンス最適化には、基礎となる原理を理解することが依然として重要である。例えば、適切なバッチサイズ、学習率スケジュール、通信バックエンドの選択は、トレーニング効率に大きな影響を与える可能性がある。モデルが成長し続けるにつれて、分散トレーニングアルゴリズムとハードウェアの革新は、人工知能分野の進歩に引き続き不可欠である。
課題と将来の方向性
その成功にもかかわらず、分散トレーニングはいくつかの課題に直面している。通信オーバーヘッドは、特に数十億のパラメータを持つモデルでは、依然として主要なボトルネックである。勾配圧縮や低精度通信などの技術が、転送されるデータ量を削減するために探求されている。フォールトトレランスも別の問題であり、デバイスやネットワークリンクの障害がトレーニングを中断させる可能性がある。チェックポイントと、デバイス数を動的に調整する弾性トレーニングは、活発な研究分野である。
将来的には、OpenAIやGoogle DeepMindが開発しているような大規模モデルへの傾向が、より効率的な分散トレーニング手法の必要性を引き続き促進するだろう。データ、モデル、パイプライン並列性を組み合わせた3D並列性の出現は、すでに大規模トレーニング実行で使用されている。さらに、AWS TrainiumやGoogle Cloud TPUなどの専用ハードウェアの開発は、分散トレーニングをよりコスト効果が高く、アクセスしやすくすることを目指している。分野が進化するにつれて、分散トレーニングの原理は人工知能の進歩の基盤であり続けるだろう。