Cache KV Condensado por Camadas para Inferência Eficiente em Modelos de Linguagem de Grande Escala

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

May 17, 2024

Autores: Haoyi Wu, Kewei Tu

cs.AI

Resumo

O alto consumo de memória tem sido um grande gargalo para a implantação de modelos de linguagem de grande escala com alta taxa de transferência em aplicações do mundo real. Além do grande número de parâmetros, o cache de chave-valor (KV) para o mecanismo de atenção na arquitetura transformer consome uma quantidade significativa de memória, especialmente quando o número de camadas é grande em modelos de linguagem profundos. Neste artigo, propomos um método inovador que calcula e armazena em cache os KVs de apenas um pequeno número de camadas, economizando significativamente o consumo de memória e melhorando a taxa de transferência na inferência. Nossos experimentos com modelos de linguagem de grande escala mostram que nosso método alcança até 26 vezes mais taxa de transferência do que transformers padrão e desempenho competitivo em modelagem de linguagem e tarefas subsequentes. Além disso, nosso método é ortogonal às técnicas existentes de economia de memória em transformers, portanto, é simples integrá-las ao nosso modelo, alcançando uma melhoria adicional na eficiência da inferência. Nosso código está disponível em 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.