PyTorch 2.2 大更新：集成 FlashAttention-2，性能提升 2 倍

图片地址：https://img.ithome.com/newsuploadfiles/2024/2/3c2e20ff-9e8e-4054-9194-de21ff3a313f.png?x-bce-process=image/format,f_auto

新的一年，PyTorch 也迎来了重大更新，PyTorch 2.2 集成了 FlashAttention-2 和 AOTInductor 等新特性，计算性能翻倍。继去年十月份的 PyTorch 大会发布了 2.1 版本之后，全世界各地的 521 位开发者贡献了 3628 个提交，由此形成了最新的 PyTorch 2.2 版本。新的版本集成了 FlashAttention-2，使得 scaled_dot_product_attention（SDPA）相较于之前的版本有了约 2 倍的性能提升。PyTorch 2.2 还引入了一个新的 TorchInductor 提前扩展，称为 AOTInductor，旨在为非 python 服务器端编译和部署 PyTorch 程序。PyTorch 中的 torch.distributed支持了一个叫作 device_mesh 的新抽象，用于初始化和表示 ProcessGroups。另外，PyTorc