PyTorch 2.0：我们下一代更快、更 Pythonic 和更动态的发布

由 PyTorch 团队

我们非常高兴地宣布 PyTorch® 2.0 的发布，我们曾在 2022 年 12 月 2 日的 PyTorch 大会上重点介绍！PyTorch 2.0 提供了相同的急切模式开发和用户体验，同时在底层编译器级别对 PyTorch 的操作进行了根本性的改变和增强，实现了更快的性能和对动态形状以及分布式支持。

这一代的发布包括加速 Transformer 的稳定版本（之前称为 Better Transformers）；Beta 版本包括 torch.compile 作为 PyTorch 2.0 的主要 API，torch.nn.functional 中的 scaled_dot_product_attention 函数，MPS 后端，torch.func 模块中的 functorch API；以及跨各种推理、性能和训练优化功能在 GPU 和 CPU 上的其他 Beta/原型改进。有关 torch.compile 的全面介绍和技术概述，请访问 2.0 入门页面。

随着 2.0 版本的发布，我们还推出了一系列 PyTorch 领域库的 beta 更新，包括树内库以及独立的库，如 TorchAudio、TorchVision 和 TorchText。TorchX 的更新也在发布中，因为它将转向社区支持模式。更多详细信息可以在该库博客中找到。

本次发布自 1.13.1 以来，共有超过 4,541 个提交和 428 位贡献者。我们衷心感谢我们忠诚的社区成员的贡献。一如既往，我们鼓励大家尝试使用这些更新，并在我们改进 2.0 版本以及今年整个 2 系列的过程中报告任何问题。

摘要：

• torch.compile 是 PyTorch 2.0 的主要 API，它封装您的模型并返回一个编译后的模型。这是一个完全可添加的（且可选的）功能，因此 2.0 版本在定义上 100%向后兼容。
• 作为 torch.compile 的基础技术，TorchInductor 与 Nvidia 和 AMD GPU 将依赖于 OpenAI Triton 深度学习编译器来生成高性能代码并隐藏底层硬件细节。OpenAI Triton 生成的内核在性能上与手动编写的内核和专门的 cuda 库（如 cublas）相当。
• 加速 Transformer 引入了对使用自定义内核架构的缩放点积注意力（SPDA）的高性能支持。API 已集成 torch.compile()，模型开发者还可以通过调用新的 scaled_dot_product_attention()运算符直接使用缩放点积注意力内核。
• 金属性能着色器（MPS）后端为 Mac 平台上的 GPU 加速 PyTorch 训练提供支持，并增加了对 Top 60 个最常用操作的支持，覆盖范围超过 300 个操作符。
• 亚马逊 AWS 优化了基于 AWS Graviton3 的 C7g 实例上的 PyTorch CPU 推理。PyTorch 2.0 相比之前的版本，在 Graviton 上的推理性能得到提升，包括 Resnet50 和 Bert 的改进。
• 涵盖 TensorParallel、DTensor、2D 并行、TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 等新的原型功能和新技术。

点击此处查看 2.0、1.13 和 1.12 版本所有公共功能提交的完整列表。

稳定功能

[稳定版] 加速 PyTorch 2 Transformers

PyTorch 2.0 版本发布了一个 PyTorch Transformer API 的新高性能实现。在发布加速 PT2 Transformers 时，我们的目标是使行业内的最先进 Transformer 模型的训练和部署变得经济实惠。本版本引入了对使用自定义内核架构进行缩放点积注意力（SPDA）训练和推理的高性能支持，扩展了之前称为“Better Transformer”的推理“fastpath”架构。

与“fastpath”架构类似，自定义内核已完全集成到 PyTorch Transformer API 中，因此，使用本地的 Transformer 和 MultiHeadAttention API 将使用户能够：

• 可见显著的速度提升；
• 支持更多用例，包括使用交叉注意力、Transformer 解码器和训练模型的模型；
• 继续使用固定和可变序列长度的 Transformer 编码器和自注意力用例的 fastpath 推理；

为了充分利用不同的硬件模型和 Transformer 用例，支持多个 SDPA 自定义内核（见下文），具有自定义内核选择逻辑，该逻辑将选择给定模型和硬件类型的最高性能内核。除了现有的 Transformer API，模型开发者还可以通过调用新的 scaled_dot_product_attention() 操作符直接使用缩放点积注意力内核。加速 PyTorch 2 Transformers 与 torch.compile() 集成。为了在使用模型的同时受益于 PT2 编译的额外加速（推理或训练），请使用 model = torch.compile(model) 预处理模型。

我们使用自定义内核和 torch.compile() 的组合，在加速 PyTorch 2 Transformer 中实现了训练 Transformer 模型，特别是大型语言模型的重大加速。

图：使用缩放点积注意力机制和自定义内核以及 torch.compile()，为训练大型语言模型（如此处所示的 nanoGPT）提供了显著的加速。

测试功能

[Beta] torch.compile

torch.compile 是 PyTorch 2.0 的主要 API，它包装您的模型并返回一个编译后的模型。这是一个完全可添加的（可选）功能，因此按照定义，2.0 是 100% 向后兼容的。

torch.compile 的基础是新技术——TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor：

• TorchDynamo 使用 Python 框架评估钩子安全地捕获 PyTorch 程序，这是我们 5 年研发成果的一个重大创新，专注于安全的图捕获。
• AOTAutograd 将 PyTorch 的 autograd 引擎作为跟踪自动微分，用于生成提前时间反向跟踪。
• PrimTorch 将 2000 多个 PyTorch 算子规范化为约 250 个基本算子的封闭集合，开发者可以针对这些算子构建完整的 PyTorch 后端。这大大降低了编写 PyTorch 功能或后端的门槛。
• TorchInductor 是一个深度学习编译器，为多个加速器和后端生成快速代码。对于 NVIDIA 和 AMD GPU，它使用 OpenAI Triton 作为关键构建块。对于英特尔 CPU，我们使用多线程、向量化指令，并在可能的情况下将适当的操作卸载到 mkldnn，以生成 C++代码。

在所有新技术的基础上，torch.compile 能够在 93%的时间内运行 165 个开源模型，平均在 float32 精度下运行速度快 20%，在 AMP 精度下平均快 36%。

如需更多信息，请参阅 https://maskerprc.github.io/get-started/pytorch-2.0/ 以及 Intel 的 TorchInductor CPU。

[Beta] PyTorch MPS 后端

MPS 后端为 Mac 平台提供 GPU 加速的 PyTorch 训练。本版本带来了改进的正确性、稳定性和算子覆盖范围。

MPS 后端现在包括对前 60 个最常用算子的支持，以及社区最常请求的操作，覆盖范围超过 300 个算子。本版本发布的主要重点是启用基于 OpInfo 的完整正向和梯度模式测试，以解决静默正确性问题。这些更改导致 MPS 后端被第三方网络（如 Stable Diffusion、YoloV5、WhisperAI）更广泛采用，并为 Torchbench 网络和基础教程提供了更全面的覆盖。我们鼓励开发者更新到最新的 macOS 版本，以在 MPS 后端上获得最佳性能和稳定性。

链接

1. MPS 后端
2. 开发者信息
3. 在 Mac 上加速 PyTorch 训练
4. 金属、Metal Performance Shaders 和 Metal Performance Shaders Graph

[Beta] 比例点积注意力 2.0

我们非常高兴地宣布 PyTorch 2.0 的发布，它引入了一个强大的比例点积注意力函数，作为 torch.nn.functional 的一部分。此函数包含多个实现，可以根据输入和使用的硬件无缝应用。

在 PyTorch 的早期版本中，您必须依赖第三方实现并安装单独的包，才能利用像 FlashAttention 这样的内存优化算法。在 PyTorch 2.0 中，所有这些实现都默认可用。

这些实现包括来自 HazyResearch 的 FlashAttention、来自 xFormers 项目的内存高效注意力以及适用于非 CUDA 设备或需要高精度的本地 C++实现。

PyTorch 2.0 将自动为您选择最佳实现，但您也可以单独切换它们以获得更细粒度的控制。此外，缩放点积注意力函数可用于构建常见的 Transformer 架构组件。

通过文档和本教程了解更多信息。

[测试版] functorch -> torch.func

受 Google JAX 的启发，functorch 是一个提供可组合的 vmap（向量化）和自动微分变换的库。它使得在 PyTorch 中难以表达的高级自动微分用例成为可能。例如包括：

• 模型集成
• 高效计算雅可比矩阵和海森矩阵
• 计算每个样本的梯度（或其他每个样本的量）

我们非常高兴地宣布，作为将 functorch 上游化和集成到 PyTorch 的最后一步，functorch API 现在已在 torch.func 模块中提供。我们的函数变换 API 与之前相同，但我们已经改变了与 NN 模块的交互方式。请参阅文档和迁移指南以获取更多详细信息。

此外，我们增加了对 torch.autograd.Function 的支持：现在可以应用函数变换（例如 vmap、grad、jvp）到 torch.autograd.Function 上。

[测试版] 可调度集体操作

可调度集体是对现有 init_process_group() API 的改进，将后端改为可选参数。对于用户来说，这个特性的主要优势是它将允许他们编写可以在 GPU 和 CPU 机器上运行的代码，而无需更改后端指定。可调度特性也将使用户更容易支持 GPU 和 CPU 集体，因为他们将不再需要手动指定后端（例如，“NCCL”或“GLOO”）。用户现有的后端指定将被尊重，无需更改。

使用示例：

import torch.distributed.dist
…
# old
dist.init_process_group(backend=”nccl”, ...)
dist.all_reduce(...) # with CUDA tensors works
dist.all_reduce(...) # with CPU tensors does not work

# new
dist.init_process_group(...) # backend is optional
dist.all_reduce(...) # with CUDA tensors works
dist.all_reduce(...) # with CPU tensors works

在这里了解更多。

[Beta] torch.set_default_device 和 torch.device 作为上下文管理器

torch.set_default_device 允许用户更改 PyTorch 工厂函数分配的默认设备。例如，如果您 torch.set_default_device(‘cuda’)，则对 torch.empty(2) 的调用将在 CUDA 上分配（而不是在 CPU 上）。您还可以使用 torch.device 作为上下文管理器在局部基础上更改默认设备。这解决了 PyTorch 初始发布以来长期存在的功能请求，即如何实现这一点。

在这里了解更多。

[Beta] “X86”作为新的默认量化后端，适用于 x86 CPU

新的 X86 量化后端，利用 FBGEMM 和 oneDNN 内核库，取代了 FBGEMM 成为 x86 CPU 平台的默认量化后端，与原始 FBGEMM 后端相比，在 40 多个深度学习模型上实现了 1.3 倍至 2 倍的推理性能提升，充分利用了两个库的优势。新的后端在功能上与原始 FBGEMM 后端兼容。

表：X86 量化后端与 FBGEMM 后端几何平均加速比

默认情况下，x86 平台上的用户将使用 x86 量化后端，并且在使用默认后端时，他们的 PyTorch 程序将保持不变。或者，用户可以选择明确指定“X86”作为量化后端。以下是一个示例代码：

import torch
from torch.ao.quantization import get_default_qconfig_mappingfrom torch.quantization.quantize_fx
import prepare_fx, convert_fx
 
# get default configuration
qconfig_mapping = get_default_qconfig_mapping()
 
# or explicitly specify the backend
# qengine = 'x86'
# torch.backends.quantized.engine = qengine
# qconfig_mapping = get_default_qconfig_mapping(qengine)
 
# construct fp32 model
model_fp32 = ...
 
# prepare
prepared_model = prepare_fx(model_fp32, qconfig_mapping, example_inputs=x)
 
# calibrate
...
 
# convert
quantized_model = convert_fx(prepared_model)

查找更多信息：https://github.com/pytorch/pytorch/issues/83888 和 https://www.intel.com/content/www/us/en/developer/articles/technical/accelerate-pytorch-int8-inf-with-new-x86-backend.html。

[测试版] CPU 上的 GNN 推理和训练优化

PyTorch 2.0 包含几个关键优化，以提高 CPU 上的 GNN 推理和训练性能。在 2.0 之前，由于缺乏对几个关键内核（scatter/gather 等）的性能调优以及缺乏 GNN 相关的稀疏矩阵乘法操作，PyG 的 GNN 模型在 CPU 上效率低下。具体来说，优化包括：

• scatter_reduce：当边索引以坐标格式（COO）存储时，在消息传递中的性能热点。
• gather：scatter_reduce 的逆过程，特别针对索引为展开张量的 GNN 计算进行了优化。
• torch.sparse.mm 带 reduce 标志：当边索引以压缩稀疏行（CSR）存储时，在消息传递中的性能热点。支持的 reduce 标志为：sum、mean、amax、amin。

在 PyG 基准/示例、OGB 基准测试中，单节点推理和训练的测量性能提升了 1.12 倍至 4.07 倍（与 2.0 相比为 1.13.1）。

了解更多：PyG CPU 性能优化。

[测试版] 利用 oneDNN Graph 加速 CPU 上的 PyTorch 推理

oneDNN Graph API 扩展了 oneDNN，通过灵活的图 API 最大化优化机会，以在 AI 硬件上生成高效的代码。

• 它自动识别需要通过融合加速的图分区。
• 融合模式专注于融合计算密集型操作，如卷积、矩阵乘法及其相邻操作，适用于推理和训练用例。
• 虽然正在进行将 oneDNN Graph 与 TorchDynamo 集成的相关工作，但其与 PyTorch JIT Fuser 的集成已在 PyTorch 2.0 中达到 Beta 状态，支持 Float32 和 BFloat16 推理（在支持 AVX512_BF16 ISA 的机器上）。

从开发者/研究人员的角度来看，使用非常简单直观，代码中唯一的变化是 API 调用：

• 利用 oneDNN Graph，通过 JIT-tracing，使用示例输入对模型进行性能分析。
• 使用 torch.jit.fuser("fuser3")的上下文管理器：也可以用来代替调用 torch.jit.enable_onednn_fusion(True)。
• 为了加速 BFloat16 推理，我们依赖 PyTorch 的 eager-mode AMP（自动混合精度）支持并禁用 JIT 模式的 AMP，因为它们目前存在分歧：

# Assuming we have a model of the name 'model'
 
example_input = torch.rand(1, 3, 224, 224)
 
# enable oneDNN Graph
torch.jit.enable_onednn_fusion(True)
# Disable AMP for JIT
torch._C._jit_set_autocast_mode(False)
with torch.no_grad(), torch.cpu.amp.autocast():
	model = torch.jit.trace(model, (example_input))
	model = torch.jit.freeze(model)
 	# 2 warm-ups (2 for tracing/scripting with an example, 3 without an example)
	model(example_input)
	model(example_input)
 
	# speedup would be observed in subsequent runs.
	model(example_input)

在这里了解更多。

原型功能

分布式 API

[原型] DTensor

PyTorch 分布式 Tensor (DTensor) 是一个用于分布式张量原语的原型设计，旨在简化在 SPMD（单程序多设备）范式中的分布式计算编写。这些原语简单但功能强大，当用于表达具有分片和复制并行策略的张量分布时。PyTorch DTensor 增强了 PyTorch 张量并行性以及其他高级并行性探索。此外，它还提供了一种统一的方式来保存/加载 state_dict，以便进行分布式检查点，即使在存在复杂张量分布策略的情况下，例如在 FSDP 中将张量并行性与参数分片相结合。更多详细信息请参阅此 RFC 和 DTensor 示例笔记本。

[原型] TensorParallel

现在支持基于 DTensor 的 Tensor Parallel，用户可以将模型参数分布到不同的 GPU 设备上。我们还支持成对并行，将两个连接的线性层分别以列向和行向的方式分片，最终只需要一个集体操作（all-reduce/reduce-scatter）。

[原型] 2D 并行

我们实现了上述 TP 与 FullyShardedDataParallel（FSDP）的集成，作为 2D 并行以进一步扩展大型模型训练。更多详情请参阅此幻灯片。

[原型] torch.compile(dynamic=True)

本版本支持 PT2 编译的动态形状实验性支持。对于简单模型，支持使用电感进行推理编译，但存在许多限制：

• 训练功能将在未来版本中提供（夜版本中部分修复！）
• 未来版本中将提供代码压缩工具。
• 很容易陷入一个情况，即你想要动态的维度仍然被特殊化。其中一些问题已在夜版本中修复，但还有一些没有修复。
• 我们没有适当地将电感器保护传播到顶层，这个问题在#96296 中被跟踪。
• 像非零这样的数据相关操作仍然需要断开图。
• 动态模式与非标准模式（如 reduce-overhead 或 max-autotune）不兼容。
• Inductor 编译中存在许多错误。要跟踪已知错误，请检查 PyTorch 问题跟踪器上的动态形状标签。

关于在 master 上动态形状支持的最新和最伟大新闻，请查看我们的状态报告。

突出/性能改进

PyTorch 2.0 对 CUDA 11.6 和 Python 3.7 支持的弃用

如果您仍在使用或依赖于 CUDA 11.6 或 Python 3.7 构建，我们强烈建议升级到至少 CUDA 11.7 和 Python 3.8，因为这将是为 PyTorch 2.0 所需的最小版本。更多详情，请参阅 PyTorch 版本发布兼容性矩阵。

Anaconda 平台对 Python 3.11 的支持

由于 Anaconda 平台上的 NumPy、SciPy、SymPy、Pillow 等依赖包不支持 Python 3.11，我们将不会发布与 PyTorch 2.0 版本兼容的 Python 3.11 编译的 Conda 二进制文件。将发布支持 Python 3.11 的 Pip 包，因此如果您打算使用 Python 3.11 与 PyTorch 2.0 一起使用，请使用我们的 Pip 包。请注意：支持 Python 3.11 的 Conda 包将在我们的夜间频道上提供。我们还计划在 Anaconda 提供这些关键依赖项后，作为未来版本的发布内容发布 Conda Python 3.11 二进制文件。有关如何下载 Pip 包的更多信息及说明，请参阅此处。

利用 AWS Graviton 处理器优化 PyTorch 推理

优化主要集中在三个关键领域：GEMM 内核、bfloat16 支持、原语缓存和内存分配器。对于 aarch64 平台，PyTorch 通过 Mkldnn(OneDNN)后端支持 Arm Compute Library (ACL) GEMM 内核。ACL 库为 fp32 和 bfloat16 格式提供 Neon/SVE GEMM 内核。在 c7g 上对 bfloat16 的支持允许高效部署 bfloat16 训练的、AMP（自动混合精度）训练的，甚至是标准 fp32 训练的模型。标准 fp32 模型通过 OneDNN 快速数学模式利用 bfloat16 内核，无需任何模型量化。接下来，我们实现了 conv、matmul 和内积运算符的原语缓存。有关即将发布的 2.0 版本改进和 TorchBench 基准测试细节的更多信息，请在此处查看。