分类： 未分类

近年来，大型语言模型（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)