Laag-gecondenseerde KV-cache voor efficiënte inferentie van grote taalmodellen

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

May 17, 2024

Auteurs: Haoyi Wu, Kewei Tu

cs.AI

Samenvatting

Het enorme geheugengebruik is een belangrijke beperking geweest bij het inzetten van hoogdoorvoerende grote taalmodel(len) in real-world toepassingen. Naast het grote aantal parameters verbruikt de key-value (KV)-cache voor het aandachtmechanisme in de transformer-architectuur een aanzienlijke hoeveelheid geheugen, vooral wanneer het aantal lagen groot is voor diepe taalmodel(len). In dit artikel stellen we een nieuwe methode voor die alleen de KVs van een klein aantal lagen berekent en cachet, waardoor het geheugengebruik aanzienlijk wordt bespaard en de inferentie-doorvoer wordt verbeterd. Onze experimenten met grote taalmodel(len) laten zien dat onze methode tot 26 keer hogere doorvoer bereikt dan standaard transformers en competitieve prestaties levert bij taalmodellering en downstream taken. Bovendien is onze methode orthogonaal aan bestaande geheugenbesparende technieken voor transformers, waardoor het eenvoudig is om deze te integreren met ons model, wat verdere verbetering in de inferentie-efficiëntie oplevert. Onze code is beschikbaar op 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.