从「根」上找出模型的瓶颈！从第一原理出发剖析深度学习

如果想提升模型的性能，你的第一直觉是问搜索引擎吗？

通常情况下你得到的建议只能是一些技巧性的操作，比如使用in-place operation，把梯度设置为None，或者是把PyTorch版本从1.10.1退回到稳定版1.10.0等等。

这些临时找到的骚操作虽然可以一时地解决当下问题，但要是用了以后性能还没提升到满意的程度，那可能就有点「抓瞎」了。

虽然深度学习本身就是一个积木类的黑盒模型，但这种调试方法仿佛深度学习真的变成了炼丹术，而非科学。

比如你的模型在训练集上的loss远远低于测试时的loss，说明模型已经「过拟合」了，如果这个时候再盲目增大模型的参数量，那就纯粹是浪费时间了。再比如模型的训练loss和验证loss一样的时候，如果再对模型加入正则化，那也是浪费时间。

所以为了让AI从业者在遇到问题之后，能从根上解决，最近康奈尔大学人工智能（CUAI）的一位联合创始人Horace He发表了一篇博客，把深度学习模型的时间损耗拆分成三部分：计算、内存和其他开销overhead，从「第一原理」出发来了解和改进深度学习模型。

其中计算（Compute）指的是GPU在计算浮点操作时所消耗的时间，也就是FLOPS；内存（Memory）指的是把tensors写到GPU里消耗的时间。

如果模型把大部分的时间都花在了内存传输上，那么增加GPU的FLOPS是没有用的。又或者如果你把所有的时间都花在执行大块的数学运算上，那么把你的模型逻辑改写成C++来减少开销也没有用。

了解你所处的状态可以让你缩小优化的范围，节省下来的时间就可以愉快地摸鱼了。

通常深度学习模型运算速度不够快的原因都是显卡性能不够，加卡解千愁啊！

但现实很骨感，越强的卡，价格也更美丽。所以为了钱花的更值，需要尽可能地提升显卡的运行效率，不断地让显卡进行矩阵运行。

并且计算比内存带宽更重要的原因还有一个，就是模型训练过程中所需的计算量不管通过何种手段，基本都不会降低，所以最大限度提升计算能力才能提升效率。

但计算量如果增长速度过快，也会加剧最大化计算利用率的难度。就拿这个关于CPU FLOPS翻倍时间与内存带宽翻倍时间的表格来说。

一种思考计算的方式是把CPU当作一个工厂。用户向工厂发送指令（开销）和原材料（内存带宽），所有这些都是为了保持工厂高效运行（计算）。

如果工厂提高效率的速度超过了为其提供原材料的速度，那么工厂就更难达到其峰值效率。即使工厂的规模（FLOPS）增加了一倍，如果带宽不能同步提升，那性能也不会增加一倍。

关于FLOPS还有一个补充。现代机器学习加速硬件都有专门用于矩阵乘法的硬件，比如Nvidia的Tensor Cores。

也就是说，如果你不做矩阵乘法，你就只能获得19.5 teraflops，而非宣传的312。并且这并非是GPU所独有的缺陷，TPU甚至比GPU更不通用。

事实上，GPU在所有非矩阵乘法的操作上都很慢，乍一看可能影响很大，但实际上神经网络模型里基本都是矩阵乘法。

在一篇关于BERT模型的flop研究中可以发现，BERT中99.8%都是矩阵乘法（Tensor Contraction）操作，所以虽然非矩阵乘法的速度要慢15倍，但也无伤大雅。

但在这种情况下，归一化和点式运算实际上比矩阵乘法运算少了250倍的FLOPS和700倍的FLOPS。

至于为什么非矩阵乘法的理论性能和现实相差这么多，研究人员给出的答案是：内存带宽（memory bandwidth）。

带宽成本本质上是将数据从一个地方移动到另一个地方所支付的成本，包括将数据从CPU转移到GPU，从一个节点转移到另一个节点，二者通常称为「数据传输成本」和「网络成本」。

深度学习模型优化关注的带宽成本主要是从CUDA全局内存转移到CUDA共享内存。

回到工厂那个例子，虽然工厂可以完成一些计算任务，但它并不是一个适合存储大量数据的地方。典型的做法是利用更便宜的硬件来建立一个数据仓库（DRAM），然后在仓库和工厂之间运送物资，也就是内存带宽。

GPU的DRAM大小可以通过nvidia-smi命令获得，仓库容量不够也是导致CUDA Out of Memory错误的主要原因。

需要注意的是，每次执行GPU内核时，都需要将数据从GPU的DRAM移出和移回。

现在我们就知道执行torch.cos这样的单个操作时，几乎每做一次这样的简单运算，数据都需要从内存运到GPU里，运送成本比计算成本要高很多，所以时间几乎都花在内存上了，这种情况也称为memory-bound operation。

错误的做法就是每次都把数据送到GPU计算后返回结果，再把结果送给GPU再次计算，可以看到，大量的时间都耗费在数据传输上了。

稍作调整之后，当预先把指令都放入计算时，内存的传输降为一次即可完成相同的任务。

如果换成pyTorch的代码就是把两行代码转为一行x.cos().cos()，效率能够提升两倍。

不过这种优化措施并不是在所有场景下都适用。因为GPU预先需要知道所有执行的指令，并生成CUDA代码，所以无法在eager-mode下使用。而且并非所有的运算符融合都像pointwise操作符这么简单。

如果你曾经写过CUDA内核代码的话，就可以知道任何两个PyTorch都有机会进行融合来节省全局内存的读写成本。现有的编译器如NVFuser和XLA通常只能进行一些简单的融合，肯定比不上AI工程师的设计。如果你想尝试自己编写一些定制的CUDA内核，Triton就比较适合新手入门。

运算符融合的效果就是更多的操作，时间成本相同，这也是为什么激活函数的计算成本几乎都是一样的，尽管gelu显然比relu多了很多操作。

当需要推理你的操作是否有内存带宽限制时，calculator可以发挥很大的作用。

对于简单的算子来说，可以直接推理内存带宽。例如，A100有1.5T字节/秒的全局内存带宽，可以进行19.5T FLOPS的计算。因此，如果你使用32位浮点（即4个字节），GPU可以执行20万亿次操作的相同时间内加载4000亿个数字。此外，为了执行一个简单的单项运算（如把一个tensor乘2），实际上需要将tensor写回全局内存。所以将单项运算做了大约一百次以后，才能够等到内存数据送进来。

在像NVFuser这样的融合编译器的帮助下，实际上可以很容易地测量成本。

以一个PyTorch函数为例，并用融合编译器对其进行基准测试，然后就可以计算出不同的重复值所达到的FLOPS和内存带宽。

增加重复次数是在不增加内存访问的情况下增加计算量的一个简单方法，这也被称为增加计算强度。

因为tensor的大小为N，需要将执行2*N次内存访问，以及N*repeat FLOP。因此，实现的内存带宽将是byte_per_elem * 2 * N / itrs_per_second，而实现的FLOPS将是N * repeat / itrs_per_second。

把运行时间、flops和实现的内存带宽取对数后绘制出来的结果可以看到，执行64次乘法之前，运行时间并没有明显的增加。这也意味着，在这之前，内存带宽是有限的，计算大部分是闲置的。

因此，一开始只实现了0.2 teraflops。当我们把计算强度提高一倍时，这个数字就会线性增长，直到我们接近9.75 teraflops的峰值，也就是「计算极限」。

内存带宽开始时接近峰值，随着计算强度的增加，开始下降。这也符合预期，因为实际上更多的时间花在了实际的计算上，而非访问内存。

在这种情况下可以很容易看到什么时候是计算约束，什么时候是内存约束。

对于重复次数小于32次时，内存带宽已经饱和，而计算能力却没有得到充分利用。相反，一旦重复大于64次，会发现计算量已经饱和（即达到接近峰值FLOPS），而内存带宽利用率开始下降。

对于更大的系统，通常很难说是计算约束还是内存带宽约束，因为可能同时包含了计算约束和内存约束。

衡量计算约束程度的一个常见方法是，将你实现的FLOPS作为峰值FLOPS的一个百分比作为指标。如果实现了峰值FLOPS的80%，那就说明计算资源利用的比较充分，其余的时间可能是花在内存带宽上了。

代码中没有花在传输或计算tensor的时间都称为开销（overhead），比如花在Python解释器上的时间，花在PyTorch框架上的时间，花在启动CUDA内核（但不执行）的时间都是开销。

开销之所以成为一个问题，主要原因是现代GPU的速度非常快。一个A100可以每秒进行312万亿次的浮点运算（312 TeraFLOPS）。相比之下，Python的运行速度就相当慢了，一秒钟内只能进行3200万次加法运算。

这也意味着，在Python可以执行一个FLOP的时间里，A100可以运行975万FLOPS。

像PyTorch这样的框架在进入实际内核之前也有很多层调度。如果你用PyTorch做同样的实验，每秒只能得到28万次操作。当然，执行小tensor并不是建立PyTorch的目的，但是如果确实在科学计算中使用小tensor，你就会发现PyTorch与C++相比慢得惊人。

一个更直观的图可以看到，PyTorch执行一个加法时产生的配置文件，除了一个小方块外，其余所有的都是纯开销。

现代深度学习模型通常都在进行大规模的计算操作，并且像PyTorch这样的框架是异步执行的。也就是说，当PyTorch正在运行一个CUDA内核时，它可以继续运行并在后面排起更多的CUDA内核。因此，只要PyTorch能够「提前」运行CUDA内核，大部分的框架开销就会被完全隐藏起来

由于开销通常不随问题的大小而变化（计算和内存则成比例增加），一个简单的判断方法是你的batch size规模增加一倍，但运行时间只增加了10%（预期是增加一倍的运行时间），那么就很可能是开销过大了。

另一种方法是使用PyTorch profiler。粉色线条显示了CPU内核与GPU内核的匹配情况。当GPU在等待CPU的开销时，就有很多空隙。

CPU比GPU运行得更快时空隙就少很多。

nvidia-smi中的GPU-Util就是在测量实际运行GPU内核的百分比，这也是一种衡量开销的好方法。

开销大部分都来自PyTorch等框架的灵活性，需要花费大量时间来「弄清该做什么」

比如当执行a+b时，需要三个步骤：

1. Python 需要查找 __add__ 在 a 上派发的内容

2. PyTorch需要确定张量的许多属性（如dtype、device以及是否需要Augrad）以确定调用哪个内核

3. PyTorch需要实际启动内核

每步都需要灵活性来支持不同操作，解决灵活性的一个方法是追踪，比如用jit.tract, FX或jax.jit，或者用CUDA Graphs在更低的层次实现。

提升模型效率，最重要的就是了解模型的性能瓶颈。

当然了，编写一个神经网络模型还需要考虑这么多开销问题，也可以说是这些系统、框架设计上的失败，因为这些本来应该是对用户透明的。

但懂得这些基本原理肯定是有意义的，可以帮助你从「根」上解决性能瓶颈。