MLKV: Cabezas de Clave-Valor Multi-Capa para la Decodificación Eficiente en Memoria de Transformers

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

June 13, 2024

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

cs.AI

Resumen

La inferencia autorregresiva de los transformadores se beneficia enormemente del almacenamiento en caché de Clave-Valor (KV), pero puede generar cuellos de botella significativos en la memoria a medida que el tamaño del modelo, el tamaño del lote y la longitud de la secuencia aumentan a gran escala. Presentamos el uso compartido de Clave-Valor Multi-Capa (MLKV), un enfoque novedoso que extiende el uso compartido de KV a través de las capas del transformador para reducir el uso de memoria más allá de lo posible con Atención Multi-Consulta (MQA) y Atención de Consulta Agrupada (GQA). Las evaluaciones en varios puntos de referencia de PLN y métricas de inferencia utilizando variantes ajustadas de Pythia-160M demuestran que MLKV reduce significativamente el uso de memoria con una pérdida mínima de rendimiento, disminuyendo el tamaño de la caché KV hasta un factor de 6x en comparación con MQA. Estos resultados destacan el potencial de MLKV para el despliegue eficiente de modelos de transformadores a gran escala. Proporcionamos el código en 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