跳转到主要内容

入门

选择偏好并运行命令以在本地安装 PyTorch,或通过支持的云平台快速入门。

概述

隆重推出 PyTorch 2.0,这是我们迈向下一代 PyTorch 2 系列版本的第一步。在过去几年中,我们从 PyTorch 1.0 到最新的 1.13 不断创新和迭代,并迁移到新成立的 PyTorch 基金会,该基金会是 Linux 基金会的一部分。

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

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

PyTorch 2.x:更快、更符合 Python 风格,且一如既往地动态

今天,我们宣布 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 的自动求导引擎,作为跟踪自动求导,用于生成提前反向跟踪。

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

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

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

尝试一下: 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 的 CTO**,也是 **PyTorch Lightning 的主要维护者**之一

“PyTorch 2.0 体现了深度学习框架的未来。无需用户干预即可捕获 PyTorch 程序,并开箱即用地获得设备上的大规模加速和程序操作,这为 AI 开发人员开启了一个全新的维度。”

动机

我们对 PyTorch 的理念始终是将灵活性和可修改性放在首位,性能紧随其后。我们一直致力于

  1. 高性能即时执行
  2. Python 式内部结构
  3. 分布式、自动求导、数据加载、加速器等方面的良好抽象。

自 2017 年推出 PyTorch 以来,硬件加速器(如 GPU)的计算速度提高了约 15 倍,内存访问速度提高了约 2 倍。因此,为了保持即时执行的高性能,我们不得不将 PyTorch 内部的大部分内容移至 C++。将内部内容移至 C++ 会降低其可修改性,并增加代码贡献的门槛。

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

技术概述

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

  • 图获取
  • 图降级
  • 图编译

图获取是构建 PyTorch 编译器时更难的挑战。

在过去的 5 年里,我们构建了 torch.jit.trace、TorchScript、FX tracing、Lazy Tensors。但它们都未能提供我们想要的一切。有些灵活但不快,有些快但不灵活,有些既不快也不灵活。有些用户体验很差(例如悄无声息地出错)。虽然 TorchScript 很有前景,但它需要对您的代码以及您的代码所依赖的代码进行大量更改。这种对代码进行大量更改的需求使其对许多 PyTorch 用户来说无法启动。

PyTorch 编译过程

TorchDynamo:可靠快速地获取图

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

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

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

AOTAutograd:重用 Autograd 以进行提前图生成

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

PrimTorch:稳定的基本运算符

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

PyTorch 2000 多个运算符的细分

因此,编写后端或交叉功能会变得一项耗费精力的事情。在 PrimTorch 项目中,我们正在努力定义更小、更稳定的运算符集。PyTorch 程序可以始终被降级到这些运算符集。我们的目标是定义两个运算符集

  • Prim ops 包含大约 250 个运算符,这些运算符相当底层。它们适用于编译器,因为它们足够底层,需要将它们重新融合才能获得良好的性能。
  • ATen ops 包含大约 750 个规范运算符,适用于按原样导出。它们适用于已经在 ATen 级别集成的后端,或无法通过编译从更底层的运算符集(如 Prim ops)恢复性能的后端。

我们将在下面的“开发者/供应商体验”部分讨论更多关于此主题的内容

用户体验

我们引入了一个简单的函数 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")

读取和更新属性

访问模型属性的工作方式与在 Eager 模式下相同。您可以像往常一样访问或修改模型的属性(例如 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 现场问答系列”中就此主题进行交流(更多详细信息见本文末尾)。

调试问题

编译模式不透明且难以调试。您会遇到以下问题:

  • 为什么我的程序在编译模式下崩溃了?
  • 编译模式是否与 Eager 模式一样准确?
  • 为什么我没有看到加速?

如果编译模式产生错误、崩溃或与 Eager 模式产生不同的结果(超出机器精度限制),则很可能不是您的代码的错。然而,了解导致此错误的代码片段是很有用的。

为了帮助调试和重现,我们创建了多种工具和日志记录功能,其中一个突出的是:Minifier。

精简器会自动将您看到的问题缩小到一小段代码。这段小代码片段重现了原始问题,您可以将精简后的代码提交到 GitHub issue。这将帮助 PyTorch 团队轻松快速地解决问题。

如果您没有看到预期的加速,我们有 **torch._dynamo.explain** 工具,它可以解释您的代码的哪些部分导致了我们所说的“图中断”。图中断通常会阻碍编译器加速代码,减少图中断的数量可能会加速您的代码(达到某种递减的限制)。

您可以在我们的故障排除指南中阅读更多相关信息。

动态形状

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

截至目前,对动态形状的支持是有限的,并且正在快速进展中。它将在稳定版本中全面实现。它被 dynamic=True 参数所限制,我们在一个功能分支 (symbolic-shapes) 上取得了更多进展,在该分支上我们已成功使用 TorchInductor 运行 BERT_pytorch 训练,并支持完整的符号形状。对于动态形状的推理,我们有更多的覆盖范围。例如,让我们看看动态形状有帮助的常见设置——使用语言模型进行文本生成。

我们可以看到,即使形状从 4 动态变化到 256,编译模式仍能持续优于 Eager 模式,最高可达 40%。如果没有动态形状支持,常见的解决方法是填充到最接近的二次幂。然而,正如我们从下面的图表中看到的那样,这会带来大量的性能开销,并且还会导致显著更长的编译时间。此外,填充有时很难正确完成。

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

目前的工作进展非常迅速,随着我们对基础设施进行根本性改进,我们可能会暂时导致一些模型退化。有关动态形状最新进展的更新,请点击此处

分布式

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

DistributedDataParallel (DDP) 和 FullyShardedDataParallel (FSDP) 均可在编译模式下工作,并提供相对于 Eager 模式改进的性能和内存利用率,但有一些注意事项和限制。

AMP 精度下的加速
左:FSDP 在编译模式下相对于 eager 模式的加速(AMP 精度)。
右:FSDP 在编译模式下比 eager 模式占用更少的内存。

分布式数据并行 (DDP)

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

完全分片数据并行 (FSDP)

FSDP 本身是一个“beta”版 PyTorch 功能,由于能够调整哪些子模块被包装以及通常有更多配置选项,因此其系统复杂度高于 DDP。FSDP 适用于 TorchDynamo 和 TorchInductor,适用于各种流行模型,如果配置了 use_original_params=True 标志。目前,预计与特定模型或配置存在一些兼容性问题,但将积极改进,如果提交 GitHub 问题,可以优先处理特定模型。

用户指定 auto_wrap_policy 参数,以指示模型的哪些子模块在用于状态分片的 FSDP 实例中一起包装,或手动将子模块包装到 FSDP 实例中。例如,许多 Transformer 模型在每个“Transformer 块”都包装在单独的 FSDP 实例中时效果很好,因此一次只需要实现一个 Transformer 块的完整状态。Dynamo 将在每个 FSDP 实例的边界插入图中断,以允许前向(和后向)的通信操作在图外并行于计算进行。

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

开发者/供应商体验

通过 PyTorch 2.0,我们希望简化后端(编译器)集成体验。为此,我们专注于**减少操作符数量**和**简化操作符集的语义**,以启动 PyTorch 后端。

以图形形式表示,PT2 堆栈如下所示

从图表中间开始,AOTAutograd 以提前的方式动态捕获自动求导逻辑,以 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)C++,生成多线程 CPU 代码,(2)Triton,生成高性能 GPU 代码。这些 Inductor 后端可以作为替代后端的灵感。

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

我们已经构建了工具,可以将 FX 图划分为包含后端支持的操作符的子图,并急切地执行其余部分。这些工具可以扩展以支持“后端混合”,配置图形的哪些部分由哪个后端运行。然而,目前还没有稳定的接口或合同供后端公开其操作符支持、操作符模式偏好等。这仍是正在进行的工作,我们欢迎早期采用者提供反馈。

最终想法

我们对 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% 到 2 倍。

opt_module = torch.compile(module)

torch.compile 支持任意 PyTorch 代码、控制流、突变,并提供对动态形状的实验性支持。我们对这项进展感到非常兴奋,称之为 PyTorch 2.0。

对我们来说,这次公告的不同之处在于,我们已经对一些最流行的开源 PyTorch 模型进行了基准测试,并获得了 30% 到 2 倍的显著加速 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 每夜版二进制文件中提供了所有必需的依赖项,您可以通过以下方式下载

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,享受免费的性能,如果您没有看到它,请打开一个 issue,我们将确保您的模型得到支持 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?有什么额外要求吗?

    安装最新的每夜构建版

    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 图。它通过 guards 保持 Eager 模式功能,以确保生成的图有效(了解更多
    • AOTAutograd 用于生成与 TorchDynamo 捕获的正向图对应的反向图(了解更多)。
    • PrimTorch 将复杂的 PyTorch 操作分解为更简单、更基本的 ops (阅读更多)。
    • [后端] 后端与 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 的编译模式下运行速度变慢了!
    性能下降最可能的原因是图中断过多。例如,您模型 forward 中的一个无害的 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
领英 | 推特
PT2 性能分析和调试 Bert Maher
领英 | 推特
深度解析 TorchInductor 和 PT 2.0 后端集成 Natalia Gimelshein、Bin Bao 和 Sherlock Huang
Natalia Gimelshein
领英
Sherlock Huang
领英
无需 C++ 和 Functorch 扩展 PyTorch:类似 JAX 的可组合函数转换 PyTorch Anjali Chourdia 和 Samantha Andow
Anjali Chourdia
领英 | 推特
Samantha Andow
领英 | 推特
深入探讨 TorchDynamo Michael Voznesensky
领英
使用 TorchData 重新思考数据加载:Datapipes 和 Dataloader2 Kevin Tse
领英
可组合训练(+ torcheval, torchsnapshot) Ananth Subramaniam
如何以及为何向 PyTorch 贡献代码和教程 Zain Rizvi、Svetlana Karslioglu 和 Carl Parker
Zain Rizvi
领英 | 推特
Svetlana Karslioglu
领英 | 推特
动态形状与计算最大批次大小 Edward Yang 和 Elias Ellison
Edward Yang
推特
PyTorch 2.0 导出:PyTorch 的健全全图捕获 Michael Suo 和 Yanan Cao
曹亚楠
领英
使用 DistributedTensor 和 PyTorch DistributedTensor 进行二维并行 Wanchao Liang 和 Alisson Gusatti Azzolini
Wanchao Liang
领英 | 推特
Alisson Gusatti Azzolini
领英
TorchRec 和 FSDP 在生产环境中的应用 Dennis van der Staay, Andrew Gu 和 Rohan Varma
Dennis van der Staay
领英
Rohan Varma
领英 | 推特
PyTorch 设备端的未来 Raziel Alvarez Guevara
领英 | 推特
TorchMultiModal
介绍博客
扩展博客
Kartikay Khandelwal
领英 | 推特
BetterTransformers(+ 与 Hugging Face 集成)、模型服务和优化
博客 1
GitHub
Hamid Shojanazeri 和 Mark Saroufim
Mark Saroufim
领英 | 推特
PT2 和分布式 Will Constable
领英

观看 PyTorch 大会演讲