H_2O: Оракул для выявления значимых элементов для эффективного генеративного вывода в крупных языковых моделях

Аннотация

Крупные языковые модели (LLM), несмотря на их недавние впечатляющие достижения, остаются крайне затратными для развертывания, особенно в приложениях, связанных с генерацией длинного контента, таких как диалоговые системы и написание рассказов. Часто большое количество временной информации, называемой кэшем ключей и значений (KV cache), хранится в памяти GPU вместе с параметрами модели, увеличиваясь линейно с длиной последовательности и размером пакета. В данной статье мы представляем новый подход к реализации KV cache, который значительно сокращает его объем памяти. Наш подход основан на важном наблюдении, что небольшая часть токенов вносит основной вклад при вычислении оценок внимания. Мы называем эти токены "тяжеловесами" (Heavy Hitters, H_2). В результате всестороннего исследования мы обнаружили, что (i) появление H_2 является естественным и сильно коррелирует с частым совместным появлением токенов в тексте, и (ii) их удаление приводит к значительному ухудшению производительности. На основе этих инсайтов мы предлагаем Heavy Hitter Oracle (H_2O) — стратегию вытеснения KV cache, которая динамически сохраняет баланс между недавними токенами и H_2. Мы формулируем вытеснение KV cache как динамическую субмодулярную задачу и доказываем (при умеренных допущениях) теоретическую гарантию для нашего нового алгоритма вытеснения, что может помочь в будущих исследованиях. Мы проверяем точность нашего алгоритма на моделях OPT, LLaMA и GPT-NeoX в широком спектре задач. Наша реализация H_2O с 20% "тяжеловесов" увеличивает пропускную способность по сравнению с тремя ведущими системами вывода — DeepSpeed Zero-Inference, Hugging Face Accelerate и FlexGen — до 29 раз, 29 раз и 3 раза на моделях OPT-6.7B и OPT-30B. При том же размере пакета H_2O может снизить задержку до 1.9 раз. Код доступен по адресу 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.

H_2O: Оракул для выявления значимых элементов для эффективного генеративного вывода в крупных языковых моделях

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

Аннотация

Support