Skip to content

新版本: PyTorch 2.2 版本,FlashAttention-v2 集成,AOTInductor

发布: 2024年01月30日

译者:@片刻

原文: https://pytorch.org/blog/pytorch2-2

翻译: https://pytorch.apachecn.org/LatestChanges/PyTorch_V2.2

来自 PyTorch团队

我们很高兴地宣布PyTorch® 2.2(发布说明)的发布!PyTorch 2.2通过FlashAttention-v2集成以及AOTInductor(一种为非python服务器端部署构建的新的提前编译和部署工具)提供了约2倍的性能改进,以达到scaledot_product_attention

此版本还包括对优化器的改进的torch.compile支持,一些新的电感优化,以及名为TORCH_LOGS的新日志记录机制。

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

除了2.2,我们还发布了对PyTorch域库的一系列更新。更多详细信息可以在图书馆更新博客中找到。

自PyTorch 2.1以来,此版本由3628个提交和521个贡献者组成。我们要真诚地感谢我们敬业的社区的贡献。一如既往,我们鼓励您尝试这些,并报告任何问题,因为我们改进了2.2。有关如何开始使用PyTorch 2系列的更多信息,请访问我们的入门页面。

总结:

  • scaled_dot_product_attention(SDPA)现在支持FlashAttention-2,与之前的版本相比,速度提高了约2倍。
  • PyTorch 2.2引入了TorchInductor的一个新的提前扩展,称为AOTInductor,旨在为非python服务器端编译和部署PyTorch程序。
  • torch.distributed支持一种名为device_mesh的新抽象,用于初始化和表示ProcessGroups。
  • PyTorch 2.2提供了一种标准化、可配置的日志记录机制,称为TORCH_LOGS
  • PyTorch 2.2中包含一些torch.compile改进,包括改进了对编译优化器的支持,以及改进了TorchInductor融合和布局优化。
  • 请注意,我们正在弃用macOS x86支持,PyTorch 2.2.x将是支持macOS x64的最后一个版本。
Stable Beta Performance Improvements
FlashAttention-2集成 Inductor 优化
AOTInductor aarch64 优化
TORCH_LOGS
device_mesh
优化器编译
  • 要查看公共功能提交的完整列表,请单击此处

测试版功能

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

torch.nn.functional.scaled_dot_product_attention(SDPA)现在支持FlashAttention-2,速度约为2倍(与之前版本相比),并在A100 GPU上达到理论最大FLOP/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引入了一种新的抽象,用于表示分布式并行性所涉及的进程组,称为torch.distributed.device_mesh。这种抽象允许用户通过N维数组表示节点间和节点内进程组,例如,一个维度可以在FSDP中表示数据并行性,而另一个维度可以表示FSDP中的张量并行性。

有关更多信息,请参阅 device_mesh 教程

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

对torch.compile-ing优化器进行了一些改进,包括减少开销和对cuda图的支持。

有关改进的更多技术细节可在dev-discuss上找到,并且可以在这里找到torch.compile-ing优化器的配方。

性能改进

Inductor性能优化

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

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

aarch64性能优化

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

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


我们一直在努力

apachecn/AiLearning

【布客】中文翻译组