团队之所以称它为 2.0，是因为它有一些标志性的新特性，包括：

 

　　TorchDynamo 可以从字节码分析生成 FX 图；

 

　　AOTAutograd 可以以 ahead-of-time 的方式生成反向图；

 

　　PrimTorch 引入了一个小型算子集，使后端更容易；

 

　　TorchInductor：一个由 OpenAI Triton 支持的 DL 编译器。

 

　　PyTorch 2.0 将延续 PyTorch 一贯的优势，包括 Python 集成、命令式风格、API 简单等等。此外，PyTorch 2.0 提供了相同的 eager-mode 开发和用户体验，同时从根本上改变和增强了 PyTorch 在编译器级别的运行方式。该版本能够为「Dynamic Shapes」和分布式运行提供更快的性能和更好的支持。

 

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

 

　

 

　　以下是详细内容。

 

　　PyTorch 2.X：速度更快、更加地 Python 化、一如既往地 dynamic

 

　　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 程序，这是一项重大创新，是研究团队对快速可靠地获取图进行 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 个开源模型。实验精心构建了测试基准，包括各种 CV 任务（图像分类、目标检测、图像生成等）、NLP 任务（语言建模、问答、序列分类、推荐系统等）和强化学习任务，测试模型主要有 3 个来源：

 

　　46 个来自 HuggingFace Transformers 的模型；

 

　　来自 TIMM 的 61 个模型：一系列 SOTA PyTorch 图像模型；

 

　　来自 TorchBench 的 56 个模型：包含来自 github 的精选流行代码库。

 

　　然后研究者测量加速性能并验证这些模型的准确性。加速可能取决于数据类型，研究团队选择测量 float32 和自动混合精度 (AMP) 的加速。

 

　　在 163 个开源模型中，torch.compile 在 93% 的情况下都有效，模型在 NVIDIA A100 GPU 上的训练速度提高了 43%。在 float32 精度下，它的平均运行速度提高了 21%，而在 AMP 精度下，它的运行速度平均提高了 51%。

 

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

 

　　TorchDynamo：快速可靠地获取图

 

　　TorchDynamo 是一种使用 Frame Evaluation API （PEP-0523 中引入的一种 CPython 特性）的新方法。研发团队采用数据驱动的方法来验证其在 Graph Capture 上的有效性，并使用 7000 多个用 PyTorch 编写的 Github 项目作为验证集。TorchScript 等方法大约在 50% 的时间里都难以获取图，而且通常开销很大；而 TorchDynamo 在 99% 的时间里都能获取图，方法正确、安全且开销可忽略不计（无需对原始代码进行任何更改）。这说明 TorchDynamo 突破了多年来模型权衡灵活性和速度的瓶颈。

 

　　TorchInductor：使用 define-by-run IR 快速生成代码

 

　　对于 PyTorch 2.0 的新编译器后端，研发团队从用户编写高性能自定义内核的方式中汲取灵感：越来越多地使用 Triton 语言。此外，研究者还想要一个编译器后端——使用与 PyTorch eager 类似的抽象，并且具有足够的通用性以支持 PyTorch 中广泛的功能。

 

　　TorchInductor 使用 pythonic define-by-run loop level IR 自动将 PyTorch 模型映射到 GPU 上生成的 Triton 代码和 CPU 上的 C++/OpenMP。TorchInductor 的 core loop level IR 仅包含约 50 个算子，并且是用 Python 实现的，易于破解和扩展。

 

　　AOTAutograd：将 Autograd 重用于 ahead-of-time 图

 

　　PyTorch 2.0 的主要特性之一是加速训练，因此 PyTorch 2.0 不仅要捕获用户级代码，还要捕获反向传播。此外，研发团队还想要复用现有的经过实践检验的 PyTorch autograd 系统。AOTAutograd 利用 PyTorch 的 torch_dispatch 可扩展机制来跟踪 Autograd 引擎，使其能够「ahead-of-time」捕获反向传递（backwards pass）。这使 TorchInductor 能够加速前向和反向传递。

 

　　PrimTorch：稳定的原始算子

 

　　为 PyTorch 编写后端具有挑战性。PyTorch 有 1200 多个算子，如果考虑每个算子的各种重载，则有 2000 多个。

 

　　在 PrimTorch 项目中，研发团队致力于定义更小且稳定的算子集，将 PyTorch 程序缩减到这样较小的算子集。目标是定义两个算子集：

 

　　Prim ops：约有 250 个相当低级的算子。这些算子适用于编译器，需要将它们重新融合在一起以获得良好的性能；

 

　　ATen ops：约有 750 个规范算子。这些算子适用于已经在 ATen 级别集成的后端或没有编译功能的后端（无法从较低级别的算子集（如 Prim ops）恢复性能）。