H$_2$O：面向大语言模型高效生成推理的重击令牌预言机
机器学习
2023-12-20 v3
摘要
大语言模型（LLMs）尽管近期取得了令人瞩目的成就，但其部署成本却高得令人望而却步，尤其是对于涉及长内容生成的应用，如对话系统和故事写作。通常，除了模型参数外，大量称为 KV 缓存的瞬时状态信息被存储在 GPU 内存中，其大小随序列长度和批大小线性增长。在本文中，我们提出了一种实现 KV 缓存的新方法，可显著减少其内存占用。我们的方法基于一个值得注意的观察：在计算注意力分数时，一小部分令牌贡献了绝大部分的值。我们称这些令牌为重击令牌（H）。通过全面研究，我们发现（i）H 的出现是自然的，并且与文本中令牌的频繁共现强相关；（ii）移除它们会导致显著的性能下降。基于这些洞察，我们提出了重击令牌预言机（HO），一种动态保留近期令牌与 H 令牌平衡的 KV 缓存淘汰策略。我们将 KV 缓存淘汰表述为一个动态子模问题，并（在温和假设下）证明了我们新淘汰算法的理论保证，这可帮助指导未来工作。我们在 OPT、LLaMA 和 GPT-NeoX 上跨广泛任务验证了我们算法的准确性。我们对 HO 的实现使用 20% 的重击令牌，在 OPT-6.7B 和 OPT-30B 上相较三个领先推理系统 DeepSpeed Zero-Inference、Hugging Face Accelerate 和 FlexGen 将吞吐量分别提升高达 29、29 和 3。在相同批大小下，H2O 可将延迟降低高达 1.9。代码可在 https://github.com/FMInference/H2O 获取。
引用
@article{arxiv.2306.14048,
title = {H$_2$O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models},
author = {Zhenyu Zhang and Ying Sheng and Tianyi Zhou and Tianlong Chen and Lianmin Zheng and Ruisi Cai and Zhao Song and Yuandong Tian and Christopher Ré and Clark Barrett and Zhangyang Wang and Beidi Chen},
journal= {arXiv preprint arXiv:2306.14048},
year = {2023}
}