你是否想要提升深度学习模型的性能？面对这个问题，很多人会习惯性地使用一些“秘诀”，比如“使用就地操作！将梯度设置为 None！安装 PyTorch 1.10.0，但不要安装 1.10.1！”。
这种“灵机一动”式的优化方法虽然看似有效，但实际上却缺乏理论基础。深度学习模型的性能优化并非炼金术，而是需要从基本原理出发进行分析。
理解性能瓶颈
深度学习模型的性能优化可以从三个方面进行考虑：
- 计算 (Compute): GPU 上执行浮点运算 (FLOPS) 所花费的时间。
- 内存 (Memory): 在 GPU 内传输张量所花费的时间。
- 开销 (Overhead): 其他所有时间。
就像训练机器学习模型一样，了解模型所处的性能瓶颈可以帮助我们找到最有效的优化方法。例如，如果你的模型主要受限于内存带宽，那么提升 GPU 的 FLOPS 性能将毫无用处。相反，如果你的模型主要受限于计算能力，那么将模型逻辑重写为 C++ 代码来降低开销也无济于事。
计算：GPU 的核心能力
从某种程度上来说，优化深度学习系统就是最大化模型在计算受限状态下运行的时间。你为那些强大的 GPU 付出了高昂的代价，理所当然地希望它们能够发挥出全部的计算能力。然而，为了充分利用 GPU 的矩阵乘法能力，我们需要减少其他方面的开销，例如内存传输和系统开销。
为什么我们更关注最大化计算能力，而不是内存带宽？原因很简单：我们可以降低开销或内存成本，但我们无法在不改变实际操作的情况下减少计算量。
计算与内存带宽的矛盾
更糟糕的是，计算能力的增长速度远超内存带宽。下表展示了 CPU FLOPS 和内存带宽的翻倍时间：
| 组件 | 翻倍时间 |
|---|---|
| CPU FLOPS | 18 个月 |
| 内存带宽 | 24 个月 |
我们可以将计算能力想象成一个工厂。我们向工厂发送指令（开销），并提供原材料（内存带宽），以确保工厂高效运转（计算）。
如果工厂的效率提升速度快于我们提供原材料的速度，那么工厂将难以达到峰值效率。
即使工厂的规模（FLOPS）翻倍，如果我们的带宽跟不上，那么性能也不会翻倍。
这种计算能力增长速度快于内存带宽的趋势，一方面意味着机器学习系统工程师将长期拥有稳定的工作，另一方面也使得理解性能瓶颈变得更加重要。
现代 GPU 的特殊性
现代机器学习加速器都拥有专门用于矩阵乘法的硬件，例如英伟达的“Tensor Cores”。
如果你没有进行矩阵乘法，那么你只能达到 19.5 teraflops，而不是标称的 312 teraflops。需要注意的是，这并非 GPU 独有，事实上，TPU 的通用性甚至比 GPU 更低。
GPU 在非矩阵乘法运算方面速度较慢，这似乎是一个问题。那么，其他运算，例如层归一化或激活函数呢？事实是，这些运算在 FLOPS 方面几乎可以忽略不计。例如，让我们看一下这篇论文中关于 BERT 不同运算类型 FLOPS 数量的表格，其中“张量收缩”= 矩阵乘法。
| 运算类型 | FLOPS 占比 |
|---|---|
| 张量收缩 | 99.8% |
| 归一化 | 0.1% |
| 点运算 | 0.1% |
可以看到，非矩阵乘法运算只占总 FLOPS 的 0.2%，因此 GPU 在非矩阵乘法运算方面速度慢 15 倍并不重要。
内存带宽：数据传输的瓶颈
然而，在实际应用中，非矩阵乘法运算却往往比预期花费更多时间。罪魁祸首通常是数据在工厂和仓库之间传输的时间，也就是内存带宽成本。
内存带宽成本是指将数据从一个地方移动到另一个地方所花费的成本。这可能包括将数据从 CPU 移动到 GPU，从一个节点移动到另一个节点，甚至从 CUDA 全局内存移动到 CUDA 共享内存。其中，最后一种情况通常被称为“带宽成本”或“内存带宽成本”。
我们可以再次使用工厂的比喻来理解内存带宽成本。
虽然工厂是进行实际工作的场所，但它并不适合作为大规模存储单元。主要原因是，由于我们在工厂进行实际工作，因此所有存储都针对快速使用进行了优化（SRAM），而不是拥有大量的存储空间。
那么，我们应该在哪里存储实际结果和原材料呢？典型的做法是建立一个仓库，可能位于土地便宜且空间充足的地方（DRAM）。然后，我们可以将物资运送到工厂和从工厂运出（内存带宽）。
将物资运送到计算单元和从计算单元运出的成本就是所谓的“内存带宽”成本。顺便说一下，你的 GPU 的 DRAM 会显示在 nvidia-smi 中，它是你遇到“CUDA 内存不足”错误的主要原因。
需要注意的是，每次执行 GPU 内核时，都需要将数据从 GPU 的 DRAM（即仓库）中读取出来，并在执行完内核后将结果写入 DRAM。
现在，想象一下执行像 torch.cos 这样的单目运算会发生什么。我们需要将数据从存储单元运送到仓库，然后对每条数据进行少量计算，最后将数据运回存储单元。运输数据的成本非常高。因此，我们几乎所有时间都花在了运输数据上，而不是实际的计算本身。
由于我们所有时间都花在了内存带宽上，因此这种运算被称为内存受限运算，这意味着我们没有花太多时间进行计算。
操作融合：优化内存带宽
那么，我们该如何解决这个问题呢？让我们看一下一系列运算的执行过程。
一系列点运算的执行过程可能如下所示：
这种安排非常愚蠢。为什么我们要将相同的数据反复发送到全局内存，然后再发送回计算单元？我们应该将数据保留在工厂中，执行所有计算，然后再将其发送回存储单元。
我们可以通过操作融合来解决这个问题。操作融合是深度学习编译器中最重要的优化方法。简单来说，与其将数据写入全局内存，然后再读取回来，不如将多个计算合并在一起执行，从而避免额外的内存访问。
例如，如果我们执行 x.cos().cos()，通常需要进行 4 次全局读写操作。
x1 = x.cos() # 从全局内存中读取 x，写入 x1
x2 = x1.cos() # 从全局内存中读取 x1，写入 x2但是，使用操作融合，我们只需要进行 2 次全局内存读写操作！因此，操作融合可以将速度提高 2 倍。
x2 = x.cos().cos() # 从全局内存中读取 x，写入 x2操作融合可以显著提高性能，但它也有一些限制。首先，GPU 需要在执行当前操作时知道下一步要执行的操作。因此，我们无法在 Eager 模式下执行操作融合，因为 PyTorch 在 Eager 模式下会逐个执行操作。其次，我们需要为此生成 CUDA 代码，这会带来新的挑战。
并非所有操作融合都像点运算那样简单。你可以将点运算融合到约简运算中，或者将点运算融合到矩阵乘法中。甚至矩阵乘法本身也可以看作是广播乘法和约简的融合。
如果你有兴趣编写自定义 CUDA 内核，那么操作融合将是你最受益的优化方法。任何两个 PyTorch 操作都存在融合的可能性，从而节省了在它们之间读写全局内存的内存带宽成本。此外，许多现有的编译器可以执行“简单的”融合，例如 NVFuser 和 XLA。然而，自动系统无法与人类的智慧相提并论，因此，如果你想尝试编写一些自定义 CUDA 内核，那么 Triton 是一个不错的起点。
最后，操作融合会带来一些意想不到的结果。例如，融合后的 x.cos().cos() 的执行时间几乎与单独调用 x.cos() 的时间相同。这就是为什么激活函数的成本几乎都相同，尽管 gelu 明显比 relu 包含更多的操作。
这一事实导致了重计算/激活检查点的一些有趣结果。本质上，进行额外的重计算可能会导致更少的内存带宽，从而减少运行时间。因此，我们可以通过重计算来降低内存和运行时间，我们在 AOTAutograd 中利用它构建了一个巧妙的最小割优化过程。你可以在此处阅读更多相关内容（也可能在未来的博客文章中进行介绍！）。
如何判断内存带宽受限
在判断你的操作是否内存带宽受限时，计算器可以起到很大的作用。
对于简单的操作，我们可以直接计算内存带宽。例如，A100 拥有 1.5 terabytes/second 的全局内存带宽，可以执行 19.5 teraflops/second 的计算。因此，如果你使用 32 位浮点数（即 4 字节），那么你可以在 GPU 执行 20 万亿次操作的同时加载 4000 亿个数字。此外，为了执行简单的单目运算（例如将张量乘以 2），实际上需要将张量写回全局内存。
因此，除非你的单目运算执行了大约一百次操作，否则你会花费更多时间执行内存访问，而不是实际的计算。
借助 NVFuser 这样的融合编译器，我们可以很容易地测量内存带宽。你可以在此处查看 Colab 代码。
如果我们使用 PyTorch 函数，例如：
def f(x: Tensor[N]):
for _ in range(repeat):
x = x * 2
return x并使用融合编译器对其进行基准测试，那么我们可以计算出不同 repeat 值下达成的 FLOPS 和内存带宽。增加 repeat 是增加计算量而不增加内存访问次数的一种简单方法，这也被称为增加计算强度。
具体来说，假设我们对这段代码进行基准测试，并找到每秒执行的迭代次数。那么，作为 N（张量大小）的函数，我们将执行 2*N 次内存访问，以及 N * repeat 次 FLOP。因此，达成的内存带宽将为 bytes_per_elem * 2 * N * itrs_per_second，达成的 FLOPS 将为 N * repeat * itrs_per_second。
现在，让我们绘制运行时间、FLOPS 和达成的内存带宽作为计算强度的函数。注意，所有内容都以对数对数刻度显示。
首先，请注意，直到我们执行 64 次乘法运算，运行时间才明显增加。这意味着，在此之前，我们主要受限于内存带宽，我们的计算能力大部分处于闲置状态。
因此，我们最初达成的 FLOPS 只有 0.2 teraflops。随着我们使计算强度翻倍，这个数字线性增长，直到我们接近 9.75 teraflops 的峰值 [1]。一旦我们接近峰值 teraflops，我们就认为是“计算受限”。
最后，你可以看到，我们达成的内存带宽最初接近峰值，随着我们增加计算强度，它开始下降。这正是我们所期望的，因为我们花费了越来越多的时间执行实际计算，而不是访问内存。
在这种情况下，很容易看出我们何时是计算受限的，何时是内存带宽受限的。对于 repeat < 32，我们饱和了内存带宽，而我们的计算能力未得到充分利用。相反，一旦 repeat > 64，我们看到我们的计算能力已经饱和（即接近峰值 FLOPS），而我们利用的内存带宽开始下降。
对于更大的系统，我们通常很难判断是计算受限还是内存带宽受限，因为它们通常包含计算受限和内存带宽受限的组件。
衡量计算受限程度的一种常见方法是测量达成的 FLOPS 占峰值 FLOPS 的百分比。例如，如果你达成了峰值 FLOPS 的 80%，那么你就知道你至少是 80% 的计算受限，这已经相当不错了！你剩下的时间可能都花在了执行内存带宽操作上。[2]
然而，除了内存带宽成本之外，还有一件事可能会导致你的 GPU 无法“Brrr”起来。
开销：系统瓶颈
开销是指你的代码花费在传输张量或计算之外的所有时间。例如，在 Python 解释器中花费的时间？开销。在 PyTorch 框架中花费的时间？开销。启动 CUDA 内核（但没有执行它们）花费的时间？也是开销。
开销之所以是一个棘手的问题，主要是因为现代 GPU 非常快。A100 每秒可以执行 312 万亿次浮点运算（312 TeraFLOPS）。相比之下，Python 非常慢。在本地进行基准测试，Python 每秒可以执行 3200 万次加法运算。
这意味着，在 Python 执行一次 FLOP 的时间里，A100 可以执行 975 万次 FLOPS。
更糟糕的是，Python 解释器并不是开销的唯一来源，像 PyTorch 这样的框架在到达实际内核之前也有很多层调度。如果你使用 PyTorch 执行相同的实验，我们每秒只能执行 28 万次操作。当然，PyTorch 的设计目的不是处理微小的张量，但是，如果你使用的是微小的张量（例如在科学计算中），你可能会发现 PyTorch 比 C++ 慢得多。
例如，看一下 PyTorch 执行一次加法的火焰图配置文件。那个方框？那是执行实际计算的部分。其他所有部分都是纯粹的开销。
鉴于此，你可能会惊讶于为什么有人会使用 PyTorch，但请记住，现代深度学习模型通常会执行大规模运算。此外，像 PyTorch 这样的框架是异步执行的。也就是说，当 PyTorch 运行 CUDA 内核时，它可以继续运行并在其后面排队更多 CUDA 内核。因此，只要 PyTorch 可以“领先于”CUDA 内核，大部分框架开销就会完全隐藏起来！
如果我们的 GPU 操作足够大，那么我们的 CPU 可以领先于 GPU（因此 CPU 开销无关紧要）。另一方面，如果我们的 GPU 操作太小，那么我们的 GPU 将大部分时间都浪费在昂贵的“纸镇”上。
那么，如何判断你是否处于这种状态呢？由于开销通常不会随着问题规模的增加而增加（而计算和内存会增加），因此判断开销受限的最简单方法是增加数据的规模。如果运行时间没有按比例增加，那么你就是开销受限的。例如，如果你将批次大小增加一倍，但运行时间只增加了 10%，那么你很可能是开销受限的。[3]
另一种方法是使用 PyTorch 分析器。在这里，粉红色的线实际上显示了 CPU 内核与 GPU 内核的匹配情况。
GPU 上有很多间隙，因为它正在等待 CPU 开销
我们的 CPU 运行速度远远超过 GPU
另一个补充说明 - nvidia-smi 中的“GPU-Util”（不是“Volatile GPU-Util”）条目基本上是测量底部的行中实际运行 GPU 内核的百分比。因此，这也是评估开销的一个好方法。
开销存在的主要原因是像 PyTorch 这样的框架具有很大的灵活性。本质上，需要花费大量时间来“弄清楚要做什么”。
这可能来自 Python（查找属性或调度到正确的函数）或 PyTorch 中的代码（PyTorch 的所有调度器）。例如，当你执行 a + b 时，需要执行以下步骤。
Python 需要查找 a 上的 add 调度到什么。
PyTorch 需要确定张量的许多属性（例如数据类型、设备以及是否需要自动微分）以确定要调用哪个内核。
PyTorch 需要实际启动内核。
从根本上说，这种开销来自于在每一步都能够执行不同操作的灵活性。如果你不需要这种灵活性，那么解决这种灵活性的方法之一是将其追踪出来，例如使用 jit.trace、FX 或 jax.jit。或者，你也可以在更低级别使用像 CUDA Graphs 这样的工具来实现。
不幸的是，这样做会失去灵活性。我期待的一种能够兼顾两者的方法是编写类似于“真正的”JIT 的工具，通过在 VM 级别进行内省来实现。有关详细信息，请参阅 TorchDynamo。
总结
如果你想加速深度学习系统，最重要的是要了解模型的瓶颈是什么。瓶颈决定了加速系统的最佳方法。
我经常看到研究人员和其他对加速 PyTorch 代码感兴趣的人盲目地尝试各种方法，而没有了解模型所处的性能瓶颈。
| 性能瓶颈 | 可行的解决方案 |
|---|---|
| 开销受限 | 追踪、操作融合、不要使用 Python、真正的 JIT :^) |
| 带宽受限 | 操作融合 |
| 计算受限 | 使用 Tensor Cores、给英伟达更多钱 |
当然，可以说，用户需要考虑这些问题本身就反映了框架的失败。PyTorch 的编译器或分析 API 并不总是那么容易使用，尽管这是一个积极的关注领域。
无论如何，我发现理解系统的基本原理总是很有用，希望这对你也有帮助。
PS：如果你喜欢这篇文章，我将在 thonking.ai 上发布我未来的大部分文章。
致谢
感谢 Emily Shen、Qian Huang 和 EleutherAI 的成员阅读这篇博文的早期草稿并提供反馈。
BibTeX 引用
@article{he2022brrrrfromfirstprinciples,
author={Horace He},
title={Making Deep Learning Go Brrrr From First Principles},
year={2022},
url={https://horace.io/brrr_intro.html},
}[1] 这可能与规格说明书上显示的 19.5 teraflops 不同。原因是 GPU 拥有更多专门用于融合乘加 (FMA) 指令的硬件。因此，对于完全通用的计算，A100 实际上只能达到 9.75 teraflops。 ↩︎
[2] 现在在 PyTorch 中以一种很好的方式计算 FLOPS 的方法有很多，请参阅 https://dev-discuss.pytorch.org/t/the-ideal-pytorch-flop-counter-with-torch-dispatch/505 ↩︎