FlashAttentionアルゴリズムの深い探求-パート3

フラッシュアテンションアルゴリズムの深い探求-パート3

Flash Attentionシリーズの第3部へようこそ！このセグメントでは、FlashAttention V1アルゴリズムの内部動作を詳しく解説し、そのコアコンセプトと原則を分解していきます。トピックについて初めて知った方や、GPUやFlashAttentionが高レベルでどのように機能するかを学びたい方は、このシリーズのGPUの理解とGPUアクセラレーションの進歩をチェックしてください。

まず、FlashAttentionの最適化と速度向上は主にGPUを対象としています。論文ではL1キャッシュやL2キャッシュについても言及されていますが、これらの最適化は基本的にGPUのパフォーマンスを中心に行われており、RAMやその他のメモリコンポーネントには関連していません。

早送り

典型的なGPUアーキテクチャでは、データはハードディスクに格納されますが、意味のある計算を行うためにはデータをRAMに移動する必要があります。そこから、データはさまざまなメモリ階層を経てGPUに到達します。FlashAttentionアルゴリズムは、現代のGPUのテンソルコアの能力を最大限に活用するように微調整されています。これは特に重要であり、GPT-3などのモデルのトレーニング中、テンソルコアは約50%の時間がアイドル状態であることがわかっています。

FlashAttentionは、主に2つの理由で注目されるアルゴリズムです：タイリングと再計算です。

• 見逃せない7つの機械学習アルゴリズム
• 人工知能の進歩：成均館大学の革新的なメモリシステム「Memoria」が長いシーケンスの複雑なタスクにおけるトランスフォーマーのパフォーマンスを向上させます
• シュナイダーエレクトリックは、SageMakerでのリトリーバルアグメントドLLMsを活用して、ERPシステムのリアルタイムの更新を確実にしています

タイリングは、Q、K、V行列をより小さなブロックに分割する技術です。この分割により、アルゴリズムはこれらの行列を一度にGPUのメモリに読み込むのではなく、ブロックごとに読み込みや処理を行うことができます。

再計算は、モデルのトレーニングの重要な側面である逆伝播に対応します。Flash Attentionは、高帯域幅メモリ（HBM）に値を格納し、このメモリに繰り返しアクセスするのではなく、必要な時に値を再計算します。再計算は浮動小数点演算の数を増やすことになりますが、メモリアクセスにかかる時間を大幅に削減します。

それでは、論文で言及されているアルゴリズムの詳細に入ってみましょう：

ビジュアルな説明をお好みですか？FlashAttention V1アルゴリズムのビデオをご覧ください

FlashAttentionアルゴリズムとSoftmaxのオンライン正規化計算

FlashAttentionアルゴリズム

タイリングコンセプト

FlashAttentionの最も重要なコンセプトの1つは、タイリングです。トランスフォーマーモデルの各トークンには、Q、K、Vの行列が関連付けられています。タイリングプロセスでは、これらの行列を処理しやすいブロックに分割します。ブロックサイズは、著者によって128で設定されることが一般的です。

まず、Q、K、V行列のブロックサイズを決定する必要があります。さらに、結果を格納するために「l」や「m」といった中間変数を初期化します。最終的な出力は、すべての中間変数の積です。これらの結果を効率的に組み合わせるために、ブロックに分割して中間結果を格納することが重要です。

セーフソフトマックス

FlashAttentionの中心には、Softmax関数の実装があります。通常のSoftmax関数は、指数値が極端に大きくなりすぎたり小さくなりすぎたりするというオーバーフローやアンダーフローの問題に直面することがあります。FlashAttentionでは、これらの問題を軽減するために「セーフソフトマックス」を採用しています。

セーフソフトマックスは、入力配列内の最大値を見つけ、各要素からこの最大値を減算した後に指数計算を行うことで動作します。この調整により、オーバーフローやアンダーフローの問題が回避され、計算が数値的に安定します。

セーフソフトマックスの式：

セーフソフトマックスとオンライン正規化計算

FlashAttentionのセーフソフトマックスは、NVIDIAの論文で説明されている「オンライン正規化計算」という概念から着想を得ています。この方法は、Softmaxのためのオンライン正規化計算というNVIDIAの論文で詳しく説明されており、冗長なメモリアクセスを行わずにSoftmaxを計算する方法を提供します。このアプローチは、メモリ操作の回数を減らすことでアルゴリズムをよりメモリ効率的にします。

この手法では、アルゴリズムは実行中の合計を保持し、計算中にオンザフライで計算された中間値を使用してSoftmaxの計算を修正します。これにより、以前の反復で行われた変更を元に戻し、新しい要素が処理される際に必要な調整を適用します。このアプローチにより、入力配列のすべての要素を再訪することなくSoftmaxの計算が可能となり、メモリアクセスを大幅に削減できます。

タイリングとセーフソフトマックスの組み合わせ

FlashAttentionは、タイリングとセーフソフトマックスの概念を巧みに組み合わせることで、アテンションメカニズムの効率を最大限に引き出します。入力を管理可能なブロックに分割し、ブロックレベルでセーフソフトマックスを適用することで、FlashAttentionはメモリアクセスを最小限に抑えつつ数値の安定性を保ちます。このアプローチにより、ディープラーニング/言語モデルのトレーニング中に頻繁にアンダーユーティライズされることが多いGPUテンソルコアとのシームレスな連携が実現されます。

結論

要約すると、FlashAttention V1は、GPUとテンソルコアのパワーを活用した高効率で数値的に安定したアテンションメカニズムです。タイリングとセーフソフトマックスの革新的な使用により、メモリアクセスのボトルネックが最小限に抑えられ、GPT-3などのLLMがより効果的にトレーニングできるようになります。

これらの技術の統合は、数学的なコンセプト、アルゴリズムの巧妙さ、ディープラーニングの領域におけるハードウェアの最適化との素晴らしい相乗効果を示しています。

参考文献

1. FlashAttention論文
2. 注釈付きFlashAttention論文
3. Softmaxのためのオンライン正規化計算​
4. NVIDIAディープラーニングGPUパフォーマンスガイド​

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