现在，只需添加一行代码，PyTorch2.0就能让你在训练变形金刚模型时实现1.5 —2倍的速度提升！

日前，PyTorch 2.0正式发布！

此次更新不仅将PyTorch的性能推向了一个新的高度，还增加了对动态形状和分布的支持。

此外，2.0系列还将PyTorch的部分代码从C++移回Python。

目前PyTorch2.0还在测试阶段，预计2023年3月初会有第一个稳定版本。

PyTorch 2.x:更快更Python！

在过去的几年里，PyTorch从1.0到最新的1.13进行了创新迭代，并转移到新成立的PyTorch基金会，成为Linux基金会的一部分。

PyTorch当前版本的挑战在于，eager—mode跟不上不断增加的GPU带宽和更加疯狂的模型架构。

PyTorch 2.0的诞生将从根本上改变和完善PyTorch在编译器层面的运行模式。

众所周知，PyTorch来自于数据科学中广泛使用的开源Python编程语言。

但是PyTorch的代码并没有完全采用Python，而是让出了一部分给C++。

但在未来的2.x系列中，PyTorch项目组计划将torch.nn相关的代码移回Python。

此外，由于PyTorch 2.0是一个完全附加的特性，2.0是100%向后兼容的。

也就是说，代码基础是一样的，API是一样的，写模型的方式也是一样的。

更多技术支持

火炬接力

使用Python框架评估钩子安全捕获PyTorch程序是团队在过去五年中在图形捕获方面开发的重大创新。

奥陶格德

超载PyTorch的自动签名引擎，作为一个跟踪自动夫，用于生成高级反向跟踪。

PrimTorch

2000多个PyTorch操作符被总结成一个约250个原始操作符的封闭集合，开发者可以为这些操作符构建一个完整的PyTorch后端大大减少了编写PyTorch函数或后端的障碍

火炬传导器

可以为多个加速器和后端生成快速代码的深度学习编译器对于NVIDIA的GPU，它使用OpenAI Triton作为关键的构建模块

值得注意的是TorchDynamo，AOTAutograd，PrimTorch和TorchInductor都是用Python写的，支持动态形状。

更快的训练速度

为了验证这些技术，团队精心构建了测试基准，包括图像分类，物体检测，图像生成等任务，以及各种NLP任务，如语言建模，问答，序列分类，推荐系统和强化学习等。其中，这些基准可以分为三类:

拥抱脸变形金刚的46个模型

来自TIMM的61个模型:罗斯·怀特曼收藏的最先进的PyTorch图像模型

来自TorchBench的56个模型:github的一组流行代码库

测试结果表明，在这163个横跨视觉，NLP等领域的开源模型上，训练速度提高了38%—76%。

NVIDIA A100 GPU对比

此外，团队还对一些流行的开源PyTorch模型进行了基准测试，得到了从30%到2倍的显著加速。

开发者Sylvain Gugger表示:PyTorch 2.0在训练变形金刚模型时，只需增加一行代码，就可以实现1.5倍到2.0倍的速度提升。这是混合精度训练问世以来最激动人心的一件事！」

技术概述

PyTorch的编译器可以分为三个部分:

图形的获取

图形的简化

图表的编制

其中，在构建PyTorch编译器时，获取图是一个比较困难的挑战。

火炬接力

今年年初，团队开始了TorchDynamo的工作该方法使用了PEP—0523中引入的CPython函数，称为框架评估API

因此，团队采用了数据驱动的方法来验证TorchDynamo在图形捕捉上的有效性——通过使用7000多个用PyTorch编写的Github项目作为验证集。

结果表明TorchDynamo可以在99%的时间内正确安全地执行图形捕获，开销可以忽略不计。

火炬传导器

对于PyTorch 2.0新的编译器后端，团队从用户如何编写高性能定制内核中获得了灵感:Triton语言被越来越多地使用。

TorchInductor通过使用Pythonic定义的逐循环IR，自动将PyTorch模型映射到CPU上的GPU和C++/OpenMP上生成的Triton代码。

TorchInductor的核心循环级IR只包含50个左右的运算符，而且是用Python实现的，很容易扩展。

奥陶格德

为了加速训练，我们不仅需要捕获用户级代码，还需要反向传播。

AOTAutograd可以使用PyTorch的torch_dispatch扩展机制来跟踪Autograd引擎，提前捕获反向传播，然后使用TorchInductor来加速前向和反向通道。

PrimTorch

PyTorch有1200多个操作符，如果考虑每个操作符的各种重载，有2000多个因此，编写后端或跨域函数成为一项非常耗费精力的任务

在PrimTorch项目中，团队定义了两个更小且更稳定的运算符集:

Prim ops有大约250个运算符，适合编译器使用因为它们足够低，所以只需要将它们融合在一起就可以获得良好的性能

ATen ops有大约750个典型的操作符，适合按原样输出这些适用于已经集成到ATen级别的后端，或者没有编译过的后端，以恢复像Prim ops这样的低级运算符集的性能

动态形状

在研究支持PyTorch代码通用性的必要条件时，一个关键要求是支持动态形状，并允许模型接受不同大小的张量，而不会导致每次形状改变时都要重新编译。

当不支持动态形状时，常见的解决方案是将它们填充到最接近的二次幂可是，正如我们从下面的图表中看到的，它带来了大量的性能开销和明显更长的编译时间

现在，由于支持动态形状，PyTorch 2.0的性能比Eager提高了40%。

最后，在PyTorch 2.x的路线图中，团队希望在性能和可扩展性方面进一步推动编译模式的发展。

参考资料: