MLKV: Многослойные головы ключ-значение для эффективного использования памяти в декодировании трансформера

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

June 13, 2024

Авторы: Zayd Muhammad Kawakibi Zuhri, Muhammad Farid Adilazuarda, Ayu Purwarianti, Alham Fikri Aji

cs.AI

Аннотация

Авторегрессивный вывод трансформеров значительно выигрывает от кэширования ключей-значений (KV), но может привести к серьезным узким местам памяти по мере увеличения размера модели, размера пакета и длины последовательности на масштабе. Мы представляем концепцию многослойного совместного использования ключей-значений (MLKV), новый подход, расширяющий совместное использование KV на протяжении слоев трансформера для снижения использования памяти за пределами того, что было возможно с множественным запросом внимания (MQA) и групповым запросом внимания (GQA). Оценки на различных бенчмарках NLP и метриках вывода с использованием вариантов Pythia-160M с обучением показывают, что MLKV значительно снижает использование памяти с минимальной потерей производительности, уменьшая размер кэша KV в 6 раз по сравнению с MQA. Эти результаты подчеркивают потенциал MLKV для эффективного развертывания моделей трансформера в масштабе. Мы предоставляем код на https://github.com/zaydzuhri/pythia-mlkv

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