作者:PyTorch 团队

我们很高兴宣布发布 PyTorch® 2.2(发行说明)!PyTorch 2.2 通过集成 FlashAttention-v2scaled_dot_product_attention 带来了约 2 倍的性能提升,同时还引入了 AOTInductor,这是一个新的面向非 Python 服务器端部署的提前编译和部署工具。

此版本还改进了对 Optimizer 的 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 系列的更多信息可在我们的入门页面找到。

摘要

  • scaled_dot_product_attention (SDPA) 现已支持 FlashAttention-2,相比之前版本带来约 2 倍的速度提升。
  • PyTorch 2.2 引入了一个新的 TorchInductor 的提前 (ahead-of-time) 扩展,称为 AOTInductor,旨在为非 Python 服务器端编译和部署 PyTorch 程序。
  • torch.distributed 支持一种新的抽象,用于初始化和表示 ProcessGroups,称为 device_mesh
  • PyTorch 2.2 内置了一个标准化、可配置的日志记录机制,称为 TORCH_LOGS
  • PyTorch 2.2 包含多项 torch.compile 改进,包括改进了对编译 Optimizer 的支持以及改进的 TorchInductor 融合和布局优化。
  • 请注意,我们正在弃用 macOS x86 支持,PyTorch 2.2.x 将是最后一个支持 macOS x64 的版本。
稳定版 测试版 性能改进
FlashAttention-2 集成 Inductor 优化
AOTInductor aarch64 优化
TORCH_LOGS
device_mesh
Optimizer 编译

*要查看公共特性提交的完整列表,请点击此处

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 的教程,请参阅此教程

[测试版] AOTInductor:用于 torch.export 程序的提前编译和部署

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

更多信息请参阅AOTInductor 教程

[测试版] 通过 TORCH_LOGS 实现细粒度可配置日志记录

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

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

更多信息,请参阅日志文档教程

[测试版] torch.distributed.device_mesh

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

更多信息,请参阅device_mesh 教程

[测试版] torch.compile 编译 Optimizer 的改进

torch.compile 编译 Optimizer 方面进行了一些改进,包括减少开销和支持 cuda graphs。

有关这些改进的更多技术细节可在dev-discuss 上找到,有关 torch.compile 编译 optimizer 的示例可在此处找到。

性能改进

Inductor 性能优化

TorchInductor 中新增了多项性能优化,包括对 torch.concat 的水平融合支持改进的卷积布局优化以及改进的 scaled_dot_product_attention 模式匹配

有关 Inductor 优化的完整列表,请参阅发行说明

aarch64 性能优化

PyTorch 2.2 包含多项针对 aarch64 的性能增强,包括支持 mkldnn 权重预打包、改进的 ideep primitive 缓存以及通过对 OneDNN 进行固定格式内核改进来提高推理速度。

有关 aarch64 优化的完整列表,请参阅发行说明