跳转到主要内容
博客

PyTorch 2.1：自动动态形状编译，分布式检查点

作者： 2023 年 10 月 4 日2024 年 11 月 14 日暂无评论

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

此外，此版本还提供了众多性能改进（例如 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 支持并行保存和加载多个进程的模型，以及因集群拓扑变化而导致的重新分片。
  • torch.compile 现在可以通过将 NumPy 操作转换为等效的 PyTorch 操作来编译它们。
  • torch.compile 现在包括对 Python 3.11 的改进支持。
  • 新的 CPU 性能特性包括 inductor 改进（例如 bfloat16 支持和动态形状）、AVX512 内核支持和 scaled-dot-product-attention 内核。
  • torch.export，一个完善的完整图捕获机制作为原型功能引入，以及基于 torch.export 的量化。
  • torch.sparse 现在包括对 NVIDIA® GPU 上半结构化 (2:4) 稀疏性的原型支持。
稳定版Beta原型性能改进
 自动动态形状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.11torchinductor 的 cpp_wrapper 
 torch.compile + autograd.Function  
 第三方设备集成：PrivateUse1  

*要查看 2.1、2.0 和 1.13 公共功能提交的完整列表，请单击此处。

Beta 功能

(Beta) 自动动态形状

动态形状是 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 上。

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

[Beta] torch.distributed.checkpoint

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

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

[Beta] torch.compile + NumPy

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

请参阅 torch.compile FAQ 中的此部分，了解有关 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 函数的反向函数，这为更大量使用扩展机制的模型解锁了训练优化。

[Beta] 改进的第三方设备支持：PrivateUse1

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

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

原型功能

[原型] torch.export()

torch.export() 提供了一种完善的跟踪机制，可以根据 PT2.0 提供的新技术从 PyTorch 程序捕获完整图。

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

欲了解更多信息，请参阅此教程。

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

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

有关使用 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 的 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.functiona.scaled_dot_product_attention 为 Transformer 原语提供了优化的 CUDA 实现。PyTorch 2.1 包含了优化的基于 FlashAttention 的 CPU 例程。

请参阅此处的文档。

bfloat16 的 CPU 优化

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