这种在计算单元之间移动东西的成本就是所谓的「内存带宽」成本。事实上,nvidia-smi 命令中出现的那个「内存」就是 DRAM,而经常让人抓狂的「CUDA out of memory」说的就是这个 DRAM。
值得注意的是:我们每执行一次 GPU 核运算都需要把数据运出和运回到我们的仓库 ——DRAM。
现在想象一下,当我们执行一个一元运算(如 torch.cos)的时候,我们需要把数据从仓库(DRAM)运送到工厂(SRAM),然后在工厂中执行一小步计算,之后再把结果运送回仓库。运输是相当耗时的,这种情况下,我们几乎把所有的时间都花在了运输数据,而不是真正的计算上。
因为我们正把所有的时间都花费在内存带宽上,这种运算也被称作内存限制运算(memory-bound operation),它意味着我们没有把大量时间花费在计算上。
显然,这并不是我们想要的。那我们能做什么呢?让我们来看看算子序列长什么样子。
一个逐点算子序列可能的样子。
在全局内存和计算单元之间来回传输数据的做法显然不是最佳的。一种更优的方式是:在数据工厂中一次性执行完全部运算再把数据传回。
这就是算子融合(operator fusion)—— 深度学习编译器中最重要的优化。简单地说,这种方法不会为了再次读取而将数据写入全局内存,而是通过一次执行多个计算来避免额外的内存访问。
例如,执行 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 需要知道执行完当前运算后下一步会发生什么,因此无法在 PyTorch 的 Eager 模式(一次运行一个运算符)下进行此优化。其次,我们需要编写 CUDA 代码,这也不是一件简单的事。
并不是所有的算子融合都像逐点算子那样简单。你可以将逐点算子融合到归约(reduction)或矩阵乘法上。甚至矩阵乘法本身也可以被认为是一种融合了广播乘法(broadcasting multiply)和归约的运算。
任何 2 个 PyTorch 算子都可以被融合,从而节省了读取 / 写入全局内存的内存带宽成本。此外,许多现有编译器通常可以执行「简单」的融合(例如 NVFuser 和 XLA)。然而,更复杂的融合仍然需要人们手动编写,因此如果你想尝试自己编写自定义 CUDA 内核,Triton 是一个很好的起点。
令人惊讶的是,融合后的 x.cos ().cos () 运算将花费几乎与单独调用 x.cos () 相同的时间。这就是为什么激活函数的成本几乎是一样的,尽管 gelu 显然比 relu 包含更多的运算。
因此,重新实现 / 激活检查点会产生一些有趣的结果。从本质上讲,进行额外的重新计算可能会导致更少的内存带宽,从而减少运行时间。因此,我们可以通过重新实现来减少内存占用和运行时间,并在 AOTAutograd 中构建一个简洁的 min-cut 优化通道。
推理内存带宽成本
对于简单的运算,直接推理内存带宽是可行的。例如,A100 具有 1.5 TB / 秒的全局内存带宽,可以执行 19.5 teraflops / 秒的计算。因此,如果使用 32 位浮点数(即 4 字节),你可以在 GPU 执行 20 万亿次运算的同时加载 4000 亿个数字。
此外,执行简单的一元运算(例如将张量 x2)实际上需要将张量写回全局内存。
因此直到执行大约一百个一元运算之前,更多的时间是花在了内存访问而不是实际计算上。
如果你执行下面这个 PyTorch 函数:
def f(x: Tensor[N]):
for _ in range(repeat):
x = x * 2
return x
并使用融合编译器对其进行基准测试,就可以计算每个 repeat 值的 FLOPS 和内存带宽。增大 repeat 值是在不增加内存访问的情况下增加计算量的简单方法 - 这也称为增加计算强度 (compute intensity)。
具体来说,假设我们对这段代码进行基准测试,首先要找出每秒执行的迭代次数;然后执行 2N(N 是张量大小)次内存访问和 N *repeat FLOP。因此,内存带宽将是 bytes_per_elem * 2 * N /itrs_per_second,而 FLOPS 是 N * repeat /itrs_per_second。
现在,让我们绘制计算强度的 3 个函数图象:运行时间、flops 和内存带宽。
请注意,在执行 64 次乘法之前,运行时间根本不会显著增加。这意味着在此之前主要受内存带宽的限制,而计算大多处于空闲状态。
一开始 FLOPS 的值是 0.2 teraflops。当我们将计算强度加倍时,这个数字会线性增长,直到接近 9.75 teraflops 的峰值,一旦接近峰值 teraflops 就被认为是「计算受限的」。
最后,可以看到内存带宽从峰值附近开始,随着我们增加计算强度开始下降。这正是我们所期待的,因为这说明执行实际计算的时间越来越多,而不是访问内存。
在这种情况下,很容易看出何时受计算限制以及何时受内存限制。repeat< 32 时,内存带宽接近饱和,而未进行充分的计算;repeat> 64 时,计算接近饱和(即接近峰值 FLOPS),而内存带宽开始下降。
对于较大的系统,通常很难说是受计算限制还是内存带宽限制,因为它们通常包含计算限制和内存限制两方面的综合原因。衡量计算受限程度的一种常用方法是计算实际 FLOPS 与峰值 FLOPS 的百分比。
然而,除了内存带宽成本之外,还有一件事可能会导致 GPU 无法丝滑运行。
额外开销
当代码把时间花费在传输张量或计算之外的其他事情上时,额外开销(overhead)就产生了,例如在 Python 解释器中花费的时间、在 PyTorch 框架上花费的时间、启动 CUDA 内核(但不执行)所花费的时间, 这些都是间接开销。
额外开销显得重要的原因是现代 GPU 的运算速度非常快。A100 每秒可以执行 312 万亿次浮点运算(312TeraFLOPS)。相比之下 Python 实在是太慢了 ——Python 在一秒内约执行 3200 万次加法。
这意味着 Python 执行单次 FLOP 的时间,A100 可能已经运行了 975 万次 FLOPS。
更糟糕的是,Python 解释器甚至不是唯一的间接开销来源,像 PyTorch 这样的框架到达 actual kernel 之前也有很多层调度。PyTorch 每秒大约能执行 28 万次运算。如果使用微型张量(例如用于科学计算),你可能会发现 PyTorch 与 C 相比非常慢。
例如在下图中,使用 PyTorch 执行单次添加,仅有一小块图是实际执行计算的内容,其他的部分都是纯粹的额外开销。