[译]PyTorch 2.x

发表于 2026-02-07 分类于算法与模型，模型训练与优化，工具与效率，开发与运维，开发者工具（CLI/IDE/调试），性能监控与调优阅读次数：本文字数： 8.2k 阅读时长 ≈ 30 分钟

Overview

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

除了我们强大的社区之外，PyTorch 最大的优势在于它始终保持着一流的 Python 集成、命令式编程风格、简洁的 API 和丰富的选项。PyTorch 2.0 提供同样的便捷开发和用户体验，同时从根本上改变并大幅提升了 PyTorch 在编译器层面的底层运行方式。我们能够提供更快的性能，并支持动态形状和分布式计算。

以下内容将帮助您更好地了解 PyTorch 2.0 的定义、发展方向，以及更重要的——如何立即上手（例如，教程、系统要求、模型、常见问题解答）。我们仍有很多东西需要学习和开发，但我们期待社区的反馈和贡献，以使 PyTorch 2.0 系列更加完善。同时，我们也感谢所有为 PyTorch 1.0 系列的成功做出贡献的人们。

PyTorch 2.x: faster, more pythonic and as dynamic as ever

今天，我们宣布推出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 的自动微分引擎（PyTorch’s autograd engine）重载为跟踪自动微分（tracing autodiff），用于生成预先反向跟踪（ahead-of-time backward traces）。
PrimTorch将 2000 多个 PyTorch 算子（operators）规范化为一个包含约 250 个基本算子的封闭集合，开发者可以利用这些基本操作符构建完整的 PyTorch 后端。这大大降低了编写 PyTorch 功能或后端的门槛。
TorchInductor是一款深度学习编译器，可为多种加速器和后端生成快速代码。对于 NVIDIA 和 AMD GPU，它使用 OpenAI Triton 作为关键构建模块。

TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 都是用 Python 编写的，并且支持动态形状（即能够传入不同大小的张量而无需重新编译），这使得它们灵活、易于修改，并降低了开发人员和供应商的准入门槛。

为了验证这些技术，我们使用了涵盖各种机器学习领域的 163 个开源模型。我们精心构建了这个基准测试集，其中包括图像分类、目标检测、图像生成、各种自然语言处理任务（例如语言建模、问答、序列分类、推荐系统和强化学习）。我们将这些基准测试分为三类：

来自HuggingFace Transformers的46款模型
TIMM中的 61 个模型：Ross Wightman 编写的最先进的 PyTorch 图像模型集合
来自TorchBench的 56 个模型：精选自 GitHub 的热门代码库。

除了添加一个封装它们的调用（torch.compile）之外，我们不对这些开源模型进行任何修改。

然后，我们测量这些模型的加速比并验证其准确性。由于加速比可能取决于数据类型，因此我们分别测量了 float32 和自动混合精度 (AMP) 两种数据类型的加速比。我们发现 AMP 在实践中更为常见，因此我们报告的加权平均加速比为0.75 * AMP + 0.25 * float32。

在这 163 个开源模型中，93% 的时间中torch.compile都能正常工作，并且在 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 模式的加速性能

试用版：torch.compile目前仍处于早期开发阶段。从今天起，您可以在nightly版本中进行试用torch.compile。我们预计将于 2023 年 3 月初发布首个稳定的 2.0 版本。

在 PyTorch 2.x 的开发路线图中，我们希望在性能和可扩展性方面不断提升编译模式。正如我们在今天的大会上所讨论的，部分工作正在进行中。部分工作尚未启动。还有一些工作是我们希望看到的，但目前我们自身资源有限，无法完成。如果您有兴趣贡献力量，欢迎参加本月开始的Ask the Engineers: 2.0 在线问答系列活动（详情见文末），或通过 GitHub 或论坛与我们联系。

Testimonials

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

Sylvain Gugger是HuggingFace transformers的主要维护者：

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

Ross Wightman是TIMM（PyTorch 生态系统中最大的视觉模型中心之一）的主要维护者：

“它无需任何代码更改，即可与大多数 TIMM 模型开箱即用，用于推理和训练工作负载。”

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

“PyTorch 2.0 代表了深度学习框架的未来。它无需用户干预即可捕获 PyTorch 程序，并开箱即用地实现设备端速度的大幅提升和程序操控，这为人工智能开发者开启了全新的维度。”

Motivation

我们在 PyTorch 开发中始终秉持的理念是：将灵活性和可定制性放在首位，性能则紧随其后。我们力求做到：

高效积极的执行
Pythonic 内部结构
为分布式系统、自动微分、数据加载、加速器等提供了良好的抽象实现。

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

从一开始，我们就清楚动态执行（eager execution）的性能局限性。2017 年 7 月，我们启动了第一个研究项目，旨在开发 PyTorch 编译器。该编译器需要提升 PyTorch 程序的运行速度，但同时又不能牺牲 PyTorch 的使用体验。我们的关键标准是保留某些灵活性——例如支持动态形状和动态编程，以满足研究人员在探索各个阶段的需求。

Technology Overview

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

graph acquisition
graph lowering
graph compilation

在构建 PyTorch 编译器时，图的获取（graph acquisition）是一个更大的挑战。

过去五年里，我们开发了torch.jit.trace、TorchScript、FX tracing和 Lazy Tensors。但它们都无法完全满足我们的需求。有些灵活但速度慢，有些速度快但不够灵活，还有些既不快也不灵活。有些用户体验很差（比如经常默默出错）。TorchScript 虽然很有潜力，但它需要对用户代码以及依赖的代码进行大量修改。这种对代码的大幅修改使得许多 PyTorch 用户望而却步。

TorchDynamo: Acquiring Graphs reliably and fast

今年年初，我们开始着手开发 TorchDynamo，它利用了PEP-0523中引入的 CPython 特性——帧评估API（the Frame Evaluation API）。我们采用数据驱动的方法来验证其在图捕获（Graph Capture）方面的有效性。我们使用了 7000 多个用 PyTorch 编写的 GitHub 项目作为验证集。TorchScript 和其他工具甚至难以在 50% 的情况下获取到图，而且通常开销很大；而 TorchDynamo在 99% 的情况下都能正确、安全地获取图，开销几乎可以忽略不计——无需对原始代码进行任何修改。正是在那时，我们意识到，我们终于突破了多年来在灵活性和速度方面一直困扰我们的瓶颈。

TorchInductor: fast codegen using a define-by-run IR

为了开发 PyTorch 2.0 的全新编译器后端，我们从用户编写高性能自定义内核的方式中汲取灵感：他们越来越多地使用Triton语言。我们也希望编译器后端能够采用与 PyTorch eager 类似的抽象，并且足够通用，能够支持 PyTorch 的丰富特性。TorchInductor 使用a pythonic define-by-run loop level IR，自动将 PyTorch 模型映射到 GPU 上的 Triton 代码和 CPU 上的 C++/OpenMP 代码。TorchInductor 的core loop level IR 仅包含约 50 个算子，并且使用 Python 实现，因此易于修改和扩展。

AOTAutograd: reusing Autograd for ahead-of-time graphs

对于 PyTorch 2.0，我们明确了加速训练的目标。因此，至关重要的是，我们不仅要捕获用户级代码，还要捕获反向传播过程。此外，我们还希望复用久经考验的 PyTorch 自动微分系统。AOTAutograd 利用 PyTorch 的torch_dispatch扩展机制来追踪我们的自动微分引擎，从而能够“提前”捕获反向传播过程。这使得我们可以使用 TorchInductor加速正向和反向传播。

PrimTorch: Stable Primitive operators

为 PyTorch 编写后端极具挑战性。PyTorch 拥有 1200 多个运算符，如果算上每个运算符的各种重载，则超过 2000 个。

对 2000 多个 PyTorch 算子进行详细分析

因此，编写一个后端或一个横切功能会变得非常耗费精力。在 PrimTorch 项目中，我们正在致力于定义更小、更稳定的算子集。PyTorch 程序可以始终如一地降维到这些算子集。我们的目标是定义两个算子集：

包含约 250 个算子的 Prim 算子库，这些算子相当底层。它们适合编译器使用，因为它们足够底层，需要将它们重新组合起来才能获得良好的性能。
包含约 750 个标准算子的ATen算子库，适合直接导出。这些算子适用于已在 ATen 层集成的后端，或无需编译即可从底层算子集（例如 Prim 算子）中恢复性能的后端。

我们将在下文“开发者/供应商体验”部分对此主题进行更详细的讨论。

User Experience

我们引入了一个简单的函数torch.compile，它可以封装你的模型并返回一个编译后的模型。

1	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)，需要编译模型，因此运行时间会较长。后续运行速度很快。

Modes

编译器内置了一些预设模式，可以从不同方面优化编译后的模型。例如，您可能正在运行一个小型模型，但由于框架开销而运行缓慢；或者，您可能正在运行一个大型模型，它几乎无法完全加载到内存中。根据您的具体需求，您可能需要使用不同的模式。

# 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")

Reading and updating Attributes

访问模型属性的方式与在eager mode下相同。您可以像往常一样访问或修改模型的属性（例如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

Hooks

Module和Tensor钩子目前还不能完全正常工作，但随着我们完成开发，它们最终会正常工作。

Serialization

你可以序列化optimized_model或者model的state-dict。它们指向相同的参数和状态，因此是等价的。

1
2
3

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。

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

Inference and Export

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

此外，我们将引入一种名为torch.export的模式，该模式会针对需要保证且可预测延迟的环境，精心导出整个模型和防护基础架构。torch.export这需要对您的程序进行一些更改，尤其是在您拥有数据相关的控制流时。

1
2
3

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

这项功能目前还处于早期开发阶段。更多详情，请关注 PyTorch 大会上关于导出路径的演讲。您也可以参加我们本月开始的“工程师问答：2.0 在线问答系列”活动，参与讨论（更多详情请见文末）。

Debugging Issues

编译后的模式不透明且难以调试。您可能会遇到以下问题：

为什么我的程序在compiled mode下会崩溃？
compiled mode和eager mode的准确度一样吗？
为什么我没有看到加速效果？

如果编译模式导致错误、崩溃或结果与即时模式（超出机器精度限制）出现差异，则很可能不是你的代码的问题。但是，了解是哪段代码导致了错误仍然很有帮助。

为了帮助进行调试和重现，我们创建了几个工具和日志记录功能，其中一个工具尤为突出：The Minifier。

The Minifier会自动将您遇到的问题简化为一小段代码。这段代码可以重现原始问题，您可以将压缩后的代码提交到 GitHub 上。这将有助于 PyTorch 团队快速轻松地修复问题。

如果您没有看到预期的速度提升，我们提供了torch._dynamo.explain工具，它可以解释您的代码中哪些部分导致了所谓的“图断裂”（graph breaks）。图断裂通常会阻碍编译器加速代码，减少图断裂的数量可能会提高代码速度（但存在收益递减的极限）。

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

Dynamic Shapes

在考虑如何支持 PyTorch 代码的通用性时，一个关键要求是支持动态形状，并允许模型接受不同大小的张量，而无需在每次形状改变时都重新编译。

截至目前，对动态形状的支持还比较有限，并且仍在快速开发中。该功能将在稳定版本中全面实现。目前，该功能需要通过参数启用。dynamic=True我们在一个特性分支（symbolic-shapes）上取得了更多进展，已成功在该分支上使用 TorchInductor 对 BERT_pytorch 进行完整符号形状的训练。对于动态形状的推理，我们也有更全面的覆盖。例如，让我们来看一个动态形状非常有用的常见场景——使用语言模型生成文本。

我们可以看到，即使数组大小从 4 动态变化到 256，编译模式的性能也始终比即时模式高出 40%。在不支持动态数组大小的情况下，一种常见的解决方法是填充到最接近的 2 的幂次方。然而，正如我们从下面的图表中看到的，这会带来显著的性能开销，并且会导致编译时间明显延长。此外，正确填充数组有时并非易事。

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

目前的工作进展非常迅速，在基础设施进行根本性改进期间，我们可能会暂时放弃一些模型。您可以在这里找到我们动态形状研究的最新进展。

Distributed

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

DistributedDataParallel（DDP）和FullyShardedDataParallel（FSDP）均以编译模式运行，与即时模式相比，可提供更高的性能和内存利用率，但也存在一些注意事项和限制。

AMP Precision 的速度提升

左图：FSDP 在编译模式下的速度比在即时模式下更快（AMP 精度）。右图：FSDP 在编译模式下占用的内存比即时模式下少得多。

DistributedDataParallel (DDP)

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

FullyShardedDataParallel (FSDP)

FSDP 本身是 PyTorch 的一个“测试版”功能，由于可以调整封装的子模块以及通常更多的配置选项，其系统复杂度高于 DDP。如果使用相应的标志进行配置，FSDP 可以与 TorchDynamo 和 TorchInductor 配合使用，支持多种常用模型。use_original_params=True目前预计会存在一些与特定模型或配置的兼容性问题，但我们会积极改进，如果用户在 GitHub 上提交 issue，我们会优先处理特定模型的问题。

用户可以通过指定auto_wrap_policy参数来指示要将模型中的哪些子模块封装到用于状态分片的 FSDP 实例中，或者手动将子模块封装到 FSDP 实例中。例如，许多 Transformer 模型在每个“Transformer 块”封装到单独的 FSDP 实例中时都能很好地工作，这样一次只需要物化一个 Transformer 块的完整状态。Dynamo 会在每个 FSDP 实例的边界处插入图断点，以允许前向（和后向）通信操作在图外部进行，并与计算并行执行。

如果 FSDP 不使用子模块封装在单独的实例中，则会退而求其次，其运行方式类似于 DDP，但不进行分桶。因此，所有梯度都会在一次操作中得到简化，即使在 Eager 模式下也不会出现计算/通信重叠。此配置仅在 TorchDynamo 上测试过功能，尚未进行性能测试。

Developer/Vendor Experience

在 PyTorch 2.0 中，我们希望简化后端（编译器）集成体验。为此，我们专注于减少运算符的数量，并简化启动 PyTorch 后端所需的运算符集的语义。

PT2协议栈的图形化表示如下：

从图表的中间开始，AOTAutograd 以预先的方式动态捕获自动微分逻辑，生成 FX 图格式的前向和后向运算符图。

我们提供一组强化分解（即用其他算子编写的算子实现），可用于减少后端需要实现的算子数量。我们还通过称为函数化的过程，选择性地重写复杂的 PyTorch 逻辑（包括变更和视图），从而简化PyTorch 算子的语义，并保证算子元数据信息（例如形状传播公式）的完整性。这项工作正在积极进行中；我们的目标是提供一套包含约 250 个语义简化的原始且稳定的算子集，称为PrimTorch，供供应商使用（即选择加入）以简化其集成。在减少和简化算子集之后，后端可以选择在 Dynamo（即中间层，紧接在 AOTAutograd 之后）或 Inductor（底层）进行集成。我们将在下文中描述做出此选择时需要考虑的一些因素，以及未来关于混合后端的工作。

Dynamo Backend

对于拥有现有编译器栈的供应商而言，最简便的集成方式可能是将其作为 TorchDynamo 后端，并接收基于 ATen/Prims IR 的 FX 图。需要注意的是，无论是训练还是推理，集成点都应紧接在 AOTAutograd 之后，因为我们目前已将分解作为 AOTAutograd 的一部分进行处理，如果目标是推理，则只需跳过反向传播的特定步骤即可。

Inductor backend

供应商还可以将他们的后端直接集成到 Inductor 中。Inductor 接收由 AOTAutograd 生成的图，该图包含 ATen/Prim 操作，并将其进一步降维到循环级中间表示 (IR)。目前，Inductor 提供针对逐点操作、归约操作、分散/聚集操作和窗口操作的循环级 IR 降维。此外，Inductor 还创建融合组、进行索引简化、维度折叠，并调整循环迭代顺序，以支持高效的代码生成。供应商可以通过提供从循环级 IR 到特定硬件代码的映射来实现集成。目前，Inductor 有两个后端：(1) 用于生成多线程 CPU 代码的 C++ 后端；(2) 用于生成高性能 GPU 代码的 Triton 后端。这些 Inductor 后端可以为其他后端提供参考。

Mixture of Backends Interface (coming soon)

我们构建了一些实用工具，用于将外汇交易图分割成子图，每个子图包含一个后端支持的运算符，其余部分则立即执行。这些工具可以扩展以支持“混合后端”，从而配置哪些子图由哪个后端运行。然而，目前后端还没有稳定的接口或协议来公开其运算符支持、运算符模式偏好等信息。这项工作仍在进行中，我们欢迎早期用户的反馈。

Final Thoughts

我们对 PyTorch 2.0 及未来发展的方向感到无比兴奋。通往最终 2.0 版本的道路注定充满挑战，但我们诚邀您尽早加入我们，共同开启这段旅程。如果您有兴趣深入了解或为编译器做出贡献，请继续阅读以下内容，其中包含更多入门信息（例如教程、基准测试、模型、常见问题解答）以及本月即将开始的工程师问答：2.0 在线问答系列。其他资源包括：

Accelerating Hugging Face and TIMM models with PyTorch 2.0

Author: Mark Saroufim

Pytorch 2.x让您能够轻松地尝试不同的编译器后端，并通过一行装饰器（torch.compile()）来提升 PyTorch 代码的运行速度。它可以直接应用于 nn.Module，作为 torch.jit.script() 的替代方案，而无需您修改任何源代码。我们预计，对于您目前运行的绝大多数模型，这一行代码的更改将带来 30% 到 2 倍的训练速度提升。

1	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/transformers、https://github.com/huggingface/accelerate 和 https://github.com/rwightman/pytorch-image-models 这样的流行库，然后对它们运行了 torch.compile()，就完成了。

性能和便利性兼得实属难得，但这正是核心团队对 PyTorch 2.0 感到如此兴奋的原因。

Requirements

对于GPU而言（新一代GPU的性能将大幅提升）

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

For CPU

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

Optional: Verify Installation

1
2
3

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

Optional: Docker installation

我们还提供了PyTorch nightly binaries中所有必需的依赖项，您可以下载这些二进制文件。

1	docker pull ghcr.io/pytorch/pytorch-nightly

对于这个实验版本，只需确保容器可以访问所有 GPU 即可。

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

Getting Started

请阅读 Mark Saroufim 的完整博客文章，他在文章中通过教程和真实模型，指导您立即尝试 PyTorch 2.0。

我们开发 PyTorch 的目标是构建一个广度优先编译器，以加速绝大多数开源用户实际运行的模型。Hugging Face Hub 最终成为了我们极其宝贵的基准测试工具，确保我们所做的任何优化都能真正帮助用户加速他们想要运行的模型。

博客教程将详细展示如何复现这些速度提升，让您和我们一样对 PyTorch 2.0 充满期待。所以，请试用 PyTorch 2.0，享受免费的性能提升。如果您没有看到任何提升，请提交 issue，我们将确保您的模型得到支持：https://github.com/pytorch/torchdynamo/issues

毕竟，除非你们的模型运行速度确实更快，否则我们不能声称我们创建的是广度优先模型。

FAQs

What is PT 2.0?
- 2.0 是 PyTorch 的最新版本。PyTorch 2.0 保留了即时开发模式的体验，同时通过 torch.compile 添加了编译模式。这种编译模式能够显著提升模型在训练和推理过程中的速度。
Why 2.0 instead of 1.14?
- PyTorch 2.0 原本应该是 1.14 的版本。但我们发布了一些重要的全新功能，相信这些功能会从根本上改变您使用 PyTorch 的方式，所以我们决定将其命名为 2.0。

How do I install 2.0? Any additional requirements?

安装最新的每日构建版本：

CUDA 11.8

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

CUDA 11.7

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

CPU

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

Is 2.0 code backwards-compatible with 1.X?
- 是的，使用 2.0 版本无需修改您的 PyTorch 工作流程。只需一行代码model = torch.compile(model)即可优化您的模型以使用 2.0 版本，并使其与其余 PyTorch 代码流畅运行。此功能完全可选，您无需使用新的编译器。
Is 2.0 enabled by default?
- 2.0 是版本号。torch.compile 是 2.0 版本中新增的功能，你需要显式地使用 torch.compile。

How do I migrate my PT1.X code to PT2.0?

您的代码应该可以直接运行，无需任何迁移。如果您想使用 2.0 版本中引入的全新编译模式功能，可以先用一行代码优化您的模型：model = torch.compile(model)。虽然速度提升主要体现在训练阶段，但如果您的模型运行速度比即时模式更快，您也可以将其用于推理阶段。

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)

Why should I use PT2.0 instead of PT 1.X?
- 请参阅问题（2）的答案。
What is my code doing differently when running PyTorch 2.0?
- 默认情况下，PyTorch 2.0 与 PyTorch 1.x 相同，模型以 eager 模式运行，即每行 Python 代码都会逐行执行。
- 在 2.0 版本中，如果您将模型包装在 <script> 标签中model = torch.compile(model)，则模型在执行前会经历以下 3 个步骤：
  - 图获取(Graph acquisition)：首先将模型重写为子图块。TorchDynamo 可以编译的子图会被“扁平化”，而其他子图（可能包含控制流代码或其他不受支持的 Python 结构）将回退到 Eager-Mode。
  - 图降维(Graph lowering)：所有 PyTorch 操作都被分解成特定于所选后端的组成内核。
  - 图编译(Graph compilation)，其中内核调用其相应的底层设备特定操作。
What new components does PT2.0 add to PT?
- TorchDynamo从 Python 字节码生成 FX 图。它使用guard机制来维护 eager-mode 功能，以确保生成的图有效（了解更多）。
- 使用AOTAutograd生成与 TorchDynamo 捕获的前向图对应的后向图（阅读更多）。
- PrimTorch将复杂的 PyTorch 操作分解为更简单、更基本的操作（阅读更多）。
- [Backend]后端与 TorchDynamo 集成，将计算图编译成可在加速器上运行的中间表示 (IR)。例如，TorchInductor将计算图编译成Triton（用于 GPU 执行）或OpenMP（用于 CPU 执行）（了解更多）。
What compiler backends does 2.0 currently support?
- 默认且最完整的后端是TorchInductor，但 TorchDynamo 拥有一个不断增长的后端列表，可以通过调用torchdynamo.list_backends().
How does distributed training work with 2.0?
- 编译模式下的 DDP 和 FSDP 在 FP32 精度下比 Eager 模式快 15%，在 AMP 精度下快 80%。PT2.0 进行了一些额外的优化，以确保 DDP 的通信-计算重叠部分能够与 Dynamo 的部分图创建机制良好配合。请确保运行 DDP 时设置 static_graph=False。更多详情请点击此处。
How can I learn more about PT2.0 developments?
- 是直接从构建这些组件的开发者那里了解2.0组件的最佳场所。
Help my code is running slower with 2.0’s Compiled Mode!
- 性能下降最可能的原因是图中断过多。例如，模型前向传播中的一个看似无害的打印语句就可能触发图中断。我们有方法可以诊断这些问题——点击此处了解更多信息。
My previously-running code is crashing with 2.0’s Compiled Mode! How do I debug it?
- 这里有一些技巧可以帮助您找出代码可能出错的地方，并打印有用的日志：https://pytorch.org/docs/stable/torch.compiler_faq.html#why-is-my-code-crashing.