我们很高兴地宣布发布 PyTorch® 2.0，我们在2022年12月2日的 PyTorch 大会上对此进行了重点介绍！PyTorch 2.0 提供相同的 Eager Mode 开发和用户体验，同时在编译器层面从根本上改变并增强了 PyTorch 的运行方式，带来了更快的性能以及对动态形状和分布式训练的支持。
这一下一代版本包括稳定版的加速 Transformer（以前称为 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）相当的性能。
- 加速 Transformer 引入了使用自定义内核架构进行缩放点积注意力（SPDA）的高性能训练和推理支持。该 API 已与 torch.compile() 集成，模型开发人员也可以通过调用新的 scaled_dot_product_attention() 运算符直接使用 缩放点积注意力 内核。
- Metal Performance Shaders (MPS) 后端为 Mac 平台上的 PyTorch 训练提供 GPU 加速，并增加了对最常用的前 60 个操作的支持，将覆盖范围扩展到 300 多个操作。
- Amazon AWS 优化了基于 AWS Graviton3 的 C7g 实例 上的 PyTorch CPU 推理。与以前的版本相比，PyTorch 2.0 改进了 Graviton 上的推理性能，包括对 Resnet50 和 Bert 的改进。
- TensorParallel、DTensor、2D Parallel、TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 的新原型功能和技术。
*要查看 2.0、1.13 和 1.12 功能提交的完整列表，请点击 此处。
稳定功能
[稳定版] 加速 PyTorch 2 Transformer
PyTorch 2.0 版本包含 PyTorch Transformer API 的新高性能实现。在发布加速 PT2 Transformer 时，我们的目标是使最先进的 Transformer 模型的训练和部署在整个行业中都能负担得起。此版本引入了使用自定义内核架构进行缩放点积注意力 (SPDA) 的高性能训练和推理支持，扩展了以前称为“Better Transformer”的推理“快速路径”架构。
类似于“快速路径”架构，自定义内核已完全集成到 PyTorch Transformer API 中——因此，使用原生 Transformer 和 MultiHeadAttention API 将使用户能够
- 透明地看到显著的速度提升；
- 支持更多用例，包括使用交叉注意力、Transformer 解码器和训练模型的用例；以及
- 继续将快速路径推理用于固定和可变序列长度的 Transformer 编码器和自注意力用例。
为了充分利用不同的硬件模型和 Transformer 用例，支持多个 SDPA 自定义内核（参见下文），并具有自定义内核选择逻辑，该逻辑将为给定模型和硬件类型选择最高性能的内核。除了现有的 Transformer API，模型开发人员还可以通过调用新的 scaled_dot_product_attention() 运算符直接使用 缩放点积注意力 内核。加速 PyTorch 2 Transformer 已与 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 Frame Evaluation Hooks 安全地捕获 PyTorch 程序，这是一项重大创新，是我们五年研发安全图捕获的成果。
- AOTAutograd 将 PyTorch 的自动微分引擎过载为用于生成超前反向跟踪的跟踪自动微分。
- PrimTorch 将 2000 多个 PyTorch 运算符规范化为大约 250 个原语运算符的封闭集，开发人员可以以此为目标构建完整的 PyTorch 后端。这大大降低了编写 PyTorch 功能或后端的障碍。
- TorchInductor 是一个深度学习编译器，可为多个加速器和后端生成快速代码。对于 NVIDIA 和 AMD GPU，它使用 OpenAI Triton 作为关键构建块。对于 Intel CPU，我们使用多线程、矢量化指令生成 C++ 代码，并在可能的情况下将适当的操作卸载到 mkldnn。
凭借所有这些新技术，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 后端上获得最佳性能和稳定性。
链接
[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 中难以表达的高级自动微分用例。示例包括
我们很高兴地宣布，作为 functorch 上游化并集成到 PyTorch 的最后一步，functorch API 现已在 torch.func 模块中提供。我们的函数转换 API 与以前相同，但我们更改了与 NN 模块交互的方式。请参阅 文档 和 迁移指南 了解更多详细信息。
此外，我们已 添加对 torch.autograd.Function 的支持：现在可以对 torch.autograd.Function 应用函数转换（例如 vmap、grad、jvp）。
[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
了解更多信息 此处。
[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 后端改进的 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 之前，PyG 的 GNN 模型在 CPU 上效率低下，原因是缺乏对几个关键内核（scatter/gather 等）的性能调优以及缺乏 GNN 相关的稀疏矩阵乘法操作。具体来说，优化包括
- scatter_reduce：当边索引以坐标格式 (COO) 存储时，消息传递中的性能热点。
- gather：scatter_reduce 的反向，专门针对索引是扩展张量时的 GNN 计算进行调优。
- 带有 reduce 标志的 torch.sparse.mm：当边索引以压缩稀疏行 (CSR) 存储时，消息传递中的性能热点。支持的 reduce 标志有：sum、mean、amax、amin。
在 PyG 基准/示例、OGB 基准测试中，单节点推理和训练的性能加速测量到 1.12x – 4.07x（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 中达到了 Beta 状态，用于 Float32 和 BFloat16 推理（在支持 AVX512_BF16 ISA 的机器上）。
从开发人员/研究人员的角度来看，用法非常简单直观，代码中唯一的更改是 API 调用
- 利用 oneDNN Graph，通过 JIT 跟踪，使用示例输入对模型进行分析。
- 也可以使用上下文管理器 with 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 DistributedTensor (DTensor) 是一项原型工作，它提供了分布式张量原语，以便在 SPMD（单程序多设备）范例中更轻松地编写分布式计算。这些原语简单而强大，当用于表达具有分片和复制并行策略的张量分布时。PyTorch DTensor 增强了 PyTorch 张量并行性 以及其他高级并行性探索。此外，它还提供了一种统一的方式来保存/加载 state_dict，用于分布式检查点目的，即使存在复杂的张量分布策略，例如在 FSDP 中将张量并行性与参数分片结合。更多详细信息可以在此 RFC 和 DTensor 示例笔记本 中找到。
[原型] TensorParallel
我们现在支持基于 DTensor 的张量并行，用户可以将其模型参数分布到不同的 GPU 设备上。我们还支持成对并行，它将两个连接的线性层分别进行列向和行向分片，以便最终只需要一次集合操作（all-reduce/reduce-scatter）。
[原型] 2D 并行
我们实现了上述 TP 与 FullyShardedDataParallel (FSDP) 的集成，作为 2D 并行，以进一步扩展大型模型训练。更多详细信息可以在此 幻灯片 中找到。
[原型] torch.compile(dynamic=True)
此版本提供了对 PT2 动态形状编译的实验性支持。支持使用 inductor 对简单模型进行推理编译，但存在许多限制
- 未来的版本将提供训练功能（这在 Nightlies 中已部分修复！）
- 未来的版本将提供缩减器。
- 很容易出现您希望动态的维度无论如何都被专门化的情况。其中一些问题已在 Nightlies 中修复，另一些则没有。
- 我们没有适当地将 Inductor 守卫传播到顶层，这在 #96296 中跟踪。
- 像 nonzero 这样依赖数据的操作仍然需要图中断。
- 动态模式不适用于非标准模式，例如 reduce-overhead 或 max-autotune。
- Inductor 编译中存在许多错误。要跟踪已知错误，请查看 PyTorch 问题跟踪器上的 动态形状 标签。
有关主分支上动态形状支持的最新消息，请查看 我们的状态报告。
亮点/性能改进
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 fast math 模式利用 bfloat16 内核，无需任何模型量化。接下来，我们实现了 conv、matmul 和 inner product 运算符的原语缓存。有关更新的 PyTorch 用户指南（包含即将发布的 2.0 版本改进和 TorchBench 基准测试详细信息）的更多信息，请参见 此处。