H_2O: Heavy-Hitter Oracle für effiziente generative Inferenz großer Sprachmodelle

H_2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models

June 24, 2023

Autoren: Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, Beidi Chen

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs), trotz ihrer jüngsten beeindruckenden Erfolge, sind besonders kostspielig in der Bereitstellung, insbesondere für Anwendungen, die die Erzeugung von langen Inhalten beinhalten, wie Dialogsysteme und das Schreiben von Geschichten. Oft wird eine große Menge an temporären Zustandsinformationen, die als KV-Cache bezeichnet werden, zusätzlich zu den Modellparametern im GPU-Speicher gespeichert, was linear mit der Sequenzlänge und der Batch-Größe skaliert. In diesem Artikel stellen wir einen neuartigen Ansatz zur Implementierung des KV-Caches vor, der dessen Speicherbedarf erheblich reduziert. Unser Ansatz basiert auf der bemerkenswerten Beobachtung, dass ein kleiner Teil der Tokens den größten Beitrag bei der Berechnung der Aufmerksamkeitswerte leistet. Wir nennen diese Tokens Heavy Hitters (H_2). Durch eine umfassende Untersuchung stellen wir fest, dass (i) das Auftreten von H_2 natürlich ist und stark mit der häufigen gemeinsamen Auftretenshäufigkeit von Tokens im Text korreliert, und (ii) deren Entfernung zu einer erheblichen Leistungsverschlechterung führt. Basierend auf diesen Erkenntnissen schlagen wir Heavy Hitter Oracle (H_2O) vor, eine KV-Cache-Verwaltungsstrategie, die dynamisch ein Gleichgewicht zwischen aktuellen und H_2-Tokens beibehält. Wir formulieren die KV-Cache-Verwaltung als ein dynamisches submodulares Problem und beweisen (unter milden Annahmen) eine theoretische Garantie für unseren neuartigen Verwaltungsalgorithmus, der zukünftige Arbeiten leiten könnte. Wir validieren die Genauigkeit unseres Algorithmus mit OPT, LLaMA und GPT-NeoX über eine breite Palette von Aufgaben. Unsere Implementierung von H_2O mit 20 % Heavy Hitters verbessert den Durchsatz gegenüber drei führenden Inferenzsystemen – DeepSpeed Zero-Inference, Hugging Face Accelerate und FlexGen – um bis zu das 29-fache, 29-fache und 3-fache bei OPT-6.7B und OPT-30B. Bei gleicher Batch-Größe kann H2O die Latenz um bis zu das 1,9-fache reduzieren. Der Code ist verfügbar unter https://github.com/FMInference/H2O.

English

Large Language Models (LLMs), despite their recent impressive accomplishments, are notably cost-prohibitive to deploy, particularly for applications involving long-content generation, such as dialogue systems and story writing. Often, a large amount of transient state information, referred to as the KV cache, is stored in GPU memory in addition to model parameters, scaling linearly with the sequence length and batch size. In this paper, we introduce a novel approach for implementing the KV cache which significantly reduces its memory footprint. Our approach is based on the noteworthy observation that a small portion of tokens contributes most of the value when computing attention scores. We call these tokens Heavy Hitters (H_2). Through a comprehensive investigation, we find that (i) the emergence of H_2 is natural and strongly correlates with the frequent co-occurrence of tokens in the text, and (ii) removing them results in significant performance degradation. Based on these insights, we propose Heavy Hitter Oracle (H_2O), a KV cache eviction policy that dynamically retains a balance of recent and H_2 tokens. We formulate the KV cache eviction as a dynamic submodular problem and prove (under mild assumptions) a theoretical guarantee for our novel eviction algorithm which could help guide future work. We validate the accuracy of our algorithm with OPT, LLaMA, and GPT-NeoX across a wide range of tasks. Our implementation of H_2O with 20% heavy hitters improves the throughput over three leading inference systems DeepSpeed Zero-Inference, Hugging Face Accelerate, and FlexGen by up to 29times, 29times, and 3times on OPT-6.7B and OPT-30B. With the same batch size, H2O can reduce the latency by up to 1.9times. The code is available at https://github.com/FMInference/H2O.