MLKV : Têtes à clé-valeur multicouches pour un décodage efficace en mémoire des Transformers

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

June 13, 2024

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

cs.AI

Résumé

L'inférence auto-régressive des transformateurs bénéficie grandement de la mise en cache des paires Clé-Valeur (KV), mais peut entraîner d'importants goulots d'étranglement en mémoire à mesure que la taille du modèle, la taille du lot et la longueur des séquences augmentent à grande échelle. Nous introduisons le partage Multi-Couches Clé-Valeur (MLKV), une approche novatrice étendant le partage KV à travers les couches du transformateur pour réduire l'utilisation de la mémoire au-delà de ce qui était possible avec l'attention Multi-Requête (MQA) et l'attention par Groupes de Requêtes (GQA). Les évaluations sur divers benchmarks de traitement du langage naturel et métriques d'inférence utilisant des variantes de Pythia-160M réentraînées démontrent que MLKV réduit significativement l'utilisation de la mémoire avec une perte de performance minimale, réduisant la taille du cache KV jusqu'à un facteur de 6x par rapport à MQA. Ces résultats mettent en évidence le potentiel de MLKV pour un déploiement efficace des modèles de transformateurs à grande échelle. Nous fournissons le code à l'adresse suivante : 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