[論文レビュー] Ring Attention with Blockwise Transformers for Near-Infinite Context

Hao Liu, Matei Zaharia+1 more|arXiv (Cornell University)|Oct 3, 2023

Topic Modeling被引用数 12

ひとこと要約

Ring Attentionは、リングトポロジーでのデバイス間通信とブロック単位の注意機構/FFN計算を重ね合わせることにより、デバイス数に応じて文脈長を拡張できるTransformerモデルの訓練と推論を可能にし、近似なしでほぼ無限の文脈を実現します。

ABSTRACT

Transformers have emerged as the architecture of choice for many state-of-the-art AI models, showcasing exceptional performance across a wide range of AI applications. However, the memory demands imposed by Transformers limit their ability to handle long sequences, thereby posing challenges in utilizing videos, actions, and other long-form sequences and modalities in complex environments. We present a novel approach, Ring Attention with Blockwise Transformers (Ring Attention), which leverages blockwise computation of self-attention and feedforward to distribute long sequences across multiple devices while fully overlapping the communication of key-value blocks with the computation of blockwise attention. Our approach enables training and inference of sequences that are up to device count times longer than those achievable by prior memory-efficient Transformers, without resorting to approximations or incurring additional communication and computation overheads. Extensive experiments on language modeling and reinforcement learning tasks demonstrate the effectiveness of our approach in allowing millions of tokens context size and improving performance.

研究の動機と目的

• 長い文脈を持つ Transformer のメモリボトルネックを動機づけ、それに対処する。
• 長いシーケンスをデバイス間に分散するためのリングベースのブロック単位計算方式を導入する。
• キー・バリューのブロック通信と計算の重ね合わせがオーバーヘッドを除去することを示す。
• 言語モデリングと強化学習タスクにおいて、百万単位のトークンとデバイス数の拡張性を示す。

提案手法

• ブロック単位の注意機構とフィードフォワード計算が、シーケンス長を複数のデバイスに分散する。
• リングトポロジがホストを調整する。各ホストはクエリブロックを処理し、キー・バリューブロックは次または前のホストへ回転する。
• キー・バリューブロックの通信とブロック単位計算を重ね合わせ、通信遅延を隠蔽する。
• ブロック単位並列トランスフォーマーを用いて、メモリコストをブロックサイズに対して線形に、シーケンス長に依存しないように保つ。
• Algorithm 1 は FSDPとリング通信を用いたリングベースの Transformer 訓練におけるメモリ削減手順を概説する。
• 実装は、メモリ効率の高いアテンションプリミティブと近似なしのブロック単位演算を活用する。

実験結果

リサーチクエスチョン

• RQ1 Ring Attentionは、デバイス数に対して線形にTransformerの文脈長をスケールさせつつ、性能を維持できるか？
• RQ2 リング上のデバイスにブロック単位の注意を分散させたときのメモリと計算のトレードオフは？
• RQ3 Ring Attentionは、異なるモデルサイズやハードウェア（GPU/TPU）におけるモデルFLOPs利用率とスループットにどう影響するか？
• RQ4 長い文脈を活用する下流タスク、例えば強化学習や長文脈言語モデリングでRing Attentionは改善をもたらすか？

主な発見

• Ring Attentionは、従来のメモリ効率的な手法よりデバイス数倍以上長いシーケンスの訓練を可能にする。
• 数百万トークンを超える文脈サイズが、近似や追加オーバーヘッドなしで達成可能。
• MFU（モデルFLOPs利用率）は、極端に長い文脈長でも高水準を維持し、いくつかのベースラインとは異なる。
• ExoRLのRL実験で、長い軌跡/文脈を使用すると、複数タスクで平均リターンがベースラインを上回る。
• 512K-token文脈でRing Attentionを用いたLLaMA-13Bのファインチューニングは、長い文脈のライン検索タスクで高い精度を維持し、いくつかの短い文脈ベースラインを上回る。
• ハードウェア（A100 GPUsとTPUs）を跨いで、Ring Attentionは、通常の/メモリ効率的なトランスフォーマーと比較して、最小限のオーバーヘッドで顕著な文脈長のスケーリングを示す。