フラッシュアテンション：トランスフォーマーの効率性を革命的に変える

</pre>
import torch

<p>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)</p>

この実装は、上記で説明した非効率性の影響を受けます。scoresテンソルは、(batch_size, seq_len, seq_len)という形を持ち、長いシーケンスの場合、非常に大きなサイズになります。

フラッシュアテンションの登場

フラッシュアテンションは、Tri Daoと同僚によって2022年の論文で導入されました。これは、メモリ使用量を大幅に削減し、計算効率を大幅に改善するアテンションの計算方法です。フラッシュアテンションの背後にある重要なアイデアは以下のとおりです。

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

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

フラッシュアテンションは、基本的にアテンションメカニズムの計算方法を再考します。アテンション行列を一度に計算するのではなく、ブロックごとに処理し、現代のGPUのメモリ階層を活用します。

ここに、アルゴリズムの高レベルな概要が示されています。

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

このブロックごとの計算により、フラッシュアテンションは、正確なアテンションを計算しながら、メモリフットプリントを大幅に小さくすることができます。

フラッシュアテンションの数学的根拠

フラッシュアテンションを機能させる鍵は、ブロックごとにソフトマックスを計算する数学的なトリックです。論文では、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)です。

これらの式により、フラッシュアテンションは、各ブロックに対して部分的なソフトマックス結果を計算し、正しく結合して最終結果を取得することができます。

実装の詳細

ここでは、フラッシュアテンションの簡略化された実装を示し、その核心概念を説明します。

import torch

<p>def flash_attention(Q, K, V, block_size=256):
batch_size, seq_len, d_model = Q.shape</p>

<p># 出力とランニング統計の初期化
O = torch.zeros_like(Q)
L = torch.zeros((batch_size, seq_len, 1))
M = torch.full((batch_size, seq_len, 1), float(&#039;-inf&#039;))</p>

<p>for i in range(0, seq_len, block_size):
Q_block = Q[:, i:i+block_size, :]</p>

<p>for j in range(0, seq_len, block_size):
K_block = K[:, j:j+block_size, :]
V_block = V[:, j:j+block_size, :]</p>

<p># このブロックのアテンションスコアを計算
S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_model ** 0.5)</p>

<p># ランニング最大値の更新
M_new = torch.maximum(M[:, i:i+block_size], S_block.max(dim=-1, keepdim=True).values)</p>

<p># 指数関数の計算
exp_S = torch.exp(S_block - M_new)
exp_M_diff = torch.exp(M[:, i:i+block_size] - M_new)</p>

<p># ランニング合計の更新
L_new = exp_M_diff * L[:, i:i+block_size] + exp_S.sum(dim=-1, keepdim=True)</p>

<p># このブロックの出力を計算
O[:, i:i+block_size] = (
exp_M_diff * O[:, i:i+block_size] +
torch.matmul(exp_S, V_block)
) / L_new</p>

<p># ランニング統計の更新
L[:, i:i+block_size] = L_new
M[:, i:i+block_size] = M_new</p>

return O

この実装は、フラッシュアテンションの本質を捉えています。ブロックごとに処理し、ランニング統計（MとL）を保持して、ブロック全体でソフトマックスを正しく計算します。

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

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

1. メモリ使用量の削減：フラッシュアテンションは、メモリ複雑さをO(N^2)からO(N)に削減します。ここで、Nはシーケンスの長さです。これにより、同じハードウェアで、はるかに長いシーケンスを処理できます。
2. 速度の向上：データ移動を最小限に抑え、GPUのコンピュート能力をよりよく活用することで、フラッシュアテンションは、標準的な実装と比較して、最大3倍の高速化を達成します。
3. 正確な計算：他のアテンション最適化技術と異なり、フラッシュアテンションは、近似ではなく、正確なアテンションを計算します。
4. スケーラビリティ：メモリフットプリントの削減により、はるかに長いシーケンス、つまり数百万トークンにスケールアップできます。

現実世界への影響

フラッシュアテンションの影響は、研究の枠を超えています。多くの人気のある機械学習ライブラリやモデルで採用されています。

• ハギングフェイストランスフォーマー：トランスフォーマーライブラリでは、フラッシュアテンションが統合されており、ユーザーが簡単にその利点を利用できます。
• GPT-4 以降：GPT-4のような先進的な言語モデルが、フラッシュアテンションに似た技術を使用して長いコンテキストを処理している可能性がありますが、確認はされていません。
• 長いコンテキストモデル：フラッシュアテンションにより、全書や長い動画など、非常に長いコンテキストを処理できる新しいモデルが開発されました。

フラッシュアテンション：最新の開発

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

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

フラッシュアテンションの成功を受けて、同チームは2023年にフラッシュアテンション2を発表しました。この更新バージョンには、いくつかの改善が含まれています。

1. さらに最適化：フラッシュアテンション2は、A100 GPUで理論上のピークFLOPSの70%に達するなど、さらにGPUの活用を改善しています。
2. 後方パスの改善：後方パスは、前方パスとほぼ同等の速度で実行できるように最適化されています。これにより、トレーニングが大幅に高速化します。
3. さまざまなアテンションバリアントのサポート：フラッシュアテンション2では、グループ化クエリアテンションやマルチクエリアテンションなどのさまざまなアテンションバリアントがサポートされています。

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

2024年にリリースされたフラッシュアテンション3は、この研究ラインの最新の進歩を表しています。これには、パフォーマンスをさらに向上させるいくつかの新しいテクニックが導入されています。

1. 非同期計算：新しいGPU命令の非同期性を利用して、異なる計算を重ねることができます。
2. FP8サポート：低精度FP8計算を利用して、さらに高速な処理を実現します。
3. 非同期処理：低精度フォーマットを使用する際の量化誤差を削減するための手法です。

ここでは、フラッシュアテンション3が非同期計算をどのように活用するかを示す簡略化された例を見てみましょう。

import torch
from torch.cuda.amp import autocast

<p>def flash_attention_3(Q, K, V, block_size=256):
with autocast(dtype=torch.float8): # FP8を使用して計算
# ... (前と同様の実装)</p>

<p># 非同期計算の例
with torch.cuda.stream(torch.cuda.Stream()):
# GEMMを非同期に計算
S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_model ** 0.5)</p>

<p># その間、デフォルトストリームで：
# ソフトマックス計算の準備</p>

<p># ストリームの同期
torch.cuda.synchronize()</p>

<p># ソフトマックスと出力の計算を続ける
# ...</p>

return O

このコードスニペットは、フラッシュアテンション3が非同期計算とFP8精度をどのように活用するかを示しています。実際の実装は、ハードウェアに依存してはるかに複雑になります。

プロジェクトでのフラッシュアテンションの実装

フラッシュアテンションを自分のプロジェクトで利用したい場合は、以下の選択肢があります。

1. 既存のライブラリを使用：人気のあるライブラリの多く、例えばハギングフェイストランスフォーマーでは、フラッシュアテンションが実装されています。最新バージョンに更新し、適切なフラグを有効にするだけで十分です。
2. カスタム実装：より高度な制御や特殊なユースケースの場合、自分でフラッシュアテンションを実装することができます。xformersライブラリは、参考実装として役立ちます。
3. ハードウェア固有の最適化：特定のハードウェア（例：NVIDIA H100 GPU）で作業している場合、ハードウェア固有の機能を利用して最大のパフォーマンスを引き出します。

ここでは、ハギングフェイストランスフォーマーライブラリでフラッシュアテンションを使用する例を見てみましょう。

from transformers import AutoModel, AutoConfig

<p># フラッシュアテンションを有効にする
config = AutoConfig.from_pretrained(&quot;bert-base-uncased&quot;)
config.use_flash_attention = True</p>

<p># フラッシュアテンションを使用したモデルをロード
model = AutoModel.from_pretrained(&quot;bert-base-uncased&quot;, config=config)</p>

<p># 通常のモデルと同じように使用
# ...

課題と将来の方向性

フラッシュアテンションは、アテンションメカニズムの効率性を大幅に改善しましたが、まだ課題と将来の研究分野があります。

1. ハードウェア依存性：現在の実装は特定のGPUアーキテクチャに最適化されており、異なるハードウェアにわたる最適化を一般化することが課題です。
2. 他の技術との統合：フラッシュアテンションを、プルーニング、量化、モデル圧縮などの他の最適化技術と組み合わせることが、活発な研究分野です。
3. 他のドメインへの拡張：フラッシュアテンションはNLPで大きな成功を収めていますが、コンピュータビジョンやマルチモーダルモデルなどの他のドメインへの拡張も、進行中の取り組みです。
4. 理論的理解の深化：フラッシュアテンションがなぜこれほど効果的であるかについて、より深い理論的理解を得ることが、さらに強力な最適化につながる可能性があります。

結論

フラッシュアテンションは、GPUのメモリ階層を巧みに活用し、数学的なトリックを採用することで、速度とメモリ使用量の両面で大幅な改善を実現しています。

この記事で探究したように、フラッシュアテンションの影響は、単なる最適化技術を超えています。より強力で効率的なモデルを開発することを可能にしました。