A Unified Sequence Parallelism Approach for Long Context Generative AI

要約

入力テンソルのシーケンス次元を複数の計算デバイスに分割するシーケンス並列処理 (SP) は、生成 AI モデルのロングコンテキスト機能を解き放つ鍵となりつつあります。
この論文では、最先端の SP アプローチ、つまり DeepSpeed-Ulysses と Ring-Attendant を調査し、トランス モデル アーキテクチャとネットワーク ハードウェア トポロジに対してより堅牢な統合 SP アプローチを提案します。
このペーパーでは、SP と既存の並列処理 (データ/テンソル/ゼロ/エキスパート/パイプライン並列処理を含む) の通信コストとメモリ コストを比較し、SP を含むハイブリッド 4D 並列処理を設計するためのベスト プラクティスについて説明します。
LLAMA3-8B モデルのシーケンス長 208K の SP を使用して、2 つの 8xA800 ノードで 86\% MFU を達成しました。
私たちのコードは \url{https://github.com/feifeibear/long-context-attention} で公開されています。

要約(オリジナル)

Sequence parallelism (SP), which divides the sequence dimension of input tensors across multiple computational devices, is becoming key to unlocking the long-context capabilities of generative AI models. This paper investigates the state-of-the-art SP approaches, i.e. DeepSpeed-Ulysses and Ring-Attention, and proposes a unified SP approach, which is more robust to transformer model architectures and network hardware topology. This paper compares the communication and memory cost of SP and existing parallelism, including data/tensor/zero/expert/pipeline parallelism, and discusses the best practices for designing hybrid 4D parallelism involving SP. We achieved 86\% MFU on two 8xA800 nodes using SP for sequence length 208K for the LLAMA3-8B model. Our code is publicly available on \url{https://github.com/feifeibear/long-context-attention}.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google