在上一篇文章中，我们深入探讨了 KV 缓存优化。现在，我们将转向探索影响机器学习模型速度的不同性能瓶颈。本文中详细介绍的概念广泛适用于任何机器学习模型，无论用于训练还是推理。但是，示例将专门针对大型语言模型（LLM）推理设置。
在开始之前，我强烈推荐这篇博文 [1]，本文在很大程度上受其启发。
四种性能瓶颈
如果您对模型的性能感到不满意，并准备投入时间进行改进，那么第一步是确定瓶颈类型。主要有四类性能瓶颈——三类与硬件限制有关，一类与软件有关。
让我们首先检查硬件瓶颈。每个瓶颈对应于特定的操作模式：
1. 计算绑定模式： 大部分处理时间，即延迟，都用于执行算术运算（图 1）。与其他模式相比，由于您主要为计算付费，因此计算绑定模式是最具成本效益的，因此我们应该努力实现这种模式。
图 1 - 计算绑定过程，计算时间和数据传输时间分别以黄色和蓝色突出显示
2. 内存带宽绑定模式： 大部分处理时间都用于在芯片内存和处理器之间移动数据，例如权重矩阵、中间计算等（图 2）。
图 2 - 内存带宽绑定过程，计算时间和数据传输时间分别以黄色和蓝色突出显示
3. 通信绑定模式（[1] 中未涵盖）： 仅在计算和数据分布在多个芯片之间时适用。大部分处理时间都用于芯片之间网络数据传输（图 3）。
图 3 - 通信绑定过程，计算时间、数据传输时间和网络通信时间分别以黄色、蓝色和绿色突出显示
注意： 我使用“芯片”一词，因为这些概念适用于任何类型的芯片：CPU、GPU、定制硅（Google TPU、AWS Neuron 核心等）等等。
请注意，现代硬件和框架经过高度优化，计算和数据传输任务通常会部分重叠（图 4）。为了简单起见，在本文中，我们将继续假设它们按顺序发生。
图 4 - 通信绑定过程，数据传输重叠
4. 开销绑定模式： 最后一种模式称为开销绑定模式，它与软件引起的限制有关。在这种模式下，大部分处理时间都用于调度工作并将其提交给硬件——本质上，我们花费更多时间来弄清楚该做什么，而不是在硬件上执行实际操作（图 5）。当使用非常灵活的语言（例如 Python）或框架（例如 PyTorch）时，更有可能出现开销绑定模式，这些语言或框架不需要在运行时显式指定所有必要的信息（例如张量数据类型、目标设备、要调用的内核等）。这种缺失的信息必须在运行时推断，相应的 CPU 周期称为开销。与 CPU 相比，加速硬件现在非常快，因此开销会影响硬件利用率，进而影响成本效益的情况很可能发生——本质上，有时硬件会保持空闲状态，等待软件提交下一个工作项。
图 5 - 开销绑定过程，计算时间、数据传输时间和软件开销时间分别以黄色、蓝色和紫色突出显示
执行模型的正向或反向传递涉及运行多个内核执行（约等于 GPU 函数调用）。所有内核不太可能在同一模式下运行。重要的是要确定大部分执行时间花费在哪种模式下。然后，优先级就变成了针对主要瓶颈进行优化，确定下一个最重要的瓶颈，依此类推。
准确识别瓶颈类型至关重要。每个问题都需要不同的解决方案。如果您误诊，您可能会浪费时间实施优化，而这种优化即使有用，也会让您失望。
诊断限制因素
我们不会在这里深入探讨细节，但 [1] 强调，当出现开销绑定模式时，运行时间不会随着计算或数据传输的增加而按比例缩放。换句话说，如果您将计算或数据传输能力翻倍，但运行时间没有相应增加，那么您的程序很可能是开销绑定模式。否则，您就是硬件绑定模式，但要区分计算瓶颈和内存带宽瓶颈，需要访问指标，例如 FLOP 计数和数据传输量，即使用探查器。
回到大型语言模型，请记住，训练和推理预填充阶段通常是计算绑定模式，而推理解码阶段通常在大多数硬件上是内存带宽绑定模式。因此，主要针对训练的优化（例如低精度矩阵乘法）如果应用于减少以解码延迟为主的总推理延迟，则可能没有那么有用。
基于瓶颈类型优化以降低延迟
让我们看看如何针对每种类型的瓶颈进行优化。
1. 计算绑定模式：
- 升级到功能更强大、价格更高的芯片，具有更高的峰值 FLOPS。
- 对于特定的操作，例如矩阵乘法，利用专门的、更快的内核，例如 NVIDIA Tensor 核心。例如，NVIDIA H100 PCIe [2] 使用通用 CUDA 核心具有 51 TFLOPS 峰值计算，而使用专门的 Tensor 核心（全精度）具有 378 TFLOPS 峰值计算。
- 减少所需操作的数量。更具体地说，对于机器学习模型，这意味着可以使用更少的参数来实现相同的结果。修剪或知识蒸馏等技术可以提供帮助。
- 使用较低精度和更快的类型进行计算。例如，对于相同的 NVIDIA H100 PCIe，8 位 Tensor 核心峰值 FLOPS（1 513 TFLOPS）是 16 位峰值 FLOPS（756 TFLOPS）的两倍，而 16 位峰值 FLOPS 又是 32 位峰值 FLOPS（378 TFLOPS）的两倍。但是，这需要对所有输入进行量化（例如，权重矩阵和激活，例如，参见 LLM.int8() [3] 或 SmoothQuant [4] 量化算法）并使用专用的低精度内核。
2. 内存带宽绑定模式：
- 升级到功能更强大、价格更高的芯片，具有更高的内存带宽。
- 使用模型压缩技术（例如量化或不太流行的修剪和知识蒸馏）来减少移动的数据量。关于大型语言模型，数据大小问题主要通过仅使用权重量化技术（例如，参见 GTPQ [5] 和 AWQ [6] 量化算法）以及 KV 缓存量化来解决。
- 减少内存操作的数量。在 GPU 上运行任务归结为执行内核的定向图（约等于 GPU 函数调用）。对于每个内核，必须从内存中获取输入并将输出写入内存。融合内核，即最初将分散在多个内核中的操作作为一个内核调用执行，可以减少内存操作的数量。运算符融合（图 6）可以由编译器自动执行，也可以通过编写自己的内核（更难，但对于复杂的融合来说是必要的）手动执行。
在 Transformer 模型的情况下，为注意力层开发高效的融合内核仍然是一个活跃的研究领域。许多优化的内核基于流行的 FlashAttention 算法 [7]。Transformer 融合内核库包括 FlashAttention、Meta 的 xFormers 和现已弃用的 NVIDIA FasterTransformer（已合并到 NVIDIA TensorRT-LLM 中）。
图 6 - 应用于 CNN 的水平和垂直层（操作）融合示例，初始状态（上）和最终状态（下）[8]
3. 通信绑定模式：
- 升级到功能更强大、价格更高的芯片，具有更高的网络带宽。
- 通过选择更有效的划分和集体通信策略来减少通信量。例如，[9] 通过引入新的张量并行策略来扩展 Transformer 模型的流行张量并行布局 [10]，这些策略允许通信时间更好地扩展（即防止通信时间成为大型芯片数量和/或批处理大小的瓶颈）。
例如，[10] 中的张量并行布局使权重分片保持静止，而激活分片在芯片之间移动。例如，在预填充阶段，对于非常大的序列批次，[9] 指出激活可能超过权重。因此，从通信的角度来看，使激活保持静止并改为移动权重分片更有效，如他们的“权重收集”划分策略中所述。
4. 开销绑定模式：
- 通过使用更不灵活但更高效的语言（例如 C++）来换取灵活性，以减少开销。
- 将内核分组提交，以便在多个内核之间摊销提交开销，而不是为每个内核支付开销。当您需要多次提交相同的一组短暂内核时，这尤其有用（例如，在迭代工作负载中）。CUDA 图（自 PyTorch 1.10 版本 [11] 开始集成）通过提供实用程序来捕获代码块中所有 GPU 活动（以一个内核启动的定向图形式）以进行一次性提交，从而实现此目的。
- 预先（AOT）提取计算图，并将其放入可部署的工件（模型跟踪）中。例如，PyTorch torch.jit.trace 会跟踪 PyTorch 程序并将其打包到可部署的 TorchScript 程序中。
您可以使用模型编译器进一步优化图。
在任何情况下，您都将灵活性换取更少的开销，因为跟踪/编译需要张量大小、类型等参数保持静态，因此在运行时保持不变。控制流结构（例如 if-else）通常也会在此过程中丢失。
对于需要与 AOT 编译不兼容的灵活性的情况（例如动态张量形状、控制流等），即时（JIT）编译器通过在代码执行之前动态优化模型代码来提供帮助（尽管不如 AOT 编译器那样彻底）。例如，PyTorch 2.x 提供了一个名为 TorchDynamo 的 JIT 编译器。由于您不需要修改 PyTorch 程序来使用它，因此您可以获得使用 JIT 编译器带来的降低的开销，同时保持 Python 开发体验的可用性。
旁注： 模型优化器和（AOT）编译器之间有什么区别？在我看来，这种区别有点模糊。以下是我在概念上区分这两个术语的方式。
首先，两者都在预先工作。典型的 AOT 编译器工作流程是：从支持的框架（PyTorch、TensorFlow、MXNet 等）跟踪代码，以将计算图提取到中间表示（IR）中，应用与硬件无关的优化（代数重写、循环展开、运算符融合等）生成优化的图，最后为目标硬件创建可部署的工件，包括硬件感知优化（选择最合适的内核、数据移动优化等）。AOT 模型编译器的示例包括 PyTorch 的 TorchInductor、XLA 和 Meta 的 Glow。
模型优化器是包含 AOT 编译的工具，但通常针对特定硬件（例如，英特尔硬件用于 OpenVINO，英伟达硬件用于 TensorRT 和 TensorRT-LLM），并且能够执行额外的训练后优化，例如量化或修剪。
到目前为止，我们只关注延迟（处理单个请求所需的时间），但让我们通过更深入地研究计算和内存带宽绑定模式来将吞吐量（每单位时间可以处理的请求数量）重新引入框架。
瓶颈 = f(硬件，算术强度)
有趣的是，相同的算法处理相同的输入，可以是计算绑定模式，也可以是内存带宽绑定模式，具体取决于使用的硬件。适用的模式由算法的算术强度决定——每访问一个字节内存执行的算术运算次数。
我们希望强度值使我们处于或更接近更具成本效益的计算绑定模式。正如我们将看到的，更高的强度与更高的吞吐量和成本效益相关。但是，一些强度驱动因素可能会降低延迟。延迟和吞吐量之间的权衡几乎是不可避免的。
令 b 为每次运行从内存传输到内存的数据字节数，令 p 为每次运行执行的 FLOP（浮点运算）数。令 BW_mem（以 TB/s 为单位）为硬件的内存带宽，令 BW_math（以 TFLOPS 为单位）为数学带宽，也称为峰值 FLOPS。令 t_mem 为移动数据字节所花费的时间，令 t_math 为执行算术运算所花费的时间。
计算绑定模式意味着在算术运算上花费的时间比传输数据的时间更多（图 7）。
图 7 - 计算与内存带宽绑定模式。计算时间和数据传输时间分别以黄色和蓝色突出显示。
因此，当以下情况成立时，我们处于计算绑定模式：
A 是算法的算术强度，其维度是每字节的 FLOP。对于每个传输的字节，算术运算越多，算术强度就越高。
如公式所示，为了使算法处于计算绑定模式，其算术强度必须超过硬件相关的峰值 FLOPS 与内存带宽的比率。相反，内存带宽绑定模式意味着以低于相同带宽比率的强度运行（图 8）。
图 8 - 内存带宽/计算绑定边界
让我们看看 NVIDIA 硬件上半精度矩阵乘法（即使用 Tensor 核心）的一些实际数字（表 1）：
表 1 - 常用于训练和/或为大型语言模型提供服务的 NVIDIA 数据中心 GPU 的规格
这意味着什么？以 NVIDIA A10 为例，带宽比率为 208 意味着在该特定硬件上移动一个字节数据与执行 208 个 FLOP 一样快。因此，如果在 NVIDIA A10 上运行的算法每传输一个字节至少执行 208 个 FLOP（或等效地每传输一个半精度数字执行 416 个 FLOP），那么它不可避免地会花费更多时间来移动数据而不是进行计算，即它处于内存带宽绑定模式。换句话说，算术强度低于硬件带宽比率的算法是内存带宽绑定模式。QED。
考虑到大型语言模型推理过程的解码阶段具有较低的算术强度（将在下一篇博文中详细介绍），因此它在大多数硬件上处于内存带宽绑定模式。与 NVIDIA H100 相比，NVIDIA H200 针对这种低强度工作负载具有更优越的带宽比率。这解释了 NVIDIA 将 H200 推销为“加速生成式 AI 推理”，因为其硬件设计针对这种内存绑定问题。
现在让我们将算术强度与延迟和吞吐量联系起来：
注意： 吞吐量在这里以 TFLOPS 表示，而不是每秒请求，但两者成正比。此外，吞吐量以 TFLOPS 表示突出了它与硬件利用率的联系，因此也与成本效益相关。为了使联系更加明显，吞吐量更准确地说是每芯片秒的请求数，每请求的芯片秒数越少，即吞吐量越高，成本效益就越高（参见 [9] 第 4 节）。
如果我们将算术强度绘制在 x 轴上，并将（可实现的最大）吞吐量作为 y 轴上的因变量，我们将得到所谓的（朴素）屋顶线模型 [12]（图 9）。
图 9 - 屋顶线模型
让我们做一个简单的思想实验，以更好地理解为什么此图上的吞吐量值是可实现的最大值。在计算绑定模式下，这是显而易见的：没有任何东西可以阻止我们利用全部计算能力，我们只受硬件峰值能力的限制。在内存带宽绑定模式下，我们可以在 1 秒内获取的最大数据量由硬件的内存带宽 BW_mem 决定。考虑到算术强度为 A 的算法，我们可以在 1 秒内实现的最大 FLOP 数因此为 BW_mem.A。QED。
增加算法的算术强度会有什么影响？我们可以检查三种情况（图 10）：
图 10 - 算术强度增加的三种情况
情况 1： 算术强度的增加太小，无法逃脱内存带宽绑定模式，但涉及吞吐量的比例增加。系统仍然是内存带宽绑定模式，因此对延迟的影响取决于更高的强度如何影响特定算法的数据传输。
情况 2： 算术强度的增加使系统切换到计算绑定模式。吞吐量上升到硬件峰值吞吐量。现在是计算绑定模式，延迟影响取决于更高的强度如何改变特定算法的总操作数。
情况 3： 由于已经处于计算绑定模式并处于峰值吞吐量，因此更高的强度不会带来吞吐量增益。延迟影响仍然取决于更高的强度如何影响特定算法的总计算量。
我们如何具体地增加算术强度？这完全取决于算法的具体情况。在下一篇文章中，我们将检查控制 Transformer 解码器块算术强度的关键参数。我们将看到，例如，提高批处理大小如何可以提高某些操作的算术强度。
我们已经讨论过的一些优化也提高了算术强度，从而提高了吞吐量和资源利用率。对于 Transformer 模型（其解码阶段是内存带宽绑定模式），算术强度主要通过减少数据传输的数量和大小来提高，方法是使用操作融合和数据（权重矩阵、KV 缓存）量化。
到目前为止，我们假设算法实现以最佳方式利用了硬件资源。例如，在传输数据时，假设算法实现使用了硬件理论内存带宽的 100%。这在实践中显然不是这种情况（尽管一些实现实现了接近最佳的资源使用率），那么次优资源使用如何影响分析？
这很简单：上面的带宽数字必须替换为实际实现的数字。次优系统在其自身的屋顶线曲线上，位于最佳屋顶线曲线下方（图 11）。现在有两个自由度可以提高吞吐量：增加算术强度和/或增强算法实现以更好地利用硬件资源。
图 11 - 具有次优资源利用率的屋顶线模型
真实示例：FlashDecoding
让我们提供一个实现改进的真实示例。在 2.2 版本之前，FlashAttention 内核的实现当应用于推理的解码阶段时，可能会变得非常次优。以前的数据加载实现使内核本质上在解码阶段利用内存带宽效率较低。更糟糕的是，带宽利用率实际上随着批处理大小的增加而进一步降低；因此，性能在需要更小批次的较长序列中表现最差，因为内存限制。FlashAttention 团队解决了这个问题（主要是通过在 KV 缓存序列长度维度上并行化数据加载），并发布了一个名为 FlashDecoding 的优化解码阶段内核，该内核在长序列长度方面实现了显着的延迟改进 [13]。
总结
在本文中，我们了解了可能影响模型延迟的四种类型的瓶颈。识别主要影响模型延迟的瓶颈类型至关重要，因为每种类型都需要特定的缓解策略。
实际硬件操作是计算绑定模式或内存带宽绑定模式。内核的算术强度决定了绑定模式。在较低强度的内存带宽绑定模式下，可实现的最大吞吐量与算术强度成线性关系。相反，在计算绑定模式下，吞吐量受硬件峰值 FLOPS 的限制。根据影响强度的因素，我们可能能够提高强度以提高最大吞吐量，甚至可能达到计算绑定模式的性能。但是，这种强度增益可能会对延迟产生负面影响。
在下一篇文章中，我们将运用这些新知识来分析大型语言模型，更详细地研究 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)