翻译和改编自Horace的博客
同样的一个算法模型代码,A同学实现的GPU利用率80%,而B同学实现的GPU利用率只有50%。
很多人希望在不改变模型效果的情况下去提升模型的性能,殊不知最快的提升方式就是在实现模型代码时按照第一性的原则进行编码。
从第一性原理出发,可以在编码时就消除大量的方法,从而使问题变的简单而更容易解决。同样的,如果你在构建模型之初就引入了导致性能瓶颈的代码,无论如何调节都不会带来较大的改变。
要了解深度学习的训练过程,首先可以将深度学习的训练效率拆解为三个不同的组成部分:
Compute:GPU 计算在实际浮点运算 (FLOPS) 上所花费的时间;
Memory:在 GPU 内传输Tensor上所花费的时间;
Overhead:其他一切;
就像训练调节 DL 模型的指标一样,实现高性能算法模型首要的工作就是了解模型所处的状态,方便缩小训练的优化调节范围。
例如,如果模型将所有时间都花在内存传输上(即处于内存带宽瓶颈状态),那么增加 GPU 的 FLOPS 将无济于事。
另一方面,如果模型将所有时间都花在执行大型的 chonky matmuls(即受计算限制的瓶颈)上,那么即使将模型逻辑重写为 C++ 来减少开销也将无济于事。
因此,如果您想让 GPU 保持正常高效运转,就需要讨论下系统可能花费时间的三个部分:Compute、memory bandwidth和overhead。
Compute
优化深度学习系统提升GPU利用率和运行效率,其根本是希望最大化GPU Compute的时间。
例如,你在平台上支付了312 teraflops算力费用,理想情况下,您会获得这 312 teraflops算力。
但是,事实上可能更多的时间花费在其他上面,为了让昂贵的矩阵乘法物有所值,你需要减少模型在其他部分花费的时间。
但为什么要关注最大化计算时间而不是内存带宽呢?
原因很简单 ,你可以减少开销或内存成本,但你(大多数情况下)无法在不更改正在执行的实际操作的情况下减少所需的计算。
想象一下,我们将Compute计算视为工厂。
我们向工厂发送指令(overhead),发送材料(memory bandwidth),所有这些都是为了保持工厂(Compute)的高效运行。

因此,如果我们工厂提高效率的速度快于我们为其供应材料的速度,我们工厂就很难达到最高效率。

如果我们的材料(带宽)跟不上,即使我们工厂的规模(FLOPS)翻倍,那么我们的性能也不会翻倍。
这就是为什么有时模型性能快了或者更好的GPU机器后,GPU利用率反而下降的原因。
现代机器学习加速器都拥有专门用于矩阵乘法的硬件,例如 Nvidia 的“Tensor Cores”。下面是其计算性能参数的图片:

从上面可以看出,如果你不进行矩阵乘法,则其只能实现 19.5 teraflops算力,而不是规定的 312算力。请注意,这并不是 GPU 所独有的,TPU 甚至比GPU更是这样。
GPU 在除矩阵乘法之外的所有操作上都慢得多。这个事实一开始可能看起来有问题。
例如,那么我们的其他运算符(如layer norm or activation)呢?
事实上,这些运算符只是 FLOPS 方面的舍入误差。
例如,让我们看看BERT论文中不同运算符类型的 FLOP 计数表,其中“Tensor contraction”= matmuls就是矩阵乘运算。

从表格中可以看出,我们的非矩阵乘操作总共只占到了FLOPS的0.2%,所以GPU计算非矩阵乘操作慢15倍并不重要。
但是,在这种情况下,归一化和逐点操作的FLOPS实际上分别比我们的矩阵相乘少了250倍和700倍。
那么为什么我们的非矩阵乘操作花费的时间比应有的多呢?
回到我们的比喻,罪魁祸首通常是物料在工厂之间运输所需的时间。换句话说,就是内存带宽。
Memory Bandwidth
带宽成本本质上是将数据从一个地方移动到另一个地方所支付的成本。
这可能会将数据从 CPU 移动到 GPU、从一个节点移动到另一个节点,甚至从 CUDA 全局内存移动到 CUDA 共享内存。
我们把最后一项称之为”内存带宽成本“。
要了解内存带宽成本是多少,让我们回到工厂的类比。
虽然我们的工厂是我们进行实际工作的地方,但它不适合作为散装存储单元。其中很大一部分原因是,由于我们在这里进行实际工作,因此所有存储都经过优化,以便能够快速实际使用(SRAM),而不是拥有大量存储。
那么,我们将实际结果和材料存储在哪里呢?典型的方法是建立一个仓库,可能是在土地便宜且有大量空间(DRAM)的地方。然后,我们可以将物资运送到我们的工厂或从我们的工厂运送物资(内存带宽)。

将数据移入和移出计算单元的成本就是所谓的“内存带宽”成本,通过nvidia-smi就可以看到GPU卡所占用的DRAM存储。
需要注意的一件事是,每次执行 GPU 内核时,我们都需要将数据从 GPU 的 DRAM(即我们的仓库)移出或移回。
当我们执行一个最简单torch.cos(…)模型命令时,我们需要将数据从存储发送到仓库,然后对每条数据执行少量计算,然后将该存储发送回。运送东西的费用相当昂贵。
因此,我们几乎所有的时间都花在了传输数据上,而不是实际的计算本身。
由于我们将所有时间都花在内存带宽上,因此此类操作称为内存限制操作,这意味着我们不会在计算上花费大量时间。
好的,所以这并不理想。我们对于它可以做些什么呢?让我们看一下运算符序列的外观。
这是一个非常愚蠢的安排。为什么我们一遍又一遍地将相同的数据发送到全局内存,然后返回到计算单元?我们应该将数据保留在工厂,执行所有计算,然后将其发回!

这就是算子融合——深度学习编译器中最重要的优化。
简而言之,我们不是将数据写入全局内存以便再次读取,而是通过一次执行多项计算来消除额外的内存访问。
例如,如果我们执行x.cos().cos(),通常我们需要执行 4 次全局读写。
x1 = x.cos() # Read from x in global memory, write to x1
x2 = x1.cos() # Read from x1 in global memory, write to x2
但是,通过运算符融合,我们只需要 2 次全局内存读取和写入!因此算子融合会将速度提高 2 倍。
x2 = x.cos().cos() # Read from x in global memory, write to x2
然而使用自动的算子融合依然存在些前提。
首先,GPU需要知道执行当前操作时接下来会发生什么。因此,您无法在eager模式下进行这种优化。在eager这种模式下,PyTorch 一次运行一个运算符。
其次,我们实际上需要为此生成 CUDA 代码,这会带来全新的麻烦。
如果您对编写自定义 CUDA 内核感兴趣,那么这可能是您受益最大的地方。
任何 2 个 PyTorch 运算符都提供了融合的机会,从而节省了它们之间读取/写入全局内存的内存带宽成本。
此外,许多现有编译器通常可以执行“简单”融合 ,例如NVFuser 和 XLA 就是两个例子。
然而,自动化系统无法与人类的聪明才智相媲美,因此,如果您想尝试自己编写一些自定义 CUDA 内核,Triton是一个很好的起点。
最后,算子融合会导致一些令人惊讶的结果。
对于一个合并的 x.cos().cos(),与单独调用 x.cos() 几乎花费相同的时间。这就是为什么激活函数几乎都具有相同的成本,尽管 gelu 明显包含了比 relu 更多的操作。
这个事实为重材料化/激活检查点提供了一些有趣的结果。基本上,做额外的重新计算可能会导致更少的内存带宽,从而减少运行时间。
因此,我们可以通过重建材料来降低内存和整体运行时间。
如何估算内存带宽成本
在考虑操作是否受内存带宽限制时,如何进行计算呢?
对于简单的算子,可以直接计算内存带宽。例如,A100 全局内存带宽为 1.5 TB/s,可以执行 19.5 TFlops/s 的计算。
因此,如果你使用 32 位浮点数(即 4 个字节),则在 GPU 执行 2000 亿次操作的同时,你可以在相同的时间内加载 4000 亿个数字。
此外,要执行一个简单的一元运算符(如将张量乘以 2),我们实际上需要将张量写回全局内存。
因此,在你的一元运算中进行约一百个操作之前,你花费的时间执行内存访问将比实际计算时间更长。
借助像 NVFuser 这样的融合编译器,其实非常容易自行测量!你可以在这里查看在Colab中的代码。
如果你拿一个 PyTorch 函数:
def f(x: Tensor[N]):
for _ in range(repeat):
x = x * 2
return x
使用融合编译器对其进行基准测试后,我们可以计算不同重复值下的 FLOPS 和内存带宽。增加重复是一种增加计算量而不增加内存访问量的简单方法,这也被称为增加计算密度。
具体来说,假设我们对这段代码进行基准测试并找到每秒执行的迭代次数。然后,作为 N(张量的大小)的一个函数,我们将执行 2*N 次内存访问和 N * repeat 次 FLOP。
因此,实现的内存带宽将是 bytes_per_elem * 2 * N * itrs_per_second,实现的 FLOPS 将是 N * repeat * itrs_per_second。
现在,让我们将运行时间、FLOPS 和实现的内存带宽作为计算密度的函数进行绘制。请注意,所有内容都在对数-对数比例上。

首先,注意到在执行 64 次乘法之前,运行时间几乎没有明显增加。这意味着在那之前,我们主要受到内存带宽的限制,计算资源大部分处于空闲状态。
其次,一开始我们只能达到0.2 Teraflops的低效率。随着计算强度的提高,这个数字呈线性增长,直到接近我们的峰值 9.75 Teraflops。一旦接近峰值,我们就被认为是"计算受限"。
最后,你可以看到我们实现的内存带宽一开始接近峰值,随着计算强度的增加开始下降。这正是我们所预期的,因为我们花费越来越多的时间进行实际的计算,而不是访问内存。
在这种情况下,很容易看出何时受到计算限制和何时受到内存带宽限制。对于 repeat < 32,我们饱和了内存带宽,但计算资源利用率不高。
相反,一旦 repeat > 64,我们发现我们正在饱和计算资源(即接近峰值 FLOPS),而利用的内存带宽开始下降。
对于较大的系统,通常更难确定是计算受限还是内存带宽受限,因为它们包含了计算受限和内存带宽受限的组件混合。
衡量计算受限程度的常见方法是将实现的 FLOPS 与峰值 FLOPS 的百分比进行对比。
例如,如果你的实际 FLOPS 达到峰值 FLOPS 的 80%,那么你就知道至少有 80% 的计算资源被利用,这是相当不错的!其余的时间可能主要用于内存带宽操作。
然而,除了内存带宽成本之外,还有一件事可能导致 GPU 的性能下降。
Overhead
Overhead是指你的代码在传输Tensor和计算之外的事情上所花费的时间。
例如,在Python解释器中花费的时间,在 PyTorch 框架上花费的时间,启动 CUDA 内核(但不执行它们)花费的时间等…。
Overhead如此严重的问题的主要原因是现代 GPU 的速度非常快。A100每秒可执行 312万亿次浮点运算 (312 TeraFLOPS)。相比之下,Python确实是非非非常慢。
在本地进行基准测试,Python 可以在一秒钟内执行 3200 万次加法。
这意味着当 Python 可以执行单次FLOPS 时,A100 可以执行975 万次 FLOPS。
更糟糕的是,Python 解释器甚至不是唯一的开销来源 - 像 PyTorch 这样的框架在进入实际内核之前也有许多层的调度。如果用 PyTorch 进行同样的实验,我们每秒只能得到 28 万次操作。
当然,微小张量不是 PyTorch 的用途,但是……如果您使用微小张量(例如在科学计算中),您可能会发现 PyTorch 与 C++ 相比慢得令人难以置信。
例如,查看 PyTorch 执行单次加法的火焰图配置文件。那个盒子在吗?这就是执行实际计算的内容。其他一切都是纯粹的开销。

鉴于此,您可能会对居然有人使用 PyTorch 感到震惊,但请记住,现代深度学习模型通常会执行大量操作。
此外,像 PyTorch 这样的框架是异步执行的。
也就是说,当 PyTorch 运行 CUDA 内核时,它可以继续并在其后面排队更多 CUDA 内核。因此,只要 PyTorch 能够“领先于”CUDA 内核,大部分框架开销就完全隐藏了!

那么,如何确定你处于这种情况下呢?
由于开销通常不随问题规模增加而增加(而计算和内存则会增加),最简单的方法是增加数据的大小。
如果运行时间没有与问题规模成比例增加,那么你就是受到开销限制了。例如,如果你将批处理大小加倍,但运行时间只增加了10%,你很可能是受到开销限制。
另一种方法是使用timeline分析器。在这里,粉色线条实际上展示了CPU内核与GPU内核的匹配情况。

另外,nvidia-smi中的"GPU-Util"(而不是"Volatile GPU-Util")条目基本上是衡量底部行实际运行GPU内核的百分比。这也是另一种很好的估算开销的方法。
存在这种开销的主要原因是PyTorch等框架具有的所有灵活性。基本上,需要花费大量时间来"弄清楚该做什么"。
这可能来自Python(查找属性或分派到正确的函数)或PyTorch中的代码(PyTorch的调度器)。例如,当你执行 a + b 时,需要执行以下步骤。
Python 需要查找 a 上的
__add__分派到什么地方。PyTorch 需要确定张量的许多属性(例如dtype、device和是否需要自动微分)以确定调用哪个内核。
PyTorch 需要实际启动内核。
从根本上讲,这种开销来自于能够在每个步骤上做不同事情的灵活性。
如果你不需要这种灵活性,解决这种灵活性的一种方法是通过跟踪它,比如使用jit.trace、FX或jax.jit。或者,你可以通过类似CUDA图的方式在更低的级别上执行。
不幸的是,这样做的代价是失去灵活性。我对一种方法感到兴奋,它可以让我们兼得两者优点,那就是在VM级别进行内省,编写更接近"真实"JIT的内容。可以参考TorchDynamo项目。
总结
如果你想加速你的深度学习系统,最重要的是要了解你模型中的瓶颈在哪里,因为这个瓶颈决定了加速系统的方法。
通常,我看到很多人在缺乏对模型算法所处状况的的了解下,盲目尝试各种方法来进行性能调节。下面总结简单的解决思路。

当然,可以说,用户需要考虑这些东西根本反映了框架的失败。然而,很多系统设计在设计之初就并非让”阿猫阿狗“写的代码也能高效运行。
总之,高效的算法代码是一切的前提 - 希望这对您也有用。