由 PyTorch 团队

我们非常高兴地宣布 PyTorch® 2.1（发布说明）的发布！PyTorch 2.1 在 torch.compile 中提供了自动动态形状支持，torch.distributed.checkpoint 支持在多个 rank 上并行保存/加载分布式训练作业，以及 torch.compile 对 NumPy API 的支持。

此外，本版本还提供了许多性能改进（例如 CPU 感应器改进、AVX512 支持、scaled-dot-product-attention 支持），以及 torch.export 的原型发布，这是一个完整的全图捕获机制，以及基于 torch.export 的量化。

与 2.1 版本一起，我们还发布了一系列 PyTorch 领域库的更新。更多详细信息请参阅库更新博客。

本版本由自 2.0 以来的 6,682 个提交和 784 位贡献者组成。我们衷心感谢我们专注的社区成员的贡献。一如既往，我们鼓励大家尝试这些更新，并在我们改进 2.1 版本时报告任何问题。有关如何开始使用 PyTorch 2 系列的信息，请参阅我们的入门页面。

摘要：

  • torch.compile 现在包含自动检测和最小化由于张量形状变化导致的重新编译的支持。
  • torch.distributed.checkpoint 允许并行保存和加载来自多个排名的模型，以及由于集群拓扑结构变化而进行的重新分片。
  • torch.compile 现在可以通过将它们转换为 PyTorch 等效操作来编译 NumPy 操作。
  • torch.compile 现在包含对 Python 3.11 的改进支持。
  • 新型 CPU 性能特性包括电感改进（例如 bfloat16 支持和动态形状）、AVX512 内核支持以及缩放点积注意力内核。
  • 引入了一个原型功能，即 torch.export，一个完整的图捕获机制，以及基于 torch.export 的量化。
  • torch.sparse 现在在 NVIDIA® GPU 上包括半结构化（2:4）稀疏的原型支持。
稳定 测试版 原型 性能提升
  自动动态形状 torch.export() AVX512 内核支持
  torch.distributed.checkpoint 基于 Torch.export 的量化 毫米波雷达的 CPU 优化
  torch.compile + NumPy 半结构化（2:4）稀疏性 bfloat16 的 CPU 优化
  torch.compile + Python 3.11 torchinductor 的 cpp 包装器  
  torch.compile + autograd.Function    
  第三方设备集成：PrivateUse1    

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

测试功能

(测试版) 自动动态形状

动态形状是 torch.compile 内置的功能，可以通过跟踪和根据张量的符号形状而不是静态形状（例如[B, 128, 4]而不是[64, 128, 4]）生成代码来最小化重新编译。这允许 torch.compile 生成一个适用于许多尺寸的单个内核，而效率的损失很小。在 PyTorch 2.1 中，动态形状得到了极大的稳定，如果 torch.compile 检测到由于输入形状变化而导致的重新编译，则会自动启用。您可以通过将 dynamic=False 传递给 torch.compile 或设置 torch._dynamo.config.automatic_dynamic_shapes = False 来禁用自动动态形状。

在 PyTorch 2.1 中，我们已启用动态形状，在包括大型语言模型在内的多种模型类型上，在 CUDA 和 CPU 上均表现出良好的性能。

关于动态形状的更多信息，请参阅此文档。

[测试版] torch.distributed.checkpoint

torch.distributed.checkpoint 允许并行保存和加载来自多个 rank 的模型。此外，检查点自动处理模型和优化器之间的完全限定名（FQN）映射，支持在不同集群拓扑结构下的负载时重新分配资源。

欲了解更多信息，请参阅 torch.distributed.checkpoint 文档和教程。

[Beta] torch.compile + NumPy

torch.compile 现在能够通过将它们转换为 PyTorch 等效操作来编译 NumPy 操作。由于这种集成以设备无关的方式进行，您现在可以通过使用 torch.compile 来加速 NumPy 程序——甚至混合 NumPy/PyTorch 程序。

请参阅 torch.compile 常见问题解答中的本节，以获取有关 torch.compile + NumPy 交互的更多信息，并关注 PyTorch 博客，以获取关于此功能的即将发布的博客文章。

[Beta] torch.compile + Python 3.11

torch.compile 之前仅支持 Python 版本 3.8-3.10。用户现在可以使用 torch.compile 在 Python 3.11 中优化模型。

[Beta] torch.compile + autograd.Function

torch.compile 现在可以追踪和优化用户定义的 autograd Functions 的逆向函数，这为使用扩展机制更频繁的模型解锁了训练优化。

[Beta] 改进第三方设备支持：私有用途 1

现在可以使用私有用途 1 调度键将第三方设备类型注册到 PyTorch。这允许设备扩展向 PyTorch 注册新的内核，并将它们与新的键关联起来，使用户代码可以与内置设备类型等效工作。例如，要注册“my_hardware_device”，可以执行以下操作：

torch.rename_privateuse1_backend("my_hardware_device")
torch.utils.generate_methods_for_privateuse1_backend()
x = torch.randn((2, 3), device='my_hardware_device')
y = x + x # run add kernel on 'my_hardware_device'

为验证此功能，Ascend NPU 的 OSS 团队已成功将 torch_npu 集成到 PyTorch 中，作为通过私有用途 1 功能实现的插件。

更多信息，请参阅此处私有用途 1 教程。

原型功能

[原型] torch.export()

torch.export() 提供了一种可靠的追踪机制，用于从基于 PT2.0 提供的新技术的 PyTorch 程序中捕获完整的图。

用户可以将 PyTorch 程序提取为数据流图形式的干净表示（导出 IR），该图主要由对 PyTorch 操作符的直线调用组成。然后可以将导出 IR 转换、序列化、保存到文件、传输，并在具有或没有 Python 的环境中重新加载以执行。

如需更多信息，请参阅此处教程。

[原型] 基于 torch.export 的量化

torch.ao.quantization 现在支持在 PyTorch 2 torch.export-based 流程中进行量化。这包括对内置的 XNNPACK 和 X64Inductor 量化器以及指定自己的量化器的支持。

关于使用 torch.export 进行训练后静态量化的说明，请参阅此教程；关于使用 torch.export 进行静态量化时的量化感知训练，请参阅此教程。

想要了解如何编写自己的量化器，请参阅这篇教程。

[原型] NVIDIA® GPU 的半结构化（2:4）稀疏性

torch.sparse 现在支持创建和加速对半结构化稀疏（2:4）张量的计算。有关格式的更多信息，请参阅 NVIDIA 的这篇博客。以下是一个介绍半结构化稀疏性的最小示例：

from torch.sparse import to_sparse_semi_structured
 
x = torch.rand(64, 64).half().cuda()
mask = torch.tensor([0, 0, 1, 1]).tile((64, 16)).cuda().bool()
linear = nn.Linear(64, 64).half().cuda()

linear.weight = nn.Parameter(to_sparse_semi_structured(linear.weight.masked_fill(~mask, 0)))
linear(x)

想要了解更多信息，请参阅文档和配套教程。

[原型] torchinductor 的 C++ 包装器

cpp_wrapper 可以通过生成 C++ 核函数包装代码来减少 torchinductor 中调用核函数的 Python 开销。此功能目前处于原型阶段；它不支持今天在 PT2 中成功编译的所有程序。如果您在使用中发现限制，请提交问题以帮助我们确定优先级。

启用此功能的 API 是：

import torch
import torch._inductor.config as config
config.cpp_wrapper = True

更多信息，请参阅教程。

性能提升

支持 AVX512 内核

在 PyTorch 2.0 中，即使 CPU 支持 AVX512 指令，也会使用 AVX2 内核。现在，PyTorch 默认使用支持 AVX512 指令的 CPU 内核，相当于之前版本中设置 ATEN_CPU_CAPABILITY=avx512。可以通过设置 ATEN_CPU_CAPABILITY=avx2 来启用之前的行为。

对缩放点积注意力（SDPA）的 CPU 优化

PyTorch 的早期版本通过 torch.nn.functional.scaled_dot_product_attention 提供了针对 transformer 原语优化的 CUDA 实现。PyTorch 2.1 包括基于 FlashAttention 的 CPU 优化例程。

请参阅此处文档。

bfloat16 的 CPU 优化

PyTorch 2.1 包括对 bfloat16 的 CPU 优化，包括改进的向量化支持和 torchinductor 代码生成。