MLKV: Mehrschichtige Schlüssel-Wert-Köpfe für speicher-effizientes Transformer-Decodieren

MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding

June 13, 2024

Autoren: Zayd Muhammad Kawakibi Zuhri, Muhammad Farid Adilazuarda, Ayu Purwarianti, Alham Fikri Aji

cs.AI

Zusammenfassung

Die autoregressive Inferenz von Transformatoren profitiert erheblich von Key-Value (KV)-Caching, kann jedoch bei wachsender Modellgröße, Batchgröße und Sequenzlänge zu erheblichen Speicherengpässen führen. Wir stellen Multi-Layer Key-Value (MLKV)-Sharing vor, einen innovativen Ansatz, der das KV-Sharing über Transformer-Schichten hinweg erweitert, um den Speicherverbrauch über das hinaus zu reduzieren, was mit Multi-Query Attention (MQA) und Grouped-Query Attention (GQA) möglich war. Evaluationen an verschiedenen NLP-Benchmarks und Inferenzmetriken unter Verwendung von trainierten Pythia-160M-Varianten zeigen, dass MLKV den Speicherverbrauch signifikant reduziert, bei minimalem Leistungsverlust, wodurch die KV-Cache-Größe im Vergleich zu MQA um den Faktor 6x verringert wird. Diese Ergebnisse unterstreichen das Potenzial von MLKV für eine effiziente Bereitstellung von Transformer-Modellen im großen Maßstab. Wir stellen den Code unter https://github.com/zaydzuhri/pythia-mlkv zur Verfügung.

English

Auto-regressive inference of transformers benefit greatly from Key-Value (KV) caching, but can lead to major memory bottlenecks as model size, batch size, and sequence length grow at scale. We introduce Multi-Layer Key-Value (MLKV) sharing, a novel approach extending KV sharing across transformer layers to reduce memory usage beyond what was possible with Multi-Query Attention (MQA) and Grouped-Query Attention (GQA). Evaluations on various NLP benchmarks and inference metrics using uptrained Pythia-160M variants demonstrate that MLKV significantly reduces memory usage with minimal performance loss, reducing KV cache size down to a factor of 6x compared to MQA. These results highlight MLKV's potential for efficient deployment of transformer models at scale. We provide code at https://github.com/zaydzuhri/pythia-mlkv