H_2O : Oracle des éléments dominants pour l'inférence générative efficace des grands modèles de langage

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

June 24, 2023

Auteurs: 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

Résumé

Les grands modèles de langage (LLMs), malgré leurs récents accomplissements impressionnants, sont particulièrement coûteux à déployer, notamment pour des applications impliquant la génération de contenus longs, tels que les systèmes de dialogue et l'écriture de récits. Souvent, une grande quantité d'informations d'état transitoires, appelée cache KV, est stockée dans la mémoire GPU en plus des paramètres du modèle, augmentant linéairement avec la longueur de la séquence et la taille du lot. Dans cet article, nous introduisons une nouvelle approche pour implémenter le cache KV qui réduit considérablement son empreinte mémoire. Notre approche repose sur l'observation notable qu'une petite partie des tokens contribue majoritairement à la valeur lors du calcul des scores d'attention. Nous appelons ces tokens les "Heavy Hitters" (H_2). À travers une investigation approfondie, nous constatons que (i) l'émergence des H_2 est naturelle et fortement corrélée à la co-occurrence fréquente de tokens dans le texte, et (ii) leur suppression entraîne une dégradation significative des performances. Sur la base de ces insights, nous proposons Heavy Hitter Oracle (H_2O), une politique d'éviction du cache KV qui conserve dynamiquement un équilibre entre les tokens récents et les H_2. Nous formulons l'éviction du cache KV comme un problème de sous-modularité dynamique et démontrons (sous des hypothèses modérées) une garantie théorique pour notre nouvel algorithme d'éviction, qui pourrait guider les travaux futurs. Nous validons la précision de notre algorithme avec OPT, LLaMA et GPT-NeoX sur une large gamme de tâches. Notre implémentation de H_2O avec 20 % de heavy hitters améliore le débit par rapport à trois systèmes d'inférence leaders : DeepSpeed Zero-Inference, Hugging Face Accelerate et FlexGen, jusqu'à 29 fois, 29 fois et 3 fois sur OPT-6.7B et OPT-30B. Avec la même taille de lot, H2O peut réduire la latence jusqu'à 1,9 fois. Le code est disponible à l'adresse 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.