なぜアテンションメモリはホームAIモデルのパラメーター数を上回って増大するのか？

エヴァ・ウォン

IceWhale author

エヴァ・ウォン は テクニカルライター であり ZimaSpaceの常駐ティンカーでもあります。 生涯のオタクであり、 ホームラボとオープンソースソフトウェアに情熱を持っています。彼女は複雑な技術的概念をわかりやすく、 実践的なガイドに翻訳することを専門としています。エヴァはセルフホスティングは楽しくあるべきで、怖がるものではないと信じています。彼女のチュートリアルを通じて、コミュニティが ハードウェアのセットアップを解明する手助けをしています。初めてのNAS構築からDockerコンテナの習得まで。

Why Does Attention Memory Outgrow Home AI Model Parameters? - Zima Store Online

アテンションに必要なメモリは、モデルのパラメータ数を上回ることがあります。重みは固定されたままですが、トークン数に依存するキャッシュ、アクティベーション、ワークスペースはコンテキスト長や同時実行数に応じて増えるためです。

量子化モデルはホームサーバーのRAMやVRAMに余裕を持って収まる場合があります。しかし、長いプロンプト、家庭内の複数ユーザー、マルチモーダル入力、大きなバッチサイズによって、メモリ不足が発生することがあります。モデルファイルが示すのは永続的な重みであり、推論に必要なワーキングセット全体ではありません。アテンションは保持された各トークンについてリクエスト固有の状態を作成し、ランタイムによってはピーク使用量が異なる一時バッファも必要になります。以下では、固定されたパラメータメモリとシーケンス依存のメモリを分け、後者が実際の容量上限になるケースを説明します。

モデルパラメータは読み込み後に固定された基準値になる

モデルの重みは、パラメータ数と数値形式が分かれば、必要なメモリ容量を予測できます。4ビットモデルは8ビット版や浮動小数点版よりパラメータ保存に必要な容量が少なくなりますが、ランタイムのメタデータやスケールによって多少のオーバーヘッドが加わります。

Transformerアーキテクチャは、短いプロンプトでも長いプロンプトでも同じ学習済みパラメータを使用します。ユーザーがコンテキストを追加したからといって、パラメータの使用量が単純に2倍になるわけではありません。

この固定された基準値があるため、モデルファイルのサイズは初期適合性を確認するうえで役立ちます。ただし、推論時のピークメモリを完全に見積もることはできません。

フルアテンションではシーケンス長の二乗に比例して中間データが増える

従来の自己アテンションは、トークン位置同士を比較します。実装が大きなアテンションスコア行列や確率行列をメモリ上に展開すると、そのサイズはシーケンス長の二乗に比例して増加します。

FlashAttentionは、二次的に増加するアテンションメモリを長いシーケンスにおける主要な問題として捉え、アテンションをタイル単位で計算することで行列全体の保存を回避します。

そのため、最新の最適化された推論カーネルは、単純な計算式が示すよりもはるかに少ない一時メモリで動作する場合があります。長いシーケンスではアテンション処理そのものが依然として難しくなりますが、完全なシーケンス長の二乗行列がデバイスメモリに現れるかどうかは実装方法によって決まります。

ピーク時のワークスペースは、カーネルのバージョン、バッチ形状、ヘッド次元、ランタイムがより効率の低いアテンション処理へフォールバックするかどうかによっても変化します。

KVキャッシュは保持するすべてのトークンに永続的な状態を追加する

自己回帰デコードでは、次のトークンを生成するたびにプレフィックス全体を再計算しなくて済むよう、過去のトークンから得たキーとバリューを保存します。

PagedAttentionは、KVキャッシュの増加を、推論サービスにおける主要なメモリ制約として扱います。その容量は、保持するトークン数、キャッシュを生成するレイヤー数、キー・バリューの次元数、精度、アクティブなシーケンス数に応じて増加します。

一時的なアテンション行列とは異なり、この状態はアクティブな会話の間、利用可能な状態で保持し続ける必要があります。そのため、モデルが一度に1トークンしか生成していなくても、長いコンテキストによってメモリ消費が増え続けることがあります。

ZimaSpaceのAIメモリの余裕に関する解説では、モデルの保存に必要な容量と、コンテキストや同時利用ユーザーに必要な追加容量を分けて説明しています。

-15% OFF

ZimaBoard 2 - あなたの大きなアイデアのためのミニホームサーバー

Single board computer zimaboard2

同時実行リクエストは動的なアテンション状態を倍増させる

複数のユーザーでモデルの重み1つを共有できますが、各ユーザーのプロンプト、生成トークン、KVキャッシュの分岐は通常、それぞれ別に保持されます。

vAttentionは、サービス開始時点ではリクエストの長さや完了時間が分からないため、動的な物理メモリ割り当てを使用します。

32,000トークンの会話を1件処理できるサーバーでも、同じ規模の会話を4件同時に処理できるとは限りません。モデルのパラメータ総量が変わらなくても、バッチサイズとユーザー数によってトークン依存の状態は増加します。

リクエストのキャッシュ済みプレフィックスが完全に同一であれば、プレフィックス共有によって重複を減らせます。しかし、家庭内でそれぞれ異なる会話を行う場合、続きの状態は独立して保持する必要があります。

アクティベーションとランタイムの予約領域がピーク使用量をさらに押し上げる

プロンプトのプリフィル、マルチモーダル投影、投機的デコード、グラフキャプチャ、一時的なテンソル変換、ライブラリのワークスペースによって、重みとKVキャッシュ以外にもメモリが割り当てられることがあります。

FlashAttention-2は、より優れたカーネルによって大規模な中間データの展開を削減しても、アテンションが依然として長いシーケンスのボトルネックであると指摘しています。

フレームワークのキャッシュアロケーターは、解放されたブロックを再利用するために保持することがあります。そのため、ピーク使用量に達したリクエストが完了した後も、デバイスのツールにはプロセスが大きなメモリ領域を使用しているように表示される場合があります。この予約領域は稼働中のテンソルメモリとは異なりますが、それでも別のプロセスが利用できる容量を制限します。

したがって、観測される最大メモリ使用量は、安定した1トークンずつのデコード中ではなく、プリフィル中やモデルの切り替え時に発生することがあります。

アテンション最適化は境界を移動させるが、なくすわけではない

Flash Attentionは一時的なIOと行列保存を削減し、ページ単位の割り当てはKVの断片化を抑え、キャッシュの低精度化はトークンあたりのバイト数を減らし、Grouped-Query Attentionはキー・バリューヘッドの数を減らします。

Grouped-Query Attentionは、単一のマルチクエリヘッドよりも高い容量を維持しながら、キー・バリューヘッドのメモリ使用量を削減します。

スライディングウィンドウ、キャッシュの削除、オフロード、検索によってアテンション状態の上限を設定したり配置を移動したりできますが、それぞれレイテンシ、利用可能なコンテキスト、回答の挙動に影響します。

モデルの精度、実際のプロンプト長、出力上限、バッチサイズ、ユーザー数、画像トークン、その他のローカルサービスを含め、想定するワークロード全体でベンチマークを実施してください。モデルの重みを変更せず、トークン状態や同時実行数を減らすことで安定性が回復するなら、アテンションメモリがパラメータメモリを上回っているといえます。