PyTorch 2 2 大更新！ 集成的 FlashAttention 2 将性能提高了 2 倍

编辑：alan

在新的一年里，PyTorch 也迎来了重大更新！

继 2 发布之后在版本 1 之后，全球 521 名开发人员为最新版本的 PyTorch 2 贡献了 3,628 次提交版本 2。

新版本集成了 FlashAttention-2，与之前的版本相比，性能提高了约 2 倍。

pytorch 2.2 还引入了一个新的 TorchInductor 高级扩展，称为 Aotinductor，旨在为非 Python 服务器端编译和部署 PyTorch 程序。

火炬。 在 pytorch 中分布式支持一个名为 Device Mesh 的新抽象，用于初始化和表示 ProcessGroups。

另外，pytorch 22提供标准化、可配置的日志记录机制,——火炬日志。

pytorch 2.2.也在火炬上Compile 进行了许多改进，包括改进了对编译优化器的支持，以及火炬电感融合和布局优化。

最后，值得注意的是，PyTorch 将放弃对 macOS x86、PyTorch 2 的支持2.X 是支持 macOS X64 的最后一个版本。

pytorch 2.2 个新功能

首先，请注意，如果您从源代码构建 pytorch 22. 海湾合作委员会 94 或更高版本，PyTorch 库已从 C++14 迁移到 C++ 17。

FlashAttention-2 通过优化 GPU 上不同线程块和扭曲之间的工作分区，解决了使用率低或不必要的共享内存读写问题。

FlashAttention-2 调整了算法以减少非 matmul 计算量，同时提高了注意力计算的并行性（即使是单个标头也可以跨不同的线程块以增加使用率），并且在每个线程块中，warp 之间的工作分配都经过优化，以减少通过共享内存的通信。

pytorch 2.2 将 FlashAttention 内核更新为 v2，但需要注意的是，之前的 Flash Attention 内核有 Windows 实现，Windows 用户可以强制使用 SDP 内核，并且只启用 Flash Attention 的上下文管理器。

而在 22、如果必须使用SDP内核上下文管理器，请使用内存效率或数学内核（在Windows上）。

在闪光灯的加持下-2，火炬nn.functional.Scaled Dot Product Attention 的速度提高了约 2 倍，达到 A100 GPU 理论峰值的 50%-73%。

Aotinductor 是 Torchinductor 的扩展，用于处理导出的 PyTorch 模型、优化它们以及生成共享库以及其他相关工件。

这些编译的项目可以部署在非 Python 环境中，通常用于服务器端推理。

以下示例演示如何调用 AOT Compile 将模型转换为共享库。

Aotinductor 支持与 Inductor 相同的后端，包括 CUDA、ROCM 和 CPU。

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

可以通过 torch logs 环境变量启用日志记录。 例如，通过在命令行中修改环境变量：

将 torchdynamo 的日志级别设置为日志记录错误，将 TorchInductor 的 log 级别设置为 loggingdebug。

当然，它也可以以 API 的形式使用 **：

pytorch 2.2 引入了一种新的抽象，用于表示分布式并行中涉及的进程组，称为 torchdistributed.device_mesh。

为分布式训练设置分布式通信器 （NCCL） 是一件繁琐的事情。 用户需要编写具有不同并行度的工作负载，并针对每个并行度手动设置和管理 NCCL 通信器 （ProcessGroup）。

此过程可能很复杂且容易出错。 DeviceMesh 可以简化此过程并使其更易于管理。

DeviceMesh 是用于管理进程组的更高级别的抽象。 它允许用户毫不费力地创建节点间和节点内进程组，而不必担心如何正确设置不同子进程组的层次结构。

例如，数组的一个维度可以表示 FSDP 中的数据并行性，而另一个维度可以表示 FSDP 中的张量并行性。

用户还可以通过DeviceMesh轻松管理底层进程组，实现多维并行。

DeviceMesh 在处理多维并行性（如 3D 并行性）时非常有用。 如上图所示，当您的并行解决方案需要在每个主机之间和每个主机内进行通信时，您可以创建一个 2D 网格，用于连接每个主机中的设备，并在同构设置中将每个设备连接到其他主机上的对应设备。

在 Init Device Mesh（） 的帮助下，我们只需两行即可完成上述 2D 设置：

如果我们不使用 devicemesh，我们可能需要编写以下一堆 **：

当然，如果需要，我们仍然可以访问底层的 processgroup：if need

可能有以下几种：

编译优化器提高了所有基准测试的性能：HuggingFace +18%、Torchbench +19%、TIMM +8% E2E;

编译优化器增加了对 cudagraph 的支持;

对测试套件中的所有模型进行平均，每个测试套件的平均基准编译时间增加了约 40 秒; 正在进行的优化可能会将其缩短到 30 秒以内。

在调度器中，将去中心化过程中注册的缓冲区列表压缩到 ForeachKernelSchedulerNodes（FusedSchedulerNode 的子类）中。

为了检查融合是否合法，每个内部调度节点执行的写入操作必须与使用调度节点位于同一列表索引的读取操作匹配。

此外，正常的垂直收敛规则必须允许在消费者和生产者调度器节点列表的每个索引处进行收敛。

如果满足这些条件，ForeachKernelSchedulerNode 将垂直融合到 ForeachKernelSchedulerNode 中，其中每个列表上的相应点操作将被融合。

通过实现这种融合，可以将一系列 foreach 操作融合到单个内核中，从而实现多张量优化器的完全融合。

TorchInductor 添加了许多性能优化，包括对 Torch 的支持Concat 的水平融合支持、改进的卷积布局优化以及改进的缩放点产品注意力模式匹配。

pytorch 2.2 还包括对 ARClCloud64 的许多性能增强，包括对 mkldnn 权重预打包的支持、改进的 iDeep 基元缓存，以及通过对 Onednn 的固定格式内核改进来提高推理速度。

引用：