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倍。

内容

我们提供两段代码来实现大型语言模型高效生成推理的重要命中预测器：

• h2o_flexgen：实现LLM生成的更高吞吐量，代码基于FlexGen。
• h2o_hf：在不同基准测试上测试性能，代码基于Hugging Face。提供了模拟代码（掩蔽注意力矩阵）和实际KV丢弃实现（请参考h2o_hf/utils_real_drop）。