PyTorch 团队

我们很高兴地宣布 PyTorch® 2.5 版本 (发行说明) 发布!此版本为 SDPA 提供了一个新的 cuDNN 后端,默认情况下可为在 H100 或更新 GPU 上使用 SDPA 的用户带来加速。此外,torch.compile 的区域编译功能提供了一种减少 torch.compile 冷启动时间的方法,允许用户编译重复的 nn.Module(例如 LLM 中的 transformer 层)而无需重新编译。最后,TorchInductor CPP 后端通过大量增强功能(如 FP16 支持、CPP 包装器、AOT-Inductor 模式和 max-autotune 模式)提供了显著的性能加速。

自 PyTorch 2.4 以来,此版本包含来自 504 位贡献者的 4095 次提交。我们衷心感谢我们敬业的社区所做的贡献。一如既往,我们鼓励您试用这些新功能并在我们改进 2.5 版本时报告任何问题。有关如何开始使用 PyTorch 2 系列的更多信息,请访问我们的入门页面。

同时,请查看我们新的生态系统项目版本:TorchRecTorchFix

Beta 版 原型
用于 SDPA 的 cuDNN 后端 FlexAttention
torch.compile 区域编译无需重新编译 编译式 Autograd
TorchDynamo 添加了对异常处理和 MutableMapping 类型的支持 Flight Recorder
TorchInductor CPU 后端优化 通过 GEMM 模板在 CPU 上支持 Max-autotune
Windows 上的 TorchInductor
Eager 模式和 TorchInductor CPP 后端在 CPU 路径上支持 FP16
自动加载设备扩展
增强的 Intel GPU 支持

*要查看完整的公共特性提交列表,请点击这里

BETA 特性

[Beta] 用于 SDPA 的 cuDNN 后端

cuDNN 的“Fused Flash Attention”后端已集成到 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+ 版本中默认为 True。此选项允许用户编译重复的 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 模板的 max-autotune 模式(在 2.5 中原型化),提供了进一步的性能提升。后端支持各种 FX 融合,下沉到定制内核,例如用于 Linear/Conv 操作的 oneDNN 和 SDPA。在 TorchBench、Hugging Face 和 timm 三个基准测试套件中,Inductor CPU 后端始终实现性能加速,在测试的 193 个模型中,有 97.5% 的模型性能优于 eager 模式。

原型特性

[原型] FlexAttention

我们引入了一个灵活的 API,只需几行地道的 PyTorch 代码即可实现各种注意力机制,例如 Sliding Window、Causal Mask 和 PrefixLM。此 API 利用 torch.compile 生成融合的 FlashAttention 内核,消除了额外的内存分配,并实现了与手写实现相当的性能。此外,我们使用 PyTorch 的 autograd 机制自动生成反向传播。此外,我们的 API 可以利用注意力掩码中的稀疏性,从而比标准注意力实现带来显著改进。

有关更多信息和示例,请参阅官方博客文章Attention Gym

[原型] 编译式 Autograd

编译式 Autograd 是 PT2 栈的扩展,允许捕获整个反向传播。与 AOT 调度器跟踪的反向图不同,编译式 Autograd 的跟踪延迟到反向执行时进行,这使得它不受正向传播图中断的影响,并允许它将反向钩子记录到图中。

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

[原型] Flight Recorder

Flight Recorder 是一种新的调试工具,可帮助调试卡住的任务。该工具通过持续捕获集体操作运行时的信息来工作。检测到卡住的任务后,可以使用这些信息快速识别出现问题的等级/机器以及代码堆栈跟踪。

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

[原型] 通过 GEMM 模板在 CPU 上支持 Max-autotune

torch.compile 中 Inductor CPU 后端的 Max-autotune 模式会在编译时对操作的多种实现进行性能分析,并选择性能最佳的实现。这对于 GEMM 相关操作特别有利,它使用基于 C++ 模板的 GEMM 实现作为替代 ATen 基于 oneDNN 和 MKL 库的方法。我们支持 x86 CPU 的 FP32、BF16、FP16 和 INT8 以及 epilogue 融合。在 dynamo 基准测试套件中,我们看到了高达 7% 的几何平均加速,在 LLM 推理中,下一 token 延迟提高了高达 20%。

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

[原型] Windows 上的 TorchInductor CPU

torch.compile 中的 Inductor CPU 后端现在可以在 Windows 上工作。目前我们支持 MSVC (cl)、clang (clang-cl) 和 Intel 编译器 (icx-cl) 用于 Windows inductor。

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

[原型] Eager 模式和 TorchInductor CPP 后端在 CPU 路径上支持 FP16

Float16 是一种常用的缩减浮点类型,用于提高神经网络推理/训练的性能。自此版本起,CPU 路径支持 eager 和 TorchInductor 的 float16。

[原型] 自动加载设备扩展

PyTorch 现在支持树外设备扩展的自动加载,通过消除手动导入的需要来简化集成。此功能通过 torch.backends 入口点启用,通过确保无缝的扩展加载来简化使用,同时允许用户在需要时通过环境变量禁用它。

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

[原型] 增强的 Intel GPU 支持

Intel GPU 支持增强现在可用于 Intel® Data Center GPU Max 系列和 Intel® Client GPU(内置 Intel® Arc™ 显卡的 Intel® Core™ Ultra 处理器以及用于 dGPU 部分的 Intel® Arc™ Graphics),旨在通过 PyTorch 2.5 版本更轻松地加速 Intel GPU 上的机器学习工作流程。在此版本中,我们还启用了 PyTorch 在 Windows 上对 Intel® Client GPU 的初始支持。

  • 扩展了 PyTorch 硬件后端支持矩阵,以包含 Intel Data Center 和 Client GPU。  
  • 实现了 SYCL* 内核,以增强 Aten 操作在 Intel GPU 上的覆盖范围和执行,从而提升 PyTorch eager 模式下的性能。
  • 增强了 torch.compile 的 Intel GPU 后端,以提高各种深度学习工作负载的推理和训练性能。

这些特性可通过 PyTorch 预览版和 nightly 二进制 PIP wheel 获取。有关 Intel GPU 支持的更多信息,请参阅文档