跳转到主要内容
公告

PyTorch 2.2:FlashAttention-v2 集成,AOTInductor

作者: 2024 年 1 月 30 日2025 年 4 月 30 日暂无评论

我们很高兴宣布 PyTorch® 2.2 版本发布(发布说明)!PyTorch 2.2 通过集成 FlashAttention-v2,使 scaled_dot_product_attention 的性能提升约 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 系列的更多信息,请访问我们的入门页面。

总结

  • scaled_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 改进,包括对编译 Optimizers 的支持改进以及 TorchInductor 融合和布局优化的改进。
  • 请注意,我们正在弃用 macOS x86 支持,PyTorch 2.2.x 将是支持 macOS x64 的最后一个版本。
稳定版 测试版 性能改进
  FlashAttention-2 集成 Inductor 优化
  AOTInductor aarch64 优化
  TORCH_LOGS  
  device_mesh  
  Optimizer 编译  

*要查看完整的公开功能提交列表,请点击此处

测试版特性

[测试版] 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-ed 程序的预编译和部署

AOTInductor 是 TorchInductor 的扩展,旨在处理导出的 PyTorch 模型、对其进行优化并生成共享库以及其他相关工件。这些编译后的工件可以部署在非 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-ing Optimizers 的改进

对 torch.compile-ing Optimizers 进行了多项改进,包括更少的开销和对 cuda graphs 的支持。

有关改进的更多技术细节可在 dev-discuss 上找到,有关 torch.compile-ing 优化器的使用指南可在此处找到

性能改进

Inductor 性能优化

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

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

aarch64 性能优化

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

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