LLM学習を加速：Accelerate ND-ParallelによるマルチGPU最適化

複数のGPUにわたる大規模な人工知能モデルの学習は、様々な並列化戦略の複雑な相互作用により、重大な課題を提示します。この複雑なプロセスを合理化するため、Hugging FaceのAccelerateのようなプラットフォームは、Axolotlと連携し、開発者がトレーニングスクリプト内で異なる並列化技術をシームレスに組み合わせることを可能にするメソッドを統合しました。この革新は、基本的な分散設定の構成から、広大なコンピューティングクラスター全体での洗練された多次元並列化のオーケストレーションまで、モデルのスケーリングを簡素化することを目的としています。

マルチGPU学習の中心には、いくつかの基本的な並列化戦略があり、それぞれが特定のスケーリングボトルネックに対処するように設計されています。

データ並列 (DP) は最も直接的なアプローチで、モデル全体、その勾配、およびオプティマイザの状態が各GPUに複製されます。データバッチはこれらのデバイス間で均等に分割され、モデルパラメータが更新される前に勾配が同期されます。これにより、単一デバイスでの学習と比較してスループットが大幅に向上しますが、その主な制限は、モデル全体が単一GPUのメモリに収まる必要があることです。

モデルが単一デバイスには大きすぎる場合、完全シャードデータ並列 (FSDP) が解決策を提供します。DeepSpeedのZeRO-3のような技術に触発され、FSDPはモデルの重み、勾配、およびオプティマイザの状態を複数のGPUにシャード化します。各デバイスはデータバッチの一部を受け取ります。DPとは異なり、FSDPはモデル全体を複製することを回避します。代わりに、フォワードパスまたはバックワードパスの前に、特定の層に必要な重みのみを収集し、その後再シャード化します。この方法は、通信オーバーヘッドの増加と引き換えに、大幅なメモリ節約を実現します。FSDPは多くのGPU、さらには複数のノード（複数のGPUをホストするマシン）にわたってスケーリングできますが、特にデバイスのネットワーク全体にわたってシャード化する場合、ノード間通信が遅いと効率が低下する可能性があります。

テンソル並列 (TP) は、モデル並列の一種であり、非常に大きな個々のモデル層の課題に対処します。モデル全体を複製したりシャード化したりする代わりに、TPは大きな線形層（Transformerモデルで一般的）をデバイス間で分割します。各デバイスは行列乗算の一部のみを計算し、同時に同一のデータバッチを受け取ります。これにより静的なメモリパーティションが作成され、TPグループサイズに比例して一定のメモリ削減が提供されます。TPは、高帯域幅のGPU間通信（例：NVLink）が利用可能な単一ノード内での計算とメモリの分散に非常に効果的です。しかし、頻繁なアクティベーション同期要件のため、TPは一般的に複数のノードにわたるスケーリングやPCIeのような遅い接続での使用は推奨されません。

大規模言語モデル (LLM) の台頭と、数十万または数百万トークンに達することもある、ますます長いシーケンス長への要求に伴い、新たなメモリ課題が生じます。Transformerのコアコンポーネントであるアテンションメカニズムは、シーケンス長に対して二次関数的にスケールし、アクティベーションのメモリ消費量が法外になります。コンテキスト並列 (CP) は、入力シーケンスをGPU間でシャード化することでこれに対処します。各デバイスは完全なコンテキストの一部のみを処理し、アテンション行列のより小さな部分を計算します。完全なシーケンスへのアクセスを必要とする正確なアテンション計算を確実にするために、RingAttentionのような技術が採用され、キーと値の行列がデバイス間で循環されます。これにより、各クエリはメモリと計算負荷を分散しながら、シーケンス全体に対してアテンションスコアを計算できます。

最も要求の厳しい学習シナリオ、特に複数のノードにわたるシナリオでは、開発者はこれらの戦略を「ND並列」として構成し、コンピューティングクラスターの二次元ビューを活用できます。一方の軸には高速なノード内通信、もう一方の軸には遅いノード間通信です。

ハイブリッドシャードデータ並列 (HSDP) はFSDPとDPを組み合わせたものです。各ノード内でFSDPを適用し、高速なノード内リンクをメモリ集約的なシャード化に利用し、同時にノード間でDPを使用してモデルを複製し、勾配を同期します。これにより、大規模なマルチノード設定における通信オーバーヘッドが最適化されます。

完全シャードデータ並列とテンソル並列 (FSDP + TP) の組み合わせは、FSDPを使用してノード間でモデルをシャード化し、TPを使用してノード内でシャード化することを伴います。この強力な組み合わせは、FSDPのレイテンシを削減し、単一デバイスには大きすぎる層を分散し、グローバルバッチサイズを減少させる可能性があります。同様に、完全シャードデータ並列とコンテキスト並列 (FSDP + CP) も、あまり一般的ではありませんが、極端に長いシーケンス長をFSDPと組み合わせる場合にメモリ使用量をさらに削減するために使用できます。

究極の柔軟性とスケールを実現するために、ハイブリッドシャードデータ並列とテンソル並列 (HSDP + TP) は、しばしば3D並列と呼ばれ、階層的な構造を作成します。データ並列はノードグループ間でモデルを複製し、FSDPは各グループ内でモデルをシャード化し、TPは各ノード内で個々の層を分割します。これにより、大規模な学習環境でメモリ使用量とスループットのバランスを取るための最大の適応性が提供されます。

適切な並列化戦略を選択するだけでなく、分散学習を最適化するためにはいくつかの実践的な考慮事項が重要です。FSDPの場合、CPU RAM効率的なロードとシャード化された状態辞書チェックポインティングを有効にすることは、単一デバイスメモリには大きすぎるモデルを処理するために不可欠です。学習の安定性とスループットに大きく影響する有効バッチサイズは、マイクロバッチサイズ、勾配蓄積ステップ、およびデータ並列ワールドサイズによって決定されます。並列化に伴い有効バッチサイズが増加するため、安定性を維持するために学習率は比例してスケーリングする必要があります。最後に、勾配チェックポインティングは、計算とメモリをトレードオフすることで追加のメモリ節約を提供します。これは、バックワードパス中に中間アクティベーションを選択的に再計算することで、学習時間のわずかな20-30%の増加と引き換えに、アクティベーションメモリを60-80%削減します。この技術はすべての並列化戦略とシームレスに連携するため、メモリ制約が続く場合に非常に価値のあるツールとなります。

最終的に、最適な構成は実験を伴うことが多く、メモリ、計算、通信オーバーヘッド間の理想的なバランスは、特定のモデル、データセット、およびハードウェア設定に大きく依存します。これらの高度な並列化技術は不可欠なツールであり、大規模なAIモデル開発で可能なことの限界を押し広げています。