大型语言模型 (LLM) 正在改变世界，从生成文本到翻译语言，再到编写代码，LLM 的应用范围越来越广。然而，LLM 的推理速度一直是制约其应用的关键因素。为了解决这个问题，研究人员一直在探索各种方法来加速 LLM 推理。

本文将介绍一篇名为 “FlashDecoding++: Faster Large Language Model Inference on GPUs” 的论文，该论文提出了一种新的 LLM 推理加速技术，可以在 GPU 上显著提高推理速度。

LLM 推理加速的挑战

加速 LLM 推理面临着三大挑战：

1. 同步部分 Softmax 更新： Softmax 操作需要对每个部分 Softmax 结果进行同步更新，这导致了 LLM 中注意力计算的约 20% 的开销。
2. 扁平 GEMM 的计算利用率低下： LLM 推理中执行 GEMM 的矩阵形状是扁平的，导致计算利用率低下，在之前的设计中，填充零后会导致超过 50% 的性能损失。
3. 静态数据流导致的性能损失： LLM 中的内核性能取决于不同的输入数据特征、硬件配置等。单一且静态的数据流会导致 LLM 推理中不同形状的 GEMM 出现 50.25% 的性能损失。

FlashDecoding++ 的解决方案

FlashDecoding++ 针对上述挑战提出了以下解决方案：

1. 异步 Softmax 与统一最大值： FlashDecoding++ 引入了一种统一最大值技术，用于不同的部分 Softmax 计算，从而避免同步。
2. 双缓冲的扁平 GEMM 优化： FlashDecoding++ 指出不同形状的扁平 GEMM 面临着不同的瓶颈。然后，引入了双缓冲等技术。
3. 硬件资源自适应的启发式数据流： FlashDecoding++ 使用不同的硬件资源，考虑输入动态，启发式地优化数据流。

性能提升

FlashDecoding++ 的优化策略使其在 NVIDIA 和 AMD GPU 上都取得了显著的性能提升，与 Hugging Face 实现相比，分别实现了高达 4.86 倍和 2.18 倍的加速。与主流 LLM 上最先进的 LLM 推理引擎相比，FlashDecoding++ 的平均加速比为 1.37 倍。

总结

FlashDecoding++ 提出了一套全面的 LLM 推理加速解决方案，有效地解决了 LLM 推理中的三大挑战。其在主流 LLM 和硬件平台上的出色表现，为 LLM 的广泛应用提供了强有力的支持。

参考文献

• [2311.01282] FlashDecoding++: Faster Large Language Model Inference on GPUs
• Stanford CRFM
• GitHub – opengear-project/GEAR: GEAR: An Efficient KV Cache Compression Recipefor Near-Lossless Generative Inference of LLM
• DistServe/distserve at main · LLMServe/DistServe · GitHub

注： 本文仅对 FlashDecoding++ 论文进行了简要介绍，更多细节请参考原文。

在GPU上推理大规模语言模型（LLM）的性能至关重要，而FlashDecoding++是一款针对LLM推理的快速引擎，通过解决同步部分softmax更新、未充分利用的扁平化GEMM计算和静态数据流等挑战，实现了显着的推理加速效果。

解决同步部分softmax更新的挑战：
FlashDecoding++引入了异步softmax和统一最大值的技术，避免了在计算部分softmax结果时需要同步更新的问题。每个部分softmax结果可以独立进行处理，无需进行同步操作，从而减少了计算中的开销。

解决未充分利用的扁平化GEMM计算的挑战：
FlashDecoding++通过双缓冲技术对扁平化GEMM计算进行了优化，隐藏了内存访问延迟，提高了计算利用率。它在共享内存中分配了两个独立的缓冲区，一个缓冲区用于进行GEMM计算，而另一个缓冲区则用于加载下一个GEMM操作所需的数据。通过这种方式，计算和内存访问可以同时进行，实现了计算与存储的重叠。

解决静态数据流的挑战：
FlashDecoding++采用了启发式数据流和硬件资源适应的方法。它根据输入动态和硬件配置，在不同的线性工作负载下动态优化数据流，选择最佳的实现方式。通过根据不同工作负载的特点进行灵活调整，FlashDecoding++实现了最佳的推理性能。

性能评估：
FlashDecoding++在多个硬件平台上进行了性能评估，包括NVIDIA和AMD的GPU。与Hugging Face、vLLM、DeepSpeed、TensorRT-LLM、OpenPPL和FlashDecoding等LLM推理引擎进行了比较。结果表明，FlashDecoding++相对于这些基线引擎实现了显着的加速效果，提供了高达4.86倍的推理速度提升。

近年来，大型语言模型（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 之间的配合情况，从而找到优化方向。

结论

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

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

参考文献：

• 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 和内存带宽的翻倍时间：

这种差距意味着，尽管 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 等技术。

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

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

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

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

参考文献

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