フラッシュアテンション：変圧器の効率を革新

</pre>
import torch

def standard_attention(Q, K, V):
# Q, K, V shape: (batch_size, seq_len, d_model)
d_k = K.size(-1)
scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))
attention_weights = torch.softmax(scores, dim=-1)
return torch.matmul(attention_weights, V)

この実装は単純ではあるが、前述の非効率性に悩まされている。 scores 形状が (batch_size, seq_len, seq_len) であるテンソルは、長いシーケンスでは法外に大きくなる可能性があります。

フラッシュアテンションを入力

フラッシュアテンション、 トリ・ダオと同僚によって紹介された 2022年の論文で発表されたFlash Attentionは、メモリ使用量を大幅に削減し、計算効率を向上させるコンピューティングアテンションのアプローチです。Flash Attentionの背後にある主要なアイデアは次のとおりです。

1. タイル: 大きなアテンション マトリックスを、高速オンチップ SRAM に収まる小さなタイルに分割します。
2. 再計算: 注意行列全体を保存する代わりに、後方パス中に必要に応じてその一部を再計算します。
3. IO 対応実装: GPU メモリ階層の異なるレベル間でのデータ移動を最小限に抑えるようにアルゴリズムを最適化します。

フラッシュ アテンション アルゴリズム

Flash Attention は、本質的に、アテンション メカニズムの計算方法を再考します。アテンション マトリックス全体を一度に計算するのではなく、最新の GPU のメモリ階層を活用してブロック単位で処理します。

アルゴリズムの概要は次のとおりです。

1. 入力: HBM (高帯域幅メモリ) 内の行列 Q、K、V およびサイズ M のオンチップ SRAM。
2. ブロック サイズは使用可能な SRAM に基づいて計算されます。
3. 出力行列 O と補助ベクトル l および m の初期化。
4. アルゴリズムは、入力マトリックスを SRAM に収まるようにブロックに分割します。
5. 2 つのネストされたループがこれらのブロックを処理します。
   • 外側のループはKブロックとVブロックをロードします
   • 内側のループはQブロックをロードし、計算を実行します。
6. オンチップ計算には、行列乗算、ソフトマックス、出力計算が含まれます。
7. 各ブロックの処理後、結果は HBM に書き戻されます。

このブロック単位の計算により、Flash Attention は正確な注意を計算しながら、メモリ フットプリントを大幅に小さく抑えることができます。

フラッシュアテンションの背後にある数学

Flash Attention を機能させるための鍵は、ブロック単位でソフトマックスを計算できる数学的なトリックです。この論文では、2 つの重要な公式が紹介されています。

1. ソフトマックス分解:
   softmax(x) = exp(x - m) / Σexp(x - m)

   ここで、m は x の最大値です。

2. ソフトマックス合併:
   softmax(x ∪ y) = softmax(softmax(x) * e^(m_x - m), softmax(y) * e^(m_y - m))

   ここで、m = max(m_x, m_y)

これらの式により、Flash Attention は各ブロックの部分的なソフトマックス結果を計算し、それらを正しく組み合わせて最終結果を得ることができます。

実装の詳細

Flash Attention の簡略化された実装を詳しく見て、その中核となる概念を説明しましょう。

import torch

def flash_attention(Q, K, V, block_size=256):
    batch_size, seq_len, d_model = Q.shape
    
    # Initialize output and running statistics
    O = torch.zeros_like(Q)
    L = torch.zeros((batch_size, seq_len, 1))
    M = torch.full((batch_size, seq_len, 1), float('-inf'))
    
    for i in range(0, seq_len, block_size):
        Q_block = Q[:, i:i+block_size, :]
        
        for j in range(0, seq_len, block_size):
            K_block = K[:, j:j+block_size, :]
            V_block = V[:, j:j+block_size, :]
            
            # Compute attention scores for this block
            S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_model ** 0.5)
            
            # Update running max
            M_new = torch.maximum(M[:, i:i+block_size], S_block.max(dim=-1, keepdim=True).values)
            
            # Compute exponentials
            exp_S = torch.exp(S_block - M_new)
            exp_M_diff = torch.exp(M[:, i:i+block_size] - M_new)
            
            # Update running sum
            L_new = exp_M_diff * L[:, i:i+block_size] + exp_S.sum(dim=-1, keepdim=True)
            
            # Compute output for this block
            O[:, i:i+block_size] = (
                exp_M_diff * O[:, i:i+block_size] +
                torch.matmul(exp_S, V_block)
            ) / L_new
            
            # Update running statistics
            L[:, i:i+block_size] = L_new
            M[:, i:i+block_size] = M_new
    
    return O

この実装は簡略化されていますが、Flash Attention の本質を捉えています。ブロック内の入力を処理し、実行中の統計 (M と L) を維持して、すべてのブロックにわたってソフトマックスを正しく計算します。

フラッシュアテンションの影響

Flash Attention の導入は、特に大規模な言語モデルや長いコンテキストのアプリケーションにおいて、機械学習の分野に大きな影響を与えました。主な利点は次のとおりです。

1. メモリ使用量の削減: Flash Attention は、メモリの複雑さを O(N^2) から O(N) に削減します。ここで、N はシーケンスの長さです。これにより、同じハードウェアではるかに長いシーケンスを処理できるようになります。
2. 速度の向上データの移動を最小限に抑え、GPU の計算能力をより有効に活用することで、Flash Attention は大幅な高速化を実現します。著者らは、標準実装と比較して GPT-3 のトレーニングが最大 2 倍高速であると報告しています。
3. 正確な計算: 他の注意最適化手法とは異なり、Flash Attention は近似値ではなく正確な注意を計算します。
4. 拡張性: メモリフットプリントが削減されたため、最大数百万のトークンまで、はるかに長いシーケンスにスケーリングできるようになります。

実世界への影響

Flash Attention の影響は学術研究だけにとどまりません。多くの人気の機械学習ライブラリやモデルに急速に採用されています。

• フェイストランスフォーマーを抱き締める: 人気の Transformers ライブラリには Flash Attention が統合されており、ユーザーはその利点を簡単に活用できます。
• GPT-4 以降: 確認はされていませんが、GPT-4 のような高度な言語モデルは、長いコンテキストを処理するために Flash Attention に似た手法を使用している可能性があるという推測があります。
• ロングコンテキストモデル: Flash Attention により、書籍全体や長いビデオを処理できるモデルなど、非常に長いコンテキストを処理できる新世代のモデルが可能になりました。

フラッシュ注目: 最近の動向

標準アテンションとフラッシュアテンション

フラッシュアテンション-2

オリジナルのFlash Attentionの成功を基に、同じチームが 2年にFlashAttention-2023を導入この更新バージョンでは、いくつかの改善が加えられています。

1. さらなる最適化FlashAttention-2 はさらに優れた GPU 使用率を実現し、A70 GPU の理論上のピーク FLOPS の最大 100% に達します。
2. 改良されたバックワードパス: 後方パスは前方パスとほぼ同じ速度になるように最適化されており、トレーニングの速度が大幅に向上します。
3. さまざまな注目バリエーションのサポート: FlashAttention-2 は、グループ化されたクエリ アテンションやマルチクエリ アテンションなど、さまざまなアテンション バリアントのサポートを拡張します。

フラッシュアテンション-3

2024年にリリースされたFlashAttention-3 この研究分野における最新の進歩を表しています。パフォーマンスをさらに向上させるためのいくつかの新しい手法が導入されています。

1. 非同期計算: 新しい GPU 命令の非同期性を活用して、さまざまな計算を重ね合わせます。
2. FP8 サポート: 低精度FP8演算を利用し、さらに高速な処理を実現します。
3. 一貫性のない処理: 低精度フォーマットを使用するときに量子化誤差を減らす手法。

以下は、FlashAttention-3 が非同期計算を活用する方法の簡略化された例です。

import torch
from torch.cuda.amp import autocast

def flash_attention_3(Q, K, V, block_size=256):
    with autocast(dtype=torch.float8):  # Using FP8 for computation
        # ... (similar to previous implementation)
        
        # Asynchronous computation example
        with torch.cuda.stream(torch.cuda.Stream()):
            # Compute GEMM asynchronously
            S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_model ** 0.5)
        
        # Meanwhile, on the default stream:
        # Prepare for softmax computation
        
        # Synchronize streams
        torch.cuda.synchronize()
        
        # Continue with softmax and output computation
        # ...

    return O

このコード スニペットは、FlashAttention-3 が非同期計算と FP8 精度を活用する方法を示しています。これは単純化された例であり、実際の実装ははるかに複雑でハードウェア固有になることに注意してください。

プロジェクトにフラッシュアテンションを実装する

独自のプロジェクトで Flash Attention を活用することに関心がある場合は、いくつかのオプションがあります。

1. 既存のライブラリを使用する: Hugging Face Transformers などの多くの人気ライブラリには現在、Flash Attention の実装が含まれています。最新バージョンに更新し、適切なフラグを有効にするだけで十分な場合があります。
2. カスタム実装: より高度な制御や特殊なユースケースの場合は、Flash Attention を自分で実装することをお勧めします。xformers ライブラリは、優れたリファレンス実装を提供します。
3. ハードウェア固有の最適化: 特定のハードウェア (NVIDIA H100 GPU など) を使用している場合は、パフォーマンスを最大限に高めるためにハードウェア固有の機能を活用することをお勧めします。

Hugging Face Transformers ライブラリで Flash Attention を使用する方法の例を次に示します。

from transformers import AutoModel, AutoConfig

# Enable Flash Attention
config = AutoConfig.from_pretrained("bert-base-uncased")
config.use_flash_attention = True

# Load model with Flash Attention
model = AutoModel.from_pretrained("bert-base-uncased", config=config)

# Use the model as usual
# ...

課題と今後の方向性

フラッシュ アテンションは注意メカニズムの効率性の向上に大きな進歩を遂げましたが、まだ課題と今後の研究領域が残っています。

1. ハードウェアの特異性現在の実装は、多くの場合、特定の GPU アーキテクチャ向けに最適化されています。これらの最適化をさまざまなハードウェアに一般化することは、依然として課題となっています。
2. 他の技術との統合: Flash Attention を、プルーニング、量子化、モデル圧縮などの他の最適化手法と組み合わせることは、活発に研究されている分野です。
3. 他のドメインへの拡張: Flash Attention は NLP で大きな成功を収めていますが、その利点をコンピューター ビジョンやマルチモーダル モデルなどの他の領域に拡張するための取り組みが現在も続いています。
4. 理論的理解: Flash Attention がなぜこれほどうまく機能するのかについての理論的理解を深めることで、さらに強力な最適化を実現できる可能性があります。

まとめ：

 Flash Attention は、GPU メモリ階層を巧みに活用し、数学的なトリックを採用することで、精度を犠牲にすることなく、速度とメモリ使用量の両方を大幅に改善します。

この記事で説明したように、Flash Attention の影響は単なる最適化手法をはるかに超えています。これにより、より強力で効率的なモデルの開発が可能になりました。