我们很高兴地宣布 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 类型的支持 | 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 支持的更多信息,请参阅文档。