导论

深度学习的许多部分至今仍像某种黑魔法，但模型性能优化并非如此——即使规模极其庞大也一样！从单个加速器到数以万计的加速器，相对简单的原理始终适用。理解这些原理后，你就能完成许多实用的事情：

• 粗略估计模型各部分距离理论最优性能还有多远。
• 在不同规模下，对不同并行方案（即如何把计算拆分到多个设备上）做出有依据的选择。
• 估算训练和运行大型 Transformer 模型所需的成本与时间。
• 设计能够利用特定硬件特性的算法。
• 在明确理解当前算法性能瓶颈的基础上设计硬件。

预期背景： 我们假设你对大语言模型（LLM）和 Transformer 架构有基本了解，但不一定熟悉它们如何在大规模系统上运行。你应当了解 LLM 训练的基础知识，最好还对 JAX 有一些基本认识。实用的背景读物包括介绍 Transformer 架构的这篇博客文章和原始 Transformer 论文。你还可以查看这份列表，其中收录了更多适合同时或日后阅读的资料。

目标与反馈： 读完本书后，你应当能够较有把握地估算：对于给定硬件平台上的 Transformer 模型，哪种并行方案最好，以及训练和推理大致需要多长时间。如果你仍然做不到，请给我们发邮件或留言！我们非常希望知道怎样才能把这些内容讲得更清楚。

你或许也会喜欢新增的 NVIDIA GPU 第 12 章！

为什么你应当关心这些？

三四年前，我认为大多数机器学习研究人员都不需要理解本书中的任何内容。但如今，即使是“小”模型也已如此接近硬件极限，以至于要开展创新研究，就必须考虑大规模系统的效率问题。¹ 如果基准成绩提高 20% 的代价是 Roofline 效率下降 20%，那么这项提升就毫无意义。 很有前景的模型架构常常失败，要么因为它们根本无法在大规模系统上高效运行，要么因为没有人愿意投入精力让它们高效运行。

“模型扩展”的目标，是在增加训练或推理所用芯片数量时，让吞吐量也按比例线性增加。 这称为“强扩展（strong scaling）”。增加芯片（即增加“并行度”）通常会缩短计算时间，但也会带来额外的芯片间通信。当通信耗时超过计算耗时时，系统便进入“通信受限”状态，无法继续实现强扩展。² 如果我们对硬件有足够深入的理解，能预判这些瓶颈会在哪里出现，就可以通过设计或重新配置模型来避开它们。³

本书旨在解释 TPU（以及 GPU）硬件如何工作，以及 Transformer 架构为了在当代硬件上高效运行而经历了怎样的演变。我们希望，这既能帮助设计新架构的研究人员，也能帮助致力于让这一代 LLM 高速运行的工程师。

全书概览

本书的整体结构如下：

第 1 章介绍 Roofline 分析，以及可能限制系统扩展能力的因素（通信、计算和内存）。第 2 章和第 3 章详细讨论 TPU 如何工作：既把它看作单独的芯片，也——这一点至关重要——把它看作由带宽和延迟都有限的芯片间链路连接起来的系统。我们将回答如下问题：

• 一次给定规模的矩阵乘法应当耗时多久？它在什么情况下受计算能力、内存带宽或通信带宽限制？
• TPU 如何互连成训练集群？系统各部分的带宽有多大？
• 在多个 TPU 之间聚合、分散或重新分布数组需要多长时间？
• 如何高效地将以不同方式分布在多个设备上的矩阵相乘？

详见第 2 章。

五年前，机器学习领域还有丰富多彩的架构版图——ConvNet、LSTM、MLP、Transformer——但如今基本只剩下 Transformer[transformers]。我们坚信，Transformer 架构的每个组成部分都值得深入理解：每个矩阵的确切大小、归一化发生的位置，以及各部分包含多少参数和 FLOPs⁴。第 4 章将细致讲解这些“Transformer 数学”，说明如何计算训练与推理过程中的参数量和 FLOPs。由此，我们便能知道模型会占用多少内存、计算或通信分别会消耗多少时间，以及注意力相对于前馈块会在什么时候变得重要。

详见第 4 章。

第 5 章：训练和第 7 章：推理是本书的核心，我们将在其中讨论一个根本问题：给定某种规模的模型和一定数量的芯片，如何并行化模型，才能让系统保持在“强扩展”区间？这是一个看似简单、答案却出乎意料地复杂的问题。从高层来看，把模型拆分到多个芯片上主要有四种并行技术（数据并行、张量并行、流水线并行和专家并行），此外还有许多用于降低内存需求的技术（重计算（rematerialization）、优化器/模型分片（即 ZeRO）、主机卸载、梯度累积）。这些内容都会在相应章节中讨论。

我们希望，在读完这些章节后，面对新的架构或设置，你能够自行选择合适的方法。第 6 章和第 8 章则是实践教程，将这些概念应用于广受欢迎的开源模型 LLaMA 3。

最后，第 9 章和第 10 章介绍如何在 JAX 中实现其中一些思想，以及在出现问题时如何对代码进行性能剖析和调试。第 12 章是新增内容，将深入讨论 GPU。

全书各处都安排了练习题，供你亲自思考。你完全不必有压力，不需要读完所有章节，也不必按顺序阅读。还请留下反馈。目前本书仍是一份草稿，会继续修订。谢谢！

在此特别感谢 James Bradbury 和 Blake Hechtman，他们推导出了本书中的许多思想。

闲话少说，第 1 章将从 TPU 的 Roofline 模型讲起。

章节导航

这一系列文章可能比实际需要的更长，但我们希望这不会让你望而却步。前三章是预备知识；如果你已经熟悉这些材料，可以跳过，不过其中会引入后续章节使用的记号。最后三部分或许最具实践价值，因为它们会说明如何处理真实模型。

第一部分：预备知识

• 第 1 章：Roofline 分析简介。算法受三类因素限制：计算、通信和内存。我们可以据此近似估算算法的运行速度。

• 第 2 章：如何理解 TPU。TPU 如何工作？这会怎样影响我们能够训练和部署推理服务的模型？

• 第 3 章：分片矩阵及其乘法。本章借助我们最喜欢的运算——（分片）矩阵乘法——来解释模型分片和多 TPU 并行。

第二部分：Transformer

• 第 4 章：你需要掌握的全部 Transformer 数学。Transformer 的前向传播和反向传播分别需要多少 FLOPs？你能计算它的参数量吗？KV 缓存有多大？本章将逐步推导这些计算。

• 第 5 章：如何并行化 Transformer 训练。FSDP、Megatron 分片、流水线并行：给定一定数量的芯片，如何用给定批大小尽可能高效地训练给定规模的模型？

• 第 6 章：在 TPU 上训练 LLaMA 3。如何在 TPU 上训练 LLaMA 3？需要多长时间？成本是多少？

• 第 7 章：Transformer 推理详解。训练完成后，我们还需要部署模型。推理引入了一个新的考量因素——延迟——并改变了内存使用格局。本章将讨论解耦式推理服务如何工作，以及应如何理解 KV 缓存。

• 第 8 章：在 TPU 上部署 LLaMA 3 推理服务。在 TPU v5e 上部署 LLaMA 3 推理服务需要多少成本？延迟与吞吐量之间有怎样的权衡？

第三部分：实践教程

• 第 9 章：如何对 TPU 代码进行性能剖析。真实的 LLM 从来不像前面的理论那样简单。本章介绍 JAX + XLA 软件栈，以及如何使用 JAX/TensorBoard 性能剖析器调试和修复真实问题。

• 第 10 章：使用 JAX 为 TPU 编程。JAX 提供了一系列仿佛魔法般的计算并行 API，但你需要知道如何使用它们。本章包含有趣的示例和附有解答的练习题。

第四部分：总结与附加内容

• 第 11 章：总结与延伸阅读。关于 TPU 和 LLM 的总结思考与延伸阅读。

• 第 12 章：如何理解 GPU。介绍 GPU 的附加章节，讨论 GPU 的工作方式、互连方式，以及它的 Roofline 模型与 TPU 有何不同。