Cache KV Condensato a Livelli per Inferenza Efficiente nei Modelli Linguistici di Grande Dimensione

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

May 17, 2024

Autori: Haoyi Wu, Kewei Tu

cs.AI

Abstract

L'elevato consumo di memoria è stato un importante collo di bottiglia per il dispiegamento di modelli linguistici di grandi dimensioni ad alta produttività in applicazioni reali. Oltre al gran numero di parametri, la cache chiave-valore (KV) per il meccanismo di attenzione nell'architettura transformer consuma una quantità significativa di memoria, specialmente quando il numero di strati è elevato nei modelli linguistici profondi. In questo articolo, proponiamo un metodo innovativo che calcola e memorizza nella cache le KV solo per un numero ridotto di strati, riducendo così significativamente il consumo di memoria e migliorando la produttività durante l'inferenza. I nostri esperimenti su modelli linguistici di grandi dimensioni dimostrano che il nostro metodo raggiunge una produttività fino a 26 volte superiore rispetto ai transformer standard e prestazioni competitive nella modellazione linguistica e nei task downstream. Inoltre, il nostro metodo è ortogonale alle tecniche esistenti per il risparmio di memoria nei transformer, quindi è semplice integrarlo con il nostro modello, ottenendo ulteriori miglioramenti nell'efficienza di inferenza. Il nostro codice è disponibile all'indirizzo 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.