スタンフォード大学の博士によるソロ作品！大規模モデルのトレーニング速度が再び2倍になり、彼はスタースタートアップの主任科学者として参加することを正式に発表した。

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

既存の大規模言語モデルのトレーニングと推論の速度はさらに速くなる可能性があります——

どれくらい速くなりますか? 2〜4回。

さまざまな大規模モデルで使用されているFlashAttention が本日、第 2 世代を正式にリリースし、オープンソース化しました。Transformerアーキテクチャのすべてのモデルでこれを使用して高速化できます。

写真

第一世代の方法は昨年6月にリリースされ、近似なしで注目度を高速化し、メモリ使用量を削減します。

現在、FlashAttention-2 はさらにアップグレードされ、コアアテンション操作が2 倍高速化され、Transformer のエンドツーエンドのトレーニングが 1.3 倍高速化され、NVIDIA A100 でトレーニングすると72% のモデル FLOP 使用率を達成できます(一般的なモデルは約 50%) 。

写真

大規模な言語モデルのトレーニングコストが数千万ドルにも達することを考えると、FlashAttention-2 の一連の操作により、直接的に数百万ドルを節約できます。

ネットユーザーたちは衝撃を受け、 （犬の頭）と罵倒した。

写真

現在、このプロジェクトは GitHub で 4.4k 個のスターを獲得しています。

同時に、著者の一人がスタンフォード大学で博士号を取得し、大手モデルスタートアップのTogether AIに入社したことにも気付きました。

具体的な実装

FlashAttention の第一世代は、アテンション計算を並べ替えるアルゴリズムであると報告されています。タイリングなどの古典的な方法を使用して計算を大幅に高速化し、シーケンス長のメモリ使用量を 2 次から 1 次まで削減します。

タイリング方式とは、入力ブロックを HBM (GPU メモリ)から SRAM (高速キャッシュ)にロードし、ブロックに対してアテンション操作を実行してから、HBM で出力を更新することを指します。

HBM の繰り返しの読み取りと書き込みが、最大のパフォーマンスボトルネックになっています。

写真

このアプローチにより、FlashAttention は HBM への大きな中間アテンション マトリックスの書き込みを回避することでメモリの読み取り/書き込みの量を削減し、クロック時間が 2 ～ 4 倍高速化されます。

ただし、このアルゴリズムにはまだいくつかの非効率性があり、最適化された行列乗算(GEMM)演算ほど高速ではなく、最終的には理論上の最大 FLOPs/s (たとえば、A100 では最大 124 TFLOPs/s)の 25 ～ 40% にしか達しません。

その理由は、異なるスレッド ブロック間の作業と GPU 上のラップ分割が理想的ではないためです。

ここで、FlashAttention-2 は3 つの改善を行います。

まず、基本アルゴリズムでは、非行列乗算（行列乗算）を削減します。   FLOP の数。

理由の 1 つは、最新の GPU の特殊な計算ユニットにより、matmul がはるかに高速になることです。たとえば、A100 上の FP16/BF16 matmul の理論上の最大スループットは312 TFLOPs/sですが、非 matmul FP32 の理論上のスループットは19.5 TFLOPs/sしかありません。

もう 1 つの理由は価格の考慮です。結局のところ、非マットマルチ FLOP はマットマルチ FLOP より 16 倍高価です。同時に、matmul FLOP にできるだけ多くの時間を費やすことで、スループットを高く保ちます。

この目的のために、著者らは FlashAttention のソフトマックストリックを書き直し、出力を変更せずに、再スケーリング操作、境界チェック、因果マスキング操作の数を削減しました。

2 番目に、バッチ サイズが小さい場合に、より高い占有率を実現するために並列化が使用されます。

FlashAttention の生成は、バッチ サイズとアテンション ヘッドの数に基づいて並列化されます。

1 つのアテンション ヘッドを処理するために 1 つのスレッド ブロックを使用するため、合計で (batch_size*アテンション ヘッドの数) 個のスレッド ブロックがあり、各スレッド ブロックはストリーミング マルチプロセッサ(SM)上で実行されるようにスケジュールされます。

A100 のような 108 SM プロセッサで動作する場合、スレッド ブロックの数が大きい場合 (たとえば 80 以上) は、このスケジュールが適切に機能します。

シーケンスが長い場合、つまりバッチ サイズとヘッドの数が少ない場合は、GPU 上のマルチプロセッサをより有効に活用するために、シーケンスの長さの次元を並列化する必要があります。

この改善は、FlashAttention-2 が大幅に高速化された主な理由でもあります。

最後に、作業の分割を改善します。

スレッド ブロック内で、異なる経糸間で作業をどのように分割するかを決定する必要があります。通常、ブロックごとに 4 または 8 個のワープが使用されます。現在、著者はこのアプローチを改良して、異なるワープ間の同期と通信の量を減らし、共有メモリの読み取りおよび書き込み操作を削減しています。

下の左の図に示すように、第 1 世代の FlashAttention のアプローチは、Q をすべてのワープからアクセス可能なままにしながら、K と V を 4 つのワープに分割することです。この結果、すべてのワープは中間結果を共有メモリに書き込み、同期して中間結果を合計する必要があり、これは非常に非効率的で、FlashAttention での前方伝播が遅くなります。

写真

FlashAttention-2 では、すべてのワープが K と V にアクセスできるようにしながら、Q を 4 つのワープに分割します。

各ワープが行列乗算を実行して QK^T のスライスを取得した後、対応する出力を取得するには、それを V の共有スライスと乗算するだけで済みます。つまり、ワープ同士が通信する必要がないため、共有メモリの読み取りおよび書き込み操作が大幅に減少し、速度が向上します。

これら 3 つの主要な改善点に加えて、FlashAttention-2 には 2 つの小さな変更点があります。
まず、アテンション ヘッドの数が 128 から 256 に増加しました。これは、GPT-J、CodeGen、CodeGen2、StableDiffusion 1.x などのモデルが高速化とメモリ節約のために FlashAttention-2 を使用できることを意味します。

2 つ目は、マルチクエリ アテンション(MQA)とグループ化クエリ アテンション(GQA)をサポートすることです。

実験的評価

著者らは、A100 80GB SXM4 GPU 上でさまざまな構成(因果マスクありまたはなし、ヘッド数 64 または 128)で実行時間を測定しました。

消す：

FlashAttention-2 は FlashAttention (xformers ライブラリと Triton の他の実装を含む)よりも約 2 倍高速です。つまり、以前の 8k コンテキスト モデルをトレーニングするのと同じコストで 16k コンテキストのモデルをトレーニングできます(つまり、モデルのコンテキストの長さが 2 倍になります) 。

PyTorch の標準的なアテンション実装と比較すると、FlashAttention-2 は最大 9 倍高速です。

写真

さらに、FlashAttention-2 を使用すると、同じ実装を H100 GPU で実行するだけで、最大 335TFLOPs/s のトレーニング速度を実現できます(TMA や第 4 世代 Tensor コアなどの新しいハードウェア機能を活用するための特別な命令は使用しません) 。

写真

また、GPT スタイル モデルのエンドツーエンドのトレーニングに使用すると、FlashAttention-2 は A100 で最大 225TFLOPs/s の速度を達成することもできます(モデル FLOP 使用率は 72% に達します) 。これは、すでに高度に最適化されている FlashAttention よりも 1.3 倍高速です。

写真

第一著者が大規模モデルスタートアップに参加

FlashAttention-2 論文には、著者としてTri Daoのみが示されています。彼はまた、『FlashAttention Generation』の共著者 2 人のうちの 1 人です。

写真

Tri Dao 氏の研究分野は機械学習とシステムの交差点であると理解されており、昨年は ICML 2022 Outstanding Paper Runner-up Award を受賞しました。

彼は最近、スタンフォード大学でコンピュータサイエンスの博士号を取得し、プリンストン大学の助教授に就任する予定で、生成AIのスタートアップ企業Together AI （オープンソースモデルの実行、トレーニング、微調整のためのクラウドプラットフォームの構築を主な目標とする）に主任科学者として参加することを発表した。

もう一つ

最後に、一部のネットユーザーは、FlashAttention-2 に加えて、DeepSpeed の ZeRO++ やマサチューセッツ大学の de ReLoRA など、最近同様の成果が相次いでいることを発見しました。

これらはすべて、大規模モデルの事前トレーニングと微調整を加速するために使用されます。これらの研究結果から、彼は次のように感じています。

将来的には、低 VRAM と低帯域幅のコンシューマー グラフィック カードで大規模なモデルをトレーニングすることも夢ではなくなるでしょう。

写真

どう思いますか？

論文アドレス: https://tridao.me/publications/flash2/flash2.pdf
ブログアドレス: https://princeton-nlp.github.io/flash-atttention-2/
GitHub ホームページ: https://github.com/Dao-AILab/flash-attention