Cache KV condensé par couche pour l'inférence efficace des grands modèles de langage

Layer-Condensed KV Cache for Efficient Inference of Large Language Models

May 17, 2024

papers.authors: Haoyi Wu, Kewei Tu

cs.AI

papers.abstract

La consommation mémoire élevée constitue un goulot d'étranglement majeur pour le déploiement de modèles de langage à haut débit dans des applications réelles. Outre le grand nombre de paramètres, le cache clé-valeur (KV) pour le mécanisme d'attention dans l'architecture des transformateurs consomme une quantité significative de mémoire, en particulier lorsque le nombre de couches est important pour les modèles de langage profonds. Dans cet article, nous proposons une méthode novatrice qui ne calcule et ne met en cache que les KVs d'un petit nombre de couches, réduisant ainsi considérablement la consommation mémoire et améliorant le débit d'inférence. Nos expériences sur des modèles de langage volumineux montrent que notre méthode atteint un débit jusqu'à 26 fois supérieur à celui des transformeurs standards, tout en offrant des performances compétitives en modélisation du langage et dans les tâches en aval. De plus, notre méthode est orthogonale aux techniques existantes d'économie de mémoire pour les transformeurs, ce qui permet de l'intégrer facilement à notre modèle pour obtenir une amélioration supplémentaire de l'efficacité de l'inférence. Notre code est disponible à l'adresse https://github.com/whyNLP/LCKV.

English

Huge memory consumption has been a major bottleneck for deploying high-throughput large language models in real-world applications. In addition to the large number of parameters, the key-value (KV) cache for the attention mechanism in the transformer architecture consumes a significant amount of memory, especially when the number of layers is large for deep language models. In this paper, we propose a novel method that only computes and caches the KVs of a small number of layers, thus significantly saving memory consumption and improving inference throughput. Our experiments on large language models show that our method achieves up to 26times higher throughput than standard transformers and competitive performance in language modeling and downstream tasks. In addition, our method is orthogonal to existing transformer memory-saving techniques, so it is straightforward to integrate them with our model, achieving further improvement in inference efficiency. Our code is available at https://github.com/whyNLP/LCKV.