如果你想提高模型的性能,你的第一本能是去问搜索引擎吗?通常你得到的建议只能是一些取巧的操作,比如使用原地操作,将梯度设置为None,或者将PyTorch版本从1.10.1恢复到稳定版1.10.0等等。虽然这些暂时觉得风骚操作可以暂时解决目前的问题,如果使用后性能没有提升到令人满意的程度,可能就有点“抓瞎了”。虽然深度学习本身就是一种积木式的黑盒模型,但这种调试方式似乎真的变成了炼金术而不是科学。比如你的模型在训练集上的损失远低于测试时的损失,说明模型已经“过拟合”了。如果此时一味地增加模型的参数个数,那是浪费时间。再比如,当模型的训练损失和验证损失相同时,给模型加上正则化是浪费时间。所以为了让AI从业者从根源上解决问题,康奈尔大学人工智能(CUAI)的联合创始人HoraceHe最近发表了一篇博客,将深度学习模型的时间损失拆分为三个部分:计算、内存等开销,从“第一性原理”出发,理解和改进深度学习模型。其中,Compute是指GPU在计算浮点运算时消耗的时间,即FLOPS;内存是指将张量写入GPU所消耗的时间。如果模型大部分时间都花在内存传输上,那么增加GPU的FLOPS是没有用的。或者,如果您将所有时间都花在执行大量数学运算上,那么用C++重写模型逻辑来减少开销是没有用的。了解自己所处的状态可以让你缩小优化的范围,省下来的时间可以愉快的钓鱼。计算通常,深度学习模型的计算速度不够快的原因是显卡的性能不够。加个卡就能解决问题!但现实很骨感,卡越强,价格越漂亮。所以,要想获得更多的性价比,就要尽可能的提高显卡的运行效率,不断让显卡进行矩阵运算。而计算比内存带宽更重要的还有一个原因,那就是模型训练过程中需要的计算量无论用什么手段都基本不会减少,所以最大化计算能力可以提高效率。但是,如果计算量增长过快,也会增加计算利用率最大化的难度。拿这个时间表来加倍CPUFLOPS与加倍内存带宽的时间。思考计算的一种方法是将CPU视为工厂。用户向工厂发送指令(开销)和原材料(内存带宽),所有这些都是为了保持工厂高效运行(计算)。如果一家工厂提高效率的速度快于为其提供原材料的速度,那么工厂将更难达到其最高效率。即使工厂规模(FLOPS)翻倍,如果不能同时增加带宽,性能也不会翻倍。关于FLOPS的补充。现代机器学习加速硬件都有专门用于矩阵乘法的硬件,比如Nvidia的TensorCores。也就是说,如果你不做矩阵乘法,你只会得到19.5teraflops,而不是标榜的312。而且这不是GPU独有的缺陷,TPU的通用性甚至不如GPU。事实上,GPU在所有非矩阵乘法运算中都非常慢。乍一看可能影响很大,但实际上神经网络模型基本上就是矩阵乘法。在对BERT模型的flop研究中,可以发现99.8%的BERT都是矩阵乘法(TensorContraction)运算,所以非矩阵乘法虽然慢15倍,但无害。但在这种情况下,归一化和逐点运算实际上比矩阵乘法运算少了250倍和700倍的FLOPS。至于为什么非矩阵乘法的理论性能与实际相差如此之大,研究人员给出的答案是:内存带宽。内存带宽成本本质上是将数据从一个地方移动到另一个地方所支付的成本,包括将数据从CPU移动到GPU以及从一个节点移动到另一个节点,这两者通常被称为“数据传输成本”和“网络成本”。深度学习模型优化的带宽成本主要从CUDA全局内存转移到CUDA共享内存。回到工厂的例子,虽然工厂可以完成一些计算任务,但是它并不是一个适合存放大量数据的地方。典型的做法是使用更便宜的硬件来构建数据仓库(DRAM),然后在仓库和工厂之间传输物料,即内存带宽。GPU的DRAM大小可以通过nvidia-smi命令获取。仓库容量不足也是CUDAOutofMemory错误的主要原因。重要的是要注意,每次执行GPU内核时,都需要将数据从GPU的DRAM移出和移回。现在我们知道,在执行torch.cos等单个操作时,几乎每次执行这种简单的操作时,都需要将数据从内存传输到GPU。运输成本远高于计算成本,所以几乎所有的时间都花在了内存上。是的,这种情况也称为内存绑定操作。错误的做法是每次都将数据发送给GPU计算并返回结果,然后再将结果发送给GPU进行计算。可以看出,大量的时间花在了数据传输上。稍微调整一下,当指令被预加载到计算中时,内存传输减少到一个以完成相同的任务。如果把pyTorch的代码改成一行x.cos().cos(),效率可以翻倍。但是,这种优化措施并不适用于所有场景。因为GPU需要提前知道所有执行的指令并生成CUDA代码,所以不能用eager-mode。并非所有运算符融合都像逐点运算符一样简单。如果你写过CUDA内核代码,就可以知道任意两个PyTorch都有机会进行融合,从而节省读写全局内存的开销。现有的NVFuser、XLA等编译器通常只能做一些简单的融合,肯定不如AI工程师的设计。如果你想尝试自己编写一些自定义的CUDA内核,Triton更适合初学者。operatorfusion的效果是操作更多,时间成本是一样的,这也是为什么激活函数的计算成本几乎一样的原因,虽然gelu明显比relu多了很多操作。当您需要推断您的操作是否受内存带宽限制时,计算器会非常有用。对于简单的算子,可以直接推断内存带宽。例如A100的全局内存带宽为1.5Tbytes/s,可以进行19.5TFLOPS的计算。因此,如果您使用32位浮点(即4个字节),GPU可以在执行20万亿次操作所需的相同时间内加载4000亿个数字。此外,为了执行简单的单向操作(例如将张量乘以2),实际上需要将张量写回全局内存。因此,在执行了大约一百次的单个操作之后,可能会等到内存数据发送进来。借助NVFuser这样的融合编译器,其实很容易衡量成本。以PyTorch函数为例,并使用FusionCompiler对其进行基准测试,然后可以计算出不同重复值所实现的FLOPS和内存带宽。增加重复次数是一种在不增加内存访问的情况下增加计算量的简单方法,这也称为增加计算强度。因为张量的大小为N,需要执行2*N次内存访问和N*repeatFLOPs。因此,实现的内存带宽将为byte_per_elem*2*N/itrs_per_second,实现的FLOPS将为N*repeat/itrs_per_second。绘制运行时间、触发器和实现的内存带宽的对数表明,在执行64次乘法之前,运行时间没有显着增加。这也意味着,在此之前,内存带宽是有限的,计算大多处于空闲状态。因此,一开始只实现了0.2teraflops。当我们将计算强度加倍时,这个数字会线性增长,直到我们接近9.75teraflops的峰值,即“计算极限”。内存带宽开始接近其峰值,并随着计算强度的增加而开始下降。这也是意料之中的,因为更多的时间实际上花在了实际计算上,而不是访问内存上。在这种情况下,很容易看出它何时受计算限制以及何时受内存限制。对于少于32次的重复,内存带宽饱和,计算能力未得到充分利用。相反,一旦重复次数大于64次,就会发现计算已经饱和(即接近峰值FLOPS),内存带宽利用率开始下降。对于较大的系统,通常很难区分是计算约束还是内存带宽约束,因为可能同时涉及计算和内存约束。衡量计算受限程度的一种常用方法是查看您实现的FLOPS占峰值FLOPS的百分比作为指标。如果达到了峰值FLOPS的80%,就意味着计算资源得到了充分利用,剩下的时间可能会花在内存带宽上。其他没有花在传输或计算张量上的开销代码所花的时间称为开销,例如花在Python解释器上的时间,花在PyTorch框架上的时间,花在启动CUDA内核上的时间(但不是执行它)时间是一种开销。开销成为问题的主要原因是现代GPU非常快。A100每秒可以执行312万亿次浮点运算(312TeraFLOPS)。相比之下,Python相当慢,每秒仅执行3200万次加法。这也意味着,A100在Python执行一次FLOP的时间内可以运行975万次FLOPS。像PyTorch这样的框架在进入实际内核之前也有多层调度。如果你用PyTorch做同样的实验,你每秒只能得到280,000次操作。当然,执行小张量并不是PyTorch的初衷,但如果你确实在科学计算中使用小张量,你会发现PyTorch与C++相比慢得惊人。更直观的一张图是PyTorch做加法时生成的配置文件,除了一个小方块,其他都是纯开销。现代深度学习模型通常是计算密集型的,而像PyTorch这样的框架是异步进行的。也就是说,当PyTorch正在运行一个CUDA内核时,它可以继续运行并在它后面排队更多的CUDA内核。因此,只要PyTorch能够“提前”运行CUDA内核,大部分框架开销就完全隐藏了。由于开销通常不会随着问题的大小而增加(计算和内存按比例增加),一个简单的判断方法是如果批量大小增加一倍,但运行时间只增加10%(预计运行时间增加一倍),那么开销可能太大了。另一种方法是使用PyTorch分析器。粉色线条显示了CPU内核与GPU内核的匹配情况。当GPU在等待CPU的开销时,有很多间隙。CPU的运行速度比GPU快,差距小得多。nvidia-smi中的GPU-Util正在测量实际运行的GPU核心的百分比,这也是衡量开销的好方法。大部分开销来自于PyTorch等框架的灵活性,需要花费大量时间“弄清楚要做什么”。例如执行a+b时,需要三步:1.Python需要找到a上__add__调度的内容2.PyTorch需要判断tensor的很多属性(比如dtype,device,是否需要Augrad)确定调用哪个内核。3.PyTorch需要真正启动内核。每一步都需要灵活性来支持不同的操作。解决灵活性的一种方法是跟踪,例如使用jit.tract、FX或jax.jit,或者在较低级别使用CUDAGraphs。要提高模型效率,最重要的是了解模型的性能瓶颈。当然,写一个神经网络模型还是要考虑这么多的开销问题,也可以说是这些系统和框架设计的失败,因为这些应该对用户透明。但是了解这些基本原理绝对是有意义的,可以帮助你从“根源”上解决性能瓶颈。
