大家都想让模型训练的更快，但是你真的找到对的方法了吗？在康奈尔大学本科生、PyTorch团队实习生HoraceHe看来，这个问题应该分几个步骤来解决：首先，你要知道你的训练为什么慢，也就是说瓶颈在哪里，二是寻找相应的解决方案。在不了解基本原理（第一原则）的情况下尝试一些事情是浪费时间。在这篇文章中，HoraceHe从计算、内存带宽和开销三个角度分析了可能存在的瓶颈，并提供了一些判断当前处于哪个瓶颈的方法，有助于我们更有针对性地对系统进行加速。这篇文章得到了陈天奇等众多资深研发人员的赞赏。以下为原文内容：如何提升深度学习模型的性能？大多数人会随意选择一些网上博客总结的技术，比如“使用系统内置的算子，设置梯度为0，使用PyTorch1.10.0版本而不是1.10.1版本……”在这个领域，当代（尤其是深度学习）系统感觉不太像科学而更像炼金术，所以不难理解为什么用户倾向于采用这种随机的方法。尽管如此，在这个领域还是有一些可以遵循的首要原则，我们可以排除大量的方法，从而使问题更容易解决。例如，如果你的训练损失远低于你的测试损失，那么你可能正在经历“过度拟合”并且试图增加模型容量是浪费时间。再举一个例子，如果你的训练损失与验证损失相同，那么对模型进行正则化是不明智的。同样，您可以将高效深度学习的问题分为以下三个不同的部分：计算：GPU计算实际浮点运算(FLOPS)所花费的时间；内存：在GPU内传输张量所花费的时间；开销：花在其他部分的时间。在训练机器学习模型时，了解您遇到的问题类型至关重要，这与提高模型效率一样重要。例如，当模型花费大量时间从内存传输到GPU时（即当内存带宽紧张时），增加GPU的FLOPS并没有帮助。另一方面，如果您正在运行大量矩阵乘法（即，当计算量很大时），用C++重写您的程序以减轻开销将无济于事。因此，要想GPU顺利运行，以上三个方面的讨论和研究必不可少。艰苦的教训背后有大量的工程师让GPU高效运行。注：本博客大部分内容基于GPU和PyTorch示例，但原理基本跨硬件和框架通用。计算优化的深度学习系统的一个方面是我们希望最大化花在计算上的时间。您支付了312teraflops的费用，并且希望将这些用于计算。但是，为了让昂贵的矩阵乘法物有所值，您需要在其他部分花费更少的时间。但是为什么这里的重点是最大化计算而不是最大化内存带宽呢？原因很简单——你可以减少开销或内存消耗，但你很难在不改变实际计算的情况下减少计算。与内存带宽相比，计算的增长速度使得计算利用率难以最大化。下表显示了使CPU的FLOPS翻倍和内存带宽翻倍所需的时间（重点关注黄色列）。考虑计算的一种方法是将其视为工厂。我们向我们的工厂传达指令（额外成本），为其提供原材料（内存带宽），所有这些都是为了使工厂更有效地运行（计算）。所以，如果工厂的产能扩张速度快于我们提供原材料的速度，它就很难达到峰值效率。即使我们的工厂容量（FLOPs）翻倍，但是带宽跟不上，我们的性能也无法翻倍。关于FLOPS还有一点要说的是，越来越多的机器学习加速器都有专用于矩阵乘法的硬件配置，比如Nvidia的“TensorCores”。所以，如果不做矩阵乘法，只能达到19.5万亿次运算，而不是312万亿次运算。请注意，GPU并不是唯一如此特殊的。事实上，TPU是比GPU更专业的计算模块。除了矩阵乘法，GPU处理其他操作的速度很慢，乍一看似乎有问题：其他操作符（如层归一化或激活函数）呢？实际上，这些运算符就像FLOPS中矩阵乘法的舍入误差一样。例如，BERT中不同算子类型占用的FLOP数见下表，其中“TensorContract”指的是矩阵乘法。如您所见，非矩阵乘法运算仅占所有运算的0.2%，因此如果它们只有矩阵乘法速度的1/15，则没有问题。事实上，归一化和逐点运算仅使用矩阵乘法FLOPS的1/250和1/700。那么，为什么非矩阵乘法运算的运行时间比应有的多得多呢？回到上面的“工厂”类比，罪魁祸首往往是原材料运往和运出工厂的方式，换句话说，就是“内存带宽”。带宽带宽消耗本质上是将数据从一个地方移动到另一个地方的成本，这可能意味着将数据从CPU移动到GPU，从一个节点移动到另一个节点，甚至从CUDA全局内存移动到CUDA共享内存。最后一个是本文的重点，我们一般称之为“带宽消耗”或者“内存带宽消耗”。前两者一般称为“数据传输消费”或“网络消费”，不在本文讨论范围之内。还是回到“工厂”的类比。虽然我们在工厂做真正的工作，但它不适合大规模存储。我们需要确保它的存储足够高效并且可以快速使用（SRAM），而不是以量取胜。那么我们将实际结果和“原材料”存放在哪里呢？通常我们有一个仓库，土地足够便宜并且有很多空间（DRAM）。然后我们可以在它和工厂之间运送东西（内存带宽）。这种在计算单元之间移动事物的成本被称为“内存带宽”成本。其实nvidia-smi命令中出现的“memory”就是DRAM，经常让人抓狂的“CUDAoutofmemory”指的就是这个DRAM。值得注意的是，每次执行GPU核心操作时，我们都需要将数据运出并运回我们的仓库-DRAM。现在想象一下，当我们进行一元运算时（比如torch.cos），我们需要将数据从仓库（DRAM）运送到工厂（SRAM），然后在工厂里进行一小步计算，然后运送结果回到仓库。运输非常耗时，在这种情况下，我们几乎所有的时间都花在运输数据上，而不是真正的计算上。因为我们把所有的时间都花在了内存带宽上，所以这个操作也被称为内存限制操作，这意味着我们不会在计算上花费很多时间。显然，这不是我们想要的。所以，我们能做些什么？让我们看看操作符序列是什么样的。逐点运算符序列可能是什么样子。在全局内存和计算单元之间来回传输数据的做法显然不是最优的。更好的方法是一次执行数据工厂中的所有操作，然后将数据发回。这就是运算符融合——深度学习编译器中最重要的优化。简单地说，这种方法不会将数据写入全局内存以便再次读取，而是通过一次执行多个计算来避免额外的内存访问。比如执行x.cos().cos()操作，写入内存的方式需要4次全局读写。x1=x.cos()#从全局内存中的x读取，写入x1x2=x1.cos()#从全局内存中的x1读取，写入x2和operatorfusion只需要2次全局内存读写，所以实现了一个2倍加速。x2=x.cos().cos()#从全局内存中的x读取，写入x2但是这个方法并不简单，需要一些条件。首先，GPU需要知道在执行完当前操作之后接下来会发生什么，所以这个优化不能在PyTorch的eager模式下完成（一次运行一个操作符）。其次，我们需要编写CUDA代码，这也不是一件容易的事。并非所有运算符融合都像逐点运算符一样简单。您可以将逐点运算符融合到归约或矩阵乘法中。甚至矩阵乘法本身也可以被认为是广播乘法和归约运算的融合。可以融合任意2个PyTorch运算符，从而节省读取/写入全局内存的内存带宽成本。此外，许多现有的编译器通常可以执行“简单”融合（例如NVFuser和XLA）。然而，更复杂的融合仍然需要人们手工编写，因此如果您想尝试自己编写自定义CUDA内核，Triton是一个不错的起点。令人惊讶的是，融合的x.cos().cos()操作将花费与单独调用x.cos()几乎相同的时间。这就是为什么激活函数的成本几乎相同，虽然gelu显然比relu涉及更多的操作。因此，重新实现/激活检查点会产生一些有趣的结果。本质上，进行额外的重新计算可能会导致更少的内存带宽，从而减少运行时间。因此，我们可以通过重新实现来减少内存占用和运行时间，并在AOTAutograd中构建一个简洁的min-cut优化pass。推断内存带宽成本对于简单的操作，直接推断内存带宽是可行的。例如，A100的全局内存带宽为1.5TB/秒，可以执行19.5teraflops/秒的计算。因此，对于32位浮点数（即4个字节），您可以在GPU执行20万亿次操作时加载4000亿个数字。此外，执行一个简单的一元运算（例如x2一个张量）实际上需要将张量写回全局内存。因此，在执行大约一百个一元运算之前，花在内存访问上的时间比实际计算要多。如果您实现以下PyTorch函数：deff(x:Tensor[N]):for_inrange(repeat):x=x*2returnx并使用融合编译器对其进行基准测试，您可以计算每个FLOPS和内存重复值的带宽。增加重复值是在不增加内存访问的情况下增加计算的简单方法——这也称为增加计算强度。具体来说，假设我们通过首先找到每秒执行的迭代次数来对这段代码进行基准测试；然后执行2N（N是张量大小）内存访问和N*重复FLOPs。因此，内存带宽将为bytes_per_elem*2*N/itrs_per_second，FLOPS为N*repeat/itrs_per_second。现在，让我们将计算强度绘制为3个函数的函数：运行时、触发器和内存带宽。请注意，在执行64次乘法之前，运行时间根本不会显着增加。这意味着在此之前，它主要受到内存带宽的限制，而计算主要是闲置的。初始FLOPS值为0.2teraflops。当我们将计算强度加倍时，这个数字会线性增长，直到它接近9.75teraflops的峰值，一旦我们接近峰值teraflops，这就被认为是“计算受限”。最后，可以看出内存带宽从峰值附近开始，并随着计算强度的增加而开始下降。这正是我们所期望的，因为这意味着越来越多的时间花在执行实际计算上，而不是访问内存上。在这种情况下，很容易看出您何时受计算限制以及何时受内存限制。当repeat<32时，内存带宽接近饱和，但未完全计算；当repeat>64时，计算接近饱和（即接近峰值FLOPS），内存带宽开始下降。对于较大的系统，通常很难判断它们是受计算限制还是受内存带宽限制，因为它们通常包含两者的组合。衡量计算受限程度的一种常用方法是计算实际FLOPS与峰值FLOPS的百分比。然而，除了内存带宽成本之外，还有一件事可以阻止GPU运行如丝般顺畅。开销发生在代码花时间做传递张量或计算以外的事情时，比如花在Python解释器上的时间，花在PyTorch框架上的时间，花在启动CUDA上的时间花在内核上（但不执行），这些都是开销。开销很重要的原因是现代GPU非常快。A100每秒可执行312万亿次浮点运算(312TeraFLOPS)。相比之下，Python慢得可怕——Python每秒执行大约3200万次加法。这意味着Python执行单个FLOP的时间，A100可能已经运行了975万次FLOPS。更糟糕的是，Python解释器甚至不是唯一的开销来源，像PyTorch这样的框架在到达实际内核之前还有很多层调度。PyTorch每秒可以执行大约280,000次操作。如果您使用微型张量（例如用于科学计算），您可能会发现PyTorch与C++相比非常慢。例如在下图中，使用PyTorch执行单个加法，只有一小部分图形是实际执行计算的部分，其余部分是纯粹的开销。鉴于此，您可能会对PyTorch成为主流框架这一事实感到困惑，因为现代深度学习模型通常执行大量计算。此外，像PyTorch这样的框架是异步执行的。因此，大部分框架开销可以完全忽略。如果我们的GPU操作足够大，CPU可以先于GPU运行（因此CPU开销无关紧要）。另一方面，如果GPU操作太小，那么GPU将把大部分时间花在镇纸上。那么，如何判断你是否遇到了这个问题呢？由于开销通常不会随着问题的大小而增加（而计算和内存会随着问题的大小而增加），因此最简单的判断方法就是简单地增加数据的大小。如果运行时间没有按比例增加，您应该可以说您达到了开销限制。例如，如果您将批处理大小加倍但仅将运行时间增加10%，您可能会受到开销的限制。另一种方法是使用PyTorch分析器。在下图中，粉红色块显示了CPU内核与GPU内核的匹配情况。CPU运行在GPU之前。另一方面，nvidia-smi中的“GPU-Util”（不是“VolatileGPU-Util”）条目测量实际运行的GPU内核的百分比，因此这是查看开销限制是否达到良好的另一种方法主意。这种开销是PyTorch等所有灵活框架的一部分，它们本质上需要花费大量时间“弄清楚要做什么”。这可能来自Python中的代码（查找属性或分派到正确的函数）或PyTorch。例如，当您执行a+b时，需要执行以下步骤：Python需要查找__add__在a上分派的内容。PyTorch需要判断tensor的很多属性（比如dtype，device，是否需要autograd）来决定调用哪个kernel。PyTorch需要实际启动内核。从根本上说，这种开销来自能够在每个步骤中执行不同操作的灵活性。如果您不需要这种灵活性，一种解决方法是跟踪它，例如使用jit.trace、FX或jax.jit。或者，可以使用诸如CUDAGraphs之类的东西来代替在较低级别执行此操作。不幸的是，这是以失去灵活性为代价的。两全其美的一种方法是通过在VM级别进行内省来编写更符合“真实”JIT的内容。有关详细信息，请参阅TorchDynamo(https://dev-discuss.pytorch.org/t/torchdynamo-an-experiment-in-dynamic-python-bytecode-transformation/361)。总结如果你想加速一个深度学习系统，最重要的是了解模型中的瓶颈是什么，因为瓶颈决定了用什么方法来加速系统是合适的。很多时候，我看到研究人员和其他对加速PyTorch代码感兴趣的人在不理解他们所面临的问题的情况下盲目尝试。当然，另一方面，如果用户需要考虑这些东西，也反映了框架的局部失效。尽管PyTorch是一个活跃的兴趣领域，但PyTorch的编译器或配置文件API并不是最容易使用的。总而言之，我发现了解系统的基本原理几乎总是有用的，希望这对您也有用。
