尽管大语言模型（LLMs）作为通用 AI 工具取得了成功，但其对计算资源的高需求使其在许多真实场景中的部署具有挑战性。模型的大小和对话状态受到可用高带宽内存的限制，限制了可服务的用户数量和最大对话长度。
目前，两种最热门的 LLM 架构是 Transformers 和选择性状态空间模型 (SSMs)，它们位于频谱的另一端：
- Transformer：对话状态由序列中每个元素的不同表示组成，其大小会迅速爆炸。
- SSMs：将整个序列压缩为单个表示形式，由于容量有限，可能会忘记过去的信息。对话状态的压缩可释放内存，对于在相同内存限制下运行更大的模型、一次处理更多的 tokens 或降低延迟至关重要。
为此，NVIDIA 的研究人员开发了一项称为动态内存压缩（Dynamic Memory Compression，DMC）的新技术，该技术可以大幅提高 LLMs 部署的效率，并在不会耗尽内存的情况下将其视野扩展到更长的序列。
DMC 开启了第三种方式，即训练 Transformer 模型以自适应压缩对话状态并实现所需的压缩率。这样可以在不替换熟悉的 Transformer 架构的情况下显著减少对话状态大小。DMC 不需要从头开始训练，因为可以通过少量额外训练来改装现有模型，这比容易出错的无训练方法更可靠。
哪些因素会影响 LLM 推理性能？
LLM 推理包含两个阶段：
- 预填充：用户查询被输入。
- 自动回归生成：响应一次生成一个 token。
在生成过程中，为执行自注意力操作，Transformers 会为缓存中的每个 token 附加一对表示 (键值对或 KVP)。系统会为每个层和注意力头存储不同的 KVP。
因此，KVP 缓存会随序列长度按比例增长。由于 KVP 缓存必须与 LLM 权重一起装入 GPU 内存，因此 KVP 缓存可能会占用显存的大部分空间，甚至耗尽显存。
此外，KVP 缓存越大，执行单个推理步骤所需的时间越长。这是因为计算注意力分数是受内存限制的运算。每个查询都有自己要加载的 KVP 缓存。
注意力层或 FFN 层的线性投影的情况则不同，如果 GPU 同时并行处理许多查询，则所有查询都必须将 每个权重矩阵从 HBM 加载到 SRAM 中一次 。
过去的研究试图通过 量化 KVP 缓存的表征、 共享 注意力头或从其中 移除令牌 来减小 KVP 缓存的大小。但是，这些方法会降低原始性能，因为它们会从内存中删除信息，而不会改变原始 LLM 的行为。
动态显存压缩
动态内存压缩 (DMC) 是一种在推理期间压缩 KV 缓存的简单方法，不会导致性能下降。模型使用以下公式分别决定每个 token、层和头，以及是否应将新的键值对附加到 KV 缓存中 (就像在普通 Transformer 中一样)，还是累加到最后一个：
这个方程是 DMC 的核心，可将密钥的子序列转换为特定的前缀和，这让人联想到 xLSTM 或 RWKV 等热门 SSM。
在推理过程中，alpha 的值严格是二进制的。当 
在普通模型中，缓存一次扩展一个 KVP。借助 DMC，决策变量可确定是否应扩展缓存，还是应将新的缓存对与 KVP 缓存中的最后一个缓存对合并。
DMC 的一个关键组成部分是 改造 过程，该过程教会了 LLM 这种行为：
- 使用 2 – 8%的原始训练数据混合训练预先存在的 LLM (例如 Llama 系列的 LLM)。
- 初始化决策模块，以便始终将 KVP 附加到 KVP 缓存中
- 通过对新配对与追踪配对的平均值施加压力，逐渐过渡到 DMC。在改装过程中，目标压缩率从 1x 提升到所需水平。
- 达到目标压缩率后，进行修复，以完成改装和整合的最后步骤。
决定是否附加或合并是独立的。要使用梯度下降训练 LLMs，您可以通过 Gumbel-Sigmoid 分布持续松此决策，从而在训练期间生成部分附加和部分合并的内存元素。
为了模拟 token 无法进入中间合并阶段的推理时间 DMC (图 2)，您可以增强 Transformers 的因果关系掩码，以越来越多地阻止对这些状态的访问。
总体而言，此过程可让您在继续预训练期间避免出现困惑峰值。事实上，我们发现 像通用查询注意力（general query-attention） 这样的方法在继续预训练之前会显著改变模型架构，很难恢复原始性能。
成果
对于我们进行基准测试的所有下游应用，包括 真实性 (MMLU)、 常识问答和编码 (HumanEval)，DMC 的性能可与普通模型相比。
至关重要的是，DMC 可以与量化和 GQA 等其他 KVP 缓存缩减方法结合使用。
| DMC 压缩率 | MMLU (真实性) | 常识问答 | HumanEval (编码) | |
| Llama-2 -7B | 1 倍 | 44.6 | 70.5 | 14.0 |
| Llama-2 -7B | 4 倍 | 44.2 | 70.2 | 16.5 |
| Llama-2 -7B | 8 倍 | 41.8 | 70.1 | 16.5 |
| Llama-2 -13B | 1 倍 | 54.5 | 73.5% | 17.5 |
| Llama-2 -13B | 4 倍 | 54.2 | 73.2 | 22.0 |
| Llama-2 -13B | 8 倍 | 52.1 | 73.3% | 21.3% |
| Llama-2 -70B+GQA | 8 倍 | 68.8 | 78.0 | 29.6 |
| Llama-2 -70B *GQA |
16 倍 | 68.8 | 77.9 | 29.9% |
| Llama-2 -7B+DMC 4x | Llama-2-7B + DMC 4x (8 位量化) | Llama-2 -7B+DMC 8x | Llama-2-7B + DMC 8x (8 位量化) (8-bit quantization) |
|
| MMLU (真实性) | 44.2 | 44.6 | 41.8 | 41.7 |
在某些情况下，性能甚至会得到提升。我们推测这是由于缓解压缩的现象，例如 中间丢失 的 tokens，不过还需要更多证据来证实这一点。
对话状态的压缩释放了内存，而内存可用于容纳更大的批量大小。由于模型受内存限制，这通常几乎直接转化为同等的吞吐量增加。
例如，选择适合内存的最大批量大小 (在 NVIDIA H100 GPU 上压缩 8 倍) 意味着每秒生成的令牌比 Vanilla 模型多 700%。
在达到一定的压缩率后，计算注意力分数会受到计算限制，因此延迟可能会增加，同时提供更多的查询。
在图 4 中，x 轴显示了适用于 Vanilla LLM 以及 DMC 2x、4x、6x 和 8x 模型的单个 GPU (7B 和 13B) 或两个具有张量并行度 (70B) 的 GPU 上显存的最大批量大小。
如果批量大小保持不变，延迟会降低。在实践中，调整批量大小有助于优化延迟吞吐量权衡。
在图 5 中，在特定批量大小下，模型受内存限制，可以相同延迟处理其他示例。除此之外，模型的计算能力也越来越受限。我们在 NVIDIA H100 GPU 上对 4K 令牌序列的最后 1K 令牌的生成进行了测量。
由于能够以不同的速率灵活地压缩不同的注意力头和层，DMC 还揭示了有关 LLMs 内部功能的新见解。例如，顶层受到的压缩最多，因为其 KVP 是最冗余的。
在图 6 中，头部按最高压缩级别到最低自上而下级别排列。
结束语
DMC 有望进一步推动 LLM 的发展。
它提供了一种自适应显存机制，该机制介于 Transformer 中 的线性显存和 SSM 的常量显存之间，并能更好地在容量和空间之间进行权衡。这在不牺牲质量的情况下提高了 Transformer LLM 的吞吐量，并且能够在相同的硬件限制条件下适应更长的上下文。这对于通过思维链实现推理时扩展至关重要，因为这需要生成长序列的 token。
DMC 还开辟了新的途径，通过重新利用现有架构，以经济实惠的方式创建原型和生产级 LLM。 使用现成的代码 改造您自己的 Transformer LLM，或者阅读 论文 以深入了解 Dynamic Memory Compression。
有关更多信息，请参阅以下资源：
致谢
在此，我们要感谢 Mostofa Patwary 和 Szymon Migacz 所提供的帮助，并感谢 Przemysław Strzelczyk、Daniel Korzekwa 和 Bryan Catanzaro 在发表此论文时提供的有益讨论和支持。这项工作得到了 UKRI Centre for Doctoral Training in Natural Language Processing 的部分支持，该中心由 UKRI (grant EP/S022481/1) 和爱丁堡大学信息学学院以及哲学、心理和语言科学学院提供资金。
