PyTorch 2.0 发布，一行代码将训练提速 76%！

在 PyTorch Conference 2022 上，PyTorch 官方正式发布了 PyTorch 2.0，与先前的 1.x 版本相比，2.0 中有了跨越式的变化。

译者 | 刘春霖    责编 | 杨紫艳

出品 | CSDN （ID：csdnnews）

在11月初，PyTorch 团队在官方博客宣布 Pytorch 1.13 发布。据官方介绍，PyTorch 1.13 中包括了 BetterTransformer 的稳定版，且不再支持 CUDA 10.2 及 11.3，并完成了向 CUDA 11.6 及 11.7 的迁移。此外 Beta 版还增加了对 Apple M1 芯片及 functorch 的支持。

令人意外的是，12月2日，PyTorch 2.0 正式发布！与先前的 1.x 版本相比，2.0 中有了跨越式的变化。

PyTorch 2.0 中发布了大量足以改变 PyTorch 使用方式的新功能，它提供了相同的 eager mode 和用户体验，同时通过 torch.compile 增加了一个编译模式，在训练和推理过程中可以对模型进行加速，从而提供更佳的性能和对 Dynamic Shapes 及分布式运行的支持。

PyTorch 团队表示，PyTorch 2.0 是他们向 2.x 系列迈出的第一步，其稳定版预计在 2023 年 3 月初发布。

PyTorch 2.x：更快、更 Python！

PyTorch 2.0 官宣了一个重要特性 —— torch.compile，这一特性将 PyTorch 的性能推向了新的高度，并将 PyTorch 的部分内容从 C++ 移回 Python。torch.compile 是一个完全附加的（可选的）特性，因此 PyTorch 2.0 是 100% 向后兼容的。

支撑 torch.compile 的技术包括研发团队新推出的 TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。

• TorchDynamo：使用 Python Frame Evaluation Hooks 安全地捕获 PyTorch 程序，这项重大创新是 PyTorch 过去 5 年来在安全图结构捕获方面的研发成果汇总；

• AOTAutograd：重载 PyTorch 的 autograd 引擎，作为一个跟踪 autodiff，用于生成 ahead-of-time 向后跟踪；

• PrimTorch：将约 2000 多个 PyTorch 算子归纳为一组约 250 个原始算子的闭集，开发人员可以将其作为构建完整 PyTorch 后端的目标。这大大降低了编写 PyTorch 功能或后端的流程；

• TorchInductor：是一种深度学习编译器，可为多个加速器和后端生成快速代码。对于 NVIDIA GPU，它使用 OpenAI Triton 作为关键构建块。

TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 是用 Python 编写的，并支持 dynamic shapes（无需重新编译就能发送不同大小的向量），这使得它们具备灵活、易于破解的特性，降低了开发人员和供应商的使用门槛。

为了验证这些技术，研发团队在机器学习领域测试了 163 个开源模型，包括图像分类、目标检测、图像生成等、NLP 任务，如语言建模、问答、序列分类、推荐系统和强化学习任务，测试模型主要有 3 个来源：

• 46 个来自 HuggingFace Transformers 的模型；

• 来自 TIMM 的 61 个模型：由 Ross Wightman 收集的SOTA PyTorch 图像模型；

• 来自 TorchBench 的 56 个模型：包含来自 Github 上收集的一组流行代码库。

对于开源模型，PyTorch 官方没有进行修改，只是增加了一个 torch.compile 调用来进行封装

接下来 PyTorch 工程师在这些模型中测量速度并验证精度，由于加速可能取决于数据类型，因此研究团队选择测量 Float32 和自动混合精度 (AMP) 的加速。

在 163 个开源模型中，该团队发现使用 2.0 可以将训练速度提高 38-76%。torch.compile 在 93% 的情况下都有效，模型在 NVIDIA A100 GPU 上的训练速度提高了 43%。在 Float32 精度下，它的平均运行速度提高了 21%，而在 AMP 精度下，它的运行速度平均提高了 51%。

目前，torch.compile 还处于早期开发阶段，预计 2023 年 3 月上旬将发布第一个稳定的 2.0 版本。

在官方博客中，PyTorch 团队还公布了他们对于整个 2.0 系列的展望：

开发背景

PyTorch 的开发理念自始至终都是灵活性和 hackability 第一，性能则是第二，致力于：

1. 高性能的 eager execution

2. 不断 Python 化内部结构

3. 分布式、自动比较、数据加载、加速器等的良好抽象

PyTorch 自 2017 年面世以来，硬件加速器（如 GPU）的计算速度提高了约 15倍，内存访问速度提高了约 2 倍。

为了保持高性能的 eager execution，PyTorch 内部的大部分内容不得不转移到 C++ 中，这使得 PyTorch hackability 下降，也增加了开发者参与代码贡献的门槛。

从第一天起，PyTorch 官方就意识到了 eager execution 的性能局限。2017 年 7 月，官方开始致力于为 PyTorch 开发一个编译器。该编译器需要在不牺牲 PyTorch 体验的前提下，加速 PyTorch 程序的运行，其关键标准是保持某种程度上的灵活性 (flexibility)：支持开发者广泛使用的 dynamic shapes 以及 dynamic programs。

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

技术概述

多年来，研究者们在 PyTorch 中建立过好几个编译器项目，这些编译器可以分为 3 类：

• 图结构的获取
• 图结构的降低
• 图结构的编译

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

过去 5 年中，官方尝试了 torch.jit.trace、TorchScript、FX tracing 以及 Lazy Tensors，但它们有些够灵活但不够快，有些够快但不灵活，有些既不快也不灵活，有些用户体验不好。

虽然 TorchScript 很有前途，但它需要大量修改代码和依赖，可行性并不高。

PyTorch 编译流程示意图

TorchDynamo：可靠快速地获取图结构 

TorchDynamo 使用了 PEP-0523 中引入的CPython 功能，称为框架评估 API (Frame Evaluation API)。为此，官方采取了一种数据驱动的方法来验证其在 Graph Capture 上的有效性，使用 7000 多个用 PyTorch 编写的 Github 项目作为验证集。

结果显示，TorchDynamo 在 99% 的时间里都能正确、安全地获取图结构，而且开销可以忽略不计，因为它无需对原始代码做任何修改。

TorchInductor：用 define-by-run IR 进行更迅速的 codegen

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

此外，对于 PyTorch 2.0 全新的编译器后端，官方还希望能够使用与 PyTorch eager 类似的抽象，并且具有足够的通用性能支持 PyTorch 中广泛的功能。

TorchInductor 使用 Pythonic define-by-run loop level IR，自动将 PyTorch 模型映射到 GPU 上生成的 Triton 代码以及 CPU 上的 C++/OpenMP。

TorchInductor 的 core loop level IR 只包含大约 50 个算子，而且是用 Python 实现的，这使得它具有很强的 hackability 和扩展性。

AOTAutograd：对于 ahead-of-time graph，重用 Autograd

PyTorch 2.0 要想加速训练，不仅要捕获用户级代码，而且要捕获反向传播算法 (backpropagation)。

AOTAutograd 利用 PyTorch torch_dispatch 扩展机制来追踪 Autograd 引擎，使开发者得以提前捕获反向传播 (backwards pas)，从而使开发者得以使用 TorchInductor 加速前向和后向通道。

PrimTorch：稳定的 Primitive operator

PyTorch 有 1200多个运算符，如果考虑到每个运算符的各种重载，数量高达 2000+。

2000+ PyTorch 算子的分类概况

因此，编写后端或交叉功能 (cross-cutting feature) 成为一项耗费精力的工作。PrimTorch 致力于定义更小更稳定的运算符集。PyTorch 程序可以持续降级到这些运算符集。

官方的目标是定义两个运算符集：

• Prim ops 包含约 250 个相对底层的运算符，因为足够底层，所以这些运算符更适用于编译器，开发者需要将这些算子进行融合，才能获得良好的性能。

• ATen ops 包含约 750 个典型运算符 (canonical operator)，适合于直接输出。这些运算符适用于已经在 ATen 级别上集成的后端，或者没有经过编译的后端，才能从底层运算符集（如 Prim ops) 恢复性能。

  动态形状

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

目前为止，对动态形状的支持有限，并且正在进行中。它将在稳定版本中具有完整的功能。

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

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

简而言之，就是 PyTorch 2.0 stable 版本预计明年 3 月发布，PyTorch 2.0 在保留原有优势的同时，大力支持编译，torch.compile 为可选功能，只需一行代码即可运行编译，还有 4 项重要技术：TorchDynamo、AOTAutograd、PrimTorch 以及 TorchInductor，PyTorch 1.x 代码无需向 2.0 迁移，更多用户体验以及Q&A，大家可以查看下方链接中原文呦~

参考链接：

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