我们很高兴地宣布发布 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 系列的更多信息，请访问我们的入门页面。
此外，请查看我们与TorchRec和TorchFix一起发布的新生态系统项目。
| Beta | 原型 |
| SDPA 的 cuDNN 后端 | FlexAttention |
| torch.compile 区域编译无需重新编译 | 编译 Autograd |
| TorchDynamo 添加了对异常处理和 MutableMapping 类型的支持 | 飞行记录仪 |
| TorchInductor CPU 后端优化 | CPU 上支持 GEMM 模板的 Max-autotune |
| Windows 上的 TorchInductor | |
| 针对 eager 模式和 TorchInductor CPP 后端在 CPU 路径上支持 FP16 | |
| 自动加载设备扩展 | |
| 增强的 Intel GPU 支持 |
*要查看完整的公开功能提交列表，请点击此处。
BETA 功能
[Beta] SDPA 的 cuDNN 后端
cuDNN“融合 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 OS 兼容，并支持默认的 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 代码即可实现各种注意力机制，例如滑动窗口、因果掩码和 PrefixLM。此 API 利用 torch.compile 生成融合的 FlashAttention 内核，消除了额外的内存分配，并实现了与手写实现相当的性能。此外，我们使用 PyTorch 的 autograd 机制自动生成反向传播。此外，我们的 API 可以利用注意力掩码中的稀疏性，从而显着优于标准注意力实现。
有关更多信息和示例，请参阅官方博客文章和Attention Gym。
[原型] 编译 Autograd
编译 Autograd 是 PT2 栈的扩展，允许捕获整个反向传播。与 AOT 调度器追踪的反向图不同，编译 Autograd 追踪被推迟到反向执行时，这使其不受前向传播图中断的影响，并允许它将反向钩子记录到图中。
有关更多信息，请参阅教程。
[原型] 飞行记录仪
飞行记录仪是一种新的调试工具，可帮助调试卡住的任务。该工具通过在集体运行时持续捕获有关集体的信息来工作。一旦检测到卡住的任务，这些信息可用于快速识别行为不当的排名/机器以及代码堆栈跟踪。
有关更多信息，请参阅以下教程。
[原型] CPU 上支持 GEMM 模板的 Max-autotune
torch.compile 中 Inductor CPU 后端的 Max-autotune 模式在编译时分析操作的多种实现并选择性能最佳的实现。这对于 GEMM 相关操作特别有用，它使用基于 C++ 模板的 GEMM 实现作为使用 oneDNN 和 MKL 库的基于 ATen 方法的替代方案。我们支持 FP32、BF16、FP16 和 INT8，并为 x86 CPU 提供尾声融合。我们发现 dynamo 基准测试套件的几何平均速度提升高达 7%，LLM 推理的下一个 token 延迟提升高达 20%。
有关更多信息，请参阅教程。
[原型] Windows 上的 TorchInductor CPU
torch.compile 中的 Inductor CPU 后端现在可以在 Windows 上运行。我们目前支持适用于 Windows inductor 的 MSVC (cl)、clang (clang-cl) 和 Intel 编译器 (icx-cl)。
有关更多详细信息，请参阅教程。
[原型] 针对 eager 模式和 TorchInductor CPP 后端在 CPU 路径上支持 FP16
Float16 是一种常用的缩减浮点类型，用于提高神经网络推理/训练的性能。自此版本以来，CPU 路径上对 eager 和 TorchInductor 都支持 float16。
[原型] 自动加载设备扩展
PyTorch 现在支持对非树设备扩展的自动加载，通过消除手动导入的需要来简化集成。此功能通过 torch.backends 入口点启用，通过确保无缝扩展加载来简化使用，同时允许用户在需要时通过环境变量禁用它。
有关更多信息，请参阅教程。
[原型] 增强的 Intel GPU 支持
Intel GPU 支持增强现在可用于 Intel® 数据中心 GPU Max 系列和 Intel® 客户端 GPU（内置 Intel® Arc™ 显卡的 Intel® 酷睿™ Ultra 处理器以及用于 dGPU 部件的 Intel® Arc™ 显卡），这将使您在 PyTorch 2.5 版本中更轻松地加速 Intel GPU 上的机器学习工作流。我们在此版本中还启用了 PyTorch 在 Windows 上对 Intel® 客户端 GPU 的初始支持。
- 扩展了 PyTorch 硬件后端支持矩阵，包括 Intel 数据中心和客户端 GPU。
- SYCL* 内核的实现，以增强 Intel GPU 上 Aten 操作符的覆盖范围和执行，从而提高 PyTorch eager 模式的性能。
- 增强了 torch.compile 的 Intel GPU 后端，以提高各种深度学习工作负载的推理和训练性能。
这些功能可通过 PyTorch 预览版和每晚二进制 PIP 轮子获得。有关 Intel GPU 支持的更多信息，请参阅文档。