PyTorch 2.0：更快、更符合 Python 习惯、更动态的下一代版本

我们很高兴地宣布发布 PyTorch® 2.0，我们已在 2022 年 12 月 2 日的 PyTorch 大会上强调了这一点！PyTorch 2.0 提供了相同的即时模式开发和用户体验，同时从根本上改变并增强了 PyTorch 在编译器层面的内部操作，带来了更快的性能以及对动态形状和分布式训练的支持。

此下一代版本包括稳定版的加速 Transformers（以前称为 Better Transformers）；Beta 版包括作为 PyTorch 2.0 主要 API 的 torch.compile、作为 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）相当的性能。
• 加速 Transformers 引入了使用自定义内核架构进行缩放点积注意力（SPDA）的高性能训练和推理支持。该 API 与 torch.compile() 集成，模型开发人员还可以通过调用新的 scaled_dot_product_attention() 运算符直接使用 缩放点积注意力 内核。
• Metal Performance Shaders (MPS) 后端在 Mac 平台上提供 GPU 加速的 PyTorch 训练，并增加了对最常用的前 60 个操作的支持，使覆盖范围达到 300 多个操作。
• Amazon AWS 在基于 AWS Graviton3 的 C7g 实例上优化了 PyTorch CPU 推理。与之前的版本相比，PyTorch 2.0 改进了 Graviton 上的推理性能，包括 Resnet50 和 Bert 的改进。
• TensorParallel、DTensor、2D Parallel、TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 的新原型功能和技术。

稳定版	Beta	原型	性能改进
加速 PT 2 Transformers	torch.compile	DTensor	CUDA 11.7 和 11.8 支持（废弃 CUDA 11.6）
	PyTorch MPS 后端	TensorParallel	Python 3.8（废弃 Python 3.7）
	缩放点积注意力	2D Parallel	AWS Graviton3
	functorch	Torch.compile (dynamic=True)
	可调度集合操作
	Torch.set_default 和 torch.device
	X86 量化后端
	GNN 推理和训练性能

*要查看 2.0、1.13 和 1.12 功能提交的完整列表，请点击 此处。

稳定功能

[稳定版] 加速 PyTorch 2 Transformers

PyTorch 2.0 版本包含 PyTorch Transformer API 的新高性能实现。发布加速 PT2 Transformers 的目标是使行业内最先进的 Transformer 模型的训练和部署变得经济实惠。此版本引入了使用缩放点积注意力 (SPDA) 的自定义内核架构的高性能训练和推理支持，扩展了推理“快速路径”架构，以前称为“Better Transformer”。

与“快速路径”架构类似，自定义内核完全集成到 PyTorch Transformer API 中——因此，使用原生 Transformer 和 MultiHeadAttention API 将使用户能够

• 透明地看到显著的速度改进；
• 支持更多用例，包括使用交叉注意力、Transformer 解码器和训练模型的模型；以及
• 继续将快速路径推理用于固定和可变序列长度的 Transformer 编码器和自注意力用例。

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

我们通过结合自定义内核和 `torch.compile()`，使用加速 PyTorch 2 Transformers 在训练 Transformer 模型，特别是大型语言模型方面，取得了显著的加速。

 图：将缩放点积注意力与自定义内核和 torch.compile 结合使用，可显著加速大型语言模型（例如此处所示的 nanoGPT）的训练。

Beta 功能

[Beta] torch.compile

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

支撑 torch.compile 的是新技术——TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor

• TorchDynamo 使用 Python 帧评估钩子安全地捕获 PyTorch 程序，这是一项重大创新，是我们 5 年安全图捕获研发的成果。
• AOTAutograd 将 PyTorch 的自动微分引擎重载为一种用于生成提前反向追踪的追踪自动微分。
• PrimTorch 将 2000 多个 PyTorch 运算符规范化为一组封闭的 250 个原始运算符，开发人员可以以此为目标构建完整的 PyTorch 后端。这大大降低了编写 PyTorch 功能或后端的障碍。
• TorchInductor 是一种深度学习编译器，可为多种加速器和后端生成快速代码。对于 NVIDIA 和 AMD GPU，它使用 OpenAI Triton 作为关键构建块。对于 Intel CPU，我们使用多线程、矢量化指令并尽可能将适当的操作卸载到 mkldnn 来生成 C++ 代码。

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

有关更多信息，请参阅 https://pytorch.ac.cn/get-started/pytorch-2.0/，有关使用 Intel 的 TorchInductor CPU，请参阅 此处。

[Beta] PyTorch MPS 后端

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

MPS 后端现在包括对最常用的前 60 个操作以及社区最常请求的操作的支持，将覆盖范围扩展到 300 多个操作。此版本的主要重点是启用基于 OpInfo 的完整前向和梯度模式测试，以解决隐式正确性问题。这些更改导致 Stable Diffusion、YoloV5、WhisperAI 等第三方网络更广泛地采用 MPS 后端，并增加了 Torchbench 网络和基本教程的覆盖范围。我们鼓励开发人员更新到最新的 macOS 版本，以在 MPS 后端获得最佳性能和稳定性。

链接

1. MPS 后端
2. 开发者信息
3. Mac 上加速 PyTorch 训练
4. Metal、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 架构组件。

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

[Beta] functorch -> torch.func

受 Google JAX 启发，functorch 是一个提供可组合 vmap（矢量化）和自动微分变换的库。它支持在 PyTorch 中难以表达的高级自动微分用例。示例包括

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

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

此外，我们还 添加了对 torch.autograd.Function 的支持：现在可以将函数转换（例如 vmap、grad、jvp）应用于 torch.autograd.Function。

[Beta] 可调度集合操作

可调度集合操作是对现有 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

在此处了解更多信息：https://pytorch.ac.cn/docs/master/distributed.html#torch.distributed.init_process_group。

[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 最初发布以来长期存在的功能请求。

在此处了解更多信息：https://pytorch.ac.cn/tutorials/recipes/recipes/changing_default_device.html。

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

新的 X86 量化后端利用 FBGEMM 和 oneDNN 内核库，取代 FBGEMM 成为 x86 CPU 平台的默认量化后端，与原始 FBGEMM 后端相比，它利用两个库的优势，提供了改进的 int8 推理性能，在 40 多个深度学习模型上测量到 1.3 倍 – 2 倍的推理性能加速。新后端在功能上与原始 FBGEMM 后端兼容。

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

	1 核/实例	2 核/实例	4 核/实例	1 插槽（32 核）/实例
Intel(R) Xeon(R) Platinum 8358 CPU @ 2.60GHz	1.76X	1.80X	2.04X	1.34X

默认情况下，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。

[Beta] 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 with reduce flag：当边缘索引以压缩稀疏行 (CSR) 存储时，消息传递中的性能热点。支持的 reduce 标志为：sum、mean、amax、amin。

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

模型-数据集	选项	加速比
GCN-Reddit（推理）	512-2-64-dense	1.22x
	1024-3-128-dense	1.25x
	512-2-64-sparse	1.31x
	1024-3-128-sparse	1.68x
	512-2-64-dense	1.22x
GraphSage-ogbn-products（推理）	1024-3-128-dense	1.15x
	512-2-64-sparse	1.20x
	1024-3-128-sparse	1.33x
	full-batch-sparse	4.07x
GCN-PROTEINS（训练）	3-32	1.67x
GCN-REDDIT-BINARY（训练）	3-32	1.67x
GCN-Reddit（训练）	512-2-64-dense	1.20x
	1024-3-128-dense	1.12x

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

[Beta] 利用 oneDNN Graph 加速 PyTorch 在 CPU 上的推理

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

• 它自动识别要通过融合加速的图分区。
• 融合模式 侧重于融合卷积、矩阵乘法及其邻近操作等计算密集型操作，适用于推理和训练用例。
• 尽管将 oneDNN Graph 与 TorchDynamo 集成的工作仍在进行中，但它与 PyTorch JIT Fuser 的集成在 PyTorch 2.0 中获得了 Float32 和 BFloat16 推理的 Beta 状态（在支持 AVX512_BF16 ISA 的机器上）。

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

• 利用 oneDNN Graph，通过 JIT 追踪，模型将使用示例输入进行分析。
• 也可以使用上下文管理器 with torch.jit.fuser("fuser3"): 而不是调用 torch.jit.enable_onednn_fusion(True)。
• 为了加速 BFloat16 推理，我们依赖 PyTorch 中的即时模式 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)

在此处了解更多信息：https://pytorch.ac.cn/tutorials/recipes/recipes/tuning_guide.html#use-onednn-graph-with-torchscript-for-inference。

原型功能

分布式 API

[原型] DTensor

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

[原型] TensorParallel

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

[原型] 2D Parallel

我们实现了上述 TP 与 FullyShardedDataParallel (FSDP) 的集成，作为 2D 并行，以进一步扩展大型模型训练。更多详细信息可以在此 幻灯片 中找到。

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

此版本提供了对具有动态形状的 PT2 编译的实验性支持。支持使用 inductor 编译简单模型的推理，但存在许多限制

• 训练将在未来版本中提供（这在夜间构建中已部分修复！）
• 精简器将在未来版本中提供。
• 很容易出现您希望动态的维度无论如何都被专业化的情况。其中一些问题已在夜间构建中修复，另一些则没有。
• 我们没有适当地将 Inductor 守卫传播到顶层，这在 #96296 中进行了跟踪。
• 数据相关操作（如 nonzero）仍然需要图中断。
• 动态与非标准模式（如 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 平台上 PyTorch 依赖的软件包（包括 NumPy、SciPy、SymPy、Pillow 等）缺乏 Python 3.11 支持，我们不会为 PyTorch 2.0 发布使用 Python 3.11 编译的 Conda 二进制文件。将发布支持 Python 3.11 的 Pip 包，因此如果您打算将 PyTorch 2.0 与 Python 3.11 一起使用，请使用我们的 Pip 包。请注意：支持 Python 3.11 的 Conda 包将在我们的 nightly 频道上提供。此外，我们计划在 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 和 inner product 运算符实现了原语缓存。有关更新的 PyTorch 用户指南（包含即将发布的 2.0 版本的改进和 TorchBench 基准测试详细信息）的更多信息可以在 此处 找到。