FlashAttention-2: より優れた並列処理と作業分割でさらに高速化したAttention
(crfm.stanford.edu)- GPT-4(32k)、MPT(65k)、Calude(100k) など、より長いコンテキストを持つ言語モデルが登場
- Transformer のコンテキスト長を拡張するのは、実行時およびメモリ要件が 4 乗で増加するため難しい
- 昨年リリースされた FlashAttention は、メモリ使用量を減らし、Attention の速度を向上させることで、さまざまな場面で利用されている
- リリース当時ですでに 2〜4 倍高速だったが、まだ改善の余地がある。最適化された行列積演算 (GEMM) と比べると依然として十分速くなく、理論上の最大 FLOPs/s の 25〜40% にとどまっていた (A100 GPU で最大 124 TFLOPs/s)
- FlashAttention-2 は前バージョンより 2 倍高速で、A100 GPU で最大 230 TFLOP/s の性能を提供
- GPT 形式の言語モデルの学習では最大 225 TFLOPS に到達した (モデル FLOP 利用率 72%)
- アルゴリズムを調整して non-matmul FLOPs を削減
- より優れた並列化と、各スレッドブロック内での作業分割方法の変更
- Head Dimensions の数を 128 から 256 に拡張
まだコメントはありません。