Caché KV Condensado por Capas para Inferencia Eficiente en Modelos de Lenguaje a Gran Escala

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

May 17, 2024

Autores: Haoyi Wu, Kewei Tu

cs.AI

Resumen

El alto consumo de memoria ha sido un cuello de botella importante para implementar modelos de lenguaje grandes de alto rendimiento en aplicaciones del mundo real. Además del gran número de parámetros, la caché de clave-valor (KV) para el mecanismo de atención en la arquitectura transformer consume una cantidad significativa de memoria, especialmente cuando el número de capas es grande en modelos de lenguaje profundos. En este artículo, proponemos un método novedoso que solo calcula y almacena en caché los KVs de un número reducido de capas, lo que reduce significativamente el consumo de memoria y mejora el rendimiento de la inferencia. Nuestros experimentos con modelos de lenguaje grandes muestran que nuestro método logra un rendimiento hasta 26 veces mayor que los transformers estándar, con un desempeño competitivo en modelado de lenguaje y tareas posteriores. Además, nuestro método es ortogonal a las técnicas existentes de ahorro de memoria en transformers, por lo que es sencillo integrarlas con nuestro modelo, logrando una mejora adicional en la eficiencia de la inferencia. Nuestro código está disponible en 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.