我们非常高兴地宣布发布 PyTorch® 2.5（发布说明）！此版本引入了用于 SDPA 的全新 cuDNN 后端，使得在 H100 或更新的 GPU 上使用 SDPA 的用户默认即可获得加速。此外，torch.compile 的区域编译功能通过允许用户编译重复的 nn.Module（例如 LLM 中的 Transformer 层）而无需重新编译，从而提供了一种减少 torch.compile 冷启动时间的方法。最后，TorchInductor CPP 后端通过 FP16 支持、CPP 包装器、AOT-Inductor 模式和 max-autotune 模式等多项增强功能，提供了稳健的性能加速。
自 PyTorch 2.4 发布以来，该版本共包含来自 504 位贡献者的 4095 次提交。我们衷心感谢社区成员的无私贡献。一如既往，我们鼓励您尝试这些新功能并报告遇到的任何问题，以便我们在 2.5 版本中不断改进。关于如何开始使用 PyTorch 2 系列的更多信息，请访问我们的入门指南页面。
此外，也请查看我们生态系统项目的新版本：TorchRec 和 TorchFix。
| Beta | 原型 |
| 用于 SDPA 的 cuDNN 后端 | FlexAttention |
| 无需重新编译的 torch.compile 区域编译 | 编译后的自动求导 (Compiled Autograd) |
| TorchDynamo 增加了对异常处理和 MutableMapping 类型的支持 | 飞行记录器 (Flight Recorder) |
| TorchInductor CPU 后端优化 | CPU 上支持基于 GEMM 模板的 Max-autotune |
| Windows 上的 TorchInductor | |
| 在 Eager 模式和 TorchInductor CPP 后端下的 CPU 路径上支持 FP16 | |
| 自动加载设备扩展 (Autoload Device Extension) | |
| 增强的 Intel GPU 支持 |
*要查看完整的公开功能提交列表，请点击此处。
测试版功能
[测试版] 用于 SDPA 的 cuDNN 后端
cuDNN “Fused Flash Attention” 后端已正式适配 torch.nn.functional.scaled_dot_product_attention。在 NVIDIA H100 GPU 上，这比 FlashAttentionV2 最高可带来 75% 的速度提升。对于在 H100 或更新 GPU 上使用 SDPA 的所有用户，该加速功能均默认开启。
[测试版] 无需重新编译的 torch.compile 区域编译
通过 torch._dynamo.config.inline_inbuilt_nn_modules 实现无需重新编译的区域编译，该选项在 2.5+ 版本中默认为 True。此选项允许用户编译重复的 nn.Module（例如 LLM 中的 Transformer 层）而无需重新编译。与编译整个模型相比，此选项可以减少编译延迟，且性能仅下降 1%-5%。
更多信息请参见教程。
[测试版] 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 操作和 SDPA 的 oneDNN。Inductor CPU 后端在三个基准测试套件（TorchBench、Hugging Face 和 timms）中持续实现性能加速，在 193 个测试模型中，有 97.5% 的模型表现优于 Eager 模式。
原型功能
[原型] FlexAttention
我们引入了一个灵活的 API，只需几行地道的 PyTorch 代码即可实现各种注意力机制，如滑动窗口 (Sliding Window)、因果掩码 (Causal Mask) 和 PrefixLM。此 API 利用 torch.compile 生成融合的 FlashAttention 内核，消除了额外的内存分配，并实现了与手写实现相当的性能。此外，我们利用 PyTorch 的自动求导机制自动生成反向传播代码。该 API 还能利用注意力掩码的稀疏性，从而比标准注意力实现有显著改进。
更多信息和示例，请参阅官方博客文章和 Attention Gym。
[原型] 编译后的自动求导 (Compiled Autograd)
编译后的自动求导是 PT2 栈的一个扩展，允许捕获整个反向传播过程。与 AOT 分发器追踪的反向图不同，编译后的自动求导追踪延迟到反向执行时进行，这使其不受前向传递图断点的影响，并允许将反向钩子 (backward hooks) 记录到图中。
更多信息请参见教程。
[原型] 飞行记录器 (Flight Recorder)
飞行记录器是一种新的调试工具，旨在帮助调试卡死的作业。该工具通过在集合通信运行过程中持续捕获相关信息来工作。一旦检测到作业卡死，即可利用这些信息快速识别异常的 rank/机器，以及相应的代码调用栈。
更多信息请参阅以下教程。
[原型] CPU 上支持基于 GEMM 模板的 Max-autotune
torch.compile 中 Inductor CPU 后端的 Max-autotune 模式会在编译时对操作的多种实现进行性能分析，并选择表现最优的一种。这对于 GEMM 相关操作特别有益，它使用基于 C++ 模板的 GEMM 实现，作为基于 oneDNN 和 MKL 库的 ATen 方法的替代方案。我们为 x86 CPU 支持 FP32、BF16、FP16 和 INT8 以及尾声融合 (epilogue fusions)。我们在 dynamo 基准套件中观察到了高达 7% 的几何平均速度提升，在 LLM 推理的 next-token 延迟中提升高达 20%。
更多信息请参阅教程。
[原型] Windows 上的 TorchInductor CPU
torch.compile 中的 Inductor CPU 后端现已支持 Windows。目前 Inductor 在 Windows 上支持 MSVC (cl)、clang (clang-cl) 和 Intel 编译器 (icx-cl)。
更多详情请查看教程。
[原型] 在 Eager 模式和 TorchInductor CPP 后端下的 CPU 路径上支持 FP16
Float16 是一种常用的精简浮点类型，可用于提升神经网络推理/训练的性能。从本版本开始，CPU 路径上的 Eager 模式和 TorchInductor 均支持 float16。
[原型] 自动加载设备扩展 (Autoload Device Extension)
PyTorch 现在支持外部设备扩展的自动加载，无需手动导入，从而简化了集成流程。此功能通过 torch.backends 入口点启用，确保扩展能无缝加载，同时也允许用户根据需要通过环境变量禁用它。
更多信息请参见教程。
[原型] 增强的 Intel GPU 支持
针对 Intel® Data Center GPU Max 系列和 Intel® Client GPU（内置 Intel® Arc™ 显卡的 Intel® Core™ Ultra 处理器以及 dGPU 部件的 Intel® Arc™ Graphics）的 Intel GPU 支持增强现已上线，旨在让您在 PyTorch 2.5 版本中更轻松地加速 Intel GPU 上的机器学习工作流。在此版本中，我们还启用了 PyTorch 在 Windows 上对 Intel® Client GPU 的初步支持。
- 扩展了 PyTorch 硬件后端支持矩阵，现已包含 Intel 数据中心 GPU 和客户端 GPU。
- 实现了 SYCL* 内核，以增强 Aten 算子在 Intel GPU 上的覆盖范围和执行效率，从而提升 PyTorch Eager 模式的性能。
- 增强了 torch.compile 的 Intel GPU 后端，以改善各种深度学习工作负载的推理和训练性能。
这些功能可通过 PyTorch 预览版和 nightly 二进制 PIP 包获取。有关 Intel GPU 支持的更多信息，请参阅文档。