PyTorch 2.0发布了！一行代码提速76%

Posted 2022-12-15 Datawhale

 Datawhale干货 

内容：PyTorch，编辑：机器之心

PyTorch 官方：我们这次的新特性太好用了，所以就直接叫 2.0 了。

前段时间，PyTorch 团队在官方博客宣布 Pytorch 1.13 发布，包含 BetterTransformer 稳定版等多项更新。在体验新特性的同时，不少人也在期待下一个版本的推出。

出乎意料的是，这个新版本很快就来了，而且是跨越式的 2.0 版！

新版本的重要进步体现在速度和可用性，而且完全向后兼容。PyTorch 团队表示，PyTorch 2.0 是他们向 2.x 系列迈出的第一步，其稳定版预计在 2023 年 3 月初发布。

首先，PyTorch 2.0 引入了 torch.compile，这是一种编译模式，可以在不更改模型代码的情况下加速模型。在 163 个涵盖视觉、NLP 和其他领域的开源模型中，该团队发现使用 2.0 可以将训练速度提高 38-76%。

其次，PyTorch 2.0 是 100% 向后兼容的：代码库一样，API 一样，写模型的方式也一样。团队之所以称它为 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）恢复性能）。

用户体验

PyTorch 2.0 引入了一个简单的函数 torch.compile，它会返回一个编译后的模型。

compiled_model *=* torch.compile(model)

compiled_model 保存对模型的引用，并将 forward 函数编译为一个更优化的版本。在编译模型时，PyTorch 2.0 给了几项设置来调整它：

def torch.compile(model: Callable,


*,


mode: Optional[str] = "default",


dynamic: bool = False,


fullgraph:bool = False,


backend: Union[str, Callable] = "inductor",


# advanced backend options go here as kwargs


**kwargs


) -> torch._dynamo.NNOptimizedModule

• 「mode」指定编译器在编译时应该优化的内容。

• default 是一种预设模式，它试图在不花费太长时间或使用额外内存的情况下高效编译。

• 其他模式，如 reduce-overhead，可以大大降低框架开销，但要消耗少量额外内存。max-autotune 编译很长时间，试图为你提供它所能生成的最快的代码。

• 「dynamic」模式指定是否为 Dynamic Shapes 启用代码路径。某些编译器优化不能应用于动态形状的程序。明确你想要一个带有动态形状还是静态形状的编译程序，将有助于编译器提供更好的优化代码。

• 「fullgraph」类似于 Numba 的 nopython。它将整个程序编译成一个图，或者给出一个错误提示，解释为什么它不能这样做。大多数用户不需要使用这种模式。如果你非常注重性能，那么你可以尝试使用它。

• 「backend 」指定使用哪个编译器后端。默认情况下使用 TorchInductor，但还有其他一些可用的工具。

编译体验想要在默认模式中提供最大的好处和最大的灵活性。

常见问答

1、什么是 PT 2.0？

2.0 是最新的 PyTorch 版本。PyTorch 2.0 提供相同的 eager 模式开发体验，同时通过 torch.compile 添加编译模式。这种编译模式有可能在训练和推理期间加速模型。

2、为什么是 2.0 而不是 1.14？

PyTorch 2.0 就是计划中的 1.14 。我们发布了大量新特性，我们相信这些新特性会改变用户使用 PyTorch 的方式，因此称其为 PyTorch 2.0。

3、如何安装 2.0？有什么额外要求吗？

安装最新的 nightlies：

CUDA 11.7

pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

CUDA 11.6

pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu116

CPU

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cpu

4、2.0 代码是否向后兼容 1.X？

是的，使用 2.0 不需要修改 PyTorch 工作流程。一行代码 model = torch.compile(model) 就可以优化模型，以使用 2.0 堆栈，并与其余 PyTorch 代码一起顺利运行。这是一个可选择项，不需要使用新的编译器。

5、默认启用 2.0 吗？

不，必须通过使用单个函数调用优化模型从而在 PyTorch 代码中显式启用 2.0。

6、如何将 PT1.X 代码迁移到 PT2.0？

代码应该按原样工作，无需任何迁移。如果想使用 2.0 中引入的新编译模式特性，那么可以从优化模型开始：

model = torch.compile(model)

虽然加速主要是在训练期间观察到的，但如果你的模型运行速度比 eager 模式快，也可以将它用于推理。

import torch


  


def train(model, dataloader):


model = torch.compile(model)


for batch in dataloader:


run_epoch(model, batch)


def infer(model, input):


model = torch.compile(model)


return model(\\*\\*input)

7、是否有任何不该使用 PT 2.0 的应用程序？

当前版本的 PT 2.0 仍处于实验阶段，并且处于 nightlies 版本。其中的动态形状支持还处于早期阶段，所以可以等到 2023 年 3 月稳定版发布后再使用该功能。

8、运行 PyTorch 2.0 时，代码有何不同？ 

开箱即用，PyTorch 2.0 与 PyTorch 1.x 相同，模型以 eager 模式运行，即 Python 的每一行都逐个执行。

在 2.0 中，如果用 model = torch.compile(model) 将模型打包，则模型在执行之前会经过 3 个步骤：

1）图获取：首先将模型重写为子图块。可由 TorchDynamo 编译的子图被「压平」，其他子图（可能包含控制流代码或其他不受支持的 Python 结构）将回退到 Eager 模式。

2）Graph lowering：所有 PyTorch 操作都被分解为特定于所选后端的组成内核。

3）图编译，内核调用其相应的低级设备专用操作。

9、2.0 目前支持哪些编译器后端？

默认和最完整的后端是 TorchInductor，但是 TorchDynamo 有一个不断增长的后端列表，可以通过调用 torchdynamo.list_backends(). 找到

10、2.0 版本的分布式训练能力如何？

Compiled 模式下的 DDP 和 FSDP ，比 FP32 中的 Eager 模式快 15%、AMP 精度快 80%。PT2.0 做了一些额外的优化，以确保 DDP 的通信 - 计算 overlap 与 Dynamo 的部分图创建良好协作。想要确保使用 static_graph=False 运行 DDP，更多细节参见：https://dev-discuss.pytorch.org/t/torchdynamo-update-9-making-ddp-work-with-torchdynamo/860

11、为什么我的代码用 2.0 的 Compiled Model 运行变慢？

性能下降最可能的原因是 graph break 太多。例如，类似模型前向 trigger 中的输出语句这样的东西会触发 graph break。详见：https://pytorch.org/docs/master/dynamo/faq.html#why-am-i-not-seeing-speedups

12、以前运行的代码在 2.0 中崩溃了，该如何调试？

参见：https ://pytorch.org/docs/master/dynamo/faq.html#why-is-my-code-crashing

参考链接：https://pytorch.org/get-started/pytorch-2.0/

整理不易，点赞三连↓