概述
隆重推出 PyTorch 2.0,这是我们迈向下一代 2 系列 PyTorch 版本的第一步。在过去几年中,我们从 PyTorch 1.0 不断创新迭代到最新的 1.13,并迁移到新成立的 PyTorch Foundation (Linux Foundation 的一部分)。
除了我们出色的社区之外,PyTorch 最大的优势在于我们始终如一的一流 Python 集成、命令式风格、简洁的 API 和丰富的选项。PyTorch 2.0 提供了相同的 eager 模式开发和用户体验,同时从根本上改变并极大提升了 PyTorch 在底层编译器层面的运作方式。我们能够提供更快的性能,并支持动态形状 (Dynamic Shapes) 和分布式 (Distributed)。
下面您将找到更好地理解 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 引擎重载为跟踪自动微分 (tracing autodiff),用于生成提前 (ahead-of-time) 的反向传播跟踪。
- PrimTorch 将 2000 多个 PyTorch 运算符规范化为约 250 个基本运算符的封闭集合,开发者可以针对这些基本运算符构建一个完整的 PyTorch 后端。这显著降低了编写 PyTorch 功能或后端的门槛。
- TorchInductor 是一个深度学习编译器,可以为多种加速器和后端生成快速代码。对于 NVIDIA 和 AMD GPU,它使用 OpenAI Triton 作为关键构建模块。
TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 使用 Python 编写,并支持动态形状 (即能够输入不同大小的张量而无需重新编译),这使得它们非常灵活、易于修改,并降低了开发者和供应商的入门门槛。
为了验证这些技术,我们使用了来自各种机器学习领域的 163 个不同的开源模型作为数据集。我们精心构建了这个基准测试,包含了诸如图像分类、目标检测、图像生成、各种 NLP 任务 (如语言建模、问答、序列分类)、推荐系统和强化学习等任务。我们将基准测试分为三类:
- 46 个来自 HuggingFace Transformers 的模型
- 61 个来自 TIMM 的模型:Ross Wightman 收集的最新 PyTorch 图像模型集
- 56 个来自 TorchBench 的模型:精选的来自 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 就能在训练 Transformer 模型时带来 1.5 倍到 2 倍的加速。这是自混合精度训练引入以来最令人兴奋的事情!”
Ross Wightman,TIMM (PyTorch 生态系统中最大的视觉模型中心之一) 的主要维护者
“它对于大多数 TIMM 模型的推理和训练工作负载开箱即用,无需任何代码更改。”
Luca Antiga,Lightning AI 的 CTO 和 PyTorch Lightning 的主要维护者之一
“PyTorch 2.0 体现了深度学习框架的未来。几乎无需用户干预即可捕获 PyTorch 程序并获得设备端显著提速和程序操作的能力,为 AI 开发者开启了一个全新的维度。”
动机
我们关于 PyTorch 的理念一直是将灵活性和可修改性放在首位,性能紧随其后。我们力求:
- 高性能 eager 执行
- Python 风格的内部实现
- 良好的分布式、自动微分、数据加载、加速器等的抽象
自 2017 年推出 PyTorch 以来,硬件加速器 (如 GPU) 在计算方面速度提高了约 15 倍,在内存访问速度方面提高了约 2 倍。因此,为了保持 eager 执行的高性能,我们不得不将 PyTorch 内部的大部分内容移至 C++。将内部内容移至 C++ 会降低其可修改性,并增加了代码贡献的门槛。
从第一天起,我们就知道 eager 执行的性能局限性。2017 年 7 月,我们启动了第一个研究项目,开发用于 PyTorch 的编译器。编译器需要使 PyTorch 程序快速运行,但不能以牺牲 PyTorch 体验为代价。我们的关键标准是保留某些类型的灵活性——支持研究人员在探索的各个阶段使用的动态形状和动态程序。
技术概述
多年来,我们在 PyTorch 内部构建了多个编译器项目。让我们将编译器分解为三个部分:
- 图捕获
- 图降低
- 图编译
在构建 PyTorch 编译器时,图捕获是更具挑战性的任务。
在过去的 5 年里,我们构建了 torch.jit.trace
、TorchScript、FX 跟踪、Lazy Tensors。但它们都没有完全达到我们想要的效果。有些灵活但不快,有些快但不灵活,有些既不快也不灵活。有些用户体验很差 (比如默默地给出错误结果)。虽然 TorchScript 很有前景,但它需要对你的代码以及你的代码所依赖的代码进行大量更改。这种需要大量代码更改的要求使得它对许多 PyTorch 用户来说无法使用。
TorchDynamo:可靠且快速地获取图
今年早些时候,我们开始研究 TorchDynamo,这是一种利用 PEP-0523 中引入的 CPython 功能 (称为 Frame Evaluation 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 以生成提前 (ahead-of-time) 图
对于 PyTorch 2.0,我们知道我们想加速训练。因此,不仅捕获用户级代码,还捕获反向传播代码至关重要。此外,我们知道我们想重用现有的、经过实战检验的 PyTorch autograd 系统。AOTAutograd 利用 PyTorch 的 torch_dispatch 扩展机制来跟踪我们的 Autograd 引擎,从而使我们能够“提前”捕获反向传播过程。这使得我们可以使用 TorchInductor 加速前向 和 反向传播过程。
PrimTorch:稳定的基本运算符
为 PyTorch 编写后端具有挑战性。PyTorch 有 1200 多个运算符,如果考虑每个运算符的各种重载,则超过 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 会在代码中插入守卫 (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
Hooks (钩子)
Module 和 Tensor 的钩子目前尚未完全工作,但随着开发的完成,它们最终将正常工作。
序列化
您可以序列化 optimized_model
或 model
的状态字典 (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 线上问答系列”中就此话题与我们互动交流 (详情见本文末尾)。
调试问题
编译模式是不透明且难以调试的。您可能会有这样的问题:
- 我的程序为什么在编译模式下崩溃?
- 编译模式的准确性是否与 eager 模式相同?
- 我为什么没有看到加速?
如果编译模式产生错误、崩溃或与 eager 模式结果不一致 (超出机器精度限制),则很可能不是您代码的错误。但是,了解是哪段代码导致了 bug 会很有用。
为了帮助调试和重现问题,我们创建了一些工具和日志记录功能,其中一个尤为突出:Minifier (精简器)。
Minifier 会自动将您遇到的问题简化为一小段代码。这小段代码能够重现原始问题,您可以将精简后的代码提交到 GitHub issue。这将帮助 PyTorch 团队轻松快速地修复问题。
如果您没有看到预期的加速,我们提供了 torch._dynamo.explain 工具,该工具会解释您代码的哪些部分导致了我们所称的“图断裂 (graph breaks)”。图断裂通常会阻碍编译器加速代码,减少图断裂的数量可能会加速您的代码 (在边际效益递减的某个限度内)。
您可以在我们的故障排除指南中阅读有关这些内容以及更多信息。
动态形状
在考察支持 PyTorch 代码通用性所需的内容时,一个关键要求是支持动态形状,并允许模型接受不同大小的张量而无需在形状每次改变时都重新编译。
截至今日,对动态形状的支持有限,且正在快速开发中。在稳定版本发布时将完全具备此功能。它通过 dynamic=True
参数启用,我们在一个功能分支 (symbolic-shapes) 上取得了更多进展,在该分支上我们已成功使用 TorchInductor 在训练中运行了具有完整符号形状 (symbolic shapes) 的 BERT_pytorch。对于动态形状的推理,我们的覆盖范围更广。例如,让我们看看动态形状有用的一个常见场景——使用语言模型进行文本生成。
我们可以看到,即使形状动态地从 4 变化到 256,编译模式仍能持续优于 eager 模式,加速高达 40%。如果没有对动态形状的支持,一个常见的变通方法是填充到最近的 2 的幂。然而,正如我们从下图中看到的,这会带来显著的性能开销,并且导致编译时间显著增加。此外,正确进行填充有时也并非易事。
通过在 PyTorch 2.0 的编译模式中支持动态形状,我们可以同时获得最佳性能 和 易用性。


当前的工作正在迅速发展,我们可能会在基础设施的基础改进落地时暂时让某些模型性能下降。关于我们在动态形状方面最新进展的更新可以在此处找到。
分布式
总之,torch.distributed 的两个主要分布式封装器在编译模式下工作良好。
DistributedDataParallel
(DDP) 和 FullyShardedDataParallel
(FSDP) 都在编译模式下工作,与 eager 模式相比,提供了更好的性能和内存利用率,但也存在一些注意事项和限制。

右图:编译模式下的 FSDP 比 eager 模式消耗的内存显著更少。


DistributedDataParallel (DDP)
DDP 依赖于将 AllReduce 通信与反向传播计算重叠,并将较小的逐层 AllReduce 操作分组到“桶 (buckets)”中以提高效率。TorchDynamo 编译的 AOTAutograd 函数在与 DDP 天真组合时会阻止通信重叠,但通过为每个“桶”编译单独的子图,并允许通信操作在子图之外和之间发生,性能得到了恢复。编译模式下的 DDP 支持目前还需要 static_graph=False
。有关 DDP + TorchDynamo 方法和结果的更多详情,请参阅此帖子。
FullyShardedDataParallel (FSDP)
FSDP 本身是一个“beta”阶段的 PyTorch 功能,由于能够调整哪些子模块被封装以及通常有更多配置选项,其系统复杂度高于 DDP。FSDP 在使用 use_original_params=True
标志配置时,可与 TorchDynamo 和 TorchInductor 配合适用于各种流行模型。目前预计会存在与特定模型或配置相关的兼容性问题,但我们将积极改进,如果提交 GitHub issue,特定模型可以优先处理。
用户通过指定 auto_wrap_policy
参数来指示模型的哪些子模块应该一起封装在一个用于状态分片 (state sharding) 的 FSDP 实例中,或者手动在 FSDP 实例中封装子模块。例如,许多 transformer 模型在每个“transformer block”被封装在单独的 FSDP 实例中时工作良好,因此只需要在某一时刻实例化一个 transformer block 的完整状态。Dynamo 会在每个 FSDP 实例的边界插入图断裂 (graph breaks),以允许前向 (和后向) 中的通信操作在图外部并行于计算发生。
如果 FSDP 在不将子模块封装在单独实例中的情况下使用,它会回退到类似于 DDP 的操作方式,但不进行分桶 (bucketing)。因此,所有梯度都在一个操作中归约 (reduce),即使在 Eager 模式下也无法实现计算/通信重叠。此配置仅在使用 TorchDynamo 时针对功能性进行了测试,但未针对性能进行测试。
开发者/供应商体验
借助 PyTorch 2.0,我们希望简化后端 (编译器) 集成体验。为此,我们专注于减少运算符数量和简化运算符集合的语义,这对于实现 PyTorch 后端是必需的。
图示形式展示了 PT2 栈的结构如下:
从图的中间部分开始,AOTAutograd 以提前 (ahead-of-time) 的方式动态捕获自动微分逻辑,生成一个以 FX 图格式表示的前向和后向运算符图。
我们提供了一套强化分解(即用其他算子实现的算子实现),可以利用这些分解来减少后端需要实现的算子数量。我们还通过一个称为函数化 (functionalization) 的过程,有选择地重写了包括变异 (mutations) 和视图 (views) 在内的复杂 PyTorch 逻辑,从而简化了 PyTorch 算子的语义,并保证了形状传播公式等算子元数据信息。这项工作正在积极进行中;我们的目标是提供一套具有简化语义的大约 250 个原始 (primitive) 且稳定 (stable) 的算子集,称为 PrimTorch,供应商可以利用(即选择启用)它们来简化其集成。
在精简和简化算子集之后,后端可以选择集成在 Dynamo(即中间层,紧随 AOTAutograd 之后)或 Inductor(较低层)上。下文我们将描述一些在做出此选择时需要考虑的因素,以及关于混合后端的未来工作。
Dynamo 后端
拥有现有编译器栈的供应商可能会发现最容易以 TorchDynamo 后端的形式进行集成,接收以 ATen/Prims IR 表示的 FX Graph。请注意,无论是训练还是推理,集成点都将在 AOTAutograd 之后,因为我们目前将分解作为 AOTAutograd 的一部分应用,如果针对推理,则仅跳过特定于反向传播的步骤。
Inductor 后端
供应商也可以将他们的后端直接集成到 Inductor 中。Inductor 接收由 AOTAutograd 生成的 Graph,该 Graph 由 ATen/Prim 操作组成,并将其进一步降低到循环级 IR。目前,Inductor 为逐点、归约、分散/聚集和窗口操作提供了到其循环级 IR 的降低。此外,Inductor 创建融合组,进行索引简化,维度折叠,并调整循环迭代顺序以支持高效的代码生成。然后,供应商可以通过提供从循环级 IR 到硬件特定代码的映射来进行集成。目前,Inductor 有两种后端:(1)C++ 后端,生成多线程 CPU 代码;(2)Triton 后端,生成高性能 GPU 代码。这些 Inductor 后端可以作为其他后端的参考。
混合后端接口(即将推出)
我们构建了用于将 FX Graph 划分为包含后端支持的算子的子图,并急切执行剩余部分的工具。这些工具可以扩展以支持“混合后端”,配置 Graph 的哪些部分由哪个后端运行。然而,目前还没有一个稳定的接口或约定供后端暴露他们支持的算子、偏好的算子模式等信息。这仍然是正在进行的工作,我们欢迎早期用户的反馈。
总结
我们对为 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/transformers、https://github.com/huggingface/accelerate 和 https://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,享受免费的性能提升,如果您没有看到(性能提升),请提交一个 Issue,我们会确保您的模型得到支持 https://github.com/pytorch/torchdynamo/issues
毕竟,除非您的模型确实运行得更快,否则我们不能声称我们创造了一个广度优先的(编译器)。
常见问题
-
什么是 PT 2.0?
2.0 是最新的 PyTorch 版本。PyTorch 2.0 提供与 eager-mode 相同的开发体验,同时通过torch.compile
添加了编译模式。这种编译模式有可能在训练和推理期间加速您的模型。 -
为什么是 2.0 而不是 1.14?
PyTorch 2.0 本来会是 1.14。我们发布了我们认为会显著改变您使用 PyTorch 方式的重大新功能,因此我们将其命名为 2.0。 -
如何安装 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
-
2.0 代码是否向后兼容 1.X 版本?
是的,使用 2.0 不需要您修改 PyTorch 工作流程。一行代码model = torch.compile(model)
就可以优化您的模型以使用 2.0 技术栈,并与您其余的 PyTorch 代码顺利运行。这是完全选择加入的,您无需使用新的编译器。 -
2.0 默认启用吗?
2.0 是版本的名称。torch.compile
是在 2.0 中发布的功能,您需要显式使用torch.compile
。 - 如何将我的 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)
-
为什么我应该使用 PT2.0 而不是 PT 1.X?
请参阅问题 (2) 的答案。 - 当运行 PyTorch 2.0 时,我的代码有何不同之处?
开箱即用,PyTorch 2.0 与 PyTorch 1.x 相同,您的模型在 eager-mode 下运行,即每一行 Python 代码都逐行执行。
在 2.0 中,如果您将模型包装在model = torch.compile(model)
中,您的模型在执行前会经历 3 个步骤:- 图获取:首先,模型被重写为子图块。可以由 TorchDynamo 编译的子图会被“展平”,而其他子图(可能包含控制流代码或其他不受支持的 Python 结构)将回退到 Eager-Mode。
- 图降低:所有 PyTorch 操作都被分解为其选定后端特有的组成内核。
- 图编译:其中内核调用其相应的低级设备特定操作。
- PT2.0 向 PT 添加了哪些新组件?
-
2.0 目前支持哪些编译器后端?
默认且最完整的后端是 TorchInductor,但 TorchDynamo 支持越来越多的后端,可以通过调用torchdynamo.list_backends()
来查看。 -
分布式训练如何与 2.0 配合工作?
编译模式下的 DDP 和 FSDP 在 FP32 中运行速度可比 Eager-Mode 快 15%,在 AMP 精度下可快 80%。PT2.0 做了一些额外优化,以确保 DDP 的通信计算重叠与 Dynamo 的部分图创建良好配合。请确保您运行 DDP 时设置static_graph=False
。更多详细信息在此。 -
如何了解更多关于 PT2.0 的进展?
PyTorch 开发者论坛是直接从构建 2.0 组件的开发者那里了解相关信息的最佳场所。 -
求助!我的代码在 2.0 的编译模式下运行变慢了!
性能下降最可能的原因是过多的 graph breaks。例如,模型 forward 方法中一个看似无害的 print 语句就会触发 graph break。我们有诊断这些问题的方法——在此阅读更多。 - 我之前可以正常运行的代码在 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 和 Sherlock Huang Natalia Gimelshein Sherlock Huang |
不使用 C++ 和 functorch 扩展 PyTorch:PyTorch 的类 JAX 可组合函数变换 | Anjali Chourdia 和 Samantha Andow Anjali Chourdia LinkedIn | Twitter Samantha Andow LinkedIn | Twitter |
深入了解 TorchDynamo | Michael Voznesensky |
使用 TorchData 重新思考数据加载:Datapipes 和 Dataloader2 | Kevin Tse |
可组合训练(+ 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 |
PyTorch 2.0 Export:PyTorch 的健全完整图捕获 | Michael Suo 和 Yanan Cao Yanan Cao |
使用 DistributedTensor 和 PyTorch DistributedTensor 实现 2-D 并行 | Wanchao Liang 和 Alisson Gusatti Azzolini Wanchao Liang LinkedIn | Twitter Alisson Gusatti Azzolini |
TorchRec 和 FSDP 在生产环境中的应用 | Dennis van der Staay, Andrew Gu 和 Rohan Varma Dennis van der Staay Rohan Varma LinkedIn | Twitter |
PyTorch On-Device 的未来 | 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 |