Schichtkondensierter KV-Cache für effiziente Inferenz großer Sprachmodelle

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

May 17, 2024

Autoren: Haoyi Wu, Kewei Tu

cs.AI

Zusammenfassung

Der hohe Speicherverbrauch war ein wesentlicher Engpass bei der Bereitstellung von hochdurchsatzfähigen großen Sprachmodellen in realen Anwendungen. Neben der großen Anzahl von Parametern verbraucht der Schlüssel-Wert (KV)-Cache für den Aufmerksamkeitsmechanismus in der Transformer-Architektur eine erhebliche Menge an Speicher, insbesondere wenn die Anzahl der Schichten bei tiefen Sprachmodellen groß ist. In diesem Artikel schlagen wir eine neuartige Methode vor, die nur die KV-Werte einer kleinen Anzahl von Schichten berechnet und zwischenspeichert, wodurch der Speicherverbrauch erheblich reduziert und die Inferenzdurchsatz verbessert wird. Unsere Experimente mit großen Sprachmodellen zeigen, dass unsere Methode eine bis zu 26-mal höhere Durchsatzrate als Standard-Transformer erreicht und eine wettbewerbsfähige Leistung bei der Sprachmodellierung und bei nachgelagerten Aufgaben erzielt. Darüber hinaus ist unsere Methode orthogonal zu bestehenden Transformer-Speicherspartechniken, sodass es einfach ist, sie mit unserem Modell zu integrieren und so eine weitere Verbesserung der Inferenzeffizienz zu erzielen. Unser Code ist verfügbar unter 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.