如何扩展你的模型

深度学习的许多方面仍然可以归结为一种黑魔法，但优化你的模型性能并非如此——即使是在巨大规模下！相对简单的原则无处不在——从处理单个加速器到数万个加速器——理解这些原则可以让你做许多有用的事情：

• 粗略估算你的模型各部分与理论最优值的接近程度。
• 在不同规模下，对不同的并行性方案（如何将计算分配到多个设备上）做出明智的选择。
• 估算训练和运行大型Transformer模型所需的成本和时间。
• 设计能够利用特定硬件功能的算法。
• 基于对当前算法性能限制的明确理解来设计硬件。

预期背景：我们假设你对LLM和Transformer架构有基本的了解，但不一定了解它们如何大规模运行。你应该了解LLM训练的基础知识，最好对JAX有一些基本熟悉。一些有用的背景阅读可能包括关于Transformer架构的这篇博客文章和原始的Transformer论文。还可以查看这个列表以获取更多有用的同步和未来阅读材料。

目标与反馈：读完本书后，你应该能够轻松地为给定硬件平台上的Transformer模型估算出最佳的并行性方案，以及大致的训练和推理时间。如果你做不到，请给我们发邮件或留言！我们很想知道如何能把内容讲得更清楚。

你可能也会喜欢阅读关于NVIDIA GPU的新的第12节！

你为什么应该关心？

三四年前，我认为大多数机器学习研究人员不需要理解这本书中的任何内容。但如今，即使是“小”模型也运行得非常接近硬件极限，以至于进行创新性研究需要你考虑大规模下的效率。历史上，机器学习研究遵循着一种系统创新和软件改进之间的“嘀嗒”循环。Alex Krizhevsky不得不编写极其复杂的CUDA代码来使CNNs变快，但几年内，像Theano和TensorFlow这样的库意味着你不再需要这样做。也许这种情况也会在这里发生，这本书中的所有内容几年后都会被抽象掉。但是，扩展定律已经将我们的模型永久地推向了硬件的最前沿，而且似乎在不久的将来，进行前沿研究将与理解如何有效地将模型扩展到大型硬件拓扑结构密不可分。 在基准测试上取得20%的提升，如果以牺牲20%的Roofline效率为代价，那是无关紧要的。 有前景的模型架构之所以失败，通常要么是因为它们无法在规模上高效运行，要么是因为没有人投入精力去实现这一点。

“模型扩展”的目标是能够增加用于训练或推理的芯片数量，同时实现吞吐量的成比例线性增长。这被称为“强扩展”。尽管增加额外的芯片（“并行性”）通常会减少计算时间，但它也会带来芯片间通信增加的成本。当通信时间超过计算时间时，我们就会变得“受通信限制”，无法实现强扩展。随着计算时间的减少，你通常也会在单个芯片层面遇到瓶颈。你闪亮的新TPU或GPU可能额定每秒执行500万亿次操作，但如果你不小心，它很可能只做到十分之一，因为它被在内存中移动参数所拖累。单芯片计算、内存带宽和总内存之间的相互作用对扩展的故事至关重要。 如果我们对硬件有足够的了解，能够预见这些瓶颈将在何处出现，我们就可以设计或重新配置我们的模型来避免它们。硬件设计者面临着相反的问题：构建能够为我们的算法提供恰到好处的计算、带宽和内存，同时最小化成本的硬件。你可以想象这个“协同设计”问题有多么令人紧张：你必须赌上第一批芯片实际可用时算法会是什么样子，这通常是2到3年之后的事情。TPU的故事就是这场博弈中的一个巨大成功。矩阵乘法是一种独特的算法，因为它每字节内存使用的FLOPs远多于几乎任何其他算法（每字节N FLOPs），早期的TPU及其脉动阵列架构在当时实现了比GPU好得多的性价比。TPU是为ML工作负载设计的，而带有TensorCore的GPU也在迅速改变以填补这一细分市场。但你可以想象，如果神经网络没有兴起，或者以TPU（本质上比GPU灵活性差）无法处理的某种根本方式发生了变化，那将是多么昂贵的代价。

我们在这本书中的目标是解释TPU（和GPU）硬件是如何工作的，以及Transformer架构是如何演变以在当前硬件上表现良好的。我们希望这对于设计新架构的研究人员和致力于让当前一代LLM快速运行的工程师都有用。

高层大纲

本书的总体结构如下：

第1节解释了Roofline分析以及哪些因素会限制我们的扩展能力（通信、计算和内存）。第2节和第3节详细讨论了TPU的工作原理，既包括作为单个芯片，也包括——至关重要的——作为一个具有有限带宽和延迟的互连芯片链接的系统。我们将回答以下问题：

• 一个特定大小的矩阵乘法应该花费多长时间？在什么情况下它会受到计算、内存或通信带宽的限制？
• TPU是如何连接在一起形成训练集群的？系统的每个部分有多少带宽？
• 在多个TPU之间收集、分散或重新分配数组需要多长时间？
• 我们如何高效地乘以在不同设备上分布不同的矩阵？

五年前，机器学习的架构景观丰富多彩——ConvNets、LSTMs、MLPs、Transformers——但现在我们主要只有Transformer。我们坚信，理解Transformer架构的每一个部分都是值得的：每个矩阵的确切大小，归一化发生在哪里，每个部分有多少参数和FLOPs浮点运算（FLoating point OPs），基本上是所需加法和乘法运算的总数。虽然许多资料将FLOPs解释为“每秒操作数”，但我们使用FLOPs/s来明确表示这一点。。第4节仔细讲解了这种“Transformer数学”，展示了如何计算训练和推理的参数和FLOPs。这告诉我们模型将使用多少内存，我们将在计算或通信上花费多少时间，以及注意力何时会相对于前馈块变得重要。

第5节：训练和第7节：推理是本文的核心，我们在这里讨论一个根本问题：给定一个特定大小的模型和一定数量的芯片，我该如何并行化我的模型以保持在“强扩展”状态？这是一个简单的问题，但答案却出人意料地复杂。从高层次来看，有4种主要的并行性技术用于在多个芯片上拆分模型（数据、张量、流水线和专家），以及一些其他技术来减少内存需求（重计算、优化器/模型分片（又名ZeRO）、主机卸载、梯度累积）。我们在这里讨论其中的许多技术。

我们希望在这些章节结束时，你应该能够为新的架构或设置自己选择合适的并行方案。第6节和第8节是将这些概念应用于流行的开源模型LLaMA-3的实践教程。

最后，第9节和第10节探讨了如何在JAX中实现其中一些想法，以及在出现问题时如何分析和调试代码。第12节是一个新章节，也深入探讨了GPU。

在整个过程中，我们尝试给你提供一些问题让你自己解决。请不要有压力去阅读所有章节或按顺序阅读。也请留下反馈。目前，这是一个草稿，并将继续修订。谢谢！

我们想感谢James Bradbury和Blake Hechtman，他们推导出了本文档中的许多想法。

话不多说，这是关于TPURoofline模型的第一节。

各章节链接

这个系列可能比它需要的要长，但我们希望这不会让你望而却步。前三章是预备知识，如果熟悉可以跳过，尽管它们介绍了后面使用的符号。最后三个部分可能是最实用的，因为它们解释了如何处理真实模型。

第1部分：预备知识

• 第1章：Roofline分析简介。算法受到三件事的限制：计算、通信和内存。我们可以用这些来近似我们的算法将运行多快。

• 第2章：如何理解TPU。TPU是如何工作的？这如何影响我们可以训练和服务的模型？

• 第3章：分片矩阵及其乘法。在这里，我们通过我们最喜欢的操作：（分片）矩阵乘法来解释模型分片和多TPU并行性。

第2部分：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的成本是多少？延迟/吞吐量的权衡是什么？

第3部分：实践教程

• 第9章：如何分析TPU代码。真实的LLM从不像上面的理论那么简单。在这里，我们解释JAX + XLA技术栈以及如何使用JAX/TensorBoard分析器来调试和修复实际问题。

• 第10章：用JAX编程TPU。JAX提供了一系列神奇的API用于并行化计算，但你需要知道如何使用它们。包含有趣的例子和已解决的问题。

第4部分：结论和附加内容

• 第11章：结论和进一步阅读。关于TPU和LLM的总结思考和进一步阅读材料。

• 第12章：如何理解GPU。一个关于GPU的附加章节，介绍它们如何工作，如何联网，以及它们的Roofline与TPU有何不同。