我们很高兴地宣布发布 PyTorch® 2.5(release note)!此版本引入了一个新的 CuDNN 后端用于 SDPA,默认情况下为使用 H100 或更新 GPU 的 SDPA 用户提供加速。此外,torch.compile 的区域编译提供了一种减少 torch.compile 冷启动时间的方法,允许用户在不重新编译的情况下编译重复的 nn.Module(例如 LLM 中的 transformer 层)。最后,TorchInductor CPP 后端通过诸如 FP16 支持、CPP 包装器、AOT-Inductor 模式和最大自动调优模式等众多增强功能提供了显著的性能提升。
此版本由 504 位贡献者的 4095 次提交组成,自 PyTorch 2.4 以来。我们衷心感谢我们敬业的社区的贡献。和往常一样,我们鼓励您尝试这些功能并报告任何问题,以便我们改进 2.5。有关如何开始使用 PyTorch 2 系列的更多信息,请访问我们的入门页面。
此外,请查看我们的新生态系统项目发布,包括 TorchRec 和 TorchFix。
Beta | 原型 |
用于 SDPA 的 CuDNN 后端 | FlexAttention |
torch.compile 区域编译,无需重新编译 | Compiled Autograd |
TorchDynamo 增加了对异常处理和 MutableMapping 类型的支持 | Flight Recorder |
TorchInductor CPU 后端优化 | CPU 上使用 GEMM 模板的最大自动调优支持 |
TorchInductor支持 Windows | |
在 CPU 上对 eager 模式和 TorchInductor CPP 后端的 FP16 支持 | |
设备自动加载扩展 | |
增强 Intel GPU 支持 |
*查看完整的功能提交列表,请点击这里.
BETA 功能
[Beta] SDPA 的 CuDNN 后端
cuDNN “融合闪存注意力” 后端已应用于 torch.nn.functional.scaled_dot_product_attention。在 NVIDIA H100 GPU 上,这可以比 FlashAttentionV2 提供高达 75% 的速度提升。对于所有在 H100 或更新 GPU 上使用 SDPA 的用户,此速度提升默认启用。
[Beta] torch.compile 区域编译无需重新编译
通过 torch._dynamo.config.inline_inbuilt_nn_modules 实现的区域编译无需重新编译,该选项在 2.5+ 版本中默认启用。此选项允许用户编译重复的 nn.Module(例如 LLM 中的 transformer 层)而无需重新编译。与编译整个模型相比,此选项可以减少编译延迟,性能下降在 1%-5% 之间。
有关更多信息,请参阅教程。
[Beta] TorchInductor CPU 后端优化
此功能改进了 Inductor 的 CPU 后端优化,包括 CPP 后端代码生成和带有自定义 CPU 内核的 FX 融合。Inductor CPU 后端支持常见数据类型的向量化和所有 Inductor IR 操作,以及静态和符号形状。它兼容 Linux 和 Windows 操作系统,并支持默认的 Python 包装器、CPP 包装器和 AOT-Inductor 模式。
此外,它扩展了 GEMM 模板的最大自动调优模式(在 2.5 中原型化),提供了进一步的性能提升。后端支持各种 FX 融合,降低到自定义内核,如用于线性/卷积操作和 SDPA 的 oneDNN。 Inductor CPU 后端在三个基准测试套件——TorchBench、Hugging Face 和 timms 中始终实现性能提升,在测试的 193 个模型中有 97.5% 超过了 eager 模式。
原型功能
[Prototype] FlexAttention
我们引入了一个灵活的API,使得只需几行惯用的PyTorch代码即可实现各种注意力机制,如滑动窗口、因果掩码和PrefixLM。该API利用torch.compile生成融合的FlashAttention内核,消除了额外的内存分配,并实现了与手写实现相当的性能。此外,我们使用PyTorch的autograd机制自动生成反向传播。此外,我们的API可以利用注意力掩码中的稀疏性,与标准注意力实现相比显著提高了性能。
有关更多信息和示例,请参阅官方博客文章和Attention Gym。
[Prototype] 编译的Autograd
编译的Autograd是PT2堆栈的扩展,允许捕获整个反向传播过程。与AOT调度器跟踪的反向图不同,编译的Autograd跟踪推迟到反向执行时间,这使其不受前向传递图中断的影响,并允许将反向钩子记录到图中。
有关更多信息,请参阅教程。
[Prototype] Flight Recorder
Flight recorder 是一种新的调试工具,有助于调试卡住的作业。该工具通过持续捕获集体运行时的信息来工作。在检测到卡住的作业时,可以使用这些信息快速识别行为异常的排名/机器以及代码堆栈跟踪。
有关更多信息,请参阅以下教程。
[Prototype] CPU上的GEMM模板最大自动调优支持
torch.compile中的Inductor CPU后端的最大自动调优模式在编译时分析操作的多种实现,并选择性能最佳的实现。这对于GEMM相关操作特别有益,使用基于C++模板的GEMM实现作为基于ATen的方法与oneDNN和MKL库的替代方案。我们支持x86 CPU的FP32、BF16、FP16和INT8以及尾融合。我们在dynamo基准测试套件上看到了高达7%的几何平均加速,并在LLM推理的下一个令牌延迟上看到了高达20%的提升。
有关更多信息,请参阅教程。
[Prototype] Windows上的TorchInductor CPU
torch.compile中的Inductor CPU后端现在可以在Windows上运行。我们目前支持Windows上的MSVC(cl)、clang(clang-cl)和Intel编译器(icx-cl)。
有关更多详细信息,请参阅教程。
[Prototype] CPU 上的 FP16 支持,包括 eager 模式和TorchInductor CPP后端
Float16是一种常用的减少浮点类型,用于提高神经网络推理/训练的性能。从此版本开始,CPU上的 eager 模式和TorchInductor都支持float16。
[Prototype] 自动加载设备扩展
PyTorch现在支持自动加载树外设备扩展,通过去除手动导入模块的代码简化集成。该功能通过torch.backends入口点启用,确保无缝扩展加载,同时用户可以通过环境变量禁用它。
有关更多信息,请参阅教程。
[Prototype] 增强的Intel GPU支持
Intel GPU支持增强现已适用于Intel®数据中心GPU Max系列和Intel®客户端GPU(带有内置Intel® Arc™图形的Intel® Core™ Ultra处理器和用于dGPU部件的Intel® Arc™图形),这使得在PyTorch 2.5版本中更容易在Intel GPU上加速您的机器学习工作流程。我们还在此版本中启用了对Intel®客户端GPU在Windows上的初步支持。
- 扩展的PyTorch硬件后端支持矩阵,包括Intel数据中心和客户端GPU。
- 实现SYCL*内核以增强在Intel GPU上执行Aten操作的覆盖范围和性能,以提高PyTorch急切模式的性能。
- 增强的torch.compile的Intel GPU后端,以提高广泛的深度学习工作负载的推理和训练性能。
这些功能通过PyTorch预览和nightly二进制PIP wheels提供。有关Intel GPU支持的更多信息,请参阅文档。