- FlashAttention-2は、FlashAttentionの計算結果を近似せず、GPU内の仕事の分担を変えて注意演算を高速化する。
- 論文のA100上の演算ベンチマークでは、前版に比べておおむね2倍の速度を報告する。
- 注意演算だけの倍率と、モデル学習全体の倍率は異なる。GPUや系列長が変われば結果も変わる。
何を確かめた研究か
Transformerの注意計算は、入力が長くなるほど計算とメモリの負担が増える。従来のFlashAttentionは、GPUの高速だが小さいメモリと、大きいが移動に時間のかかるメモリの間を行き来する回数を減らした。注意の答えを近似せずに、読み書きの順序を変えて効率化した方法だ。
それでもGPUの行列積の能力を十分に使い切れていなかった。FlashAttention-2の論文は、前版で残った負担を、GPU上のブロックとワープへの仕事の割り振りから分析した。単にメモリ使用量を減らすだけでなく、利用できる計算器を働かせ、不要な共有メモリの読み書きを減らすことが目的である。
どう検証したか
FlashAttention-2は、行列積以外の演算を減らし、系列の方向にも処理を分割する。さらに、一つのブロック内のワープ間での分担を見直し、共有メモリを介する通信を減らす。論文は、この変更を前版やPyTorch、xFormersの注意計算と比較した。
演算単体の測定はNVIDIA A100 80GBで行い、因果マスクの有無や注意ヘッドの大きさを変えた。加えて、1.3Bまたは2.7BパラメータのGPT型モデルを8台のA100で学習し、学習全体の処理量も測った。この二つの測定は対象が違うため、数字を混ぜずに読む必要がある。
何が分かったか
注意演算の比較では、FlashAttention-2は前版のおおむね2倍の速度を示した。論文はA100の理論上の演算能力に対して、設定により50〜73%に達すると報告する。これはGPUの特定の演算で到達できる理論値との比であり、アプリケーション全体のGPU利用率を示す値ではない。
学習全体の表1では、2.7Bモデルで8,000トークンの設定において、A100 1台当たりの処理量はFlashAttentionの175 TFLOPs/sからFlashAttention-2の225 TFLOPs/sへ増えた。同じ表でFlashAttentionを使わない基準は80 TFLOPs/sである。注意演算単体では約2倍でも、この設定の学習全体では前版比約1.3倍となる。改善箇所以外の処理にも時間がかかるためだ。
限界は何か
論文第1版の主な比較はA100上で行われた。H100でも既存の実装をそのまま動かした測定を示す一方、H100固有の新しい命令を使う最適化、AMD製GPU、新しいデータ型への対応を今後の課題として挙げている。異なるGPUやモデル、短い系列、推論時の実サービスで同じ倍率が出ると示したものではない。
また、長文脈のモデルを作りやすくすることと、そのモデルが長い文書を正しく理解することは別だ。この論文は注意演算と学習処理量を調べた研究であり、長い入力での回答品質を直接測った研究ではない。演算の効率改善を、品質向上と読み替えないことが重要だ。
仕事との接点
長い文書を扱うモデルの計算基盤を比べるときは、まず使用するGPU、系列長、バッチサイズを固定し、注意演算の速度とエンドツーエンドの学習・推論時間を別々に測る。さらに回答品質も独立に評価する。これは論文の測定条件を実務の比較に移したTRAIL編集部の提案であり、特定サービスの改善率を論文が保証したものではない。