作者：PyTorch 团队

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

此外，此版本提供了众多性能改进（例如，CPU inductor 改进、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 支持在多个 rank 上并行保存和加载模型，并能处理因集群拓扑变化导致的重新分片（resharding）。
  • torch.compile 现在可以通过将 NumPy 操作转换为 PyTorch 等效操作来编译 NumPy 操作。
  • torch.compile 现在改进了对 Python 3.11 的支持。
  • 新的 CPU 性能特性包括 inductor 改进（例如，bfloat16 支持和动态形状）、AVX512 内核支持和 scaled-dot-product-attention 内核。
  • 引入了 torch.export（一种健全的完整图捕获机制）作为原型特性，以及基于 torch.export 的量化。
  • torch.sparse 现在包含对 NVIDIA® GPU 上半结构化 (2:4) 稀疏性的原型支持。
稳定版 测试版 原型版 性能改进
  自动动态形状 torch.export() AVX512 内核支持
  torch.distributed.checkpoint 基于 Torch.export 的量化 针对 scaled-dot-product-attention (SPDA) 的 CPU 优化
  torch.compile + NumPy 半结构化 (2:4) 稀疏性 针对 bfloat16 的 CPU 优化
  torch.compile + Python 3.11 torchinductor 的 cpp_wrapper  
  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 检测到因输入形状变化导致的重新编译，则会自动启用该功能。您可以通过向 torch.compile 传递 dynamic=False 或设置 torch._dynamo.config.automatic_dynamic_shapes = False 来禁用自动动态形状。

在 PyTorch 2.1 中，我们展示了在启用动态形状后，在各种模型类型（包括大型语言模型）上，无论是在 CUDA 还是 CPU 上都取得了良好的性能。

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

【测试版】torch.distributed.checkpoint

torch.distributed.checkpoint 支持在多个 rank 上并行保存和加载模型。此外，检查点功能会自动处理模型和优化器之间的完全限定名 (FQN) 映射，从而支持在不同的集群拓扑结构下进行加载时重新分片（resharding）。

有关更多信息，请参阅 torch.distributed.checkpoint 文档和教程。

【测试版】torch.compile + NumPy

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

请参阅 torch.compile 常见问题解答中的此部分，了解有关 torch.compile + NumPy 交互的更多信息，并关注PyTorch 博客，了解有关此特性的即将发布的博客。

【测试版】torch.compile + Python 3.11

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

【测试版】torch.compile + autograd.Function

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

【测试版】改进的第三方设备支持：PrivateUse1

第三方设备类型现在可以使用 privateuse1 dispatch 键注册到 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 的开源团队已通过 PrivateUse1 功能成功将 torch_npu 集成到 PyTorch 中作为插件。

有关更多信息，请参阅 PrivateUse1 教程此处。

原型特性

【原型版】torch.export()

torch.export() 提供了一种健全的追踪机制，可根据 PT2.0 提供的新技术从 PyTorch 程序中捕获完整的计算图。

用户可以提取 PyTorch 程序的清晰表示（Export IR），其形式为数据流图，主要由对 PyTorch 操作符的直线调用组成。然后可以将 Export IR 进行转换、序列化、保存到文件、传输、加载回来以便在有或没有 Python 的环境中执行。

有关更多信息，请参阅教程此处。

【原型版】基于 torch.export 的量化

torch.ao.quantization 现在支持在基于 PyTorch 2 torch.export 的流程中进行量化。  这包括对内置的 XNNPACK 和 X64Inductor Quantizer 的支持，以及指定自己的 Quantizer 的能力。

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

有关如何编写自己的 Quantizer 的说明，请参阅此教程。

【原型版】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 的 cpp_wrapper

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 内核。  现在，如果 CPU 支持 AVX512 指令，PyTorch 默认使用 AVX512 CPU 内核，这相当于在以前的版本中设置 ATEN_CPU_CAPABILITY=avx512。  可以通过设置 ATEN_CPU_CAPABILITY=avx2 来启用先前的行为。

针对 scaled-dot-product-attention (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 代码生成。