Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

以下は、論文「Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference」の詳細な技術的サマリーです。

1. 問題定義

大規模言語モデル（LLM）の学習と推論は、主に 2 つのメモリボトルネックに直面しています。

1. パラメータメモリ: モデルの重み、勾配、オプティマイザ状態の保存に必要なメモリ。
2. アクティベーションメモリ: 中間アクティベーションの保存に必要なメモリ。アテンション機構では、シーケンス長に対して二次的に（$S^2$）スケーリングします。

既存の並列化戦略は、これらの問題を個別に、あるいはトレードオフを伴って解決しています。

• テンソル並列化（TP）: モデルの重みをデバイス間で分割（シャード）することでパラメータメモリを削減しますが、アクティベーションメモリは完全に複製されたままになります。高いアクティベーションメモリと通信量により、長いシーケンス長では非効率になります。
• シーケンス並列化（SP）: 入力シーケンスをデバイス間で分割することでアクティベーションメモリを削減しますが、モデルの重みは完全に複製されたままになります。パラメータメモリがボトルネックとなる場合には効果がありません。
• 従来の TP+SP: 両者を組み合わせ、デバイスメッシュの直交軸（例：$T \times \Sigma$ グリッド）に割り当てます。効果的ではありますが、モデル並列化に多数のデバイス（$T \times \Sigma$）を消費し、データ並列化に割り当てられるデバイスが減少します。さらに、並列グループが低速なノード間リンクにまたがることを余儀なくされ、NVLink や AMD Infinity Fabric のような高帯域幅のノード内相互接続が十分に活用されない傾向があります。

核心的な課題: 単一のデバイス軸上でパラメータメモリとアクティベーションメモリの両方を同時に削減し、メモリ効率を最大化してデータ並列化のためのデバイスを解放しつつ、過大な通信コストを回避する方法は何か。

2. 手法：テンソルおよびシーケンス並列化（TSP）

著者は、テンソル並列化とシーケンス並列化を単一のデバイス軸に「折りたたむ」戦略であるTSPを提案します。TP と SP を直交する次元に割り当てるのではなく、TSP は各ランク（GPU）にモデル重みのシャードとシーケンスのシャードの両方を割り当てます。

主要な設計要素:

• 統合シャディング: 各デバイスはモデルパラメータの $1/D$ とシーケンストークンの $1/D$ を保持します。これにより、単一の軸上でパラメータに比例するメモリとアクティベーションメモリの両方が $D$ 倍削減されます。
• ランタイムスケジューリング: 重みとアクティベーションが異なってシャードされているという事実に対処するため、TSP はレイヤーごとに特定の通信スケジューリングを採用します。
  • アテンションレイヤー: ブロードキャストスケジューリングを使用します。
    • ランクは重みシャード（パックされた $W_Q, W_K, W_V, W_O$）を反復処理します。
    • 各反復で、1 つのランクが自身の重みシャードを他のすべてのランクにブロードキャストします。
    • すべてのランクが自身のシーケンスシャード上で局所的な投影を計算します。
    • All-Gather操作により、シーケンス次元全体にわたって完全なキー/バリューコンテキストを再構築し、FlashAttention を実行します。
    • ジグザグ分割方式を用いて、因果アテンションのワークロードをバランスさせます（後続のトークンが不均衡に多くの作業を負担しないようにします）。
  • ゲート付き MLP レイヤー: リング通信スケジューリングを使用します。
    • 重みシャード（$W_1, W_2, W_3$）がリング周りでポイントツーポイントに循環します。
    • 各ランクは部分的な出力を局所的に蓄積します。
    • シーケンスは局所的なまま重みのみが移動するため、TP で標準的な All-Reduce 操作の必要性が排除されます。
• ハードウェア認識: この設計は、参加する GPU の数に応じて集合帯域幅がスケーリングする AMD の MI300X（Infinity Fabric）などのトポロジーに最適化されています。TSP は、通信プリミティブ（ブロードキャストとリング）がグループ全体を関与させるか、または計算とオーバーラップさせて効率を最大化するように設計されています。

3. 主要な貢献

1. 新規並列化戦略: 2 つの直交する並列化軸を 1 つに縮退させる TSP の導入により、パラメータメモリとアクティベーションメモリの同時削減を可能にしました。
2. 理論的解析: DP、TP、SP、TP+SP と TSP を比較する厳密な数学モデルの提示。
   • メモリ: $D^2$ 個のデバイスを必要とせずに、重みに比例するメモリとアクティベーションメモリの両方を $1/D$ 削減する唯一の方式が TSP であることを示しました。
   • 通信: TSP は TP+SP が回避する重み移動の通信を導入しますが、著者は長いシーケンスや中程度のバッチサイズの場合、アクティベーションに比例する通信が支配的となり、TSP が競争力を持つことを証明しました。
   • FLOPs: TSP は、TP や SP と同様のデバイスあたりの計算効率を維持します。
3. 実装とオーバーラップ: GEMM 計算の背後で重み転送をパイプライン化するなど、明示的な通信/計算オーバーラップ戦略を備えた TSP アテンションおよび MLP ブロックの詳細な実装。
4. トポロジー最適化: TSP により、モデルのレプリカ全体を高帯域幅の単一ノード（ノード内）に収容でき、2 次元 TP+SP メッシュでしばしば必要とされる低速なノード間リンクにまたがることによる性能ペナルティを回避できることを実証しました。

4. 結果

著者は、7B パラメータの密なトランスフォーマーモデルを用いて、最大1,024 個の AMD MI300X GPU（128 ノード）のクラスター上で TSP を評価しました。

• メモリ効率:
  • TSP は、テストされたすべてのシーケンス長（16K から 128K トークン）において、一貫して最小のピークメモリを達成しました。
  • 短いコンテキストでは、TSP は TP（パラメータメモリが支配的）と同等の性能を示しました。
  • 長いコンテキストでは、TSP は TP および TP+SP（アクティベーションメモリが支配的）を大幅に上回り、他の戦略が失敗する Out-Of-Memory（OOM）エラーを防止しました。
• スループット:
  • 重み移動による理論的な通信量の増加にもかかわらず、TSP は TP+SP と競争力のあるスループットを維持しました。
  • 追加の通信は、パイプライン化による計算の背後に隠蔽することに成功しました。
  • スループットの優位性は、並列度の増加（$D$）およびシーケンス長の増加に伴って拡大しました。
• スケーラビリティ:
  • TSP は、メモリフットプリントの削減によりより大きなマイクロバッチを可能にし、スループットをさらに向上させました。
  • モデルレプリカを単一ノード（8 GPU）内に完全に収容することを可能にし、大規模クラスターにおけるデータ並列化（DP）のためのランクを解放しました。

5. 意義

• ハードウェア認識型スケーリング: TSP は、ノード内帯域幅が高いがノード間帯域幅がボトルネックとなる現代の HPC クラスターに対する実用的な解決策を提供します。高価なノード内相互接続の利用率を最大化します。
• 長コンテキストの妥当性: モデルがより長いコンテキストウィンドウ（10 万トークン以上）へと移行するにつれ、アクティベーションメモリが主要なボトルネックとなります。TSP は、2 次元並列化の莫大なデバイスオーバーヘッドなしにこれに対処できるという点で、独自の位置を占めています。
• 柔軟性: TSP は、パイプライン並列化、エキスパート並列化、データ並列化などの他の並列化方式と直交的に組み合わせることができ、システム設計者が特定のハードウェアトポロジーに最適化するための柔軟な「軸」を提供します。
• コスト効率: モデル並列化に必要なデバイスの数を削減することで、TSP はより多くのデバイスをデータ並列化に使用することを可能にし、特定の学習ジョブに必要なクラスター全体のサイズを削減するか、有効なバッチサイズを増加させる可能性があります。

結論として、TSP は分散学習戦略における重要な進化を表しており、管理可能な通信量の増加と引き換えに、メモリ効率とハードウェア利用率の大幅な向上をもたらします。これは、次世代の長コンテキスト LLM の学習にとって不可欠なツールとなります。