入门

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

概述

隆重推出 PyTorch 2.0,这是我们迈向下一代 PyTorch 2 系列版本的首要步骤。在过去几年中,我们从 PyTorch 1.0 不断创新迭代至最近的 1.13 版本,并加入新成立的 PyTorch 基金会(Linux 基金会的一部分)。

除了出色的社区之外,PyTorch 最大的优势在于我们始终保持一流的 Python 集成、命令式风格、简洁的 API 和多样化的选择。PyTorch 2.0 提供了相同的即时(eager-mode)开发模式和用户体验,同时在编译器底层根本性地改变并增强了 PyTorch 的运作方式。我们得以提供更快的性能,并支持动态形状(Dynamic Shapes)和分布式训练。

您将在下方找到了解 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 帧评估钩子(Frame Evaluation Hooks)安全地捕获 PyTorch 程序,这是一项重大创新,也是我们 5 年来在安全图捕获方面研发的成果。

  • AOTAutograd 重载了 PyTorch 的自动微分(autograd)引擎,作为一种用于生成提前(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。

在 NVIDIA A100 GPU 上 torch.compile 相较于即时模式(eager mode)的加速比

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

在 PyTorch 2.x 的路线图中,我们希望在性能和可扩展性方面进一步推动编译模式。正如我们今天在大会上谈到的那样,其中一些工作正在进行中。有些工作尚未开始。有些是我们希望看到但目前尚无人力去完成的工作。如果您有兴趣参与贡献,欢迎在从本月开始的“咨询工程师:2.0 直播问答系列”中与我们交流(详情见本文末尾),或者通过 Github/论坛联系我们。

推荐语

以下是部分 PyTorch 用户对我们新方向的评价:

Sylvain GuggerHuggingFace Transformers 的主要维护者

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

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

“它无需任何代码更改,就能直接与大多数 TIMM 模型进行推理和训练工作负载。”

Luca AntigaLightning AI 的 CTOPyTorch 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 内部构建了多个编译器项目。让我们将编译器分为三个部分:

  • 图捕获(graph acquisition)
  • 图降低(graph lowering)
  • 图编译(graph compilation)

在构建 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 即时模式抽象的编译器后端,并且足够通用,以支持 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 程序可以一致地降低到这些算子集。我们的目标是定义两个算子集:

  • 包含约 250 个算子的原始(Prim)算子集,它们非常底层。这些算子适用于编译器,因为它们足够底层,需要您将它们重新融合在一起以获得良好的性能。
  • 包含约 750 个规范算子的 ATen 算子集,适用于按原样导出。这些适用于已经在 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 在代码中插入保护(guards)来检查其假设是否成立。如果属性以某种方式发生变化,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 的 state-dict。它们指向相同的参数和状态,因此是等效的。

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 大会上关于 Export Path 的演讲以了解更多详细信息。您还可以参与本月开始的“咨询工程师:2.0 直播问答系列”讨论此主题(详情见本文末尾)。

调试问题

编译模式是不透明的,难以调试。您可能会有以下疑问:

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

如果编译模式产生错误、崩溃或结果与即时模式不一致(超出机器精度限制),那么这极大概率不是您代码的错。然而,了解哪段代码是导致 Bug 的原因是有用的。

为了辅助调试和重现,我们创建了几个工具和日志功能,其中一个脱颖而出:Minifier(简化器)。

Minifier 会自动将您看到的问题简化为一小段代码。这一小段代码重现了原始问题,您可以提交一个带有简化代码的 GitHub Issue。这将帮助 PyTorch 团队轻松快捷地修复问题。

如果您没有看到预期的加速效果,我们提供了 torch._dynamo.explain 工具,它可以解释您代码中哪些部分导致了所谓的“图中断(graph breaks)”。图中断通常会阻碍编译器加速代码,减少图中断的数量可能会加速您的代码(直到出现边际收益递减的极限)。

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

动态形状

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

截至今天,对动态形状的支持仍然有限,且正在快速改进中。到稳定版发布时,它将功能齐全。它被限制在 dynamic=True 参数之后,我们在功能分支(symbolic-shapes)上取得了更多进展,在该分支上我们已成功使用 TorchInductor 以完整的符号形状训练了 BERT_pytorch。对于具有动态形状的推理,我们有了更多的覆盖范围。例如,让我们看看动态形状有用的常见设置——语言模型的文本生成。

我们可以看到,即使形状从 4 动态变化到 256,编译模式也能持续超过即时模式高达 40%。在没有动态形状支持的情况下,一种常见的变通方法是填充(pad)到最接近的 2 的幂。然而,正如我们从下面的图表中看到的那样,这会产生大量的性能开销,并导致显著更长的编译时间。此外,正确执行填充有时并非易事。

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

当前的工作正在飞速演进,随着我们对基础设施进行根本性改进,可能会暂时让某些模型出现倒退。关于动态形状进展的最新更新可以在 这里 找到。

分布式

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

DistributedDataParallel (DDP) 和 FullyShardedDataParallel (FSDP) 在编译模式下均可工作,并相较于即时模式提供了更高的性能和内存利用率,但也存在一些警告和限制。

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

DistributedDataParallel (DDP)

DDP 依赖于将 AllReduce 通信与反向计算重叠,并将较小的逐层 AllReduce 操作分组为“存储桶(buckets)”以提高效率。当与 DDP 朴素结合时,TorchDynamo 编译的 AOTAutograd 函数会阻止通信重叠,但通过为每个“存储桶”编译单独的子图并允许通信操作在子图之外和之间进行,可以恢复性能。编译模式下的 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 的操作,但没有存储桶。因此,所有梯度在一次操作中被归约,即使在即时模式下也没有计算/通信重叠。此配置仅通过 TorchDynamo 进行了功能测试,未进行性能测试。

开发者/供应商体验

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

用图表形式表示,PT2 堆栈看起来像这样:

从图表的中间开始,AOTAutograd 以提前方式动态捕获自动微分逻辑,生成一个以 FX 图格式呈现的前向和反向算子图。

我们提供了一组强化后的分解(即以其他算子编写的算子实现),可以用来减少后端需要实现的算子数量。我们还通过称为函数化(functionalization)的过程选择性地重写复杂的 PyTorch 逻辑(包括变异和视图),以及保证形状传播公式等算子元数据信息,来简化 PyTorch 算子的语义。这项工作正在积极进行中;我们的目标是提供一组包含约 250 个具有简化语义的原始(primitive)稳定(stable)的算子,称为 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,享受免费的性能提升。如果您没有看到提升,请提交一个 Issue,我们将确保您的模型得到支持 https://github.com/pytorch/torchdynamo/issues

毕竟,除非您的模型运行得更快,否则我们不能声称我们创建了一个宽度优先的解决方案。

常见问题

  1. 什么是 PT 2.0?
    2.0 是最新的 PyTorch 版本。PyTorch 2.0 提供了相同的即时开发体验,同时通过 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)
    虽然加速主要在训练期间观察到,但如果您的模型运行速度比即时模式快,您也可以将其用于推理。
     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 相同,您的模型以即时模式运行,即每一行 Python 代码都依次执行。
    在 2.0 中,如果您用 model = torch.compile(model) 包裹模型,您的模型在执行前会经过 3 个步骤:
    1. 图获取:首先将模型重写为子图块。可以被 TorchDynamo 编译的子图会被“扁平化”,而其他子图(可能包含控制流代码或其他不支持的 Python 结构)将回退到即时模式。
    2. 图降低:所有 PyTorch 操作都被分解为特定于所选后端的组成内核。
    3. 图编译,内核调用其相应的底层设备特定操作。
  9. PT2.0 向 PT 添加了哪些新组件?
    • TorchDynamo 从 Python 字节码生成 FX 图。它使用 保护机制(guards) 保持即时模式功能,以确保生成的图有效 (阅读更多)。
    • AOTAutograd 用于生成对应于 TorchDynamo 捕获的前向图的反向图 (阅读更多)。
    • PrimTorch 将复杂的 PyTorch 操作分解为更简单和更基本的算子 (阅读更多)。
    • [后端] 后端与 TorchDynamo 集成,将图编译为可在加速器上运行的 IR。例如,TorchInductor 将图编译为用于 GPU 执行的 Triton 或用于 CPU 执行的 OpenMP (阅读更多)。
  10. 2.0 目前支持哪些编译器后端?
    默认且最完整的后端是 TorchInductor,但 TorchDynamo 有一个不断增长的后端列表,可以通过调用 torchdynamo.list_backends() 找到。

  11. 分布式训练在 2.0 中是如何工作的?
    编译模式下的 DDP 和 FSDP 在 FP32 中可以比即时模式快 15%,在 AMP 精度下快 80%。PT2.0 做了一些额外的优化,以确保 DDP 的通信-计算重叠与 Dynamo 的部分图创建良好配合。确保您以 static_graph=False 运行 DDP。更多详细信息 在这一处

  12. 我该如何了解更多关于 PT2.0 的发展?
    PyTorch 开发者论坛 是直接从构建它们的开发者那里了解 2.0 组件的最佳场所。

  13. 我的代码在 2.0 的编译模式下运行变慢了,帮帮我!
    性能下降最可能的原因是过多的图中断。例如,模型前向传播中的打印语句这样无害的东西也会触发图中断。我们有诊断这些问题的方法——阅读更多 这里

  14. 我之前运行的代码在 2.0 的编译模式下崩溃了!我该如何调试它?
    这里有一些技术可以排查您的代码可能在哪里失败,并打印有用的日志:https://pytorch.ac.cn/docs/stable/torch.compiler_faq.html#why-is-my-code-crashing

咨询工程师:2.0 直播问答系列

我们将为社区举办一系列直播问答会议,以便提出更深层次的问题并与专家进行对话。请回来查看全年完整的主题日历。如果您无法参加:1) 它们将被录制以供将来观看;2) 您可以在每周五太平洋标准时间上午 10 点参加我们的开发基础设施办公时间(Dev Infra Office Hours) @ 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 and Sherlock Huang
Natalia Gimelshein
LinkedIn
Sherlock Huang
LinkedIn
无需 C++ 扩展 PyTorch 和 functorch:类似 JAX 的可组合函数变换 Anjali Chourdia and 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 and Carl Parker
Zain Rizvi
LinkedIn | Twitter
Svetlana Karslioglu
LinkedIn | Twitter
动态形状和计算最大批处理大小 Edward Yang and Elias Ellison
Edward Yang
Twitter
PyTorch 2.0 导出:PyTorch 的健全全图捕获 Michael Suo and Yanan Cao
Yanan Cao
LinkedIn
使用 DistributedTensor 和 PyTorch DistributedTensor 进行 2-D 并行 Wanchao Liang and Alisson Gusatti Azzolini
Wanchao Liang
LinkedIn | Twitter
Alisson Gusatti Azzolini
LinkedIn
生产环境中的 TorchRec 和 FSDP Dennis van der Staay, Andrew Gu and 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 and Mark Saroufim
Mark Saroufim
LinkedIn | Twitter
PT2 和分布式 Will Constable
LinkedIn

观看 PyTorch 大会的演讲