スタンフォード研究所がFlashAttention-2を発表：長い文脈の言語モデルの速度と効率の飛躍

Stanford Research Institute announces FlashAttention-2 a leap in the speed and efficiency of language models with long context.

過去の1年間、自然言語処理は著しい進歩を遂げており、長い文脈を備えた言語モデルが登場しました。これらのモデルには、32kの文脈長を持つGPT-4、65kの文脈を持つMosaicMLのMPT、そして驚異的な100kの文脈長を誇るAnthropicのClaudeなどがあります。長いドキュメントのクエリやストーリー作成などのアプリケーションが成長し続ける中、拡張された文脈を持つ言語モデルの需要が明らかになってきました。ただし、課題は、トランスフォーマーの文脈長を拡大することであり、その注意層は入力シーケンスの長さと二乗的に増加する計算とメモリの要件があります。

この課題に取り組むために、わずか1年前にリリースされた革新的なアルゴリズムであるFlashAttentionは、さまざまな組織や研究所で急速に採用されました。このアルゴリズムは、精度を損なうことなく注意の計算を加速し、そのメモリの使用量を削減することに成功しました。初期リリース時に最適化されたベースラインと比較して2〜4倍高速なパフォーマンスを発揮したFlashAttentionは、画期的な進歩を証明しました。しかし、A100 GPUで最大124 TFLOPs/sを達成した高速最適化マトリックス乗算（GEMM）操作には及びませんでした。

次なる飛躍を遂げたFlashAttentionの開発者は、FlashAttention-2という前作を大幅に上回る再発明版を発表しました。NvidiaのCUTLASS 3.xとCuTeコアライブラリを活用したFlashAttention-2は、A100 GPU上で最大230 TFLOPs/sの驚異的なスピードアップを実現しています。さらに、GPTスタイルの言語モデルのエンドツーエンドトレーニングでは、FlashAttention-2は最大225 TFLOPs/sのトレーニング速度を実現し、驚異的な72%のモデルFLOP利用率を達成しています。

• 「Google DeepMindの最新研究、ICML 2023にて」
• GoogleのAI研究者は、HyperDreamBoothを紹介しましたこれは、人の単一の画像から個別の重みを効率的に生成するAIアプローチであり、DreamBoothよりも小さく、25倍高速です
• KAISTとGoogleの研究者は、コラボレーションスコア蒸留（CSD）を導入しましたこれは、一貫した視覚合成のためのテキストから画像への拡散モデルの単数形を拡張したAIの手法です

FlashAttention-2の主な改良点は、より優れた並列性と作業の分割戦略にあります。最初に、FlashAttentionはバッチサイズとヘッド数を並列化し、GPU上の計算リソースを効果的に活用しました。しかし、バッチサイズが小さいか、ヘッド数が少ない長いシーケンスの場合、FlashAttention-2はシーケンス長の次元で並列化するようになり、これによりこれらのシナリオで大幅なスピードアップが実現されました。

もう一つの改善点は、各スレッドブロック内の異なるワープ間での効率的な作業の分割です。FlashAttentionでは、KとVを4つのワープに分割し、Qをすべてのワープでアクセス可能な状態に保つ「sliced-K」スキームは、不要な共有メモリの読み書きを引き起こし、計算を遅くしました。FlashAttention-2は異なるアプローチを取り、KとVをすべてのワープでアクセス可能な状態に保ちながら、Qを4つのワープに分割するようにしました。これにより、ワープ間の通信が不要となり、共有メモリの読み書きが大幅に削減され、パフォーマンスがさらに向上しました。

FlashAttention-2は、その適用範囲を広げ、機能を向上させるためにいくつかの新機能を導入しています。最大256のヘッド次元をサポートし、GPT-J、CodeGen、CodeGen2、およびStableDiffusion 1.xなどのモデルを収容できるようになり、より高速化とメモリの節約の機会が広がりました。さらに、FlashAttention-2はマルチクエリアテンション（MQA）およびグループ化クエリアテンション（GQA）のバリアントに対応し、クエリの複数のヘッドがキーと値の同じヘッドにアテンドできるようになり、推論のスループットとパフォーマンスが向上します。

FlashAttention-2のパフォーマンスは本当に印象的です。A100 80GB SXM4 GPUでベンチマークを行った結果、前作と比較して約2倍の高速化を実現し、PyTorchの標準的なアテンション実装と比較して最大9倍の高速化を達成しました。さらに、GPTスタイルのモデルのエンドツーエンドトレーニングに使用すると、FlashAttention-2は既に高度に最適化されたモデルに比べて1.3倍のエンドツーエンドの高速化を実現します。

FlashAttention-2の将来の応用は非常に有望です。前作の8k文脈モデルと同じ価格で16kより長い文脈のモデルをトレーニングできる能力により、この技術は長い本、レポート、高解像度の画像、音声、ビデオの分析に役立つことができます。H100 GPUやAMD GPUなどのデバイスへの広範な適用と、fp8などの新しいデータタイプへの最適化の計画も進行中です。さらに、FlashAttention-2の低レベルの最適化を高レベルのアルゴリズムの変更と組み合わせることで、前例のない長い文脈でのAIモデルのトレーニングの可能性が開かれるかもしれません。プログラム性を向上させるためのコンパイラ研究者との協力も展望されており、次世代の言語モデルに明るい未来が約束されています。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!