H2O：大型语言模型高效生成推理的重要命中预测器

论文"H2O：大型语言模型高效生成推理的重要命中预测器"的代码

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

概述

尽管大型语言模型（LLMs）近期取得了令人印象深刻的成就，但其部署成本仍然非常高，特别是对于涉及长内容生成的应用，如对话系统和故事写作。通常，除了模型参数外，还需要在GPU内存中存储大量临时状态信息，即所谓的KV缓存，其规模与序列长度和批处理大小呈线性关系。本文提出了一种新颖的KV缓存实现方法，显著降低了其内存占用。我们的方法基于一个值得注意的观察：在计算注意力分数时，只有一小部分标记贡献了大部分价值。我们称这些标记为重要命中（H2）。通过全面研究，我们发现：(i) H2的出现是自然的，并且与文本中标记的频繁共现有很强的相关性；(ii) 移除它们会导致性能显著下降。基于这些洞见，我们提出了重要命中预测器（H2O），这是一种动态平衡保留最近和H2标记的KV缓存淘汰策略。我们将KV缓存淘汰问题表述为一个动态子模问题，并在温和假设下为我们的新颖淘汰算法提供理论保证，这可能有助于指导未来的工作。我们在广泛的任务中验证了我们算法在OPT、LLaMA和GPT-NeoX上的准确性。我们实现的H2O，使用20%的重要命中，在OPT-6.7B和OPT-30B上将吞吐量相比三个领先的推理系统DeepSpeed Zero-Inference、Hugging Face Accelerate和FlexGen分别提高了最多29倍、29倍和3倍。在相同的批处理大小下，H2O可以将延迟降低最多1.9倍。