一行代码，炼丹 2 倍速！PyTorch 2.0 惊喜问世，LeCun 激情转发

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

12 月 2 日，PyTorch 2.0 正式发布！

这次的更新不仅将 PyTorch 的性能推到了新的高度，同时也加入了对动态形状和分布式的支持。

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

目前，PyTorch 2.0 还处在测试阶段，预计第一个稳定版本会在 2023 年 3 月初面世。

PyTorch 2.x：更快、更 Python！

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

当前版本的 PyTorch 所面临的挑战是，eager-mode 难以跟上不断增长的 GPU 带宽和更疯狂的模型架构。

而 PyTorch 2.0 的诞生，将从根本上改变和提升了 PyTorch 在编译器级别下的运行方式。

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

然而，PyTorch 的代码却并没有完全采用 Python，而是把一部分交给了 C++。

不过，在今后的 2.x 系列中，PyTorch 项目团队计划将与 torch.nn 有关的代码移回到 Python 中。

除此之外，由于 PyTorch 2.0 是一个完全附加的（和可选的）功能，因此 2.0 是 100% 向后兼容的。

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

更多的技术支持

• TorchDynamo

使用 Python 框架评估钩子安全地捕获 PyTorch 程序，这是团队 5 年来在 graph capture 方面研发的一项重大创新。

• AOTAutograd

重载了 PyTorch 的 autograd 引擎，作为一个追踪的 autodiff，用于生成超前的反向追踪。

• PrimTorch

将约 2000 多个 PyTorch 运算符归纳为约 250 个原始运算符的封闭集，开发人员可以针对这些运算符构建一个完整的 PyTorch 后端。大大降低了编写 PyTorch 功能或后端的障碍。

• TorchInductor

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

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

更快的训练速度

通过引入新的编译模式「torch.compile」，PyTorch 2.0 用一行代码，就可以加速模型的训练。

这里不用任何技巧，只需运行 torch.compile () 即可，仅此而已：

opt_module = torch.compile(module)

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

• 来自 HuggingFace Transformers 的 46 个模型

• 来自 TIMM 的 61 个模型：Ross Wightman 收集的最先进的 PyTorch 图像模型

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

测试结果表明，在这 163 个跨越视觉、NLP 和其他领域的开源模型上，训练速度得到了 38%-76% 的提高。

此外，团队还在一些流行的开源 PyTorch 模型上进行了基准测试，并获得了从 30% 到 2 倍的大幅加速。

开发者 Sylvain Gugger 表示：「只需添加一行代码，PyTorch 2.0 就能在训练 Transformers 模型时实现 1.5 倍到 2.0 倍的速度提升。这是自混合精度训练问世以来最令人兴奋的事情！」

技术概述

PyTorch 的编译器可以分解成三个部分：

• 图的获取

• 图的降低

• 图的编译

其中，在构建 PyTorch 编译器时，图的获取是更难的挑战。

TorchDynamo

今年年初，团队便开始了 TorchDynamo 的工作，这种方法使用了 PEP-0523 中引入的 CPython 功能，称为框架评估 API。

为此，团队采取了一种数据驱动的方法来验证 TorchDynamo 在 graph capture 上的有效性 -- 通过使用 7000 多个用 PyTorch 编写的 Github 项目，来作为验证集。

结果显示，TorchDynamo 在 99% 的时间里都能正确、安全地进行 graph capture，而且开销可以忽略不计。

TorchInductor

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

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

TorchInductor 的核心循环级 IR 只包含大约 50 个运算符，而且它是用 Python 实现的，这使得它很容易得到扩展。

AOTAutograd

想要加速训练，就不仅需要捕获用户级代码，而且还要捕获反向传播。

AOTAutograd 可以利用 PyTorch 的 torch_dispatch 扩展机制来追踪 Autograd 引擎，「提前」捕获反向传播，进而能够使用 TorchInductor 来加速前向和后向通道。

PrimTorch

PyTorch 有 1200 多个运算符，如果考虑到每个运算符的各种重载，则有 2000 多个。因此，编写后端或跨领域的功能成为一项耗费精力的工作。

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

Prim ops 有大约~250 个运算符，适合于编译器。由于足够低级，因此只需将它们融合在一起以获得良好的性能。

ATen ops 有大约~750 个典型的运算符，适合于按原样输出。这些适合于已经在 ATen 级别上集成的后端，或者没有编译的后端，从而恢复像 Prim ops 这样的低级别运算符集的性能。

动态形状

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

在不支持动态形状的情况下，一个常见的解决方法是将其填充到最接近的 2 次方。然而，正如我们从下面的图表中所看到的，它产生了大量的性能开销，同时也带来了明显更长的编译时间。

现在，有了对动态形状的支持，PyTorch 2.0 也就获得了比 Eager 高出了最多 40% 的性能。

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

参考资料：

• https://pytorch.org/get-started/pytorch-2.0/

本文来自微信公众号：新智元 （ID：AI_era），编辑：好困