Dynamic Memory Compression: Retrofitting LLMs for Accelerated Inference

要約

トランスフォーマーは、大規模言語モデル (LLM) のバックボーンとして登場しました。
ただし、過去のトークンのキーと値の表現のキャッシュをメモリに保存する必要があるため、生成は依然として非効率的であり、そのサイズは入力シーケンスの長さとバッチ サイズに比例して変化します。
解決策として、推論時にオンラインでキーと値のキャッシュを圧縮する方法である動的メモリ圧縮 (DMC) を提案します。
最も重要なのは、モデルがさまざまなヘッドとレイヤーにさまざまな圧縮率を適用することを学習することです。
Llama 2 (7B、13B、70B) などの事前トレーニング済み LLM を DMC Transformer に改造し、NVIDIA H100 GPU での自動回帰推論中に最大 7 倍のスループット向上を達成しました。
DMC は、追加のパラメーターを追加することなく、元のデータのごくわずかな割合で事前トレーニングを継続することによって適用されます。
DMC は、最大 4 倍のキャッシュ圧縮により元のダウンストリーム パフォーマンスを維持し、アップトレーニングされたグループ化クエリ アテンション (GQA) およびキー値エビクション ポリシー (H$_2$O、TOVA) を上回るパフォーマンスを発揮します。
GQA と DMC を組み合わせて、複合的な利益を得ることができます。
したがって、DMC は、既存の LLM の KV キャッシュのドロップイン代替として機能し、特定のメモリ バジェット内でより長いコンテキストやより大きなバッチに適合できます。

要約(オリジナル)

Transformers have emerged as the backbone of large language models (LLMs). However, generation remains inefficient due to the need to store in memory a cache of key-value representations for past tokens, whose size scales linearly with the input sequence length and batch size. As a solution, we propose Dynamic Memory Compression (DMC), a method for online key-value cache compression at inference time. Most importantly, the model learns to apply different compression ratios in different heads and layers. We retrofit pre-trained LLMs such as Llama 2 (7B, 13B and 70B) into DMC Transformers, achieving up to 7x throughput increase during auto-regressive inference on an NVIDIA H100 GPU. DMC is applied via continued pre-training on a negligible percentage of the original data without adding any extra parameters. DMC preserves the original downstream performance with up to 4x cache compression, outperforming up-trained grouped-query attention (GQA) and key-value eviction policies (H$_2$O, TOVA). GQA and DMC can be even combined to obtain compounded gains. Hence, DMC can serve as a drop-in replacement for KV caching in existing LLMs to fit longer contexts and larger batches within any given memory budget.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google