Dynamic Memory Compression: Retrofitting LLMs for Accelerated Inference

要約

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

要約(オリジナル)

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 on-line key-value cache compression at inference time. Most importantly, the model learns to apply different compression rates in different heads and layers. We retrofit pre-trained LLMs such as Llama 2 (7B, 13B and 70B) into DMC Transformers, achieving up to ~3.7x throughput increase in auto-regressive inference on a NVIDIA H100 GPU. DMC is applied via continued pre-training on a negligible percentage of the original data without adding any extra parameters. We find that DMC preserves the original downstream performance with up to 4x cache compression, outperforming up-trained grouped-query attention (GQA). GQA and DMC can be even combined to obtain compounded gains. As a result DMC fits longer contexts and larger batches within any given memory budget.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google