由 PyTorch 团队

我们非常高兴地宣布 PyTorch® 2.2（发布说明）的发布！PyTorch 2.2 通过 FlashAttention-v2 集成提供了约 2 倍的性能提升，以及 AOTInductor，这是一个为非 Python 服务器端部署构建的新型的即时编译和部署工具。

本次发布还包括对 Optimizers 的 torch.compile 支持的改进、多项新的 Inductor 优化以及名为 TORCH_LOGS 的新日志机制。

请注意，我们将弃用 macOS x86 支持，PyTorch 2.2.x 将是最后一个支持 macOS x64 的版本。

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

本次发布包含自 PyTorch 2.1 以来的 3,628 个提交和 521 位贡献者。我们衷心感谢我们忠诚的社区成员的贡献。一如既往，我们鼓励大家尝试这些新功能，并在我们改进 2.2 版本的过程中报告任何问题。有关如何开始使用 PyTorch 2 系列的信息，请参阅我们的入门页面。

摘要：

  • 扩展的 dot_product_attention（SDPA）现在支持 FlashAttention-2，与之前的版本相比，速度提升了约 2 倍。
  • PyTorch 2.2 引入了 TorchInductor 的新的预编译扩展，称为 AOTInductor，旨在编译和部署用于非 Python 服务器端的 PyTorch 程序。
  • torch.distributed 支持用于初始化和表示 ProcessGroups 的新抽象，称为 device_mesh。
  • PyTorch 2.2 发布了名为 TORCH_LOGS 的标准、可配置的日志机制。
  • PyTorch 2.2 包含多个 torch.compile 改进，包括对编译优化器支持的改进以及 TorchInductor 融合和布局优化。
  • 注意，我们将弃用 macOS x86 支持，PyTorch 2.2.x 将是最后一个支持 macOS x64 的版本。
稳定 测试版 性能提升
FlashAttention-2 集成 电感优化
AOT 电感 aarch64 优化
TORCH_LOGS
设备网格
优化器编译

*点击此处查看所有公开功能提交的完整列表。

测试功能

[Beta] 在 torch.nn.functional.scaled_dot_product_attention 中支持 FlashAttention-2

torch.nn.functional.scaled_dot_product_attention (SDPA) 现在支持 FlashAttention-2，速度提升约 2 倍（与之前版本相比），在 A100 GPU 上达到理论最大 FLOPs/s 的 50-73%。

本文中有关于 FlashAttention-2 的更多信息。

想了解如何使用 SDPA 的教程，请参阅此教程。

[Beta] AOTInductor：为 torch.export-ed 程序提供即时编译和部署

AOTInductor 是 TorchInductor 的扩展，旨在处理导出的 PyTorch 模型，优化它们，并生成共享库以及其他相关工件。这些编译工件可以在非 Python 环境中部署，这些环境通常用于服务器端推理。请注意，AOTInductor 支持与 Inductor 相同的后端，包括 CUDA、ROCm 和 CPU。

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

[Beta] 通过 TORCH_LOGS 进行细粒度可配置日志记录

PyTorch 现在提供了一种标准化的、可配置的日志记录机制，可用于分析各种子系统的状态，例如编译和分布式操作。

可以通过 TORCH_LOGS 环境变量启用日志记录。例如，要将 TorchDynamo 的日志级别设置为 logging.ERROR 并将 TorchInductor 的日志级别设置为 logging.DEBUG，请将 TORCH_LOGS 设置为”-dynamo,+inductor”传递给 PyTorch。

如需更多信息，请参阅日志文档和教程。

[Beta] torch.distributed.device_mesh

PyTorch 2.2 引入了一种新的抽象，用于表示参与分布式并行计算的 ProcessGroups，称为 torch.distributed.device_mesh。这种抽象允许用户通过一个 N 维数组来表示节点间和节点内进程组，例如，一个维度可以表示 FSDP 中的数据并行，另一个维度可以表示 FSDP 内部的张量并行。

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

[Beta] torch.compile-ing 优化器的改进

对 torch.compile-ing 优化器进行了多项改进，包括降低开销和支持 cuda 图。

改进的更多技术细节可在 dev-discuss 上找到，torch.compile-ing 优化器的配方请在此处查看。

性能提升

电感性能优化

已添加多项性能优化到 TorchInductor，包括对 torch.concat 的水平融合支持、改进的卷积布局优化以及改进的 scaled_dot_product_attention 模式匹配。

完整的 inductor 优化列表请参阅发布说明。

aarch64 性能优化

PyTorch 2.2 为 aarch64 提供了多项性能增强，包括对 mkldnn 权重预打包的支持、改进的 ideep 原语缓存以及通过 OneDNN 固定格式内核改进实现的推理速度提升。

关于 aarch64 优化功能的完整列表，请参阅发布说明。