月度归档： 2024 年 6 月

近年来，大型语言模型（LLM）在各个领域都取得了显著的进步，其强大的文本理解和生成能力令人惊叹。那么，LLM 是否能胜任模拟文本世界的任务呢？换句话说，能否利用 LLM 来构建虚拟环境，并准确预测行动对世界状态的影响，从而省去繁琐的人工编码？

本文将探讨 LLM 作为文本世界模拟器的潜力，并通过一个名为 BYTESIZED32-State-Prediction 的全新基准数据集进行评估。

模拟世界的挑战

模拟世界对于理解和研究世界至关重要，但传统上，构建一个复杂的模拟环境需要大量的人工投入，耗费大量时间和精力。LLM 的出现为我们提供了一种新的思路，即利用其庞大的预训练数据集，直接将其作为模拟器使用。

然而，LLM 真的能胜任模拟器的角色吗？为了回答这个问题，研究人员将目光投向了文本游戏领域。文本游戏以自然语言描述环境和动态变化，长期以来被用于决策过程、信息提取和人工智能推理等研究领域。

两种利用 LLM 进行世界建模的方法

研究人员提出两种利用 LLM 进行世界建模和模拟的方法：

1. 神经符号方法: 利用 LLM 生成符号表示的代码，以便进行形式化规划或推理。例如，REASONING VIA PLANNING (RAP) 方法利用 LLM 的先验知识构建世界模型，然后使用专门的规划算法来决定代理策略。
2. 直接模拟方法: 利用 LLM 直接生成文本描述，构建虚拟环境，并根据用户输入的行动进行模拟。例如，AI-DUNGEON 项目使用 LLM 生成文本描述，构建一个纯粹由语言模型驱动的游戏世界。

本文重点关注第二种方法，即直接模拟方法，并首次对 LLM 直接模拟虚拟环境的能力进行了量化分析。

BYTESIZED32-State-Prediction 基准数据集

为了评估 LLM 作为文本世界模拟器的能力，研究人员构建了一个名为 BYTESIZED32-State-Prediction (BYTESIZED32-SP) 的全新基准数据集。该数据集包含 76,369 个文本游戏状态转换，每个转换都由一个七元组 (S, A, T, O, R, C, D) 表示，分别对应状态空间、动作空间、转换函数、观察函数、奖励函数、上下文信息和完成指示函数。

该数据集从 BYTESIZED32 数据集派生而来，BYTESIZED32 数据集包含 32 个由人类编写的文本游戏，每个游戏模拟不同的科学或常识推理概念。研究人员通过修改每个 BYTESIZED32 游戏，使其能够在每个时间步输出游戏状态 (st, rt, dt) 和中间状态 sactt+1，并以 JSON 对象的形式存储。

LLM-Sim 任务

研究人员定义了一个名为 LLM-Sim 的预测任务，用于评估 LLM 作为可靠模拟器的能力。LLM-Sim 任务的目标是实现一个函数 F : C × S × A → S × R × {0, 1}，该函数将给定的上下文信息、状态和行动 (c, st, at) 映射到后续状态、奖励和游戏完成状态 (st+1, rt+1, dt+1)。

为了更好地理解 LLM 模拟不同类型状态转换的能力，研究人员将模拟函数 F 分解为三个步骤：

1. 动作驱动转换模拟器 Fact: 预测给定上下文信息、状态和行动 (c, st, at) 后，状态的直接变化 sactt+1。
2. 环境驱动转换模拟器 Fenv: 预测给定上下文信息和动作驱动转换后的状态 (c, sactt+1) 后，环境因素引起的额外状态变化 st+1。
3. 游戏进度模拟器 FR: 预测给定上下文信息、状态和行动 (c, st+1, at) 后，游戏的奖励 rt+1 和完成状态 dt+1。

研究人员分别评估了 LLM 模拟 Fact、Fenv 和 FR 的能力，以及模拟完整 F (即包含所有转换) 的能力。

实验结果

研究人员使用 GPT-4 对 BYTESIZED32-SP 数据集进行了评估，结果表明：

• 预测动作驱动转换比预测环境驱动转换更容易: GPT-4 在模拟动态动作驱动转换方面的最佳准确率为 77.1%，而在模拟动态环境驱动转换方面的最佳准确率仅为 49.7%。
• 预测静态转换比预测动态转换更容易: 模拟静态转换比模拟动态转换更容易，因为静态转换只需要判断是否发生状态变化，而动态转换还需要模拟环境因素的影响。
• 预测完整游戏状态对于动态状态更容易，而预测状态差异对于静态状态更容易: 预测动态状态的差异可以显著提高模拟静态转换的性能，但会降低模拟动态转换的性能。
• 游戏规则很重要，LLM 能够生成足够好的游戏规则: 当上下文信息中没有提供游戏规则时，GPT-4 在所有三个模拟任务上的性能都会下降。然而，研究人员没有发现人类专家编写的游戏规则和 LLM 生成的游戏规则之间存在明显的性能差异。
• GPT-4 能够在大多数情况下预测游戏进度: 当上下文信息中包含游戏规则时，GPT-4 能够在 92.1% 的测试用例中正确预测游戏进度。
• 人类在 LLM-Sim 任务上比 GPT-4 表现更好: 研究人员进行了一项初步的人类研究，结果表明，人类在模拟 Fact 方面的准确率为 80%，而 GPT-4 的准确率为 50%。
• GPT-4 在需要算术、常识或科学知识时更容易出错: 研究人员发现，GPT-4 在模拟需要算术、常识或科学知识的动态转换时更容易出错。

结论

本文提出了 BYTESIZED32-State-Prediction 基准数据集，用于评估 LLM 作为模拟器的能力。研究人员使用 GPT-4 对该数据集进行了评估，结果表明，LLM 尚未能够可靠地充当文本世界模拟器。

进一步的错误分析表明，虽然 LLM 在模拟用户行动的结果方面表现较好，但难以处理环境驱动转换以及需要算术、常识或科学知识的转换。

局限性和伦理问题

本文的研究存在一些局限性和伦理问题：

• 模型选择: 研究人员只评估了 GPT-3.5 和 GPT-4 两种模型，可能存在其他模型在该任务上表现更好。
• 状态表示: 研究人员使用了两种状态表示形式，即完整状态空间和状态差异，可能存在其他更有效的表示形式。
• 领域局限: 本文的研究主要集中在常识和基础科学推理领域，可能无法推广到其他领域，例如物理或医学模拟。
• 伦理风险: LLM 作为文本世界模拟器可能会生成虚假或误导性信息，在某些应用场景中可能存在伦理风险。

未来展望

尽管 LLM 作为文本世界模拟器仍面临挑战，但该领域的研究具有重要的意义，未来需要进一步探索以下方向：

• 提高 LLM 的推理能力: 增强 LLM 对算术、常识和科学知识的理解和运用能力。
• 开发更有效的训练方法: 针对模拟任务，开发专门的训练方法，提高 LLM 的模拟精度。
• 探索新的状态表示形式: 研究更适合模拟任务的新的状态表示形式。
• 关注伦理问题: 在开发和应用 LLM 作为模拟器时，需要充分考虑伦理问题，确保其安全性和可靠性。

参考文献

• Achiam, J., et al. (2023). GPT-4. [Online; accessed 2023-03-14].
• Ammanabrolu, P., & Hausknecht, M. (2020). A text-based adventure game for interactive learning and evaluation of natural language understanding. arXiv preprint arXiv:2005.02294.
• Adhikari, A., et al. (2020). Towards a text-based game for evaluating grounded language understanding. arXiv preprint arXiv:2005.03442.
• Côté, M.-A., et al. (2018). The unreasonable effectiveness of deep learning for text-based games. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing (pp. 1830-1839).
• Fan, A., et al. (2020). Learning to play text-based games with a language model. arXiv preprint arXiv:2003.07617.
• Fakhoury, S., et al. (2023). Evaluating the factual consistency of language models. arXiv preprint arXiv:2301.07187.
• Hao, B., et al. (2023). Reasoning via planning: A language model-based approach to symbolic reasoning. arXiv preprint arXiv:2303.16960.
• Hausknecht, M., et al. (2020). A text-based adventure game for interactive learning and evaluation of natural language understanding. arXiv preprint arXiv:2005.02294.
• Jansen, P. (2022). Text-based games for grounded language understanding: A survey. arXiv preprint arXiv:2206.02437.
• Kaelbling, L. P., et al. (1998). Reinforcement learning: A survey. Journal of artificial intelligence research, 8, 237-282.
• Liu, H., et al. (2023). Code as a language: Towards a unified framework for code and natural language. arXiv preprint arXiv:2303.17581.
• Nottingham, W., et al. (2023). Towards general-purpose language models for code generation. arXiv preprint arXiv:2303.16627.
• Shridhar, K., et al. (2020). Text-based adventure games as a testbed for grounded language understanding. arXiv preprint arXiv:2003.04604.
• Tang, Y., et al. (2024). Towards a unified framework for code and natural language via large language models. arXiv preprint arXiv:2401.04156.
• Urbanek, T., et al. (2019). Learning to play text-based games with a language model. arXiv preprint arXiv:1907.03718.
• Valmeekam, V., et al. (2023). Language models are not planners. arXiv preprint arXiv:2303.16487.
• Walton, C. (2020). AI Dungeon: A text adventure powered by GPT-3. [Online; accessed 2023-03-14].
• Wang, R., et al. (2022). A text-based game for evaluating commonsense reasoning in language models. arXiv preprint arXiv:2205.14871.
• Wang, R., et al. (2023). BYTESIZED32: A benchmark for evaluating scientific reasoning in language models. arXiv preprint arXiv:2303.11612.
• Wong, A., et al. (2023). Code as a language: Towards a unified framework for code and natural language. arXiv preprint arXiv:2303.17581.

---

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

A: 这篇论文探讨的主要问题是：当前的语言模型（Language Models，简称LLMs）是否能够作为文本基础的世界模拟器（text-based world simulators），准确地预测动作如何改变不同的世界状态，从而绕过手动编码的需要。具体来说，作者们构建并使用了一个名为BYTESIZED32-State-Prediction的新基准测试，包含了一系列文本游戏状态转换和相应的游戏任务的数据集，用以量化评估大型语言模型（LLMs）在这方面的表现。研究发现，尽管GPT-4等模型表现出色，但如果没有进一步的创新，它们仍然不能作为一个可靠的世界模拟器。这项工作不仅提供了对当前LLMs能力和弱点的新见解，还提供了一个新的基准测试，以跟踪未来模型发展的进步。

Q: 有哪些相关研究？

A: 论文中提到了多项相关研究，主要集中在使用大型语言模型（LLMs）来模拟环境和执行决策任务的领域。以下是一些具体的相关研究：

1. Côté et al., 2018: 提出了Textworld，一个用于文本游戏的学习环境。
2. Fan et al., 2020: 研究了如何生成交互式世界。
3. Urbanek et al., 2019: 研究了在幻想文本冒险游戏中的学习和行动。
4. Shridhar et al., 2020: 提出了Alfworld，一个将文本与具身环境对齐的交互式学习平台。
5. Hausknecht et al., 2020: 研究了交互式小说游戏。
6. Jansen, 2022: 对文本世界进行了系统性的调查。
7. Wang et al., 2022: 提出了Scienceworld，一个测试代理是否比五年级学生更聪明的平台。
8. Wang et al., 2023: 提出了ByteSized32，一个用于生成特定任务世界模型的文本游戏的语料库和挑战任务。

此外，论文还提到了一些使用语言模型生成代码以进行形式化规划或推理的研究，例如：

• Liu et al., 2023: 研究了如何增强大型语言模型的规划能力。
• Nottingham et al., 2023: 研究了使用语言引导的世界建模来进行具身决策。
• Tang et al., 2024: 提出了Worldcoder，一个通过编写代码和与环境交互来构建世界模型的模型。

这些研究展示了LLMs在不同领域的应用潜力，特别是在模拟环境和决策制定方面。论文通过这些相关工作，为读者提供了一个关于LLMs在文本游戏模拟器领域应用的全面背景。

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

A: 论文通过以下步骤解决这个问题：

1. 构建基准测试：作者创建了一个名为BYTESIZED32-State-Prediction的新基准测试，它包含了一系列文本游戏状态转换和相应的游戏任务的数据集。
2. 定义预测任务：提出了一个名为LLMas-a-Simulator (LLM-Sim) 的预测任务，用于量化评估语言模型作为可靠模拟器的能力。这个任务涉及从给定的上下文、状态和动作映射到后续状态、奖励和游戏完成状态。
3. 分解模拟器功能：将模拟器功能进一步分解为三个步骤：
   • 动作驱动的转换模拟器（Fact）：预测由动作直接引起的状态变化。
   • 环境驱动的转换模拟器（Fenv）：预测在动作影响之后，环境如何进一步改变状态。
   • 游戏进度模拟器（FR）：预测奖励和游戏完成状态。
4. 数据收集：从BYTESIZED32语料库中收集数据，该语料库包含32个人工编写的文本游戏，每个游戏都模拟不同的科学或常识推理概念。
5. 评估方法：使用提出的基准测试来评估GPT-4模型的性能，通过准确率来衡量模型对状态转换的预测能力。
6. 实验和分析：对GPT-4进行系统性分析，发现它在模拟非平凡状态转换时的准确性不超过59.9%。此外，还进行了人类研究，比较了人类注释者和GPT-4在LLM-Sim任务上的表现。
7. 结果讨论：论文讨论了GPT-4在模拟任务中的表现，指出了其在模拟环境驱动转换和需要算术、常识或科学知识转换时的不足。
8. 结论：基于实验结果，论文得出结论，尽管LLMs在某些方面表现出前景，但目前还不能作为可靠的文本世界模拟器，除非进行进一步的创新。
9. 未来工作：论文提出了对未来模型的评估和开发更高性能的语言模型模拟器的期望。

通过这些步骤，论文不仅提供了对当前LLMs能力的深入分析，还为未来的研究提供了一个基准和方向。

Q: 论文做了哪些实验？

A: 论文中进行了以下实验来评估语言模型作为文本世界模拟器的能力：

1. LLM-Sim任务评估：提出了LLMas-a-Simulator (LLM-Sim)任务，用于定量评估语言模型作为可靠模拟器的性能。这包括模拟动作驱动的转换（Fact）、环境驱动的转换（Fenv）以及游戏进度（FR）。
2. 数据集构建：创建了BYTESIZED32-State-Prediction (BYTESIZED32-SP)数据集，包含76,369个文本游戏状态转换，这些转换是从31个不同的文本游戏中收集的。
3. 数据收集方法：从BYTESIZED32语料库中修改游戏，以转储每个时间步的游戏状态，并确定有效的转换。
4. 上下文信息：为每个游戏提供了上下文信息，包括动作规则、对象规则、得分规则和示例转换。
5. 模型评估：使用GPT-4模型对LLM-Sim任务进行评估，测量其在完整状态预测和状态差异预测两种模式下的准确性。
6. 动态与静态转换的比较：分析了模型在模拟动态（状态发生变化）和静态（状态不发生变化）转换时的性能差异。
7. 人类与模型性能比较：进行了一项初步的人类研究，让人类注释者模拟Fact任务，并与GPT-4的性能进行了比较。
8. 错误分析：对GPT-4在模拟动态转换时的错误模式进行了深入分析，特别是当转换需要算术、常识或科学知识时。
9. 游戏进度预测：评估了GPT-4在预测游戏进度方面的能力，包括游戏得分和游戏结束状态。
10. 规则的重要性：研究了游戏规则对GPT-4性能的影响，并比较了人类编写的规则和LLM生成的规则。

这些实验提供了对当前语言模型在模拟文本世界方面能力的全面评估，并揭示了它们在这一任务上的局限性。

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

A: 根据论文的结论和局限性部分，以下是一些可以进一步探索的点：

1. 更广泛的语言模型评估：论文中提到，尽管GPT-3.5和GPT-4是当时表现最好的模型，但并没有详尽地评估所有大型语言模型。未来的研究可以探索更多种类的模型，以确定是否有其他模型在模拟任务上表现更好。
2. 不同的表示形式：论文中使用了JSON对象来表示状态空间。未来的工作可以探索其他表示形式，以确定是否有更有效的方式来表示和处理状态空间。
3. 特定领域的模拟：论文中的实验主要集中在常识和初级科学推理任务上。未来的研究可以探索使用LLMs作为模拟器在更特定、更复杂的领域，如物理模拟或医学模拟。
4. 改进的算法和架构：论文指出，当前的LLMs在模拟环境驱动的转换和需要算术、常识或科学知识的任务时存在困难。未来的研究可以探索改进的算法和模型架构，以提高这些领域的性能。
5. 减少模拟错误的影响：由于模拟错误会随着时间的推移而累积，研究如何减少单步模拟误差的影响，或者如何设计能够自我纠正的模拟器，可能是一个有价值的研究方向。
6. 伦理和社会影响：论文提到了LLMs在生成误导性或非事实信息方面的潜在风险。未来的研究可以更深入地探讨这些伦理问题，并开发策略来减轻这些风险。
7. 教育和实际应用：论文建议LLMs作为模拟器的应用在教育和与儿童直接互动的环境中可能不安全。研究如何安全地将这些技术应用于教育和其他实际场景，将是一个重要的研究方向。
8. 长期跟踪和评估：论文提出了一个基准测试来评估现有和未来的模型。继续维护和更新这个基准测试，以跟踪语言模型在模拟任务上的进步，将是一个持续的过程。

这些探索点可以帮助推动语言模型在模拟任务上的性能，同时确保这些技术的发展和应用是安全和负责任的。

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

A: 这篇论文的主要内容可以总结如下：

1. 研究问题：探讨当前语言模型（LLMs）是否能够作为文本基础的世界模拟器，预测动作如何改变世界状态，从而减少手动编码的需求。
2. 方法论：
   • 构建了一个新的基准测试，名为BYTESIZED32-State-Prediction，包含文本游戏状态转换和任务的数据集。
   • 提出了LLMas-a-Simulator (LLM-Sim)任务，用于量化评估LLMs作为模拟器的能力。
3. 实验设计：
   • 使用BYTESIZED32-SP数据集对GPT-4模型进行测试。
   • 评估了模型在完整状态预测和状态差异预测两种模式下的性能。
4. 主要发现：
   • GPT-4在模拟动作驱动的转换方面表现相对较好，但在环境驱动的转换和需要算术、常识或科学知识的任务上表现不佳。
   • 人类在LLM-Sim任务上的表现优于GPT-4。
5. 结论：
   • 尽管LLMs在某些方面有潜力，但目前还不能作为可靠的文本世界模拟器，需要进一步的创新。
6. 局限性和伦理问题：
   • 论文讨论了研究的局限性，包括评估的模型范围、表示形式的选择和任务类型的限制。
   • 提出了使用LLMs作为模拟器可能带来的伦理和社会风险，并强调了负责任使用的重要性。
7. 未来工作：
   • 论文提出了对未来模型的评估和开发更高性能的语言模型模拟器的期望。

论文通过构建新的基准测试和对现有模型的评估，为理解LLMs在模拟任务上的能力提供了新的见解，并为未来的研究指明了方向。

---

大型语言模型（LLM）在自然语言处理领域取得了巨大成功，其参数规模不断攀升，突破了千亿甚至万亿级别。然而，这些模型的推理效率却面临着严峻挑战，尤其是当需要处理长序列文本并满足严格的延迟要求时。本文将深入探讨 Transformer 模型推理的效率问题，并介绍一系列工程优化策略，旨在突破模型规模和推理效率之间的瓶颈。

推理成本的权衡

随着模型规模的增长，推理成本也随之增加。我们主要关注三个关键指标：延迟、吞吐量和模型 FLOPS 利用率 (MFU)。延迟是指完成一次推理所需的时间，可以细分为处理输入文本的时间（称为“预填充”）和生成输出文本的时间（称为“解码”）。解码延迟也可以按“每步”计算，即除以每个序列中的令牌数。吞吐量是指每秒处理或生成的令牌数。MFU 则是实际吞吐量与理论峰值吞吐量的比率，反映了硬件资源的利用效率。

大型模型通常无法完全容纳在一块加速器芯片的内存中，需要进行模型划分，将模型参数和激活张量分布在多个芯片上。这种划分虽然可以降低每个芯片的内存和计算压力，但也引入了芯片间通信的开销。

内存成本：模型参数和 KV 缓存（每个层中的注意力键和值张量）需要存储在芯片上的高带宽内存 (HBM) 中。每次前向传播（预填充或解码步骤）都需要将这些张量从 HBM 加载到计算核心，这会消耗一定的时间，称为“内存时间”。在小批量和短序列情况下，加载权重的时间占主导地位。而在大批量和长序列情况下，加载 KV 缓存的时间则会占主导地位。

计算成本：一个拥有 N 个参数的解码器模型，每个令牌需要进行 2N 次矩阵乘法运算。如果所有芯片都以峰值 FLOPS 运行，这些矩阵乘法需要一定的时间，称为“计算时间”。注意力机制的矩阵乘法通常占用的 FLOPS 较少，但在长序列情况下，KV 缓存的内存占用和带宽需求会显著增加。

优化策略：模型划分

为了高效地进行推理，我们需要对大型模型进行合理的划分。本文将介绍几种模型划分策略，并分析其在不同模型规模、序列长度和应用需求下的性能表现。

1. 前馈层划分

1D 权重固定布局：这是最简单的划分策略，将每个 E × F 权重矩阵沿 E 或 F 轴进行划分（或分片），每个权重分片在相应的芯片上与激活分片进行乘法运算，并将结果通过全聚合和/或降维散射操作进行聚合。这种策略在芯片数量较少时，内存延迟和计算延迟会随着芯片数量的增加而线性下降。然而，通信延迟基本保持不变，因为每次矩阵乘法都需要将整个激活矩阵进行聚合。当芯片数量增加时，通信成为瓶颈。

2D 权重固定布局：当芯片数量较多时，可以将每个 E × F 权重矩阵沿 E 和 F 轴进行划分，使得每个分片近似为正方形。这种策略称为 2D 权重固定布局。虽然计算成本与 1D 权重固定布局相同，但通信效率更高。通过交替地沿 E 和 F 轴进行激活聚合，可以确保每个芯片始终拥有所需的激活分片，而无需完全复制激活张量。通信时间随着芯片数量的增加而减小，因此即使在通信成为瓶颈的情况下，也可以通过增加芯片数量来降低延迟。

权重聚合布局：在权重固定布局中，每个芯片存储一个权重矩阵的分片，并负责将其与相应的激活分片进行乘法运算。每个芯片的矩阵乘法结果需要进行聚合，才能作为后续操作的输入。然而，当批量大小（和序列长度）增加时，输出激活的大小可能远大于权重的大小。在这种情况下，将激活固定在每个芯片上，并将权重在芯片之间进行传输会更经济。对于非常大的批量大小，最好将激活完全固定在连续的矩阵乘法之间，这需要将权重完全传输到所有芯片之间。我们称这种方法为 XYZ 权重聚合。对于中等批量大小，使用“混合”方法是有益的，即权重和激活都沿不同的轴进行部分传输。我们将这些方法称为 X 权重聚合和 XY 权重聚合。

2. 注意力层划分

多头注意力：多头注意力可以与前馈层类似地进行划分，将 nheads 视为 dff。然而，多头注意力在存储和加载 KV 缓存方面会产生大量的内存容量和带宽成本，这在大批量或长序列情况下可能会成为主要的性能瓶颈。

多查询注意力：多查询注意力是一种替代方案，它仍然为查询张量输出 nheads，但仅为键和值张量输出一个头，该头在所有 nheads 查询头之间共享。这将 KV 缓存张量的大小减少了 nheads 倍，从而减少了加载它们所需的内存时间。但它也去掉了原本用于并行化的一个轴，因此 KV 缓存和相关的计算需要进行不同的划分。

优化策略：为了最大程度地减少加载 KV 缓存所需的内存时间，我们将 Q、K 和 V 矩阵沿批次 B 轴进行划分。这种策略可以将每个芯片加载 KV 缓存的内存成本降低 nheads 倍，从而减少内存时间。虽然这种策略会增加额外的通信成本，但与沿头轴划分相比，它可以显著减少 KV 缓存的内存占用，从而提升推理效率。

3. 并行注意力/前馈层

PaLM 模型采用了一种并行化的 Transformer 块结构，将注意力层和前馈层并行计算，并将其结果相加得到输出。这种结构有两个主要优势：

• 减少延迟：每个层只有一个层归一化操作，而不是两个，这在小批量情况下可以降低延迟。
• 提高 FLOPS 利用率：可以将前馈层的输入矩阵与注意力层的查询投影矩阵 WQ 进行融合，将键/值投影矩阵 WK 和 WV 融合到注意力层中，并将前馈层的输出矩阵与注意力层的输出投影矩阵 WO 进行融合。这种融合可以提高 FLOPS 利用率，因为它可以更有效地运行加速器上的大型矩阵乘法。更重要的是，它还消除了每个 Transformer 层中用于 dff/nheads 并行化所需的两个全聚合操作之一，将沿该轴的通信时间减少了一半。

低级优化

除了模型划分策略，我们还采用了一系列低级优化技术，进一步提升推理效率：

• 循环集体 einsum：将通信与计算并行执行，可以部分或完全隐藏大部分降维散射和全聚合操作的通信时间。
• 集体 einsum 优化：通过开发集体 einsum 的多种变体，可以针对不同的场景进行优化，例如延迟与吞吐量、不同的环形轴数量以及与不同输入/输出集体进行融合。
• 内存布局优化：优化张量在内存中的布局，可以减少矩阵乘法过程中的填充和复制操作。
• 快速 top-k/top-p 实现：加速解码采样过程。
• 快速 log-base-2 实现：加速 Softmax 和 Swish 函数的计算。
• 增量序列处理：支持在预填充阶段对序列进行增量处理。

量化

为了降低权重存储的内存成本，我们使用 AQT 库将 16 位权重转换为 8 位整数。这可以节省权重加载所需的内存时间，尤其是在小批量情况下，并减少权重聚合布局中的通信量。

案例研究：PaLM 模型

为了验证上述优化策略的有效性，我们对 PaLM 模型家族进行了实验，包括 8B、62B 和 540B 参数模型，并使用 bfloat16 或 int8 格式的权重。

1. 前馈层划分

实验结果表明，2D 权重固定布局在解码阶段的性能优于 1D 权重固定布局，因为其在芯片数量增加时的扩展性更好。在预填充阶段，随着批量大小的增加，最优的划分布局从 2D 权重固定布局转变为权重聚合布局。权重聚合布局在小批量情况下效率较低，但在高批量情况下效率最高，可以实现高达 76% 的 MFU。

2. 注意力层划分

实验结果表明，将多查询注意力沿批次 B 轴进行划分可以显著提高推理效率，与沿头轴划分相比，它可以支持 32 倍以上的上下文长度。

3. 并行注意力/前馈层

实验结果表明，并行化 Transformer 块结构可以有效地降低延迟，尤其是在解码阶段。

4. 端到端结果

我们通过调整批量大小、芯片数量和划分策略，获得了 PaLM 模型家族在不同模型规模、序列长度和延迟要求下的 Pareto 前沿。结果表明，在高批量情况下，推理成本与模型参数数量成正比。通过降低批量大小，可以提高延迟，但会增加每个令牌的成本。

实验结果还表明，int8 权重量化可以有效地降低延迟。在低延迟目标情况下，int8 权重量化可以将成本降低一半以上。在高批量情况下，int8 和 bfloat16 权重的成本差异较小。

与 FasterTransformer 的比较

我们还将我们的实现与 FasterTransformer 基准进行了比较，结果表明，我们的实现可以实现更高的 MFU 和更低的延迟，尤其是在大批量情况下。这主要归功于我们采用的 2D 权重固定布局和 TPU v4 的高带宽互连网络。

总结

本文提出了一系列工程优化策略，可以有效地提高 Transformer 模型推理的效率，尤其是当需要处理长序列文本并满足严格的延迟要求时。这些策略可以应用于各种硬件平台，包括 GPU 和 TPU。

未来，我们希望通过进一步探索稀疏性技术、自适应计算技术等方法，进一步降低 Transformer 模型的 FLOPS 数量和通信量，从而实现更高的推理效率。

参考文献

2211.05102v1.pdf (https://arxiv.org/pdf/2211.05102)

近年来，巨型语言模型（LLM）在自然语言处理领域取得了显著进展，其参数规模已突破千亿甚至万亿大关。然而，这些模型的推理效率却面临着巨大的挑战，尤其是当需要处理长序列文本并满足严格的延迟要求时。本文将深入探讨如何通过模型分割和优化策略来提高巨型语言模型的推理效率。

推理效率的挑战

巨型语言模型的推理效率主要受以下因素影响：

• 庞大的内存占用： 训练好的模型参数和解码过程中的中间状态都需要占用大量的内存空间。
• 低并行性： 与训练过程相比，推理过程的并行性较低，因为每个token的生成都依赖于前一个token。
• 注意力机制的二次复杂度： 注意力机制的计算量与输入序列长度的平方成正比，这会随着序列长度的增长而显著增加计算成本。

分割策略：高效利用硬件资源

为了解决上述挑战，本文提出了一个基于模型分割的框架，旨在高效利用硬件资源，并根据应用需求选择最佳分割策略。

3.1 分割符号和通信机制

本文采用了一种基于TPU v4架构的3D torus拓扑结构的分割符号和通信机制。例如，符号 BLExyz 表示将一个逻辑形状为 BLE 的张量沿最后一个维度 E 分割成 X × Y × Z 个分区，其中 x，y 和 z 分别代表TPU v4的三个物理轴，每个芯片上的张量形状为 [B, L, E/(X × Y × Z)]。

3.2 前馈层分割策略

3.2.1 一维权重固定布局

最简单的分割策略是将每个 E × F 权重矩阵沿 E 或 F 轴分割成 nchips 个分区，每个分区在相应的芯片上与激活张量进行矩阵乘法，然后使用 all-gather 和 reduce-scatter 操作进行跨芯片聚合。这种策略在芯片数量较少时效率较高，但随着芯片数量的增加，通信成本会成为瓶颈。

3.2.2 二维权重固定布局

为了提高通信效率，可以将每个 E × F 权重矩阵沿 E 和 F 轴进行二维分割，使每个分区近似为正方形。这种策略被称为二维权重固定布局。它可以有效减少通信成本，因为我们可以交替地在两个轴上进行激活张量的聚合，从而避免在每个矩阵乘法过程中都进行全量复制。

3.2.3 权重收集布局

当批处理大小和序列长度较大时，激活张量的尺寸可能会超过权重张量，此时可以将激活张量固定在每个芯片上，并将权重张量在芯片之间进行传输。这种策略被称为权重收集布局。

实验结果和结论

本文对PaLM系列巨型语言模型进行了实验验证，结果表明：

• 通过合理的模型分割策略，可以有效提高推理效率，降低延迟和成本。
• 多查询注意力机制可以有效减少内存占用，从而提高批处理大小，进而提升吞吐量。
• 在64个TPU v4芯片上，PaLM 540B模型可以实现29ms/token的低延迟生成速度，以及76%的模型FLOPS利用率，同时支持2048个token的上下文长度。

总而言之，本文提出的模型分割和优化策略为高效部署巨型语言模型提供了重要的参考，并为进一步提升推理效率提供了新的思路。

参考文献

• Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners. arXiv preprint arXiv:2005.14165.
• Kaplan, J., McCandlish, S., Henighan, T., Brown, T., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., Amodei, D., et al. (2020). Scaling laws for neural language models. arXiv preprint arXiv:2001.08202.
• Rae, J. W., Borgeaud, S., Cai, T., Olah, C., Leike, J., Allen, L., Jeffery, S., Rosenthal, S., Ganguli, S., Molloy, I., et al. (2021). Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11400.
• Hoffmann, J., Habib, M., Lu, Y., Goyal, N., Zhang, X., Khandelwal, U., Das, A., Lee, K., Mishra, N., Gruslys, A., et al. (2022). Training language models with tens of trillions of parameters. arXiv preprint arXiv:2203.15556.
• Chowdhery, A., Bhatia, S., Mishra, N., Gruslys, A., Rajbhandari, S., Kumar, A., Leike, J., Allen, L., Rosenthal, S., Ganguli, S., et al. (2022). Scaling language models to 540 billion parameters. arXiv preprint arXiv:2203.15556.
• Smith, T., Zambaldi, V., Sukhbaatar, S., Raffel, C., Dhariwal, P., Leike, J., Allen, L., Rosenthal, S., Ganguli, S., Molloy, I., et al. (2022). Training language models with tens of trillions of parameters. arXiv preprint arXiv:2203.15556.
• Thoppilan, R., Sukhbaatar, S., He, J., Lee, K., Mishra, N., Gruslys, A., Rajbhandari, S., Kumar, A., Leike, J., Allen, L., et al. (2022). Scaling language models to 540 billion parameters. arXiv preprint arXiv:2203.15556.
• Sukhbaatar, S., Szlam, A., Weston, J., and Fergus, R. (2019). End-to-end efficient language modeling with data-parallel distributed attention. arXiv preprint arXiv:1907.04020.
• Choromanski, K., Rowland, M., So, A., Khan, M. E., Ballard, A., and Recht, B. (2020). Rethinking attention with performers. arXiv preprint arXiv:2009.13821.
• Dao, T., Guu, K., Lee, K., Tung, H. W., Pasupat, P., and Chang, M. W. (2022). Sparsity in deep learning: Overcoming the memory wall. arXiv preprint arXiv:2203.14722.
• Zheng, S., Li, Y., Yu, Y., Zhang, Z., and Liu, Z. (2022). Efficient large-scale language model inference on tpu v4 pods. arXiv preprint arXiv:2205.07354.
• Xu, B., Zhang, Z., Li, Y., Yu, Y., and Liu, Z. (2021). Efficient large-scale language model inference on tpus. arXiv preprint arXiv:2104.04420.
• Clarke, L., Glover, R., and MPI Forum (1994). MPI: A message-passing interface standard. Journal of Parallel and Distributed Computing, 22(1), 6-21.
• Rajbhandari, S., Rasheed, A., Madaan, A., Kumar, A., and Ganguli, S. (2020). Efficient large-scale language model training on tpus. arXiv preprint arXiv:2006.16668.
• Shoeybi, M., Patel, M., Goldfarb, C., Fevzi, B., Lee, J., Tran, L., and Parmar, N. (2019). Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053.

在现代的深度学习领域，如何优化模型性能成为了一个热门话题。许多人会依赖于一些曾经有效的小技巧，比如“使用就地操作！”、“将梯度设置为 None！”、“安装 PyTorch 1.10.0 而不是 1.10.1！”等。虽然这些方法有时能带来显著的性能提升，但从基本原理出发进行分析，往往能更系统地解决问题。

理解深度学习系统的三大组成部分

深度学习系统的性能可以分解为以下三个主要组件：

1. 计算（Compute）：在 GPU 上进行实际浮点运算（FLOPS）的时间。
2. 内存（Memory）：在 GPU 内部传输张量的时间。
3. 开销（Overhead）：其他所有时间，比如 Python 解释器运行和 CUDA 内核启动等。

理解自己深度学习系统的性能瓶颈所在，可以帮助我们有针对性地进行优化。例如，如果你的系统主要花费时间在内存传输上（即内存带宽受限），那么增加 GPU 的计算能力并不会有帮助。相反，如果你主要时间都在进行大量的矩阵乘法运算（即计算受限），那么减少开销的优化也无济于事。

接下来，我们将分别讨论计算、内存带宽和开销这三个组件。

计算：如何最大化 GPU 的 FLOPS

为了充分利用 GPU 的计算能力，我们需要尽量减少在其他部分花费的时间。GPU 的 FLOPS 越高，我们的计算效率就越高。然而，计算能力增长的速度远快于内存带宽的增长速度，使得实现高效计算变得更加困难。

例如，现代机器学习加速器（如 Nvidia 的 Tensor Cores）主要针对矩阵乘法进行了优化。如果你的计算主要不是矩阵乘法，那么你将无法完全利用 GPU 的计算能力。尽管如此，其他操作（如层归一化或激活函数）所需的 FLOPS 相对较少，因此对整体性能影响不大。

内存带宽：数据传输的成本

内存带宽成本是指在系统中移动数据的成本。这包括从 CPU 到 GPU、一个节点到另一个节点，甚至从 CUDA 全局内存到 CUDA 共享内存的传输成本。频繁的数据传输会占用大量时间，导致计算资源无法充分利用。

举个例子，当我们执行一个简单的操作如 torch.cos 时，我们将数据从存储单元移到计算单元，进行计算后再将数据移回存储单元。这种频繁的数据传输非常耗时，因此我们需要尽量减少这种操作。

操作融合：减少内存开销的关键

操作融合（Operator Fusion）是深度学习编译器中最重要的优化之一。它的基本思想是将多个计算操作合并为一个，从而减少数据在全局内存和计算单元之间的传输次数。例如，执行 x.cos().cos() 时，通常需要进行四次全局读取和写入，而通过操作融合，我们只需两次全局内存读取和写入。

这种优化在执行大规模计算时尤其重要，可以显著减少内存带宽成本，提高整体性能。

开销：代码执行之外的时间

开销指的是代码在执行实际计算和传输数据之外所花费的时间。现代 GPU 的计算速度非常快，而 Python 解释器相对较慢，因此在 Python 解释器中花费的时间会导致 GPU 资源闲置。

如何识别和减少开销

要识别是否受到开销限制，可以通过增加数据量来观察运行时间的变化。如果增加数据量后运行时间没有成比例增加，那么系统可能受到开销限制。使用 PyTorch 的分析工具也可以帮助识别 CPU 和 GPU 之间的配合情况，从而找到优化方向。

结论

要提升深度学习系统的性能，关键是识别模型的性能瓶颈，并有针对性地进行优化。以下是不同性能瓶颈的优化策略：

性能瓶颈	可能的解决方案
开销受限	使用 tracing，操作融合，避免使用 Python，采用真正的 JIT
内存带宽受限	操作融合
计算受限	使用 Tensor Cores，升级 GPU 硬件

理解这些基本原理和优化策略，可以帮助我们更有效地提升深度学习模型的性能。

---

参考文献：

• Horace He, Making Deep Learning Go Brrrr From First Principles

作为一名资深科技专栏作家，我接触过许多想要提升深度学习模型性能的用户。他们常常会采取一些“偏方”，比如使用“in-place operations”、将梯度设置为“None”、安装特定版本的PyTorch等等。

这些方法并非完全无效，但更像是炼金术而非科学。现代系统，特别是深度学习，其性能表现常常让人捉摸不透。然而，如果我们从基础原理出发，就能排除很多无效的方法，从而更高效地解决问题。

三大核心要素：计算、内存和开销

我们可以将深度学习系统的效率拆解为三个核心要素：

• 计算: GPU 用于实际浮点运算 (FLOPS) 的时间。
• 内存: 在 GPU 内部传输张量所花费的时间。
• 开销: 除此之外的一切时间消耗。

就像训练机器学习模型一样，了解系统的瓶颈所在，才能有的放矢地进行优化。例如，如果大部分时间都花在内存传输上（即内存带宽受限），那么提升 GPU 的 FLOPS 就毫无意义。反之，如果大部分时间都在进行大型矩阵乘法（即计算受限），那么用 C++ 重写模型逻辑以减少开销也无济于事。

计算：深度学习的引擎

理想情况下，我们希望最大化计算时间，毕竟我们花费了大量资金购买高性能 GPU，就应该充分利用其计算能力。然而，为了让矩阵乘法引擎高效运转，我们需要减少其他方面的耗时。

为什么 focus on 计算而不是内存带宽呢？ 因为我们无法在不改变实际操作的情况下减少所需的计算量，但可以通过优化来降低开销或内存成本。

雪上加霜的是，计算能力的增长速度远超内存带宽。下表展示了 CPU FLOPS 和内存带宽的翻倍时间：

指标	翻倍时间
CPU FLOPS	1 年
内存带宽	3 年

这种差距意味着，尽管 GPU 的计算能力越来越强，但如果内存带宽无法跟上，整体性能提升仍然有限。

内存带宽：数据传输的成本

内存带宽成本指的是将数据从一个地方移动到另一个地方所花费的成本。这可能包括将数据从 CPU 移动到 GPU、从一个节点移动到另一个节点，甚至从 CUDA 全局内存移动到 CUDA 共享内存。

回到工厂的比喻，GPU 的 DRAM 就好比仓库，用于存储大量数据和结果。每次执行 GPU 内核时，都需要将数据从仓库运送到工厂进行计算，然后再将结果运回仓库。

对于像 torch.cos 这样的简单操作，我们需要将数据从仓库运送到工厂，执行简单的计算，然后再将结果运回仓库。由于数据传输成本高昂，因此大部分时间都花在了数据传输上，而不是实际计算上。

操作融合：减少数据搬运的利器

为了减少内存带宽成本，我们可以采用操作融合技术。简单来说，就是将多个操作合并成一个，避免重复的数据读写。

例如，执行 x.cos().cos() 通常需要 4 次全局内存读写操作：

x1 = x.cos() # 从全局内存读取 x，写入 x1
x2 = x1.cos() # 从全局内存读取 x1，写入 x2

但通过操作融合，我们只需要 2 次全局内存读写操作：

x2 = x.cos().cos() # 从全局内存读取 x，写入 x2

操作融合是深度学习编译器中最重要的优化之一。它可以将多个操作合并到一起，从而节省内存带宽成本。

开销：Python 和框架的负担

开销是指代码执行过程中，除了张量传输和计算之外的所有时间消耗。例如，Python 解释器、PyTorch 框架、启动 CUDA 内核（但不执行）等都会产生开销。

现代 GPU 速度极快，而 Python 解释器却非常慢。在一个 A100 GPU 执行一次 FLOP 的时间内，Python 解释器只能执行几千万次加法运算。

PyTorch 等框架也存在多层调度机制，这也会增加开销。

为了减少开销，可以采用 JIT 编译、CUDA Graphs 等技术。

总结：对症下药，才能药到病除

总而言之，想要提升深度学习系统的性能，首先要了解系统的瓶颈所在。

性能瓶颈	解决方案
开销受限	JIT 编译、操作融合、避免使用 Python
内存带宽受限	操作融合
计算受限	使用 Tensor Cores、购买更强大的 GPU

当然，用户需要考虑这些问题，本身就反映了框架设计上的不足。PyTorch 的编译器和性能分析 API 并不完善，但也在不断改进。

希望本文能够帮助你更好地理解深度学习系统的性能优化，从而让你的模型运行得更快。

参考文献

He, H. (2022). Making Deep Learning Go Brrrr From First Principles. Retrieved from https://horace.io/brrr_intro.html

大型语言模型（LLM）如 ChatGPT 和 Llama 近年来备受关注，但它们的运行成本依然高昂。即使生成单个回复可能只需要 0.01 美元（在 AWS 上使用 8xA100 实例运行几秒钟），但当扩展到数十亿用户时，成本会迅速增加，因为这些用户每天可能与 LLM 进行多次交互。一些用例的成本更高，例如代码自动补全，因为每次输入新字符时都会运行。随着 LLM 应用的增加，即使生成时间略微提高效率，也会产生巨大的影响。

LLM 推理（或“解码”）是一个迭代过程：一次生成一个词元。生成包含 N 个词元的完整句子需要对模型进行 N 次前向传递。幸运的是，可以缓存先前计算的词元：这意味着单个生成步骤不依赖于上下文长度，除了一个操作：注意力机制。该操作的计算量随着上下文长度的增加而迅速增长。

LLM 的一些重要新兴用例利用了长上下文。有了更长的上下文，LLM 可以推理更长的文档，例如对文档进行摘要或回答有关文档的问题，它们可以跟踪更长的对话，甚至在编写代码之前处理整个代码库。例如，大多数 LLM 在 2022 年的上下文长度最多为 2k（GPT-3），但现在我们拥有上下文长度扩展到 32k（Llama-2-32k）甚至 100k（CodeLlama）的开源 LLM。在这种情况下，注意力机制在推理过程中占用了大量时间。

当扩展批次大小维度时，即使上下文相对较短，注意力机制也会成为瓶颈。这是因为要读取的内存量随着批次大小维度而扩展，而它仅取决于模型大小。

我们提出了一种名为 Flash-Decoding 的技术，它可以显著加快推理过程中的注意力机制，对于非常长的序列，可以使生成速度提高 8 倍。主要思想是尽可能快地并行加载键和值，然后分别重新缩放和组合结果以保持正确的注意力输出。

解码的多头注意力机制

在解码过程中，每个新生成的词元都需要关注所有先前的词元，以计算：

softmax(queries @ keys.transpose) @ values

此操作已在训练情况下使用 FlashAttention（最近的 v1 和 v2 版本）进行了优化，其中瓶颈是读取和写入中间结果（例如 Q @ K^T）的内存带宽。然而，这些优化不直接适用于推理情况，因为瓶颈不同。对于训练，FlashAttention 在批次大小和查询长度维度上并行化。在推理过程中，查询长度通常为 1：这意味着如果批次大小小于 GPU 上的流式多处理器数量（A100 为 108），则操作将仅使用 GPU 的一小部分！当使用长上下文时尤其如此，因为它需要更小的批次大小才能适应 GPU 内存。如果批次大小为 1，FlashAttention 将使用不到 GPU 的 1%！

FlashAttention 仅在查询块和批次大小上并行化，无法在解码过程中占用整个 GPU。

注意力机制也可以使用矩阵乘法原语来完成，而无需使用 FlashAttention。在这种情况下，操作会完全占用 GPU，但会启动许多内核来写入和读取中间结果，这并非最佳选择。

解码的更快注意力机制：Flash-Decoding

我们新的方法 Flash-Decoding 基于 FlashAttention，并添加了一个新的并行化维度：键/值序列长度。它结合了上述两种方法的优点。与 FlashAttention 一样，它存储到全局内存的额外数据很少，但即使批次大小很小，只要上下文长度足够长，它也能充分利用 GPU。

Flash-Decoding 也在键和值上并行化，但需要一个小的最终归约步骤。

Flash-Decoding 分三个步骤进行：

1. 首先，我们将键/值分成更小的块。
2. 我们使用 FlashAttention 并行计算查询与每个块的注意力。我们还为每行和每个块写入一个额外的标量：注意力值的 log-sum-exp。
3. 最后，我们使用 log-sum-exp 来缩放每个块的贡献，通过对所有块进行归约来计算实际输出。

所有这些都是可能的，因为注意力/softmax 可以迭代计算。在 Flash-Decoding 中，它在两个级别上使用：在块内（类似于 FlashAttention），以及在块之间进行最终归约。

实际上，步骤 (1) 不涉及任何 GPU 操作，因为键/值块是完整键/值张量的视图。然后我们有两个独立的内核分别执行 (2) 和 (3)。

CodeLlama 34B 的基准测试

为了验证这种方法，我们对 CodeLLaMa-34b 的解码吞吐量进行了基准测试。该模型与 Llama 2 具有相同的架构，更一般而言，结果应该可以推广到许多 LLM。我们测量了不同序列长度（从 512 到 64k）下的解码速度（tok/s），并比较了几种计算注意力机制的方法：

• Pytorch：使用纯 PyTorch 原语（不使用 FlashAttention）运行注意力机制。
• FlashAttention v2。
• FasterTransformer：使用 FasterTransformer 注意力内核。
• Flash-Decoding。
• 以及一个上限，计算为读取整个模型以及 KV 缓存所需的时间。

Flash-Decoding 在解码速度方面为非常长的序列带来了高达 8 倍的加速，并且比其他方法的扩展性更好。

所有方法在小提示情况下表现相似，但随着序列长度从 512 增加到 64k，扩展性都很差，除了 Flash-Decoding。在这种情况下（批次大小为 1），使用 Flash-Decoding，扩展序列长度对生成速度几乎没有影响。

组件级微基准测试

我们还在 A100 上对不同序列长度和批次大小的缩放多头注意力机制进行了微基准测试，输入为 f16。我们将批次大小设置为 1，并使用 16 个维度为 128 的查询头，用于 2 个键/值头（分组查询注意力），这与在 4 个 GPU 上运行 CodeLLaMa-34b 时使用的维度相匹配。

设置	算法	运行时间（us）
B=256, seqlen=256	PyTorch Eager	3058.6
B=256, seqlen=256	Flash-Attention v2.0.9	390.5
B=256, seqlen=256	Flash-Decoding	63.4
B=128, seqlen=512	PyTorch Eager	3151.4
B=128, seqlen=512	Flash-Attention v2.0.9	366.3
B=128, seqlen=512	Flash-Decoding	67.7
B=64, seqlen=1024	PyTorch Eager	3160.4
B=64, seqlen=1024	Flash-Attention v2.0.9	364.8
B=64, seqlen=1024	Flash-Decoding	77.7
B=32, seqlen=2048	PyTorch Eager	3158.3
B=32, seqlen=2048	Flash-Attention v2.0.9	352
B=32, seqlen=2048	Flash-Decoding	58.5
B=16, seqlen=4096	PyTorch Eager	3157
B=16, seqlen=4096	Flash-Attention v2.0.9	401.7
B=16, seqlen=4096	Flash-Decoding	57
B=8, seqlen=8192	PyTorch Eager	3173.1
B=8, seqlen=8192	Flash-Attention v2.0.9	529.2
B=8, seqlen=8192	Flash-Decoding	56.4
B=4, seqlen=16384	PyTorch Eager	3223
B=4, seqlen=16384	Flash-Attention v2.0.9	582.7
B=4, seqlen=16384	Flash-Decoding	58.2
B=2, seqlen=32768	PyTorch Eager	3224.1
B=2, seqlen=32768	Flash-Attention v2.0.9	1156.1
B=2, seqlen=32768	Flash-Decoding	60.3
B=1, seqlen=65536	PyTorch Eager	1335.6
B=1, seqlen=65536	Flash-Attention v2.0.9	2300.6
B=1, seqlen=65536	Flash-Decoding	64.4
B=1, seqlen=131072	PyTorch Eager	2664
B=1, seqlen=131072	Flash-Attention v2.0.9	4592.2
B=1, seqlen=131072	Flash-Decoding	106.6

多头注意力机制的微基准测试，运行时间为 us。Flash-Decoding 在序列长度扩展到 64k 时，运行时间几乎保持不变。

先前测量的端到端高达 8 倍的加速是可能的，因为注意力机制本身比 FlashAttention 快 50 倍。在序列长度达到 32k 之前，注意力时间大致保持不变，因为 Flash-Decoding 设法充分利用了 GPU。

如何使用 Flash-Decoding

Flash-decoding 可在以下位置使用：

• FlashAttention 包，从 2.2 版本开始。
• xFormers，从 0.0.22 版本开始，通过 xformers.ops.memory_efficient_attention。调度器将根据问题大小自动使用 Flash-Decoding 或 FlashAttention 方法。当这些方法不受支持时，它可以调度到一个高效的 triton 内核，该内核实现了 Flash-Decoding 算法。

LLaMa v2 / CodeLLaMa 的完整解码示例可在 FlashAttention 存储库 这里 和 xFormers 存储库 这里 找到。我们还提供了一个 LLaMa v1/v2 模型的最小高效解码代码示例，旨在快速、易于阅读、具有教育意义和可修改性。

致谢

感谢 Erich Elsen、Ashish Vaswani 和 Michaël Benesty 建议了将 KV 缓存加载拆分的这个想法。我们要感谢 Jeremy Reizenstein、Patrick Labatut 和 Andrew Tulloch 的宝贵讨论，以及 Quentin Carbonneaux 为 xFormers 提供高效的解码示例。我们还要感谢 Geeta Chauhan 和 Gregory Chanan 在写作方面提供的帮助，以及更广泛地为将本文发表在 PyTorch 博客上做出的贡献。

参考文献：

• Flash-Decoding for long-context inference | PyTorch

你是否想要提升深度学习模型的性能？面对这个问题，很多人会习惯性地使用一些“秘诀”，比如“使用就地操作！将梯度设置为 None！安装 PyTorch 1.10.0，但不要安装 1.10.1！”。

这种“灵机一动”式的优化方法虽然看似有效，但实际上却缺乏理论基础。深度学习模型的性能优化并非炼金术，而是需要从基本原理出发进行分析。

理解性能瓶颈

深度学习模型的性能优化可以从三个方面进行考虑：

• 计算 (Compute): GPU 上执行浮点运算 (FLOPS) 所花费的时间。
• 内存 (Memory): 在 GPU 内传输张量所花费的时间。
• 开销 (Overhead): 其他所有时间。

就像训练机器学习模型一样，了解模型所处的性能瓶颈可以帮助我们找到最有效的优化方法。例如，如果你的模型主要受限于内存带宽，那么提升 GPU 的 FLOPS 性能将毫无用处。相反，如果你的模型主要受限于计算能力，那么将模型逻辑重写为 C++ 代码来降低开销也无济于事。

计算：GPU 的核心能力

从某种程度上来说，优化深度学习系统就是最大化模型在计算受限状态下运行的时间。你为那些强大的 GPU 付出了高昂的代价，理所当然地希望它们能够发挥出全部的计算能力。然而，为了充分利用 GPU 的矩阵乘法能力，我们需要减少其他方面的开销，例如内存传输和系统开销。

为什么我们更关注最大化计算能力，而不是内存带宽？原因很简单：我们可以降低开销或内存成本，但我们无法在不改变实际操作的情况下减少计算量。

计算与内存带宽的矛盾

更糟糕的是，计算能力的增长速度远超内存带宽。下表展示了 CPU FLOPS 和内存带宽的翻倍时间：

组件	翻倍时间
CPU FLOPS	18 个月
内存带宽	24 个月

我们可以将计算能力想象成一个工厂。我们向工厂发送指令（开销），并提供原材料（内存带宽），以确保工厂高效运转（计算）。

如果工厂的效率提升速度快于我们提供原材料的速度，那么工厂将难以达到峰值效率。

即使工厂的规模（FLOPS）翻倍，如果我们的带宽跟不上，那么性能也不会翻倍。

这种计算能力增长速度快于内存带宽的趋势，一方面意味着机器学习系统工程师将长期拥有稳定的工作，另一方面也使得理解性能瓶颈变得更加重要。

现代 GPU 的特殊性

现代机器学习加速器都拥有专门用于矩阵乘法的硬件，例如英伟达的“Tensor Cores”。

如果你没有进行矩阵乘法，那么你只能达到 19.5 teraflops，而不是标称的 312 teraflops。需要注意的是，这并非 GPU 独有，事实上，TPU 的通用性甚至比 GPU 更低。

GPU 在非矩阵乘法运算方面速度较慢，这似乎是一个问题。那么，其他运算，例如层归一化或激活函数呢？事实是，这些运算在 FLOPS 方面几乎可以忽略不计。例如，让我们看一下这篇论文中关于 BERT 不同运算类型 FLOPS 数量的表格，其中“张量收缩”= 矩阵乘法。

运算类型	FLOPS 占比
张量收缩	99.8%
归一化	0.1%
点运算	0.1%

可以看到，非矩阵乘法运算只占总 FLOPS 的 0.2%，因此 GPU 在非矩阵乘法运算方面速度慢 15 倍并不重要。

内存带宽：数据传输的瓶颈

然而，在实际应用中，非矩阵乘法运算却往往比预期花费更多时间。罪魁祸首通常是数据在工厂和仓库之间传输的时间，也就是内存带宽成本。

内存带宽成本是指将数据从一个地方移动到另一个地方所花费的成本。这可能包括将数据从 CPU 移动到 GPU，从一个节点移动到另一个节点，甚至从 CUDA 全局内存移动到 CUDA 共享内存。其中，最后一种情况通常被称为“带宽成本”或“内存带宽成本”。

我们可以再次使用工厂的比喻来理解内存带宽成本。

虽然工厂是进行实际工作的场所，但它并不适合作为大规模存储单元。主要原因是，由于我们在工厂进行实际工作，因此所有存储都针对快速使用进行了优化（SRAM），而不是拥有大量的存储空间。

那么，我们应该在哪里存储实际结果和原材料呢？典型的做法是建立一个仓库，可能位于土地便宜且空间充足的地方（DRAM）。然后，我们可以将物资运送到工厂和从工厂运出（内存带宽）。

将物资运送到计算单元和从计算单元运出的成本就是所谓的“内存带宽”成本。顺便说一下，你的 GPU 的 DRAM 会显示在 nvidia-smi 中，它是你遇到“CUDA 内存不足”错误的主要原因。

需要注意的是，每次执行 GPU 内核时，都需要将数据从 GPU 的 DRAM（即仓库）中读取出来，并在执行完内核后将结果写入 DRAM。

现在，想象一下执行像 torch.cos 这样的单目运算会发生什么。我们需要将数据从存储单元运送到仓库，然后对每条数据进行少量计算，最后将数据运回存储单元。运输数据的成本非常高。因此，我们几乎所有时间都花在了运输数据上，而不是实际的计算本身。

由于我们所有时间都花在了内存带宽上，因此这种运算被称为内存受限运算，这意味着我们没有花太多时间进行计算。

操作融合：优化内存带宽

那么，我们该如何解决这个问题呢？让我们看一下一系列运算的执行过程。

一系列点运算的执行过程可能如下所示：

这种安排非常愚蠢。为什么我们要将相同的数据反复发送到全局内存，然后再发送回计算单元？我们应该将数据保留在工厂中，执行所有计算，然后再将其发送回存储单元。

我们可以通过操作融合来解决这个问题。操作融合是深度学习编译器中最重要的优化方法。简单来说，与其将数据写入全局内存，然后再读取回来，不如将多个计算合并在一起执行，从而避免额外的内存访问。

例如，如果我们执行 x.cos().cos()，通常需要进行 4 次全局读写操作。

x1 = x.cos()  # 从全局内存中读取 x，写入 x1
x2 = x1.cos()  # 从全局内存中读取 x1，写入 x2

但是，使用操作融合，我们只需要进行 2 次全局内存读写操作！因此，操作融合可以将速度提高 2 倍。

x2 = x.cos().cos()  # 从全局内存中读取 x，写入 x2

操作融合可以显著提高性能，但它也有一些限制。首先，GPU 需要在执行当前操作时知道下一步要执行的操作。因此，我们无法在 Eager 模式下执行操作融合，因为 PyTorch 在 Eager 模式下会逐个执行操作。其次，我们需要为此生成 CUDA 代码，这会带来新的挑战。

并非所有操作融合都像点运算那样简单。你可以将点运算融合到约简运算中，或者将点运算融合到矩阵乘法中。甚至矩阵乘法本身也可以看作是广播乘法和约简的融合。

如果你有兴趣编写自定义 CUDA 内核，那么操作融合将是你最受益的优化方法。任何两个 PyTorch 操作都存在融合的可能性，从而节省了在它们之间读写全局内存的内存带宽成本。此外，许多现有的编译器可以执行“简单的”融合，例如 NVFuser 和 XLA。然而，自动系统无法与人类的智慧相提并论，因此，如果你想尝试编写一些自定义 CUDA 内核，那么 Triton 是一个不错的起点。

最后，操作融合会带来一些意想不到的结果。例如，融合后的 x.cos().cos() 的执行时间几乎与单独调用 x.cos() 的时间相同。这就是为什么激活函数的成本几乎都相同，尽管 gelu 明显比 relu 包含更多的操作。

这一事实导致了重计算/激活检查点的一些有趣结果。本质上，进行额外的重计算可能会导致更少的内存带宽，从而减少运行时间。因此，我们可以通过重计算来降低内存和运行时间，我们在 AOTAutograd 中利用它构建了一个巧妙的最小割优化过程。你可以在此处阅读更多相关内容（也可能在未来的博客文章中进行介绍！）。

如何判断内存带宽受限

在判断你的操作是否内存带宽受限时，计算器可以起到很大的作用。

对于简单的操作，我们可以直接计算内存带宽。例如，A100 拥有 1.5 terabytes/second 的全局内存带宽，可以执行 19.5 teraflops/second 的计算。因此，如果你使用 32 位浮点数（即 4 字节），那么你可以在 GPU 执行 20 万亿次操作的同时加载 4000 亿个数字。此外，为了执行简单的单目运算（例如将张量乘以 2），实际上需要将张量写回全局内存。

因此，除非你的单目运算执行了大约一百次操作，否则你会花费更多时间执行内存访问，而不是实际的计算。

借助 NVFuser 这样的融合编译器，我们可以很容易地测量内存带宽。你可以在此处查看 Colab 代码。

如果我们使用 PyTorch 函数，例如：

def f(x: Tensor[N]):
    for _ in range(repeat):
        x = x * 2
    return x

并使用融合编译器对其进行基准测试，那么我们可以计算出不同 repeat 值下达成的 FLOPS 和内存带宽。增加 repeat 是增加计算量而不增加内存访问次数的一种简单方法，这也被称为增加计算强度。

具体来说，假设我们对这段代码进行基准测试，并找到每秒执行的迭代次数。那么，作为 N（张量大小）的函数，我们将执行 2*N 次内存访问，以及 N * repeat 次 FLOP。因此，达成的内存带宽将为 bytes_per_elem * 2 * N * itrs_per_second，达成的 FLOPS 将为 N * repeat * itrs_per_second。

现在，让我们绘制运行时间、FLOPS 和达成的内存带宽作为计算强度的函数。注意，所有内容都以对数对数刻度显示。

首先，请注意，直到我们执行 64 次乘法运算，运行时间才明显增加。这意味着，在此之前，我们主要受限于内存带宽，我们的计算能力大部分处于闲置状态。

因此，我们最初达成的 FLOPS 只有 0.2 teraflops。随着我们使计算强度翻倍，这个数字线性增长，直到我们接近 9.75 teraflops 的峰值 [1]。一旦我们接近峰值 teraflops，我们就认为是“计算受限”。

最后，你可以看到，我们达成的内存带宽最初接近峰值，随着我们增加计算强度，它开始下降。这正是我们所期望的，因为我们花费了越来越多的时间执行实际计算，而不是访问内存。

在这种情况下，很容易看出我们何时是计算受限的，何时是内存带宽受限的。对于 repeat < 32，我们饱和了内存带宽，而我们的计算能力未得到充分利用。相反，一旦 repeat > 64，我们看到我们的计算能力已经饱和（即接近峰值 FLOPS），而我们利用的内存带宽开始下降。

对于更大的系统，我们通常很难判断是计算受限还是内存带宽受限，因为它们通常包含计算受限和内存带宽受限的组件。

衡量计算受限程度的一种常见方法是测量达成的 FLOPS 占峰值 FLOPS 的百分比。例如，如果你达成了峰值 FLOPS 的 80%，那么你就知道你至少是 80% 的计算受限，这已经相当不错了！你剩下的时间可能都花在了执行内存带宽操作上。[2]

然而，除了内存带宽成本之外，还有一件事可能会导致你的 GPU 无法“Brrr”起来。

开销：系统瓶颈

开销是指你的代码花费在传输张量或计算之外的所有时间。例如，在 Python 解释器中花费的时间？开销。在 PyTorch 框架中花费的时间？开销。启动 CUDA 内核（但没有执行它们）花费的时间？也是开销。

开销之所以是一个棘手的问题，主要是因为现代 GPU 非常快。A100 每秒可以执行 312 万亿次浮点运算（312 TeraFLOPS）。相比之下，Python 非常慢。在本地进行基准测试，Python 每秒可以执行 3200 万次加法运算。

这意味着，在 Python 执行一次 FLOP 的时间里，A100 可以执行 975 万次 FLOPS。

更糟糕的是，Python 解释器并不是开销的唯一来源，像 PyTorch 这样的框架在到达实际内核之前也有很多层调度。如果你使用 PyTorch 执行相同的实验，我们每秒只能执行 28 万次操作。当然，PyTorch 的设计目的不是处理微小的张量，但是，如果你使用的是微小的张量（例如在科学计算中），你可能会发现 PyTorch 比 C++ 慢得多。

例如，看一下 PyTorch 执行一次加法的火焰图配置文件。那个方框？那是执行实际计算的部分。其他所有部分都是纯粹的开销。

鉴于此，你可能会惊讶于为什么有人会使用 PyTorch，但请记住，现代深度学习模型通常会执行大规模运算。此外，像 PyTorch 这样的框架是异步执行的。也就是说，当 PyTorch 运行 CUDA 内核时，它可以继续运行并在其后面排队更多 CUDA 内核。因此，只要 PyTorch 可以“领先于”CUDA 内核，大部分框架开销就会完全隐藏起来！

如果我们的 GPU 操作足够大，那么我们的 CPU 可以领先于 GPU（因此 CPU 开销无关紧要）。另一方面，如果我们的 GPU 操作太小，那么我们的 GPU 将大部分时间都浪费在昂贵的“纸镇”上。

那么，如何判断你是否处于这种状态呢？由于开销通常不会随着问题规模的增加而增加（而计算和内存会增加），因此判断开销受限的最简单方法是增加数据的规模。如果运行时间没有按比例增加，那么你就是开销受限的。例如，如果你将批次大小增加一倍，但运行时间只增加了 10%，那么你很可能是开销受限的。[3]

另一种方法是使用 PyTorch 分析器。在这里，粉红色的线实际上显示了 CPU 内核与 GPU 内核的匹配情况。

GPU 上有很多间隙，因为它正在等待 CPU 开销

我们的 CPU 运行速度远远超过 GPU

另一个补充说明 – nvidia-smi 中的“GPU-Util”（不是“Volatile GPU-Util”）条目基本上是测量底部的行中实际运行 GPU 内核的百分比。因此，这也是评估开销的一个好方法。

开销存在的主要原因是像 PyTorch 这样的框架具有很大的灵活性。本质上，需要花费大量时间来“弄清楚要做什么”。

这可能来自 Python（查找属性或调度到正确的函数）或 PyTorch 中的代码（PyTorch 的所有调度器）。例如，当你执行 a + b 时，需要执行以下步骤。

Python 需要查找 a 上的 add 调度到什么。
PyTorch 需要确定张量的许多属性（例如数据类型、设备以及是否需要自动微分）以确定要调用哪个内核。
PyTorch 需要实际启动内核。

从根本上说，这种开销来自于在每一步都能够执行不同操作的灵活性。如果你不需要这种灵活性，那么解决这种灵活性的方法之一是将其追踪出来，例如使用 jit.trace、FX 或 jax.jit。或者，你也可以在更低级别使用像 CUDA Graphs 这样的工具来实现。

不幸的是，这样做会失去灵活性。我期待的一种能够兼顾两者的方法是编写类似于“真正的”JIT 的工具，通过在 VM 级别进行内省来实现。有关详细信息，请参阅 TorchDynamo。

总结

如果你想加速深度学习系统，最重要的是要了解模型的瓶颈是什么。瓶颈决定了加速系统的最佳方法。

我经常看到研究人员和其他对加速 PyTorch 代码感兴趣的人盲目地尝试各种方法，而没有了解模型所处的性能瓶颈。

性能瓶颈	可行的解决方案
开销受限	追踪、操作融合、不要使用 Python、真正的 JIT :^)
带宽受限	操作融合
计算受限	使用 Tensor Cores、给英伟达更多钱

当然，可以说，用户需要考虑这些问题本身就反映了框架的失败。PyTorch 的编译器或分析 API 并不总是那么容易使用，尽管这是一个积极的关注领域。

无论如何，我发现理解系统的基本原理总是很有用，希望这对你也有帮助。

PS：如果你喜欢这篇文章，我将在 thonking.ai 上发布我未来的大部分文章。

致谢

感谢 Emily Shen、Qian Huang 和 EleutherAI 的成员阅读这篇博文的早期草稿并提供反馈。

BibTeX 引用

@article{he2022brrrrfromfirstprinciples,
  author={Horace He},
  title={Making Deep Learning Go Brrrr From First Principles},
  year={2022},
  url={https://horace.io/brrr_intro.html},
}

[1] 这可能与规格说明书上显示的 19.5 teraflops 不同。原因是 GPU 拥有更多专门用于融合乘加 (FMA) 指令的硬件。因此，对于完全通用的计算，A100 实际上只能达到 9.75 teraflops。 ↩︎

[2] 现在在 PyTorch 中以一种很好的方式计算 FLOPS 的方法有很多，请参阅 https://dev-discuss.pytorch.org/t/the-ideal-pytorch-flop-counter-with-torch-dispatch/505 ↩︎

在之前的文章中，我们介绍了 KV 缓存，这是一种常见的 LLM 推理过程优化方法，它可以使（自）注意力机制的计算需求随总序列长度（提示 + 生成的完成）线性增长，而不是二次增长。

更具体地说，KV 缓存通过在生成过程中计算时将过去标记的键和值张量存储（“缓存”）在 GPU 内存中，来避免在每个生成步骤中重新计算这些张量。

KV 缓存是一种折衷方案：我们用内存换取计算。在这篇文章中，我们将了解 KV 缓存可以增长到多大，它会带来什么挑战，以及解决这些挑战最常用的策略。

KV 缓存到底可以增长到多大？

这很简单：对于批次中每个序列的每个标记，我们需要为每个注意力层的每个注意力头存储两个向量张量（一个键张量和一个值张量），每个张量的尺寸为 $d_{head}$。每个张量参数所需的存储空间取决于精度：全精度（FP32）为 4 字节/参数，半精度（BF16，FP16）为 2 字节/参数，8 位数据类型（INT8，FP8）为 1 字节/参数，等等。

假设 $b$ 为批次大小，$t$ 为总序列长度（提示 + 完成），$n_{layers}$ 为解码器块/注意力层的数量，$n_{heads}$ 为每个注意力层的注意力头的数量，$d_{head}$ 为注意力层的隐藏维度，$p_a$ 为精度。多头注意力（MHA）模型的 KV 缓存的每个标记的内存消耗（以字节为单位）为：

$$
KV_cache_size_per_token = 2 \cdot n_{heads} \cdot d_{head} \cdot p_a
$$

注意： 我们提醒，在 MHA 模型中，$n_{heads} \cdot d_{head} = d_{model}$，但我们不会使用它来简化上面的公式。

因此，KV 缓存的总大小（以字节为单位）为：

$$
KV_cache_size = b \cdot t \cdot KV_cache_size_per_token
$$

KV 缓存的第一个挑战出现了：它随着批次大小线性增长，更重要的是随着总序列长度线性增长。由于它随着总序列长度增长，因此 KV 缓存的大小实际上是无界的，而我们的 GPU 内存显然是有限的。更糟糕的是，由于总序列长度无法提前知道，因此 KV 缓存的内存需求也是未知的，这使得内存管理特别具有挑战性。

让我们看看一些流行的 MHA 模型（表 1）的数字，包括 Meta 的 Llama-2 [1] 和 OPT [2]，MosaicML 的 MPT [3] 和 BigScience 的 BLOOM [4]：

模型	参数	层数	头数	隐藏维度
Llama-2–7B	7B	32	32	4096
OPT-1.3B	1.3B	24	16	768
MPT-7B	7B	24	32	4096
BLOOM-176B	176B	60	80	12288

表 1 – 流行的多头注意力 (MHA) 模型规格

假设参数存储在半精度（FP16，BF16）中，并选择一个较小的模型（Llama-2–7B）和一个较大的模型（BLOOM-176B）。对于 Llama-2–7B（分别为 BLOOM-176B），KV 缓存的内存消耗约为 0.5MB/标记（分别为 4MB/标记）。

让我们关注 Llama-2–7B。使用半精度，加载模型权重消耗约 14GB 内存，与缓存 28k 个标记的键和值相同。28k 个标记可能对应于 56 个长度为 512 的序列的批次，这并不算特别极端。

从上面的数字可以看出，KV 缓存的内存消耗会变得非常大，甚至可能超过加载大型序列模型权重所需的内存量。

现在让我们将这些数字与常用的 NVIDIA 数据中心 GPU 的内存容量进行比较（表 2）：

GPU	内存
A10	24GB
A100	40GB
H100	80GB
H200	96GB

表 2 – 常用于训练和/或服务 LLM 的 NVIDIA 数据中心 GPU 规格

让我们选择性价比比较高的 A10 GPU，坚持使用 Llama-2–7B 并计算最大 KV 缓存容量。加载模型权重后，还有 24-2×7=10 GB 可用于 KV 缓存，即总容量约为 20k 个标记，包括提示，这显然不允许在处理或生成长序列时，特别是当处理或生成长序列时，无法处理大量并发请求。

我们现在了解到，KV 缓存阻止我们处理或生成非常长的序列（即阻碍长上下文窗口）和/或处理大型批次，因此无法最大限度地提高硬件效率。

从这个角度来看，最大限度地提高处理能力意味着为 KV 缓存留出尽可能多的空间，这可以通过以下方法实现：

• 减少模型权重的内存占用（权重量化）
• 减少 KV 缓存的内存占用（见下文）
• 通过将模型在多个 GPU 上分片，以网络通信为代价，从多个设备中汇集内存（模型并行），或使用其他类型的存储，如 CPU 内存或磁盘（卸载）。

由于模型权重和不断增长的 KV 缓存必须在每次前向传递时加载，因此解码步骤涉及非常大的数据传输，正如我们将在后面的文章中看到，实际上是内存带宽受限的，即我们实际上花费更多的时间移动数据，而不是进行有用的工作，即计算。在这种情况下，延迟只能通过拥有更多内存带宽（即更好的硬件）或传输更少数据来改善。更小的模型权重和 KV 缓存释放了内存，可以容纳更多序列，因此能够提高吞吐量（和/或最大序列长度）。

在这方面，内存占用减少策略具有三重用途，因为它们允许我们提高硬件利用率，从而提高成本效益，同时降低延迟并提高吞吐量。

为什么我需要为输入标记付费？

在这一点上，你应该明白为什么你需要为输入和输出标记付费。一旦输入提示被处理，即在预填充阶段结束时，我们已经消耗了 GPU 内存（用于存储每个输入标记的键和值张量）和计算（用于将提示标记通过模型）。

让我们看看一些实际数字。假设 P 参数模型的前向传递的总 FLOPs 数量约为 2.P FLOPs/标记 [5]，使用 Llama-2-7B 处理提示消耗约 0.5 MB/标记的 GPU 内存（见上文）和约 14 GFLOPs/标记的 GPU 计算。对于一个 1000 个标记的提示（略小于两页），那就是约 500MB 的内存和 14 TFLOPs 的计算，而且我们还没有生成任何内容。

现在让我们看看所有可以减少 KV 缓存的内存占用的方法，方法是采用上面的公式，依次查看每个项：

减少批次大小？

在大多数情况下，我们不想减小批次大小，因为虽然它有助于减少 KV 缓存的内存占用，从而降低延迟，但它会降低硬件利用率，从而降低成本效益。在后面的文章中，我们将看到，相反，我们希望尽可能地增加批次大小。

减少对总序列长度的依赖？

不为序列中的所有标记存储键和值的一个原因是，我们明确选择在每次迭代时重新计算缺失的键和值，因为花费 FLOPs 比消耗 GPU 内存更划算（例如，因为我们是内存带宽受限的，在自回归阶段就是这样）。据我所知，这在实践中并不是我所知道的，所以我们不会深入研究这个方向。

另一个角度是，我们可以不必为模型不关注或很少关注的标记存储键和值。对于经过训练只关注序列一部分的模型（例如，使用 Mistral AI 的 Mistral-7B），或者作为内存消耗和模型精度之间的一种折衷方案，这可能是这种情况。让我解释一下。

像 Mistral-7B [6] 这样的模型经过训练不会关注整个序列。Mistral-7B 的注意力层实际上通过只关注最后（4096）个相邻标记来构建标记表示。这种注意力机制的变体称为滑动窗口注意力（SWA）或局部注意力。通过设计，局部注意力保证我们永远不会在 KV 缓存中为每个序列存储超过窗口大小（例如 4096）的张量对。

另一种方法是利用注意力层在序列中将注意力分散到标记上的方式中的模式。众所周知，注意力模块不成比例地且始终如一地将更多注意力分配给序列中的一小部分标记（图 1）。相反，许多标记始终对输出的贡献很小，那么为什么要费心存储它们的键和值呢？

通过丢弃这些标记，我们实际上将相应的注意力分数设置为零，并用更稀疏的矩阵来近似注意力矩阵。成功的近似将最大限度地减少近似误差，从而最大限度地减少对模型精度的影响（例如，使用困惑度来衡量）。

让我们看看过去几个月出现的一些方法，这些方法可以直接应用，无需任何重新训练或微调：StreamingLLM 框架、H2O（重击者预言机）、剪刀手和 FastGen。据我所知，目前还没有任何流行的 LLM 推理框架支持它们。

针对使用有限长度上下文窗口训练的模型，StreamingLLM 框架 [7] 基于这样一个观察结果：初始标记会收集大量的注意力。因此，该框架通过只在缓存中保留最开始的几个位置标记（“汇聚标记”）和最后一个相邻标记（局部注意力）来构建一个滑动窗口。因此，StreamingLLM KV 缓存的长度是固定的，既有固定部分（通常为 1 到 4 个标记），也有滑动部分。

类似的 H2O [8] 和剪刀手 [9] 方法明确地旨在通过设置一个最大缓存标记数（预算）并通过在每次达到缓存预算时丢弃标记来压缩 KV 缓存。H2O 算法一次只丢弃一个标记，而剪刀手则根据目标压缩率（例如，减少 30% 的 KV 缓存大小）丢弃尽可能多的标记。

这两种方法都基于这样一个观察结果：在给定步骤中具有影响力的标记（“关键标记”或“重击者”）在未来步骤中仍然具有影响力（剪刀手作者称之为“重要性持久性假设”）。换句话说，我们可以确保被丢弃的低影响力标记在未来步骤中将仍然相对被忽略，因此可以安全地丢弃。

这两种算法的一个关键方面显然是缓存驱逐策略。剪刀手只保留最新的标记和历史窗口内注意力分数最高的标记。H2O 丢弃累积注意力分数最低的标记，因此只保留在迭代过程中始终获得高注意力分数的标记。这两个作者团队都表明，他们的算法可以实现高达 80% 的 KV 缓存大小减少，而模型精度损失可以忽略不计。

FastGen 方法 [10]（不要与无关的 DeepSpeed-FastGen 混淆）仍然基于注意力模式，但采用另一种方法，即不设置缓存预算，而是设置注意力矩阵的最大近似误差，因此侧重于保持模型精度。

FastGen 是一种两步法：首先，在预填充阶段结束时对模型的注意力层进行分析，以确定允许满足误差目标的压缩策略集。与其他方法一样，它假设所识别的注意力模式将在未来的生成步骤中保持不变。压缩策略包括：保留特殊标记、保留标点符号标记、保留最后一个相邻标记（局部注意力）等等（图 2）。如果误差目标过于严格，无法满足，FastGen 将回退到常规的 KV 缓存。然后，在每个生成步骤中将所选的压缩策略应用于 KV 缓存。

图 2 – FastGen 论文中压缩策略集的示例：特殊标记（绿色）+ 标点符号标记（橙色）+ 局部注意力（蓝色）。被丢弃的标记用灰色表示。

请注意，与其他方法不同，FastGen 为每个提示构建一个量身定制的压缩策略。FastGen 作者表明，对于给定的 KV 缓存压缩率，他们比 H2O 和剪刀手更好地保持模型精度。

在任何情况下，打破对不可预测的总序列长度的依赖都是一种解脱，因为它允许为每个序列分配一个内存预算，从而大大简化内存管理。由于数据传输是延迟的主要贡献者，因此没有一个随着序列长度线性增长的 KV 缓存，尤其是在较长的序列长度下，可以带来惊人的加速。

减少层数？

这里没有太多可以获得的。较小的模型通常具有较少的层（表 4），因此如果较小的模型在你的用例中表现良好，只需选择它即可。

模型	参数	层数	头数	隐藏维度
Llama-2–7B	7B	32	32	4096
Llama-2–13B	13B	40	40	4096
Llama-2–70B	70B	60	64	4096
Llama-2–13B	13B	80	80	4096

表 4 – Llama-2 模型规格

减少注意力头的数量？

由于对于给定的模型架构，模型大小主要由层数和头数控制，因此减少头数可能意味着选择一个更小的模型（见表 4）。

但是，如果我们仔细观察，我们会注意到，我们只需要减少键和值头的数量，查询头的数量不会影响 KV 缓存的大小。这正是多查询注意力（MQA）[11] 和分组查询注意力（GQA）[12] 架构背后的理念。这些多头注意力（MHA）变体的唯一动机是减少 KV 缓存的大小。

MQA 最早于 2019 年被引入。在 MQA 中，所有查询头共享相同的单个键和值头。换句话说，所有查询头使用相同的键计算它们的注意力分数，所有头的输出使用相同的值（但不是相同的注意力分数）计算（图 3）。

图 3 – 多头注意力（上）与多查询注意力（下）（两个注意力头）

然而，剥离所有头对于较大的模型来说相对更激进。例如，从 64 个头减少到 1 个头，与从 32 个头减少到 1 个头相比，模型的表示能力的削减更大。GQA 通过提供一种折衷方案来解决这个问题：我们不是让所有查询头共享相同的唯一 KV 头，而是将它们分成 g 个查询头的组，同一个组的查询头共享相同的唯一 KV 头。换句话说，我们不是从 $n_{heads}$ 缩减到 1 个 KV 头，而是将 KV 头的数量从 $n_{heads}$ 缩减到 $1<g<n_{heads}$。

从这个角度来看，MHA 和 MQA 都是 GQA 的特例（分别为 $g=1$ 和 $g=n_{heads}$）。QGA 允许在两个极端情况（MHA 和 MQA）之间更平滑地导航模型精度/KV 缓存大小（与延迟和吞吐量都相关）的折衷方案。

考虑到这个新的参数 g，KV 缓存大小公式变为：

$$
KV_cache_size = b \cdot t \cdot 2 \cdot g \cdot d_{head} \cdot p_a
$$

在实践中，MQA/GQA 架构已被 Google Research 的 PaLM [13]，TII 的 Falcon [14] 模型，Meta 的 Llama-2 [1]（仅 70B）和 Mistral AI 的 Mistral-7B [7]（表 5）显著地实现。

模型	参数	层数	头数	隐藏维度	注意力类型
PaLM	540B	52	64	16384	MQA
Falcon	40B	40	64	5120	MQA
Llama-2–70B	70B	60	64	4096	GQA
Mistral-7B	7B	24	32	4096	GQA

表 5 – 使用 MQA 或 GQA 的模型系列

注意力头的隐藏维度？

再次强调，如果你不准备选择另一个模型，这里没有太多可以获得的。根据模型系列的不同，头隐藏维度在不同模型大小之间可能是恒定的（例如 Llama-2，Falcon），因此选择同一个系列中较小的变体不会有帮助。

使用更少的字节数/参数？

量化 KV 缓存确实是一种可以大幅减少其大小的好方法。但是，像 AWQ [15] 或 GPTQ [16] 这样的仅权重量化算法从定义上来说不会有帮助。只有量化权重和“激活”（即任何不是权重的东西）的算法，例如 LLM.int8()[17] 或 SmoothQuant [18]，才会产生量化的 KV 缓存。

请注意，在权重和激活上都起作用的量化算法的意图之一是在较低的精度下执行计算密集型矩阵乘法。如果像训练期间那样是计算受限的，这会带来性能提升，但正如我们将在后面的文章中看到，推理的自回归阶段实际上是内存带宽受限的，因此能够更快地计算不会带来太多价值。

我们实际上只对内存占用的减少感兴趣，因为它意味着更少的数据传输。从这个角度来看，像 LLM.int8() 或 SmoothQuant 这样的量化算法有点过分：在将缓存的张量移动到 GPU 内存之前对其进行量化，并在从 GPU 内存中获取相同的张量后对其进行反量化（以额外的开销为代价）应该就足够了。

一些 LLM 推理系统已经包含了这样的 KV 缓存量化功能。例如，FlexGen [19] 将 KV 缓存和模型权重都量化并存储在 4 位数据格式中。NVIDIA TensorRT-LLM 能够将 KV 缓存量化到 8 位数据格式（INT8 或 FP8）中。流行的 vLLM 框架从 0.3.0 版本开始也支持 KV 缓存（FP8）量化。由于量化是在每次迭代时动态执行的，因此不需要校准步骤。

高效内存管理的重要性

到目前为止，我们隐含地假设内存没有浪费：所有保留的内存都用于存储标记，所有可用的内存都可以被保留。在实践中，简单的内存管理策略会导致大量内存被浪费（PagedAttention 论文 [20] 表明，实际的有效内存利用率可能低至 20%，即 80% 的浪费！）：

• 由于请求的总序列长度事先未知，我们可以保留能够容纳最大序列长度的连续内存块。这种分配中很大一部分肯定永远不会被使用，并且由于无法用于其他请求，因此被浪费了（内部内存碎片）。
• 即使序列长度事先已知，由于内存是逐渐消耗的，但内存块是为请求的生命周期保留的，因此较短的请求无法使用仍然未使用的内存块。
• 如果我们使用像束搜索这样的解码策略，为每个请求生成多个序列，那么多个候选序列实际上可以部分共享它们的 KV 缓存。如果我们不考虑这种情况，我们不可避免地会通过存储本可以共享的重复 KV 条目来浪费内存。

这些缺点正是现在流行的 PagedAttention 算法旨在解决的问题。PagedAttention 分配固定大小且相对较小的内存块，称为块。每个块可以包含固定数量的标记，如果需要，可以跨不同的请求共享。按需分配和较小的块大小减轻了内部内存碎片，而相同大小的块消除了外部内存碎片。

总的来说，PagedAttention 实现了 KV 缓存内存的近乎零浪费（低于 4% [21]）。以前被浪费的内存现在可以用来容纳更多请求，从而提高吞吐量。PagedAttention 推出时，吞吐量提高的数字与当时内存浪费的程度一样惊人。

PagedAttention 最初由 vLLM 推理系统实现，但现在所有主要的推理框架都支持它（例如 HuggingFace TGI、NVIDIA TensorRT-LLM、LMDeploy TurboMind 等）。

PagedAttention 未涵盖的另一个可能的优化是在请求之间重用键值缓存。这将适用于提示共享公共前缀的情况，这种情况在像聊天和代理这样的多轮用例中很常见，或者在使用提示模板时（图 4）。

图 4 – SGLang 论文中 KV 缓存共享示例（多轮聊天），总共四个生成请求。蓝色框表示可共享的提示部分。

能够在请求之间重用 KV 缓存将能够显著降低延迟（尤其是第一个标记的延迟）和吞吐量（通过大大减少具有共享前缀的并发请求的内存占用）。

LMSYS SGLang 论文 [22] 中介绍的 RadixAttention 技术实现了这种 KV 缓存重用。

RadixAttention 算法不是在完成生成请求后丢弃 KV 缓存，而是将其保留在 GPU 内存中，并向一个专门的数据结构（基数树）添加一个新条目，该结构将标记序列映射到它们的 KV 缓存张量。当一个新的请求进来时，调度程序使用基数树进行前缀匹配。如果有缓存命中，调度程序将重用缓存的 KV 张量来满足请求。

由于 GPU 内存有限，缓存的 KV 张量不能永远保留。因此，RadixAttention 算法包含一个驱逐策略（例如，最近最少使用 (LRU) 驱逐策略）。最佳缓存重用可能与先到先得的调度程序不兼容。因此，RadixAttention 带有一个修改后的调度程序，该调度程序优先考虑与缓存的前缀匹配的请求（缓存感知调度）。

注意： PagedAttention 和 RadixAttention 的命名有点误导，因为与人们可能认为的相反，它们不是模型注意力层的优化（如 FlashAttention），而是在模型服务器级别运行（它们帮助服务应用程序更好地管理主机上的 KV 缓存）。

如果 GPU 内存不足，为什么不“只”使用多个 GPU？或者卸载到 CPU 内存甚至磁盘？

这些是两种不同的但有效的方案。

首先关于卸载到更丰富但更慢的存储器（CPU 内存和磁盘）。并非所有推理框架都支持此功能，让我们列举 HuggingFace Accelerate、DeepSpeed-Inference 和更先进的 FlexGen。由于它涉及使用更慢的存储器，因此卸载会以较大的延迟损失为代价，因此对于对延迟敏感的用例来说，显然不应优先考虑此选项。卸载系统通常用于面向吞吐量的用例，例如离线批处理。

关于使用多个 GPU（对于较大的模型来说不可避免），将模型在多个设备上分片可以释放内存压力，因为可以从聚合的内存容量和内存带宽中获益。

如果选择管道并行 [23]，模型和 KV 缓存都将在层维度上分片。如果选择张量并行 [24]（更常见于推理），KV 缓存将在头维度上分片。请注意，MQA 在这种设置下效率很低：由于我们无法跨多个设备分片单个头，因此 KV 缓存必须在所有设备上复制，从而失去了 MQA 的优势。对于实现 MQA 的模型，另一种选择是跨批次大小维度分片 KV 缓存 [25]。

在任何情况下，以上所有情况都假设单个主机，我们仍然受限于我们可以得到的最大多 GPU 实例的存储容量。据我所知，目前还没有推理框架支持多主机模型并行。如果我们能够将模型和 KV 缓存都分片到多个主机上，那么我们可以处理的可用内存量和最大序列长度将实际上变得无限。这是 Infinite-LLM 论文 [26] 旨在通过引入一种新的分布式注意力算法 (DistAttention) 和调整 Ray 框架来构建多主机分布式 KV 缓存管理和调度系统 (DistKV-LLM) 来解决的问题。

总结

在这篇文章中，我们学习了选择 KV 缓存会带来额外的挑战。多头注意力 (MHA) 模型的 KV 缓存确实会消耗大量 GPU 内存，大约为 1MB/标记，并且很容易超过模型权重的内存占用。

鉴于 GPU 内存的有限性，KV 缓存内存压力导致了许多不同方向的举措：新的注意力架构（MQA、GQA、SWA）、缓存压缩策略（H2O、剪刀手、FastGen）、高效内存管理（PagedAttention、RadixAttention）、量化和存储容量扩展（卸载系统、单主机和多主机模型并行）。

正如我们将在后面的文章中看到，减少 KV 缓存的大小至关重要，不仅因为 GPU 内存有限，而且因为数据移动量实际上是每个自回归步骤的延迟的主要贡献者，因此也是整个生成过程的延迟的主要贡献者。

在下一篇文章中，我们将看看可能影响模型延迟和吞吐量的不同类型的瓶颈。我们下次再见！

参考文献：

[1]: Llama 2: Open Foundation and Fine-Tuned Chat Models (Touvron 等人，2023)
[2]: OPT: Open Pre-trained Transformer Language Models (Zhang 等人，2022)
[3]: MPT-7B（2023 年 5 月）和 MPT-30B（2023 年 6 月）的发布博客文章
[4]: BLOOM: A 176B-Parameter Open-Access Multilingual Language Model (BigScience，2023)
[5]: Scaling Laws for Neural Language Models (Kaplan 等人，2020)
[6]: Mistral 7B (Jiang 等人，2023)
[7]: Efficient Streaming Language Models with Attention Sinks (Xiao 等人，2023) + GitHub 仓库
[8]: H_2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (Zhang 等人，2023) + GitHub 仓库
[9]: Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time (Liu 等人，2023)
[10]: Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (Ge 等人，2023)
[11]: Fast Transformer Decoding: One Write-Head is All You Need (Shazeer，2019)
[12]: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie 等人，2023)
[13]: PaLM: Scaling Language Modeling with Pathways (Chowdhery 等人，2022)
[14]: The Falcon Series of Open Language Models (Almazrouei 等人，2023)
[15]: AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (Lin 等人，2023) + GitHub 仓库
[16]: GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers (Frantar 等人，2022) + GitHub 仓库
[17]: LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (Dettmers 等人，2022) + GitHub 仓库
[18]: SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models (Xiao 等人，2022) + GitHub 仓库
[19]: FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU (Sheng 等人，2023) + GitHub 仓库
[20]: Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon 等人，2023) + GitHub 仓库
[21]: vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention (Kwon 等人，2023)
[22]: Efficiently Programming Large Language Models using SGLang (Zheng 等人，2023) + 博客文章
[23]: GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (Huang 等人，2018)
[24]: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan 等人，2021)
[25]: Efficiently Scaling Transformer Inference (Pope 等人，2022)
[26]: Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache (Lin 等人，2024)