LLM 模型的福音：QJL 量化技术实现 KV 缓存零开销压缩

作者

近年来，大型语言模型（LLM）在各个领域都取得了显著的成就，其应用范围涵盖聊天机器人、文本到图像/视频合成、代码助手等。然而，LLM 的强大能力与其庞大的模型规模密不可分，这也带来了巨大的内存消耗挑战。

在 LLM 的生成阶段，为了避免重复计算，模型需要将所有先前生成的键值（KV）嵌入缓存到内存中。然而，随着序列长度的增加，KV 缓存所需的内存空间也会急剧增长，成为内存和速度的瓶颈。因此，如何在保持 LLM 准确性的同时减少 KV 缓存的大小成为了一个亟待解决的问题。

量化技术：压缩 KV 缓存的利器

一种简单而有效的方法是使用量化技术，用更少的比特数来表示 KV 缓存中的浮点数（FPN）。现有的 KV 缓存量化方法通常将数据分组，并为每组数据计算和存储量化常数（至少一个零点和一个比例因子）。然而，这种方法会带来显著的内存开销，因为它需要存储额外的量化常数。

QJL：一种零开销的量化方法

为了解决这个问题，本文介绍了一种名为 QJL 的新型量化方法，它结合了 Johnson-Lindenstrauss (JL) 变换和符号位量化，实现了零开销的 KV 缓存压缩。

QJL 的核心思想是将 JL 变换应用于键嵌入，然后将结果量化为单个符号位（即正负号）。与传统方法不同，QJL 不需要存储量化常数，从而消除了内存开销。

研究表明，将 JL 变换应用于键嵌入并将其量化为单个符号位，同时将相同的 JL 变换应用于查询嵌入（不进行量化），仍然可以获得键和查询嵌入的内积的无偏估计。

QJL 的优势

QJL 量化方法具有以下优点：

  • 零内存开销: QJL 不需要存储量化常数，因此不会增加内存开销。
  • 数据无关性: QJL 是一种数据无关算法，不需要根据输入数据进行调整。
  • 易于并行化: QJL 可以轻松地并行化，适用于实时应用。

实验结果

实验结果表明，QJL 在各种 LLM 和 NLP 任务中都能将 KV 缓存量化为每个 FPN 仅需 3 比特，同时保持与使用 16 比特 FPN 的精确模型相当的精度。这意味着 QJL 可以将缓存内存使用量减少五倍以上，并显著提高长序列的生成速度。

结论

QJL 是一种高效且实用的 KV 缓存压缩方法，它能够在不损失精度的情况下显著减少内存使用量和运行时间。这项技术为 LLM 的部署和应用开辟了新的可能性。

Q: 这篇论文试图解决什么问题？

A: 这篇论文试图解决大型语言模型（LLMs）在部署过程中遇到的内存消耗问题，特别是在生成阶段需要存储大量先前生成的键值（KV）嵌入到缓存中以避免重复计算，这导致了显著的内存和速度瓶颈。为了解决这一挑战，论文提出了一种有效的KV缓存压缩方法，即通过量化技术减少KV缓存所需的内存。

具体来说，论文中提出了一种名为QJL（Quantized Johnson-Lindenstrauss）的新型量化方法，它结合了Johnson-Lindenstrauss（JL）变换和符号位量化。与传统量化方法相比，QJL消除了存储量化常数（至少一个零点和一个比例因子）所需的内存开销，因为这些常数需要以全精度存储在每个数据块中。通过这种方法，QJL能够在不牺牲准确性的情况下，显著减少KV缓存的内存使用量，并且提高运行速度。

Q: 有哪些相关研究？

A: 论文中提到了与KV缓存量化相关的一些研究工作，具体包括：

  1. 多查询注意力（Multi-query attention） [29] 和 多组注意力（Multi-group attention） [2]：这些方法通过减少KV缓存中的头数来解决内存问题，但需要对预训练模型进行微调或从头开始训练。
  2. 剪枝（Pruning） 或 逐出（Evicting） 不重要的令牌 [39, 21, 33, 37]：这些研究尝试通过移除不重要的令牌来减少KV缓存的大小。
  3. 系统级优化，例如 卸载（Offloading） [30] 或在注意力机制中使用 虚拟内存和分页技术 [18]。
  4. KV缓存量化 特定方法，包括 [36, 34, 10, 16, 38]。这些方法提出了不同的量化技术，以减少KV缓存中的浮点数（FPN）所需的位数。
  5. KIVI [22] 和 KVQuant [13]：这些是最近提出的量化方法，它们通过进行每通道量化（per-channel quantization）来实现更好的性能。
  6. 其他量化方法，例如 [36, 34, 10, 16, 38]，这些方法专注于为KV缓存提出特定的量化技术。
  7. 随机傅里叶特征（Random Fourier features） [35] 和 局部敏感哈希（Locality sensitive hashing） [14]：这些与JL变换相关的应用在论文中被提及，因为它们利用了JL变换的性质。
  8. 其他与大型语言模型压缩和加速相关的研究，如 [9, 20]。

这些相关工作提供了不同的视角和方法来解决大型语言模型在内存和速度方面的挑战。QJL方法则是在此基础上提出的一个新颖的量化方法，旨在实现零开销的KV缓存量化。

Q: 论文如何解决这个问题？

A: 论文通过提出一种新的量化方法，称为QJL（Quantized Johnson-Lindenstrauss）变换，来解决大型语言模型（LLMs）在KV缓存中面临的内存和速度瓶颈问题。以下是论文解决该问题的关键步骤和方法：

  1. 引入QJL量化方法：QJL由Johnson-Lindenstrauss（JL）变换和符号位量化组成。它首先对键值嵌入向量应用JL变换，然后将其量化为符号位。
  2. 消除内存开销：与传统量化方法不同，QJL不需要存储量化常数（如零点和比例因子），这些常数通常需要为每个数据块以全精度存储，导致显著的内存开销。QJL通过不存储这些常数来消除内存开销。
  3. 不对称估计器：论文提出了一种不对称的内积估计器，用于量化向量与未量化向量之间的内积。这种估计器对于KV缓存机制特别有用，因为它允许对一个向量进行量化，而另一个向量保持未量化状态。
  4. 保持内积的无偏估计：论文证明了，即使对键嵌入向量应用了JL变换并量化为符号位，同时对查询嵌入向量仅应用JL变换而不量化，仍然可以获得它们内积的无偏估计。
  5. 控制失真：论文展示了QJL变换的内积估计器具有小的失真，并且与标准JL变换的失真相当。
  6. 数据不可知算法：QJL是一个数据不可知的算法，它不需要根据输入数据进行调整或微调，可以轻松地并行化，并实时应用。
  7. 针对异常值的处理：论文还分析了大型语言模型中的异常值分布，并提出了一种处理方法，通过在提示阶段识别异常值通道，并分别对异常值和非异常值应用独立的量化器。
  8. 高效的CUDA实现：为了优化计算，论文开发了一个轻量级的CUDA内核来实现QJL变换及其对应的内积估计器。
  9. 实验验证：论文在多个LLMs和自然语言处理（NLP）任务上应用QJL，将KV缓存量化到仅3位，证明了在不牺牲准确性的情况下，显著减少了KV缓存的内存使用量，并提高了长上下文生成的速度。

通过这些方法，论文成功地减少了KV缓存的内存占用，同时保持了模型的准确性，并提高了运行速度。

Q: 论文做了哪些实验？

A: 论文中进行了以下实验来验证QJL量化方法的有效性：

  1. 异常值分析：分析了Llama-2模型中不同层的键嵌入坐标的异常值分布，观察到深层网络中某些固定坐标的键嵌入持续表现出较大的幅度。
  2. 端到端文本生成：在LongBench基准测试上评估了QJL方法，这是一个针对长范围上下文理解的多任务基准测试。使用了longchat-7b-v1.5-32k模型（微调过的Llama-2模型，具有7B参数和16384的上下文长度），并将其与KIVI [22]、KVQuant [13]等其他量化方法进行了比较。
  3. 性能比较：在LongBench的六个问答数据集上进行了评估，设置了最大序列长度为31,500，并遵循原始存储库中的提示和评估方法。
  4. 运行时间比较：报告了在单个注意力层中提示编码和解码的运行时间，输入序列长度在1k到128k之间。比较了KVQuant、KIVI和QJL方法的运行时间。
  5. 内存使用比较：测试了QJL方法在Lambada-OpenAI、HellaSwag、PIQA、MathQA和MMLU等数据集上的表现，这些数据集具有较短的序列长度。使用LM-eval框架进行了评估，并比较了Llama-2-7B和Llama-3-8B模型的准确率。
  6. 不同量化方法的比较：在表1和表2中，展示了QJL与其他量化方法在不同数据集上的F1分数和准确率的比较结果。
  7. 不同序列长度的运行时间测试：展示了不同输入序列长度下，不同量化方法在提示编码和解码阶段的运行时间。

这些实验验证了QJL方法在减少内存使用的同时，能够保持或甚至提高生成速度，并在多个NLP任务上保持了与其他量化方法相比的竞争力或优越性。

Q: 有什么可以进一步探索的点？

A: 论文提出了QJL量化方法，有效地减少了大型语言模型（LLMs）的KV缓存内存使用，同时保持了模型的准确性和提高了运行速度。尽管论文已经取得了显著的成果，但仍有一些潜在的探索方向可以进一步研究：

  1. 算法优化：虽然论文已经提供了基于CUDA的轻量级实现，但进一步优化算法，提高计算效率和减少内存占用仍然是可能的。
  2. 更广泛的模型和任务：论文主要在自然语言处理（NLP）任务上进行了实验验证。将QJL方法应用于其他类型的模型（如图像、视频处理模型）和任务上，以评估其通用性和有效性。
  3. 不同位数的量化：论文中使用了3位量化，研究不同位数（如4位、5位等）对性能和内存使用的影响，以找到最佳平衡点。
  4. 异常值处理策略：论文中提出了一种处理异常值的策略，但进一步探索更高效或更精确的异常值检测和处理方法仍然是有价值的。
  5. 模型微调：研究在量化后是否需要对模型进行微调，以及如何有效地进行微调以保持或恢复模型性能。
  6. 硬件加速：除了CUDA实现，还可以探索其他硬件加速方法，如FPGA或ASIC，以进一步提高QJL方法的性能。
  7. 量化方法的泛化：研究QJL方法是否可以泛化到其他类型的量化问题，例如权重量化或混合精度量化。
  8. 理论分析：深入研究QJL方法的理论基础，包括量化误差、失真控制和算法收敛性等。
  9. 实际部署：在实际的硬件和软件环境中测试QJL方法，评估其在现实世界应用中的性能和可扩展性。
  10. 与其他压缩技术的结合：探索将QJL方法与其他模型压缩技术（如剪枝、稀疏化等）结合使用的可能性，以实现更高效的模型压缩。

这些探索方向可以帮助进一步推动大型语言模型的量化技术和应用的发展。

Q: 总结一下论文的主要内容

A: 这篇论文介绍了一种名为QJL（Quantized Johnson-Lindenstrauss）的新型量化方法，旨在解决大型语言模型（LLMs）在存储键值（KV）缓存时面临的内存消耗问题。以下是论文的主要内容总结：

  1. 问题背景：大型语言模型（LLMs）在处理长序列数据时，需要在KV缓存中存储大量的键值嵌入，这导致了显著的内存和速度瓶颈。
  2. QJL量化方法：提出了一种结合了Johnson-Lindenstrauss（JL）变换和符号位量化的新方法。与传统量化方法相比，QJL消除了存储量化常数（零点和比例因子）的内存开销。
  3. 无偏内积估计：论文证明了通过QJL变换和量化后，可以无偏地估计两个向量之间的内积，即使其中一个向量被量化到符号位。
  4. 数据不可知算法：QJL是一个数据不可知的算法，不需要根据输入数据进行调整或微调，易于并行化和实时应用。
  5. 异常值处理：论文分析了大型语言模型中的异常值分布，并提出了一种处理方法，通过在提示阶段识别异常值通道，并分别对异常值和非异常值应用独立的量化器。
  6. 实验验证：在多个LLMs和自然语言处理（NLP）任务上应用QJL，展示了在将KV缓存量化到仅3位的情况下，显著减少了内存使用量，同时保持了模型的准确性，并提高了运行速度。
  7. CUDA实现：为了优化计算，论文开发了一个轻量级的CUDA内核来实现QJL变换及其对应的内积估计器。
  8. 性能提升：实验结果表明，QJL方法在减少内存使用的同时，能够显著提高长上下文生成的速度，并且在多个NLP任务上保持了与其他量化方法相比的竞争力或优越性。

论文的代码已在GitHub上公开，供有兴趣的研究者和开发者进一步探索和使用。

参考文献

  • [1] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
  • [2] J. Ainslie, S. Ontanon, C. Alberti, P. Pham, A. Ravula, S. Sanghai, P. Shafto, and Y. Bisk. Etc: Encoding long and diverse items with transformers for dense retrieval. arXiv preprint arXiv:2004.08922, 2020.
  • [3] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
  • [4] S. Chang, W. L. Hamilton, K. Roberts, and L. Zettlemoyer. Longbench: Towards evaluation of long-form language models. arXiv preprint arXiv:2206.08153, 2022.
  • [5] R. Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.

发表评论