跳转到主要内容

入门

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

概览

PyTorch 2.0 来了,这是 PyTorch 下一代 2.x 系列发布的第一步。在过去几年中,我们从 PyTorch 1.0 不断创新和迭代,直至最新的 1.13 版本,并已迁移到新成立的 PyTorch 基金会,该基金会是 Linux 基金会的一部分。

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

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

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 的自动求导引擎作为跟踪自动微分器,用于生成提前(ahead-of-time)反向传播跟踪。

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

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

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

在 PyTorch 2.x 的路线图中,我们希望在性能和可伸缩性方面进一步推动编译模式。其中一些工作正在进行中,正如我们今天在大会上谈到的那样。其中一些工作尚未开始。其中一些工作是我们希望看到的,但我们没有精力自己完成。如果您有兴趣做出贡献,请在本月开始的“工程师问答:2.0 现场问答系列”(详情见本文末尾)和/或通过 Github / 论坛与我们交流。

推荐语

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

Sylvain GuggerHuggingFace Transformers 的主要维护者

“只需添加一行代码,PyTorch 2.0 就能在训练 Transformer 模型时带来 1.5 倍到 2 倍的速度提升。这是自混合精度训练推出以来最令人兴奋的事情!”

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

“它开箱即用,无需更改代码即可与大多数 TIMM 模型进行推理和训练工作。”

Luca AntigaLightning AI 的 CTO,以及 PyTorch Lightning 的主要维护者之一

“PyTorch 2.0 体现了深度学习框架的未来。在几乎无需用户干预的情况下捕获 PyTorch 程序,并开箱即用地获得巨大的设备端加速和程序操作,为 AI 开发者开启了全新的维度。”

动机

我们对 PyTorch 的理念始终是把灵活性和可修改性放在首位,性能紧随其后。我们一直追求:

  1. 高性能 Eager 执行
  2. Python 风格的内部实现
  3. 对分布式、自动微分、数据加载、加速器等方面的良好抽象

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

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

技术概览

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

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

在构建 PyTorch 编译器时,图获取是更具挑战性的任务。

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

PyTorch 编译过程

TorchDynamo:可靠快速地获取图

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

TorchInductor:使用即时定义 IR 快速生成代码

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

AOTAutograd:重用自动求导以实现提前图

对于 PyTorch 2.0,我们知道我们想要加速训练。因此,不仅要捕获用户级代码,还要捕获反向传播至关重要。此外,我们知道我们想要重用现有的经过实战检验的 PyTorch 自动求导系统。AOTAutograd 利用 PyTorch 的 torch_dispatch 可扩展机制来跟踪我们的自动求导引擎,使我们能够“提前”捕获反向传播。这使我们能够使用 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")

读取和更新属性

访问模型属性的方式与在 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

钩子

Module 和 Tensor 钩子目前尚未完全发挥作用,但随着我们开发完成,它们最终将发挥作用。

序列化

您可以序列化 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。

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 模式占用更少的内存。

DistributedDataParallel (DDP)

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

FullyShardedDataParallel (FSDP)

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

用户通过 auto_wrap_policy 参数指定模型的哪些子模块应该在 FSDP 实例中一起包装以用于状态分片,或者手动将子模块包装到 FSDP 实例中。例如,许多 Transformer 模型在每个“Transformer 块”都被包装在单独的 FSDP 实例中时工作良好,因此只需要同时实例化一个 Transformer 块的完整状态。Dynamo 将在每个 FSDP 实例的边界插入图中断,以允许前向(和反向)中的通信操作在图外部并行于计算发生。

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

开发者/供应商体验

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

以图形形式表示,PT2 栈如下:

从图的中间开始,AOTAutograd 动态地以提前(ahead-of-time)方式捕获自动求导逻辑,生成 FX 图格式的前向和反向算子图。

我们提供了一组经过强化的分解(即用其他算子实现的算子),可以利用它们来减少后端需要实现的算子数量。我们还通过一种称为“功能化”(functionalization)的过程,通过选择性地重写复杂的 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 图划分为包含后端支持的算子的子图,并急切地执行剩余部分。这些工具可以扩展以支持“后端混合”,配置图形的哪些部分由哪个后端运行。然而,目前还没有一个稳定的接口或约定供后端公开其算子支持、对算子模式的偏好等。这仍然是一项正在进行的工作,我们欢迎早期采用者的反馈。

总结

我们对 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,享受免费的性能提升,如果您没有看到,请提出问题,我们将确保您的模型得到支持 https://github.com/pytorch/torchdynamo/issues

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

常见问题

  1. 什么是 PT 2.0?
    2.0 是最新的 PyTorch 版本。PyTorch 2.0 提供了相同的 Eager Mode 开发体验,同时通过 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 Mode 下运行,即每一行 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 的部分图创建良好配合。请确保您运行 DDP 时设置 static_graph=False。更多详情请参阅此处

  12. 我如何了解更多关于 PT2.0 的发展?
    PyTorch 开发者论坛是直接向构建 2.0 组件的开发者了解 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
领英 | 推特
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
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
领英 | 推特
更优的 Transformer(+与 Hugging Face 集成)、模型服务和优化
博客 1
GitHub
Hamid Shojanazeri 和 Mark Saroufim
Mark Saroufim
领英 | 推特
PT2 与分布式 Will Constable
领英

观看 PyTorch 大会演讲