PyTorch 2.0 编译器提高了模型训练速度

PyTorch基金会最近发布了PyTorch 2.0版本，这是一个 100%向后兼容的更新。该版本的主要 API 贡献是为深度学习模型提供了一个编译函数，可以加快训练速度。163 个开源人工智能项目的内部基准测试显示，在训练期间，这些模型平均运行速度提高了 43%。

PyTorch 2.0 的发布计划于 2022 年 12 月在PyTorch大会上宣布。除了新的编译特性外，该版本还包括通过缩放点积注意力（SDPA）的新实现来提高 Transformer-based 模型（如大语言模型和扩散模型）的性能。通过改进的Metal Performance Shaders（MPS）加快了在 Apple Silicon 上的训练，目前在 MPS 中实施了 300 项操作。除了核心版本外，包括 TorchAudio、TorchVision 和 TorchText 在内的域库也更新了新的测试版功能。总体而言，自 1.13.1 版本发布以来，2.0 版本包含了来自 428 名开发人员的 4500 多次提交。PyTorch 基金会的博客上写到：

能宣布 PyTorch®2.0 的发布，我们很激动，我们在 2022 年 12 月 2 日的 PyTorch 大会上强调了这一点！PyTorch 2.0 提供了相同的动态图模式（eager-mode）开发和用户体验，同时从根本上改变并加强了 PyTorch 在编译器级别的操作方式，提供了更快的性能和对动态形状和分布式的支持。

在 2022 年 PyTorch 大会的主题演讲中，PyTorch 联合创始人Soumith Chintala指出，由于 GPU 计算能力的增加，许多现有的 PyTorch 工作负载受到了内存带宽或 PyTorch 框架开销的限制。此前，PyTorch 团队通过用 C++编写一些核心组件来解决性能问题；Chintala 将 PyTorch 描述为“基本上是一个 C++代码库”，并表示他“讨厌”为 C++组件做出贡献。

新的编译特性基于四个用 Python 编写的底层组件：

• TorchDynamo——通过将表示深度学习模型的 Python 代码重写为计算图块来执行图的获取

• AOTAutograd ——为后退步骤执行“提前”自动微分

• PrimTorch——将超过 2k 个 PyTorch 操作符规范化为固定的约 250 个原始操作符

• TorchInductor——为加速器生成特定于硬件的快速后端代码

为了演示编译函数的性能改进和易用性，PyTorch 团队确定了 163 个开源深度学习项目进行基准测试。其中包括各种任务的实现，包括计算机视觉、自然语言处理和强化学习。除了对编译函数的单行调用外，该团队没有对代码进行任何更改。这一单一改动在 93%的项目中有效，在NVIDIA A100 GPU 上训练时，编译后的模型的运行速度提高了 43%。

在黑客新闻（Hacker News）关于此次发布的讨论中，一位用户指出：

通过与其他框架相比，我从 PyTorch 中学到的一条重要教训是，生产力胜过增量的性能改进。Caffe 和 MXNet 都以速度快为卖点的，但显然在此处或这里的速度都快了一些，但这并不重要。另一方面，一旦我们让一个系统运行并流行起来，社区将会以比竞争对手预期更快地缩小性能差距。另一个教训可能是老生常谈的但同样值得重复：对开源项目的投资和专业打磨很重要。

PyTorch的代码和2.0版本的发布说明可在 GitHub 上获得。

原文链接：

https://www.infoq.com/news/2023/03/pytorch-release-compile/

相关阅读：

里程碑！PyTorch 正式加入 Linux 基金会，社区治理这一核心将不会改变

深度学习为什么要选择 PyTorch