Dynamic Memory Compression: Retrofitting LLMs for Accelerated Inference

📄 arXiv: 2403.09636v2 📥 PDF

作者: Piotr Nawrot, Adrian Łańcucki, Marcin Chochowski, David Tarjan, Edoardo M. Ponti

分类: cs.CL

发布日期: 2024-03-14 (更新: 2024-07-23)

期刊: Proceedings of the 41st International Conference on Machine Learning (2024) 37396-37412

💡 一句话要点

提出动态内存压缩以解决大语言模型推理效率问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 动态内存压缩 大语言模型 推理效率 键值缓存 在线压缩 预训练模型 自然语言处理

📋 核心要点

  1. 现有的LLM在推理时需要存储大量的键值缓存，导致效率低下，尤其是在长序列和大批量情况下。
  2. 本文提出动态内存压缩（DMC），通过在线压缩键值缓存，允许模型在不同层和头中使用不同的压缩比，从而提高推理效率。
  3. 实验结果表明，DMC在NVIDIA H100 GPU上实现了高达7倍的吞吐量提升，并且在缓存压缩率达到4倍的情况下保持了原有性能。

📝 摘要（中文）

随着变换器成为大型语言模型（LLMs）的基础，生成过程的效率仍然受到限制，主要是由于需要存储与过去标记相关的键值表示的缓存，其大小与输入序列长度和批量大小线性相关。为此，本文提出了动态内存压缩（DMC），一种在推理时进行在线键值缓存压缩的方法。DMC允许模型在不同的头和层中应用不同的压缩比。通过对预训练的LLM（如Llama 2）进行改造，DMC在NVIDIA H100 GPU上实现了高达7倍的推理吞吐量提升，同时保持了原有下游性能，并在缓存压缩率达到4倍的情况下，超越了现有的分组查询注意力（GQA）和键值驱逐策略（H$_2$O，TOVA）。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在推理过程中由于键值缓存存储需求导致的效率低下问题。现有方法在处理长序列和大批量时，缓存大小呈线性增长，造成内存压力和计算延迟。

核心思路：动态内存压缩（DMC）通过在线压缩键值缓存，允许模型在不同的层和头中应用不同的压缩比，从而提高推理效率。该方法通过继续预训练来实现，且不增加额外参数。

技术框架：DMC的整体架构包括预训练的LLM模型，通过对少量原始数据的继续训练，模型学习如何在推理时动态调整缓存的压缩比。主要模块包括键值缓存的生成、压缩和应用。

关键创新：DMC的最大创新在于其动态调整压缩比的能力，使得模型能够根据不同的上下文需求灵活地优化内存使用。这一特性与传统的静态缓存策略形成鲜明对比。

关键设计：在DMC中，模型通过继续预训练来学习不同层和头的压缩策略，采用的损失函数和网络结构设计旨在最大化推理效率，同时保持下游任务的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，DMC在NVIDIA H100 GPU上实现了高达7倍的推理吞吐量提升，同时在缓存压缩率达到4倍的情况下，保持了原有的下游性能，超越了现有的分组查询注意力（GQA）和键值驱逐策略（H$_2$O，TOVA）。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等，能够显著提高大语言模型在实际应用中的推理速度和效率。未来，DMC可能会推动更大规模模型的应用，使其在资源受限的环境中也能高效运行。

📄 摘要（原文）

Transformers have emerged as the backbone of large language models (LLMs). However, generation remains inefficient due to the need to store in memory a cache of key-value representations for past tokens, whose size scales linearly with the input sequence length and batch size. As a solution, we propose Dynamic Memory Compression (DMC), a method for online key-value cache compression at inference time. Most importantly, the model learns to apply different compression ratios in different heads and layers. We retrofit pre-trained LLMs such as Llama 2 (7B, 13B and 70B) into DMC Transformers, achieving up to 7x throughput increase during auto-regressive inference on an NVIDIA H100 GPU. DMC is applied via continued pre-training on a negligible percentage of the original data without adding any extra parameters. DMC preserves the original downstream performance with up to 4x cache compression, outperforming up-trained grouped-query attention (GQA) and key-value eviction policies (H$_2$O, TOVA). GQA and DMC can be even combined to obtain compounded gains. Hence, DMC can serve as a drop-in replacement for KV caching in existing LLMs to fit longer contexts and larger batches within any given memory budget.