在本系列文章《LLM推理入门指南①：文本生成的初始化与解码阶段》中，作者对Transformer解码器的文本生成算法进行了高层次概述，着重介绍了两个阶段：提示的处理阶段和逐个生成补全词元的多步生成阶段。在上一篇文章《LLM推理入门指南②：深入解析KV Cache》中，深入探讨了KV Cache优化。
本文将转变方向，探索可能影响机器学习模型速度的不同性能瓶颈。本文所详细介绍的概念广泛适用于任何ML模型，无论是用于训练还是推理，不过提供的示例主要聚焦于LLM推理设置。
（本文作者为AWS的GenAI解决方案架构师Pierre Lienhart。以下内容由OneFlow编译发布，转载请联系授权。原文：https://medium.com/@plienhar/...）
作者 |Pierre Lienhart
OneFlow编译
翻译｜张雪聃
1 四种性能瓶颈
如果你对模型性能不太满意，并且打算继续改进，那么第一步就是确定性能瓶颈的类型。瓶颈主要有四类——其中三类与硬件限制相关，另一类与软件相关。
我们先来看看硬件瓶颈。每一个都有其对应的特定操作模式。
- 计算受限模式：大部分处理时间（即时延）被用于执行算术运算（见图1）。与其他模式相比，由于主要成本是计算，计算受限模式是最具成本效益的，因此我们应该以此为目标。
- 内存带宽受限模式：大部分处理时间花在了芯片内存和处理器之间搬运数据，如权重矩阵、中间计算等（见图2）。
- 通信受限模式（引文[1]中未涵盖）：仅适用于计算和数据分布在多个芯片之间的情况。大部分处理时间花在了芯片之间的网络数据传输（见图3）。
注意：我用“芯片”这个词是因为这些概念适用于任何类型的芯片：CPU、GPU、定制化芯片（如Google的TPU、AWS的Neuron Cores等）。
需注意的是，现代硬件和框架经过了高度优化且计算和数据传输任务常常出现部分重叠的现象（见图4）。简单起见，我们在本文中将继续假设它们按顺序进行。
最后一种模式称为计算开销受限，与软件引起的限制相关。在这种模式下，大部分处理时间花在了调度工作并将其提交给硬件上 —— 基本上，我们花更多的时间去弄清楚要做什么，而不是在硬件上执行实际操作（见图5）。计算开销受限的情况在使用非常灵活的语言（例如Python）或框架（例如PyTorch）时更有可能出现，这些语言或框架不需要在运行时明确指定全部所需的信息（如张量数据类型、目标设备、要调用的kernel等）。这些缺失的信息必须在运行时推理出来，相应的CPU周期称为计算开销。现在的加速硬件与CPU相比速度更快，以至于计算开销影响硬件利用率，进而影响成本效益的情况很可能发生 —— 基本上，有时硬件处于空闲状态，等待软件提交下一个工作项。
执行模型的反向或正向传播涉及运行多kernel的执行（~ GPU函数调用）。所有kernel都在相同的模式下运行是不可能实现的。问题的关键在于确定大部分执行时间都花在哪个模式上。由此，优先事项就变成针对主要瓶颈进行优化，找出下一个最重要的瓶颈，依此类推。
正确地判断瓶颈类型至关重要。每个问题需要不同的解决方案。如果判断错误，可能会浪费不少时间，即使实施这种优化确实有所帮助，但结果可能依旧令人失望。
2 判断限制因素
这篇文章中我们不会深入讨论细节问题，但博文[1]强调了当计算开销受限制时，运行时不会随着增加的计算或数据传输成比例地扩展。换句话说，如果你将计算或数据传输容量加倍，但运行时未相应增加，那么你的程序很可能受到计算开销限制。否则，你可能受到硬件限制，但区分计算与内存带宽等瓶颈需要访问诸如FLOP计数和数据传输量等指标，即使用性能分析器（profiler）。
当我们重新把LLM纳入讨论范围时，请记住，训练和推理的预填充阶段通常是计算受限，而推理解码阶段通常在大多数硬件上是内存带宽受限。因此，主要用于训练的优化（例如低精度矩阵乘法）如果应用于减少由解码时延为主的总推理时延，可能就不会那么有帮助。
3 基于瓶颈类型优化以降低时延
让我们看看针对每种瓶颈类型的优化。
如果处于计算受限模式下，你可以：
- 升级到更强大、更昂贵、峰值FLOPS更高的芯片。
- 针对特定操作，如矩阵乘法，可以使用专门的、更快的核心，如NVIDIA Tensor Cores。例如，NVIDIA H100 PCIe [2]使用通用CUDA核心的峰值计算性能为51 TFLOPS，而使用专门的Tensor Cores（全精度）可达到378 TFLOPS。
- 减少所需的操作数。更具体地说，对于ML模型，这可能意味着，使用更少的参数来实现相同的结果。剪枝或知识蒸馏等技术可以帮助实现这一点。
- 使用较低精度和更快速的数据类型进行计算。例如，对于相同的NVIDIA H100 PCIe，8位Tensor Cores的峰值FLOPS（1,513 TFLOPS）是16位峰值FLOPS（756 TFLOPS）的两倍，是32位峰值FLOPS（378 TFLOPS）的两倍。但这需要量化所有的输入（例如，权重矩阵和激活函数），并使用专门的低精度kernel。
如果处于内存带宽受限模式下，你可以：
- 升级到更强大、更昂贵、有更高的内存带宽的芯片
- 减少所搬运数据的大小，例如使用模型压缩技术，如量化或不太流行的剪枝和知识蒸馏。就LLM而言，数据大小问题主要通过仅包含权重的量化技术（例如GTPQ [5] 和AWQ [6] 量化算法）与KV Cache量化来解决。
- 减少内存操作次数。GPU上的运行任务可以归结为执行一系列kernel（GPU函数调用）的有向图。对于每个kernel，必须从内存中获取输入并将输出写入其中。kernel融合，即将最初分散在多个kernel中的操作作为单个kernel调用执行，可以减少内存操作次数。算子融合（见图6）可以由编译器自动执行，也可以通过编写自己的kernel来手动执行（这更难，但对于复杂的融合是必要的）。
对于Transformer模型而言，针对注意力层开发高效融合的kernel仍然是一个十分活跃的领域。许多优化的kernel都基于流行FlashAttention算法[7]。Transformer融合kernel库包括FlashAttention、Meta的xFormers和现已弃用的NVIDIA FasterTransformer（已合并到NVIDIA TensorRT-LLM中）。
如果处于通信受限模式下，你可以：
- 升级到更强大、更昂贵、具有更高的网络带宽的芯片
- 通过选择更高效的分区和集合通信策略来减少通信量。例如，通过引入新的张量并行策略，[9]扩展了来自[10]中用于Transformer模型的流行张量并行布局，使通信时间能够更好地扩展（即防止它们成为瓶颈），以适应大量芯片或批大小。
例如，[10]中的张量并行布局保持权重分片不变，而激活分片在芯片之间移动。例如，在预填充阶段和对于非常大的序列批次，例如，[9]指出激活可能比权重更大。因此，从通信的角度来看，保持激活不变并将权重分片移动会更有效，就像他们所采用的“权重聚合（weight-gathered）”分区策略那样。
如果处于计算开销受限模式下，你可以：
- 通过使用不太灵活但更高效的语言，如C++，来换取更少的计算开销。
- 将kernel分组提交以分摊跨多个kernel提交的计算开销，而不是按每个kernel都算一次计算开销。
- 这在需要多次提交相同的“短命（short-lived）”kernel组时特别有益（例如在迭代工作负载中）。CUDA Graph（自PyTorch 1.10版以来已集成）通过提供工具捕获由代码块引起的所有GPU活动，形成一个kernel启动的有向图，实现一次性提交，从而服务于这一目的。
- 提前提取计算图（AOT）并转换为可部署的工件（模型追踪）。例如，PyTorch的torch.jit.trace 可以追踪PyTorch程序，并将其打包成可部署的TorchScript程序。
你可以通过使用模型编译器进一步优化图。
在任何情况下，你再次通过牺牲灵活性来减少计算开销，因为追踪/编译需要参数（如张量大小、类型等）是静态的，因此运行时（runtime）保持不变。控制流结构，如if-else，在这一过程中通常也会丢失。
在对灵活性有需求但同时与AOT编译不兼容的情况中（例如，动态张量形状、控制流等），即时（JIT）编译器能通过在模型代码在执行之前进行动态优化来提供帮助（虽然不如AOT编译器彻底）。例如，PyTorch 2.x 引入了一个名为TorchDynamo的JIT编译器。由于无需修改你的PyTorch程序即可使用它，因此你可以使用JIT编译器减少计算开销，同时保持Python开发体验的易用性。
附注：模型优化器和（AOT）编译器之间是否有区别？在我看来，这个区分比较模糊。我在概念上对这两个术语的区分如下。
首先，它们都是提前执行的。典型的AOT编译器工作流程是：跟踪来自支持框架（如PyTorch、TensorFlow、MXNet等）的代码，将计算图提取到中间表示（IR）中，应用与硬件无关的优化（如代数重写、循环展开loop unrolling、算子融合等）生成优化后的图形，最后为目标硬件创建可部署的工件，包括硬件感知优化（选择最合适的内核、数据移动优化等）。一些AOT模型编译器的示例包括PyTorch的TorchInductor、XLA和Meta的Glow。
模型优化器是一种工具，包含AOT编译，但通常针对特定硬件（例如，OpenVINO针对英特尔硬件，TensorRT和TensorRT-LLM针对NVIDIA硬件），能够执行额外的后训练优化，如量化或剪枝。
到现在，我们只聚焦于时延（处理单个请求所需的时间），接下来让我们通过深入探讨计算和内存带宽受限模式，重新引入吞吐量（单位时间内可以处理的请求数量）。
4 瓶颈 = f（硬件，算术强度）
有趣的是，处理相同输入的相同算法可以是计算受限模式，也可以是内存带宽受限，这取决于所使用的硬件。适用的模式由算法的算术强度确定，即每个访问的字节执行的算术操作数量。
我们想要的是能使我们处于或更具成本效益的计算受限模式的强度值。正如我们将看到的那样，较高的强度与更好的吞吐量和成本效率相关。然而，一些强度驱动因素可能会降低延迟。我们不得不在时延和吞吐量之间进行权衡。
设b为每次运行从内存传输的数据字节数，p为每次运行执行的浮点运算次数（FLOPs）。设BW\_mem（以TB/s表示）为硬件的内存带宽，BW\_math（以TFLOPS表示）为数学带宽，也称为峰值FLOPS。设 t\_mem 为移动数据字节所花费的时间，t\_math 为算术运算所花费的时间。
处于计算受限模式意味着在算术运算上花费的时间比传输数据多（见图7）。
因此，当以下条件满足时，我们处于计算受限状态：
A是算法的算术强度，其维度为FLOP/byte。每字节传输的数据进行的算术运算越多，算术强度就越高。
如方程所示，要使算法处于计算受限状态，其算术强度必须超过峰值FLOPS与内存带宽的比率有关的硬件相关比值。相反，内存带宽受限意味着操作强度低于相同带宽比率（见图8）。
让我们来看一些关于在NVIDIA硬件上进行半精度矩阵乘法（即使用张量核心）的实际数据（表1）：
这意味着什么？以NVIDIA A10为例，带宽比为208意味着在该特定硬件上搬运一字节的数据与执行208 FLOPs的速度相同。因此，如果在NVIDIA A10上运行的算法每传输一字节不能执行至少208次FLOPs（或等效地每传输半精度数执行416次FLOPs），那么它必然会花费更多的时间在数据传输上而不是在计算上，即它是内存带宽受限。换句话说，算法的算术强度低于硬件带宽比率的情况下，就是内存带宽受限。证明完毕。
已知LLM推理过程的解码阶段具有较低的算术强度（在下一篇博文中详细介绍），这导致它在大多数硬件上成为内存带宽受限。与NVIDIA H100相比，NVIDIA H200针对这种低强度工作负载具有更有利的带宽比率。这解释了NVIDIA将H200宣传为“为生成式AI推理提速”的原因，因为其硬件设计针对的是这种内存受限情况。
现在我们将算术强度与时延和吞吐量联系起来：
注意：这里的吞吐量表示为TFLOPS而不是每秒请求数，但两者是直接成比例的。另外，吞吐量以TFLOPS表示突显了它与硬件利用率的联系，因此与成本效率有关。为了使这种联系更明显，吞吐量更精确地表示为每个芯片-秒（chip-second）的请求数量，每个请求的芯片秒数越低，即吞吐量越高，成本效率越高（参见[9]第4节）。
如果我们将算术强度作为自变量绘制在x轴上，将（最大可实现的）吞吐量作为因变量绘制在y轴上，我们就得到了所谓的（简单）屋顶线模型[12]（图9）。
让我们进行一个小小的思维实验，以更好地理解为什么此图中的吞吐量值是最大可实现的水平。显而易见，在计算受限模式下：没有什么可以阻止我们利用全部计算能力，我们只受到硬件峰值容量的限制。在内存带宽受限模式下，我们在1秒内可以获取的最大数据量由硬件的内存带宽BW\_mem确定。考虑到算法的算术强度A，因此我们在1秒内可以实现的最大FLOP数量是BW\_mem乘以A。
增加算法的算术强度会产生什么影响？我们可以考虑三种情况（图10）：
- 情景1：算术强度的增加太小，无法摆脱内存带宽受限的状态，但吞吐量线性增长。系统仍然受到内存带宽的限制，因此对时延的影响取决于更大强度如何影响该特定算法的数据传输。
- 情景2：算术强度的增加使系统转换到计算受限模式。吞吐量增加到硬件的峰值吞吐量。现在处于计算受限状态，时延的影响取决于更高强度如何改变该算法的总操作量。
- 情景3：由于已经处于计算受限状态并且达到了峰值吞吐量，增加强度不会提供吞吐量增益。时延的影响仍取决于更高强度如何影响该算法的总计算量。
如何具体增加算术强度？这完全取决于算法的具体情况。在下一篇文章中，我们将考查控制Transformer解码器块算术强度的关键参数。例如，我们会看到如何通过提高批大小来增加某些操作的算术强度。
一些已经讨论过的优化方法也可以增加算术强度，从而提高吞吐量和资源利用率。对于Transformer模型（其解码阶段受内存带宽限制），通过操作融合和数据（权重矩阵、KV Cache）量化来减少数据传输的数量和大小，来改善算术强度。
目前为止，我们进行了一个关键的假设，即算法的实现充分利用了硬件资源。例如，在传输数据时，假设算法的实现使用了硬件的理论内存带宽的100%。显然，在实践中并非总是如此（尽管一些实现达到了接近最佳资源利用率），那么次优的资源使用对分析有何影响呢？
很简单：上述带宽数字必须被实际实现的数字所替代。次优系统位于它自己的屋顶线曲线上，位于最优化的屋顶线曲线下方（见图11）。现在有两个提高吞吐量的自由度：增加算术强度和/或改善算法的实现以更好地利用硬件资源。
让我们通过一个实际的改进实现的示例来总结。在版本2.2之前，FlashAttention kernel的实现在应用于推理的解码阶段时可能会变得非常不理想。先前的数据加载实现使得kernel在解码阶段利用内存带宽方面的效率较低。更糟的是，随着批大小的增加，带宽利用实际上进一步降低；因此，对于由于内存限制需要较小批次的长序列，性能受到的影响最大。FlashAttention团队通过解决这个问题（主要是通过在KV Cache序列长度维度上并行加载数据）发布了一个名为FlashDecoding的优化解码阶段kernel，为长序列长度实现了显著的时延改进。[13]
5 总结
通过本文，我们了解了影响模型时延的四种瓶颈类型。识别模型时延主要成因的类别至关重要，因为每种类别都需要特定的缓解策略。
不考虑分布式设置以简化问题，实际的硬件操作要么是计算受限，要么是内存带宽受限。kernel的算术强度决定了绑定的模式。在较低强度的内存带宽受限模式下，最大可实现的吞吐量与算术强度呈线性关系。相反，在计算受限模式下，吞吐量受到峰值硬件FLOPS的限制。根据影响强度的因素，我们可能能够增加强度以提高最大吞吐量，甚至可能达到计算受限的性能。然而，这种强度增益可能会增加时延。
在下一篇文章中，我们将运用这些新知识来研究LLM，具体探讨Transformer解码器块的算术强度。
参考文献
[1]: Making Deep Learning Go Brrrr From First Principles (He, 2022)
[2]: NVIDIA H100 product specifications
[3]: LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (Dettmers et al., 2022) + GitHub repository
[4]: SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models (Xiao et al., 2022) + GitHub repository
[5]: GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers (Frantar et al., 2022) + GitHub repository
[6]: AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (Lin et al., 2023) + GitHub repository
[7]: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022), FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (Dao, 2023) + GitHub repository.
[8]: Deploying Deep Neural Networks with NVIDIA TensorRT (Gray et al., 2017)
[9]: Efficiently Scaling Transformer Inference (Pope et al., 2022)
[10]: Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism Shoeybi et al., 2019)
[11]: Blog post — Accelerating PyTorch with CUDA Graphs (Ngyuen et al., 2021)
[12]: Roofline: an insightful visual performance model for multicore architectures (Williams et al., 2009)
[13]: Blog post — Flash-Decoding for long-context inference (Dao et al., 2023)
作者：Pierre Lienhart
来源：OneFlow
推荐阅读
- LLM推理入门指南②：深入解析KV缓存
- LLM推理入门指南①：文本生成的初始化与解码阶段
- AffineQuant: 大语言模型的仿射变换量化
- 媲美Sora！可生成16s 1080视频！清华联合生数提出视频生成新模型Vidu！
欢迎大家点赞留言，更多Arm技术文章动态请关注极术社区嵌入式客栈专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
