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

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

出品 | CSDN （ID：csdnnews）

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%。

开发背景

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 很有前途，但它需要大量修改代码和依赖，可行性并不高。

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.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/