我们很高兴地宣布 PyTorch® 2.6 版本发布(发布说明)!此版本针对 PT2 有多项改进:torch.compile
现在可以与 Python 3.13 一起使用;新增了与性能相关的设置项 torch.compiler.set_stance
;还有一些 AOTInductor 的增强功能。除了 PT2 的改进之外,另一个亮点是 X86 CPU 对 FP16 的支持。
注意:从这个版本开始,我们将不再在 Conda 上发布,详情请参阅 [公告] 弃用 PyTorch 的官方 Anaconda 渠道。
在此版本中,实验性的 Linux 二进制文件附带了 CUDA 12.6.3(以及 Linux Aarch64、Linux ROCm 6.2.4 和 Linux XPU 二进制文件),它们是使用 CXX11_ABI=1 构建的,并且 采用了 Manylinux 2.28 构建平台。如果你使用自定义的 C++ 或 CUDA 扩展来构建 PyTorch 扩展,请将这些构建也更新为使用 CXX_ABI=1,并报告你遇到的任何问题。对于下一个 PyTorch 2.7 版本,我们计划将所有 Linux 构建切换到 Manylinux 2.28 和 CXX11_ABI=1,详情和讨论请参阅 [RFC] PyTorch 下一个 wheel 构建平台:manylinux - 2.28。
此外,在这个版本中,作为一项重要的安全改进措施,我们更改了 torch.load
的 weights_only
参数的默认值。这是一个向后兼容性有重大变更的改动,更多详情请参阅 此论坛帖子。
自 PyTorch 2.5 以来,此版本由 520 位贡献者的 3892 次提交组成。我们衷心感谢我们敬业的社区所做出的贡献。和往常一样,我们鼓励你尝试这些新功能,并在我们改进 PyTorch 的过程中报告任何问题。有关如何开始使用 PyTorch 2 系列的更多信息,可以在我们的 入门指南 页面找到。
测试版功能 | 原型功能 |
torch.compiler.set_stance | 英特尔 GPU 上改进的 PyTorch 用户体验 |
torch.library.triton_op | X86 CPU 上对大语言模型的 FlexAttention 支持 |
torch.compile 对 Python 3.13 的支持 | Dim.AUTO |
AOTInductor 的新打包 API | AOTInductor 的 CUTLASS 和 CK GEMM/CONV 后端 |
AOTInductor:精简器 | |
AOTInductor:ABI 兼容模式代码生成 | |
X86 CPU 对 FP16 的支持 |
*要查看公共功能提交的完整列表,请点击 此处。
BETA特性
[Beta] torch.compiler.set_stance
此功能使用户能够指定torch.compile在编译函数的不同调用之间可以采取的不同行为(“立场”)。例如,其中一个立场是 “eager_on_recompile” 指示 PyTorch 在需要重新编译时回退到 eager 而不是报错,并尽可能重用缓存的编译代码。
有关更多信息,请参阅 set_stance 文档 和 使用 torch.compiler.set_stance 进行动态编译控制 教程。
[Beta] torch.library.triton_op
torch.library.triton_op
提供了一种标准方法,用于创建由用户定义的 Triton 内核支持的自定义操作符。
当用户将用户定义的 Triton 内核转换为自定义操作符时,torch.library.triton_op
允许 torch.compile
深入了解其实现,从而使 torch.compile
能够对其中的 Triton 内核进行优化。
有关更多信息,请参阅 triton_op 文档 和 在 torch.compile 中使用用户定义的 Triton 内核 教程。
[Beta] torch.compile 对 Python 3.13 的支持
torch.compile
之前仅支持最高到 Python 3.12 版本。用户现在可以在 Python 3.13 中使用 torch.compile
来优化模型。
[Beta] AOTInductor 的新打包 API
引入了一种新的打包格式 “PT2 存档”。它本质上是一个包含 AOTInductor 需要使用的所有文件的压缩文件,允许用户将所需的所有内容发送到其他环境中。还有将多个模型打包到一个工件中的功能,以及在包内存储额外元数据的功能。
有关更多详细信息,请参阅更新后的 用于 Python 运行时的 torch.export AOTInductor 教程。
[Beta] AOTInductor:minifier
如果用户在使用 AOTInductor API 时遇到错误,AOTInductor Minifier 允许创建一个最小 nn.Module 来重现该错误。
有关更多信息,请参阅 AOTInductor 精简器文档。
[Beta] AOTInductor:ABI 兼容模式代码生成
AOTInductor 生成的模型代码依赖于 PyTorch 的 C++ 库。由于 PyTorch 发展迅速,确保之前由 AOTInductor 编译的模型能够在较新的 PyTorch 版本上继续运行,即 AOTInductor 的向后兼容性,这一点非常重要。
为了保证应用程序二进制接口(ABI)的向后兼容性,我们在 libtorch 中精心定义了一组稳定的 C 接口,并确保 AOTInductor 生成的代码仅引用这组特定的 API,而不引用 libtorch 中的其他内容。我们将在不同的 PyTorch 版本中保持这组 C API 的稳定性,从而为 AOTInductor 编译的模型提供向后兼容性保证。
[Beta] FP16 支持 X86 CPU(eager 模式和 Inductor 模式)
Float16 数据类型通常用于在 AI 推理和训练中减少内存使用量并加快计算速度。最近推出的带有 P-Cores 的 Intel® Xeon® 6 等 CPU通过原生加速器 AMX 支持 Float16 数据类型。X86 CPU 上的 Float16 支持在 PyTorch 2.5 中作为原型功能引入,现在它已针对 eager 模式和 Torch.compile + Inductor 模式进行了进一步改进,使其成为 Beta 级功能,其功能和性能均已通过广泛的验证。
原型功能
[Prototype] 英特尔 GPU 上改进的 PyTorch 用户体验
通过简化安装步骤、发布 Windows 二进制文件以及扩大支持的 GPU 型号范围(包括最新的英特尔® Arc™ B 系列独立显卡),PyTorch 在英特尔 GPU 上的用户体验得到了进一步提升。希望在 英特尔® 酷睿™ Ultra AI 电脑 和 英特尔® Arc™ 独立显卡 上使用 PyTorch 模型进行微调、推理和开发的应用程序开发人员和研究人员现在可以直接安装适用于 Windows、Linux 和 Windows 子系统 for Linux 2 的二进制版本的 PyTorch。
- 简化了英特尔 GPU 软件栈的设置,实现了一键安装 torch - xpu PIP 软件包,以便开箱即用的方式运行深度学习工作负载,消除了安装和激活英特尔 GPU 开发软件包的复杂性。
- 针对英特尔 GPU 发布了 torch 核心、torchvision 和 torchaudio 的 Windows 二进制版本,并且支持的 GPU 型号已从配备英特尔® Arc™ 显卡的英特尔® 酷睿™ Ultra 处理器、配备英特尔® Arc™ 显卡的英特尔® 酷睿™ Ultra 2 系列处理器 和 英特尔® Arc™ A 系列显卡 扩展到最新的 GPU 硬件 英特尔® Arc™ B 系列显卡。
- 进一步增强了使用 SYCL* 内核的英特尔 GPU 上 Aten 运算符的覆盖范围,以实现流畅的 eager 模式执行,以及对英特尔 GPU 上的 torch.compile 进行错误修复和性能优化。
有关英特尔 GPU 支持的更多信息,请参阅 入门指南。
[Prototype] X86 CPU 上对大语言模型的 FlexAttention 支持
FlexAttention 最初在 PyTorch 2.5 中引入,旨在为 Attention 变体提供优化的实现,并提供灵活的 API。在 PyTorch 2.6 中,通过 TorchInductor CPP 后端为 X86 CPU 添加了对 FlexAttention 的支持。此新功能利用并扩展了当前的 CPP 模板能力,基于现有的 FlexAttention API 支持广泛的 Attention 变体(例如:对大语言模型推理至关重要的 PageAttention),并在 x86 CPU 上带来了优化的性能。有了这个功能,就可以轻松地使用 FlexAttention API 在 CPU 平台上构建 Attention 解决方案并实现良好的性能。
[Prototype] Dim.AUTO
Dim.AUTO
允许在 torch.export
中使用自动动态形状。用户可以使用 Dim.AUTO
进行导出,并“发现”其模型的动态行为,其中最小/最大范围、维度之间的关系以及静态/动态行为将被自动推断出来。
与现有的用于指定动态形状的命名维度方法相比,这是一种更加用户友好的体验,因为现有的方法要求用户在导出时完全了解其模型的动态行为。Dim.AUTO
允许用户编写不依赖于模型的通用代码,从而提高了使用动态形状进行导出的易用性。
有关更多信息,请参阅 torch.export 教程。
[Prototype] AOTInductor 的 CUTLASS 和 CK GEMM/CONV 后端
CUTLASS 和 CK 后端为 Inductor 中的 GEMM 自动调优添加了内核选择。现在 AOTInductor 中也可以使用这些后端,它可以在 C++ 运行时环境中运行。这两个后端的一个重大改进是通过消除冗余的内核二进制编译和支持动态形状,提高了编译时的速度。