KVQuant: 通过KV缓存量化实现1000万上下文长度的LLM推理 [论文]
KVQuant是一种高效的KV缓存量化方法,结合多项创新实现精确的低精度量化,从而实现高效的长上下文长度推理。
TLDR: KVQuant通过将KV缓存量化到低精度来解决长上下文长度推理的内存瓶颈问题。KVQuant通过考虑不同LLMs中缓存的KV值所观察到的若干一致模式,并开发利用这些模式的方法,实现高精度的低精度KV缓存量化,包括:
- 每通道,预RoPE键量化,更好地匹配键中的异常通道
- 非均匀量化(NUQ)以更好地表征非均匀激活
- 稠密-稀疏量化以减轻数值异常对量化难度的影响
KVQuant能够在单个A100-80GB GPU上提供上下文长度为100万的LLaMA-7B模型的服务,甚至在8-GPU系统上提供上下文长度为1000万的LLaMA-7B模型 🔥
[TLDR: Twitter 主线] [论文]
具有长上下文长度推理的大型世界模型
大型世界模型(LWM) 是一项最近的工作,能够训练上下文长度最长可达100万的长上下文长度模型。然而,由于必须在整个推理过程中存储大量缓存的KV,这些模型的推理非常耗资源。使用KVQuant,我们现在可以在单个A100上高效地推理这些长上下文长度模型!
lmw/ 目录包含使用量化大型世界模型进行推理和评估的脚本。
附加方法改进
为了进一步改进支持长上下文长度推理的方法,我们做了多项改进:
GPU上的并行topK支持和并行提示处理的核
- 我们通过增加额外的内核来执行多个输入令牌的并行打包,并修改我们的推理代码在并行附加多个值令牌时利用GPU进行并行topK来增强我们的开源支持。键异常值封顶
- 我们增加了支持在校准和推理过程中为每个令牌键运行固定数量的异常值。这使我们能够设计更高效的内核,因为每个键和值的令牌有一个最大的异常值数量,并且这使我们的方法更容易进行内存分配,因为我们可以为每个键分配固定大小的内存。注意汇量化
- 基于Attention Sink 论文的见解,即模型将注意力集中在第一个令牌上,我们在校准和推理过程中增加了支持为初始的少量键和值(例如5个)保留fp16。这可以带来显著的性能提升,也是另一项并发工作IntactKV中改善量化性能的方法。
安装
代码库包含五个子文件夹,每个文件夹都有自己的README文件,提供了安装该步骤所需环境的说明。
代码结构
gradients
- 计算渔民信息的代码库 - 这对于能够量化新模型是必要的quant
- 运行模拟量化+评估实验的代码库(需要先计算渔民信息)deployment
- 使用压缩向量进行高效推理的代码库(需要先从量化步骤获得量化器)lwm
- 运行推理和评估量化大世界模型的代码benchmarking
- 内核基准测试代码(需要先从量化步骤获得量化器)
要重现论文中报告的困惑度数字,请运行gradients
然后运行quant
。
路线图:
添加部署代码优化内核- 合并优化内核与更新的端到端推理部署代码
- 添加用于运行额外长上下文长度评估的脚本
致谢
此代码重用了包括GPTQ、GPTQ-For-LLaMA和SqueezeLLM在内的多个库中的组件。
引用
KVQuant是以下论文的一部分。如果您发现该库对您的工作有用,请引用以下论文:
@article{hooper2024kvquant,
title={KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization},
author={Hooper, Coleman and Kim, Sehoon and Mohammadzadeh, Hiva and Mahoney, Michael W and Shao, Yakun Sophia and Keutzer, Kurt and Gholami, Amir},
journal={arXiv preprint arXiv:2401.18079},
year={2024}
}