概述

介绍 PyTorch 2.0,我们迈向 PyTorch 下一代 2.x 版本的第一步。在过去几年中,我们从 PyTorch 1.0 创新并迭代到最新的 1.13,并迁移到新成立的 PyTorch 基金会,该基金会是 Linux 基金会的一部分。

除了我们惊人的社区之外,PyTorch 最大的优势在于我们继续作为一流的 Python 集成,命令式风格,API 的简单性和选项。PyTorch 2.0 提供相同的渴望模式开发和用户体验,同时从根本上改变和增强 PyTorch 在幕后编译器级别的操作方式。我们能够提供更快的性能和对动态形状和分布式支持。

以下是你需要了解 PyTorch 2.0 是什么、它将走向何方以及更重要的是如何立即开始使用(例如,教程、要求、模型、常见问题解答)的所有信息。我们还有很多东西要学习和开发,但我们期待社区的反馈和贡献,使 2.x 系列变得更好,并感谢所有使 1.x 系列如此成功的人。

PyTorch 2.x:更快、更 Pythonic 且与以往一样动态

今天,我们宣布 torch.compile,这项功能将 PyTorch 的性能推向新的高度,并开始将 PyTorch 的部分从 C++ 迁移回 Python。我们相信这对于 PyTorch 来说是一个全新的方向——因此我们将其称为 2.0。 torch.compile 是一款完全增量(且可选)的功能,因此 2.0 本质上是 100% 向后兼容的。

支撑 torch.compile 的是新技术——TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。

  • TorchDynamo 使用 Python 帧评估钩子安全地捕获 PyTorch 程序,这是一项重大创新,是我们在安全图捕获方面 5 年研发工作的成果

  • AOTAutograd 将 PyTorch 的 autograd 引擎重载为跟踪自动微分,用于生成提前时间反向跟踪。

  • PrimTorch 将 ~2000+ PyTorch 运算符规范化为一组封闭的 ~250 个原始运算符,开发者可以针对这些运算符构建完整的 PyTorch 后端。这大大降低了编写 PyTorch 功能或后端的难度。
  • TorchInductor 是一种深度学习编译器,可以为多个加速器和后端生成快速代码。对于 NVIDIA 和 AMD GPU,它使用 OpenAI Triton 作为关键构建块。

TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 是用 Python 编写的,支持动态形状(即能够发送不同大小的张量而不会引起重新编译),使它们灵活、易于破解,并降低了开发人员和供应商的入门门槛。

为了验证这些技术,我们使用了跨越各种机器学习领域的 163 个开源模型的各种集合。我们精心构建了此基准测试,其中包括图像分类、目标检测、图像生成、各种 NLP 任务(如语言建模、问答、序列分类、推荐系统和强化学习)。我们将基准测试分为三类

  • 来自 HuggingFace Transformers 的 46 个模型
  • 来自 TIMM 的 61 个模型:Ross Wightman 收集的一组最先进的 PyTorch 图像模型
  • 来自 TorchBench 的 56 个模型:从整个 github 收集的一组经过精心挑选的流行代码库

我们不会修改这些开源模型,只是添加了用 torch.compile 调用包装它们的代码。

然后,我们测量加速并验证这些模型的准确性。由于加速可能会依赖于数据类型,因此我们在 float32 和自动混合精度 (AMP) 上测量加速。我们报告一个不均匀的加权平均加速,即 *0.75 * AMP + 0.25 * float32*,因为我们发现 AMP 在实践中更为常见。

在这些 163 个开源模型中,torch.compile 在 93% 的时间内有效,并且模型在 NVIDIA A100 GPU 上的训练速度提高了 43%。在 Float32 精度下,平均速度提高了 21%,在 AMP 精度下,平均速度提高了 51%。

注意事项: 在 NVIDIA 3090 等台式机级 GPU 上,我们测量的速度提升低于 A100 等服务器级 GPU 上的速度提升。截至今天,我们的默认后端 TorchInductor 支持 CPU 和 NVIDIA Volta 和 Ampere GPU。它(尚未)支持其他 GPU、xPU 或较旧的 NVIDIA GPU。

在 NVIDIA A100 GPU 上,torch.compile 相对于渴望模式的加速

试试看: torch.compile 处于开发的早期阶段。从今天开始,你可以在 nightly 二进制文件中试用 torch.compile。我们预计将在 2023 年 3 月初发布第一个稳定的 2.0 版本。

在 PyTorch 2.x 的路线图中,我们希望在性能和可扩展性方面进一步推动编译模式。我们今天在大会上谈到了一些正在进行的工作。其中一些工作尚未开始。其中一些工作是我们希望看到的,但我们没有带宽自行完成。如果你有兴趣做出贡献,请在本月开始的 **与工程师对话:2.0 直播问答系列** 中与我们聊天(有关详细信息,请参阅本文末尾)或通过 Github / 论坛与我们聊天。

推荐

以下是一些 PyTorch 用户对我们新方向的评价

Sylvain Gugger,**HuggingFace Transformers 的主要维护者**

“只需添加一行代码,PyTorch 2.0 就能使 Transformers 模型的训练速度提高 1.5 倍到 2 倍。这是自引入混合精度训练以来最激动人心的事件!”

Ross Wightman,**TIMM 的主要维护者**(PyTorch 生态系统中最大的视觉模型中心之一)

“它可以直接用于大多数 TIMM 模型的推理和训练工作负载,无需代码更改。”

Luca Antiga,**Lightning AI 的首席技术官**,也是 **PyTorch Lightning 的主要维护者之一**

“PyTorch 2.0 体现了深度学习框架的未来。能够在几乎没有用户干预的情况下捕获 PyTorch 程序,并直接获得巨大的设备速度提升和程序操作,这为 AI 开发人员打开了全新的维度。”

动机

我们对 PyTorch 的理念一直是将灵活性和可破解性作为我们的首要任务,并将性能作为紧随其后的第二要务。我们努力实现

  1. 高性能渴望执行
  2. Pythonic 内部
  3. 针对分布式、自动微分、数据加载、加速器等的良好抽象

自从我们于 2017 年推出 PyTorch 以来,硬件加速器(如 GPU)的计算速度提高了约 15 倍,内存访问速度提高了约 2 倍。因此,为了保持渴望执行的高性能,我们不得不将 PyTorch 内部的大量部分迁移到 C++。将内部迁移到 C++ 使它们的可破解性降低,并增加了代码贡献的入门门槛。

从一开始,我们就知道渴望执行的性能限制。在 2017 年 7 月,我们启动了第一个关于为 PyTorch 开发编译器的研究项目。编译器需要使 PyTorch 程序变快,但不能以牺牲 PyTorch 体验为代价。我们的关键标准是保留某些类型的灵活性——支持动态形状和动态程序,研究人员在探索的各个阶段都会使用它们。

技术概述

多年来,我们在 PyTorch 中构建了多个编译器项目。让我们将编译器分解为三个部分

  • 图获取
  • 图降低
  • 图编译

在构建 PyTorch 编译器时,图获取是最困难的挑战。

在过去的 5 年中,我们构建了 torch.jit.trace、TorchScript、FX 跟踪、延迟张量。但没有一个让我们感到满意。有些灵活但速度不快,有些速度快但不够灵活,有些既不快也不灵活。有些用户体验很糟糕(例如静默错误)。虽然 TorchScript 很 promising,但它需要对你的代码和你代码所依赖的代码进行实质性更改。这种对代码进行实质性更改的需要,让很多 PyTorch 用户望而却步。

PyTorch 编译过程

TorchDynamo:可靠快速地获取图

今年早些时候,我们开始研究 TorchDynamo,这是一种使用 CPython 在 PEP-0523 中引入的功能(称为帧评估 API)的方法。我们采用数据驱动的方法来验证其在图捕获方面的有效性。我们使用了 7,000 多个用 PyTorch 编写的 Github 项目作为验证集。虽然 TorchScript 和其他方法在大多数情况下甚至无法获取 50% 的图,而且开销很大,但 TorchDynamo 99% 的情况下都能正确、安全地获取图,并且开销可以忽略不计——无需对原始代码进行任何更改。这是我们知道,在灵活性和速度方面,我们终于突破了多年来一直困扰我们的障碍。

TorchInductor:使用 define-by-run IR 进行快速代码生成

对于 PyTorch 2.0 的一个新的编译器后端,我们从用户编写高性能自定义内核的方式中汲取了灵感:越来越多地使用 Triton 语言。我们还希望编译器后端使用类似于 PyTorch 渴望的抽象,并且通用性足以支持 PyTorch 中的各种功能。TorchInductor 使用一个类似 Python 的 define-by-run 循环级别 IR,将 PyTorch 模型自动映射到 GPU 上生成的 Triton 代码和 CPU 上的 C++/OpenMP 代码中。TorchInductor 的核心循环级别 IR 仅包含约 50 个运算符,它是在 Python 中实现的,使其易于 hack 和扩展。

AOTAutograd:重复使用 Autograd 进行预先编译的图

对于 PyTorch 2.0,我们知道我们希望加速训练。因此,至关重要的是,我们不仅要捕获用户级代码,还要捕获反向传播。此外,我们知道我们希望重复使用现有的经过实战检验的 PyTorch autograd 系统。AOTAutograd 利用 PyTorch 的 torch_dispatch 可扩展性机制来跟踪我们的 Autograd 引擎,允许我们“预先”捕获反向传播。这使我们能够使用 TorchInductor 加速我们的前向和反向传播。

PrimTorch:稳定的原始运算符

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

2000 多个 PyTorch 运算符的细分

因此,编写后端或横切功能变成了一项令人精疲力竭的工作。在 PrimTorch 项目中,我们正在努力定义更小、更稳定的运算符集。PyTorch 程序可以始终如一地降低到这些运算符集。我们的目标是定义两个运算符集

  • Prim 运算符,约有 250 个运算符,这些运算符相当低级。它们适合编译器,因为它们足够低级,以至于你需要将它们融合在一起才能获得良好的性能。
  • ATen 运算符,约有 750 个规范运算符,适合按原样导出。它们适合已经集成到 ATen 级别的后端或不会编译以从 Prim 运算符之类的低级运算符集中恢复性能的后端。

我们将在下面的开发者/供应商体验部分讨论有关此主题的更多信息。

用户体验

我们引入了简单的函数 torch.compile,它包装你的模型并返回编译后的模型。

compiled_model = torch.compile(model)

compiled_model 持有一个对你的模型的引用,并将 forward 函数编译成一个更优化的版本。在编译模型时,我们会提供一些旋钮来调整它。

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 指定编译器在编译时应优化什么。

    • 默认模式是一种预设,它试图在编译效率和编译时间或使用额外内存之间取得平衡。
    • 其他模式,例如 reduce-overhead,可以大幅减少框架开销,但会花费少量额外的内存。 max-autotune 会花费很长时间进行编译,试图为你提供它能够生成的最快代码。
  • dynamic 指定是否启用动态形状的代码路径。某些编译器优化无法应用于动态形状的程序。明确说明你是否想要一个具有动态形状或静态形状的编译程序,将有助于编译器为你提供更好的优化代码。
  • fullgraph 类似于 Numba 的 nopython。它将整个程序编译成一个单一的图,或者给出解释为什么无法这样做的错误。大多数用户不需要使用此模式。如果你非常注重性能,那么你可以尝试使用它。
  • backend 指定要使用的编译器后端。默认情况下,使用 TorchInductor,但还有一些其他后端可用。

编译体验旨在以默认模式提供最大的收益和最大的灵活性。以下是你在每种模式下能获得什么的思维模型。

现在,让我们看一个编译真实模型并运行它的完整示例(使用随机数据)。

import torch
import torchvision.models as models

model = models.resnet18().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
compiled_model = torch.compile(model)

x = torch.randn(16, 3, 224, 224).cuda()
optimizer.zero_grad()
out = compiled_model(x)
out.sum().backward()
optimizer.step()

第一次运行 compiled_model(x) 时,它会编译模型。因此,运行时间会更长。后续运行速度很快。

模式

编译器有一些预设,可以以不同的方式调整编译后的模型。你可能正在运行一个由于框架开销而速度很慢的小模型。或者,你可能正在运行一个几乎无法容纳在内存中的大型模型。根据你的需求,你可能想使用不同的模式。

# API NOT FINAL
# default: optimizes for large models, low compile-time
#          and no extra memory usage
torch.compile(model)

# reduce-overhead: optimizes to reduce the framework overhead
#                and uses some extra memory. Helps speed up small models
torch.compile(model, mode="reduce-overhead")

# max-autotune: optimizes to produce the fastest model,
#               but takes a very long time to compile
torch.compile(model, mode="max-autotune")

读取和更新属性

访问模型属性的工作方式与渴望模式相同。你可以像往常一样访问或修改模型的属性(例如 model.conv1.weight)。这在代码校正方面是完全安全和可靠的。TorchDynamo 将保护机制插入代码以检查其假设是否成立。如果属性以某种方式发生更改,那么 TorchDynamo 会知道何时需要自动重新编译。

# optimized_model works similar to model, feel free to access its attributes and modify them
optimized_model.conv1.weight.fill_(0.01)

# this change is reflected in model

钩子

模块和张量 钩子 目前还不完全有效,但最终会随着我们完成开发而有效。

序列化

你可以序列化 optimized_modelmodel 的状态字典。它们指向相同的参数和状态,因此是等效的。

torch.save(optimized_model.state_dict(), "foo.pt")
# both these lines of code do the same thing
torch.save(model.state_dict(), "foo.pt")

你目前无法序列化 optimized_model。如果你想直接保存对象,请改为保存 model

torch.save(optimized_model, "foo.pt") # Error
torch.save(model, "foo.pt")           # Works

推理和导出

对于模型推理,在使用 torch.compile 生成编译后的模型后,请在实际模型服务之前运行一些预热步骤。这有助于缓解初始服务期间的延迟峰值。

此外,我们将引入一种名为 torch.export 的模式,该模式会仔细导出整个模型和保护基础设施,以用于需要保证和可预测延迟的环境。 torch.export 需要对你的程序进行更改,尤其是在你具有数据依赖的控制流的情况下。

# API Not Final
exported_model = torch._dynamo.export(model, input)
torch.save(exported_model, "foo.pt")

这还处于开发的早期阶段。在 PyTorch 大会上观看有关导出路径的演讲,以了解更多详细信息。你也可以在本月开始的“向工程师提问:2.0 直播问答系列”中参与有关此主题的讨论(更多详细信息请参见本文末尾)。

调试问题

编译模式是不透明的,难以调试。你可能会遇到以下问题

  • 为什么我的程序在编译模式下崩溃?
  • 编译模式与渴望模式一样精确吗?
  • 为什么我没有看到加速?

如果编译模式产生错误或崩溃,或者与渴望模式产生不同的结果(超出机器精度限制),那么很可能不是你的代码的错误。但是,了解导致 bug 的代码片段非常有用。

为了帮助调试和可重复性,我们创建了一些工具和日志记录功能,其中一个脱颖而出:缩减器。

缩减器会自动将你遇到的问题缩减为一小段代码。这段代码会重现原始问题,你可以使用缩减后的代码提交 github 问题。这将有助于 PyTorch 团队轻松快速地修复问题。

如果你没有看到你期望的加速,那么我们有 torch._dynamo.explain 工具,它解释了你的代码的哪些部分导致了我们所说的“图中断”。图中断通常会阻碍编译器加速代码,减少图中断的数量很可能会加速你的代码(直到收益递减的某个极限)。

你可以在我们的 故障排除指南 中阅读有关这些内容以及更多内容的信息。

动态形状

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

截至今天,对动态形状的支持有限,并且正在快速进行中。它将在稳定版本发布之前完全具备功能。它被一个 dynamic=True 参数所控制,并且我们对一个功能分支(symbolic-shapes)有更多进展,在这个分支中,我们已经成功地使用 TorchInductor 在训练中运行 BERT_pytorch,具有完整的符号形状。对于使用动态形状进行推理,我们的覆盖范围更广。例如,让我们看一下动态形状有帮助的常见设置——使用语言模型进行文本生成。

我们可以看到,即使形状动态地从 4 变化到 256,编译模式也能够始终如一地比渴望模式快 40%。在没有动态形状支持的情况下,常见的解决方法是填充到最接近的 2 的幂。但是,正如我们从下面的图表中看到的那样,这会导致大量的性能开销,并且还会导致编译时间显著增加。此外,填充有时并非易事。

通过在 PyTorch 2.0 的编译模式中支持动态形状,我们可以获得最佳的性能和易用性。

目前的工作发展非常迅速,在我们对基础设施进行基本改进时,我们可能会暂时让一些模型出现退化。有关我们对动态形状的进展的最新更新,可以 在此处找到。

分布式

总而言之,torch.distributed 的两个主要分布式包装器在编译模式下运行良好。

DistributedDataParallel (DDP) 和 FullyShardedDataParallel (FSDP) 都在编译模式下工作,并且相对于渴望模式提供了更好的性能和内存利用率,但有一些注意事项和限制。

AMP 精度的加速
左:编译模式下 FSDP 相对于渴望模式的加速(AMP 精度)。
右:编译模式下 FSDP 比渴望模式占用的内存少得多。

DistributedDataParallel (DDP)

DDP 依赖于与反向计算重叠的 AllReduce 通信,并将较小的每层 AllReduce 操作分组到“桶”中,以提高效率。TorchDynamo 编译的 AOTAutograd 函数在与 DDP 结合使用时会阻止通信重叠,但这可以通过为每个“桶”编译单独的子图并允许通信操作在子图之外和之间发生来恢复性能。已编译模式下的 DDP 支持目前也需要 static_graph=False。有关 DDP + TorchDynamo 的方法和结果的更多详细信息,请参阅 这篇文章

全分片数据并行 (FSDP)

FSDP 本身是 PyTorch 的“beta”功能,由于能够调整哪些子模块被封装以及通常有更多的配置选项,因此比 DDP 具有更高等级的系统复杂性。如果配置了 use_original_params=True 标志,FSDP 可以与 TorchDynamo 和 TorchInductor 协同工作以应对各种流行模型。目前预计某些模型或配置存在兼容性问题,但会积极改进,如果提交 github 问题,可以优先考虑特定模型。

用户指定一个 auto_wrap_policy 参数来指示其模型的哪些子模块应一起封装在一个用于状态分片的 FSDP 实例中,或者手动将子模块封装在 FSDP 实例中。例如,许多转换器模型在每个“转换器块”都封装在一个单独的 FSDP 实例中时运行良好,因此一次只需实例化一个转换器块的完整状态。Dynamo 将在每个 FSDP 实例的边界处插入图断点,以允许正向(和反向)中的通信操作在图之外并与计算并行发生。

如果使用 FSDP 而不将子模块封装在单独的实例中,它将退回到类似于 DDP 的操作,但没有分桶。因此所有梯度都在一次操作中减少,即使在 Eager 中也无法实现计算/通信重叠。此配置仅在 TorchDynamo 中进行了功能测试,而没有进行性能测试。

开发人员/供应商体验

使用 PyTorch 2.0,我们希望简化后端(编译器)集成体验。为此,我们专注于减少操作符数量简化构建 PyTorch 后端所需操作符集的语义。

以图形形式,PT2 堆栈看起来像

从图表的中间开始,AOTAutograd 以提前的方式动态捕获 autograd 逻辑,生成一个包含 FX 图格式的前向和反向操作符的图。

我们提供了一组经过强化分解(即用其他操作符编写的操作符实现),可以用来减少后端需要实现的操作符数量。我们还通过称为功能化的过程,简化PyTorch 操作符的语义,包括有选择地重写复杂的 PyTorch 逻辑(包括变异和视图),以及保证操作符元数据信息(例如形状传播公式)。这项工作正在积极进行中;我们的目标是提供一个原始稳定的操作符集(约 250 个),称为PrimTorch,供应商可以利用(即选择加入)以简化他们的集成。
在减少和简化操作符集之后,后端可以选择在 Dynamo(即中间层,紧接在 AOTAutograd 之后)或 Inductor(底层)进行集成。我们在下面描述了做出此选择的一些考虑因素,以及围绕混合后端的未来工作。

Dynamo 后端

拥有现有编译器堆栈的供应商可能会发现将自己集成到 TorchDynamo 后端最为容易,从而获得一个以 ATen/Prims IR 表示的 FX 图。请注意,对于训练和推理,集成点都应该紧接在 AOTAutograd 之后,因为我们目前将分解应用于 AOTAutograd 的一部分,并且如果针对推理,只会跳过特定于反向的步骤。

Inductor 后端

供应商还可以将其后端直接集成到 Inductor 中。Inductor 接受由 AOTAutograd 生成的由 ATen/Prim 操作组成的图,并将其进一步降低到循环级 IR。如今,Inductor 为逐点、归约、散布/收集和窗口操作提供降低到其循环级 IR 的方法。此外,Inductor 创建融合组,进行索引简化、维度折叠,并调整循环迭代顺序以支持高效的代码生成。然后,供应商可以通过提供从循环级 IR 到特定于硬件的代码的映射来进行集成。目前,Inductor 有两个后端:(1) 生成多线程 CPU 代码的 C++,(2) 生成高效 GPU 代码的 Triton。这些 Inductor 后端可以用作替代后端的灵感来源。

混合后端接口(即将推出)

我们已经构建了用于将 FX 图划分为包含后端支持的操作符的子图并在 Eager 模式下执行其余部分的实用程序。这些实用程序可以扩展为支持“混合后端”,配置要为哪个后端运行图的哪些部分。但是,还没有用于后端公开其操作符支持、对操作符模式的偏好等的稳定接口或契约。这仍然是正在进行的工作,我们欢迎早期采用者的反馈。

最后的想法

我们对 PyTorch 2.0 及更高版本所采取的方向感到非常兴奋。通往最终 2.0 版本的道路将会很艰难,但请尽早加入我们的旅程。如果您有兴趣深入研究或贡献编译器,请继续阅读以下内容,其中包含有关如何开始的信息(例如,教程、基准测试、模型、常见问题解答)以及工程师问答:2.0 直播问答系列,本月开始。其他资源包括

使用 PyTorch 2.0 加速 Hugging Face 和 TIMM 模型

作者:Mark Saroufim

torch.compile() 使得通过一个简单的行装饰器 torch.compile() 轻松尝试不同的编译器后端,以使 PyTorch 代码更快。它可以直接在 nn.Module 上运行,作为 torch.jit.script() 的直接替代,但不需要您进行任何源代码更改。我们预计此单行代码更改将为您提供大多数您正在运行的模型的 30%-2x 训练时间加速。

opt_module = torch.compile(module)

torch.compile 支持任意 PyTorch 代码、控制流、变异,并附带对动态形状的实验支持。我们对这一发展感到非常兴奋,因此将其称为 PyTorch 2.0。

对我们来说,这个公告的不同之处在于,我们已经对一些最流行的开源 PyTorch 模型进行了基准测试,并获得了实质性的加速,范围从 30% 到 2x https://github.com/pytorch/torchdynamo/issues/681

这里没有技巧,我们已经使用 pip 安装了流行的库,例如 https://github.com/huggingface/transformershttps://github.com/huggingface/acceleratehttps://github.com/rwightman/pytorch-image-models,然后对它们运行了 torch.compile(),就是这样。

在性能和便利性方面都能取得成功非常难得,但这就是核心团队为何对 PyTorch 2.0 感到如此兴奋的原因。

需求

对于 GPU(新一代 GPU 将看到显著的性能提升)

pip3 install numpy --pre torch --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu117

对于 CPU

pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu

可选:验证安装

git clone https://github.com/pytorch/pytorch
cd tools/dynamo
python verify_dynamo.py

可选:Docker 安装

我们还提供 PyTorch nightly 二进制文件中所需的所有依赖项,您可以使用以下命令下载:

docker pull ghcr.io/pytorch/pytorch-nightly

对于临时实验,请确保您的容器可以访问所有 GPU

docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash

入门

请阅读 Mark Saroufim 的 完整博文,他将在其中引导您完成教程和真实模型,以便您今天尝试 PyTorch 2.0。

我们使用 PyTorch 的目标是构建一个广度优先编译器,它可以加速人们在开源中运行的大多数实际模型。Hugging Face Hub 最终成为我们非常宝贵的基准测试工具,确保我们进行的任何优化都能真正帮助加速人们想要运行的模型。

博客教程将向您展示如何精确地复制这些加速效果,以便您像我们一样对 PyTorch 2.0 感到兴奋。因此,请尝试 PyTorch 2.0,享受免费性能,如果您没有看到性能提升,请打开一个问题,我们将确保您的模型得到支持 https://github.com/pytorch/torchdynamo/issues

毕竟,除非您的模型确实运行得更快,否则我们不能声称我们创建了一个广度优先的编译器。

常见问题解答

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

  2. 为什么是 2.0 而不是 1.14?
    PyTorch 2.0 本来应该是 1.14。我们发布了大量我们认为会改变您使用 PyTorch 的方式的新功能,因此我们将其称为 2.0。

  3. 如何安装 2.0?还有其他要求吗?

    安装最新的 nightly 版本

    CUDA 11.8

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

    CUDA 11.7

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

    CPU

     pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --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 默认启用吗?
    2.0 是版本名称。torch.compile 是在 2.0 中发布的功能,您需要显式使用 torch.compile。

  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. 为什么我应该使用 PT2.0 而不是 PT 1.X?
    请参阅问题 (2) 的答案。

  8. 运行 PyTorch 2.0 时我的代码做了哪些不同的操作?
    开箱即用,PyTorch 2.0 与 PyTorch 1.x 相同,您的模型在 eager 模式下运行,即每行 Python 代码都按顺序执行。
    在 2.0 中,如果您将您的模型封装在 model = torch.compile(model) 中,您的模型将在执行之前经过 3 个步骤
    1. 图获取:首先,模型被重写为子图块。可以通过 TorchDynamo 编译的子图将被“扁平化”,而其他子图(可能包含控制流代码或其他不受支持的 Python 结构)将退回到 Eager 模式。
    2. 图降低:所有 PyTorch 操作都将分解为其特定于所选后端的组成内核。
    3. 图编译,内核在其中调用其相应的低级设备特定操作。
  9. PT2.0 向 PT 添加了哪些新组件?
    • TorchDynamo 从 Python 字节码生成 FX 图。它使用 保护 保持 eager 模式功能,以确保生成的图是有效的 (了解更多)
    • AOTAutograd 用于生成与 TorchDynamo 捕获的前向图相对应的反向图 (了解更多)
    • PrimTorch 用于将复杂的 PyTorch 操作分解为更简单、更基本的运算 (了解更多)
    • [后端] 后端与 TorchDynamo 集成,将图编译成可在加速器上运行的 IR。例如,TorchInductor 将图编译成 Triton 用于 GPU 执行或 OpenMP 用于 CPU 执行(了解更多)。
  10. 2.0 目前支持哪些编译器后端?
    默认且最完整的后端是 TorchInductor,但 TorchDynamo 拥有不断增长的后端列表,可以通过调用 torchdynamo.list_backends() 来查看。

  11. 2.0 如何进行分布式训练?
    在编译模式下,DDP 和 FSDP 比 FP32 中的 Eager 模式快 15%,在 AMP 精度下快 80%。PT2.0 进行了一些额外的优化,以确保 DDP 的通信-计算重叠与 Dynamo 的部分图创建良好地配合。请确保使用 static_graph=False 运行 DDP。更多详情请查看 这里

  12. 如何了解更多关于 PT2.0 的开发信息?
    PyTorch 开发者论坛 上,你可以直接从开发人员那里了解 2.0 的组件。

  13. 我的代码在 2.0 的编译模式下运行速度变慢了!
    性能下降最可能的原因是图中断过多。例如,模型前向函数中一个简单的 print 语句就会触发图中断。我们有一些方法可以诊断这些问题 - 更多内容请查看 这里

  14. 我的代码在 2.0 的编译模式下崩溃了!如何调试它?
    以下是一些技术,可以帮助你找出代码可能出错的地方并打印有用的日志:https://pytorch.ac.cn/docs/stable/torch.compiler_faq.html#why-is-my-code-crashing

向工程师提问:2.0 直播问答系列

我们将举办一系列直播问答活动,让社区成员有机会向专家提出更深入的问题并进行交流。请随时回来查看全年的主题日历。如果你无法参加:1)这些活动将被录制,以便日后观看;2)你也可以参加我们每周五上午 10 点(太平洋标准时间)的开发基础设施办公室时间 @ https://github.com/pytorch/pytorch/wiki/Dev-Infra-Office-Hours

请点击 这里 查看日期、时间、描述和链接。

免责声明:请勿在参加直播活动或提交问题时分享你的个人信息、姓氏或公司名称。

主题 主持人
使用 2.0 的全新开发者体验(安装、设置、克隆示例、使用 2.0 运行) Suraj Subramanian
LinkedIn | Twitter
PT2 分析和调试 Bert Maher
LinkedIn | Twitter
深入探讨 TorchInductor 和 PT 2.0 后端集成 Natalia Gimelshein、Bin Bao 和 Sherlock Huang
Natalia Gimelshein
LinkedIn
Sherlock Huang
LinkedIn
在没有 C++ 和 functorch 的情况下扩展 PyTorch:适用于 PyTorch 的类似 JAX 的可组合函数变换 Anjali Chourdia 和 Samantha Andow
Anjali Chourdia
LinkedIn | Twitter
Samantha Andow
LinkedIn | Twitter
深入探讨 TorchDynamo Michael Voznesensky
LinkedIn
用 TorchData 重新思考数据加载:Datapipes 和 Dataloader2 Kevin Tse
LinkedIn
可组合训练(+ torcheval、torchsnapshot) Ananth Subramaniam
如何以及为何为 PyTorch 贡献代码和教程 Zain Rizvi、Svetlana Karslioglu 和 Carl Parker
Zain Rizvi
LinkedIn | Twitter
Svetlana Karslioglu
LinkedIn | Twitter
动态形状和计算最大批次大小 Edward Yang 和 Elias Ellison
Edward Yang
Twitter
PyTorch 2.0 导出:完整的 PyTorch 图捕获 Michael Suo 和 Yanan Cao
Yanan Cao
LinkedIn
使用 DistributedTensor 和 PyTorch DistributedTensor 进行 2D 并行化 Wanchao Liang 和 Alisson Gusatti Azzolini
Wanchao Liang
LinkedIn | Twitter
Alisson Gusatti Azzolini
LinkedIn
TorchRec 和 FSDP 在生产中的应用 Dennis van der Staay、Andrew Gu 和 Rohan Varma
Dennis van der Staay
LinkedIn
Rohan Varma
LinkedIn | Twitter
PyTorch 设备端未来 Raziel Alvarez Guevara
LinkedIn | Twitter
TorchMultiModal
简介博客
扩展博客
Kartikay Khandelwal
LinkedIn | Twitter
BetterTransformers(+ 与 Hugging Face 的集成)、模型服务和优化
博客 1
Github
Hamid Shojanazeri 和 Mark Saroufim
Mark Saroufim
LinkedIn | Twitter
PT2 和分布式 Will Constable
LinkedIn

观看 PyTorch 大会上的演讲