作者:PyTorch 团队

我们很高兴宣布 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.loadweights_only 参数的默认值。这是一个破坏向后兼容性的更改,更多详情请参阅 此论坛帖子

自 PyTorch 2.5 以来,此版本包含来自 520 位贡献者的 3892 次提交。我们要衷心感谢我们专注的社区所做的贡献。与往常一样,我们鼓励您试用这些功能并报告任何问题,以便我们改进 PyTorch。有关如何开始使用 PyTorch 2 系列的更多信息,请访问我们的入门页面。

Beta Prototype
torch.compiler.set_stance 改进 Intel GPU 上的 PyTorch 用户体验
torch.library.triton_op X86 CPU 上对 LLM 的 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 在编译函数的不同调用之间可以采用的不同行为(“stances”)。例如,其中一种 stance 是

“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 文档将用户定义的 Triton 内核与 torch.compile 一起使用教程。

[Beta] torch.compile 支持 Python 3.13

torch.compile 之前仅支持最高版本为 3.12 的 Python。用户现在可以在 Python 3.13 中使用 torch.compile 优化模型。

[Beta] AOTInductor 的新打包 API

引入了一种新的包格式“PT2 归档”。这本质上包含一个 zip 文件,其中包含 AOTInductor 需要使用的所有文件,并允许用户将所有必需内容发送到其他环境。还具有将多个模型打包到一个工件中以及在包内存储其他元数据的功能。

有关更多详细信息,请参阅更新后的 torch.export AOTInductor Python 运行时教程

[Beta] AOTInductor:最小化器

如果用户在使用 AOTInductor API 时遇到错误,AOTInductor 最小化器允许创建可重现错误的最小 nn.Module。

有关更多信息,请参阅 AOTInductor 最小化器文档

[Beta] AOTInductor:ABI 兼容模式代码生成

AOTInductor 生成的模型代码依赖于 Pytorch cpp 库。由于 Pytorch 发展迅速,因此确保以前 AOTInductor 编译的模型可以继续在新版本的 Pytorch 上运行非常重要,即 AOTInductor 向后兼容。

为了保证应用程序二进制接口 (ABI) 的向后兼容性,我们在 libtorch 中仔细定义了一组稳定的 C 接口,并确保 AOTInductor 生成的代码仅引用这组特定的 API,而不引用 libtorch 中的其他任何内容。我们将保持这组 C API 在各个 Pytorch 版本中的稳定性,从而为 AOTInductor 编译的模型提供向后兼容性保证。

[Beta] X86 CPU 的 FP16 支持(eager 模式和 Inductor 模式)

Float16 数据类型通常用于减少内存使用量并加快 AI 推理和训练速度。诸如最近推出的英特尔® 至强® 6 P 核处理器等 CPU 支持带有原生加速器 AMX 的 Float16 数据类型。X86 CPU 上的 Float16 支持在 PyTorch 2.5 中作为原型功能引入,现在已针对 eager 模式和 Torch.compile + Inductor 模式进行了进一步改进,使其成为 Beta 级别的功能,其功能和性能均已通过广泛的工作负载验证。

PROTOTYPE 功能

[Prototype] 改进 Intel GPU 上的 PyTorch 用户体验

Intel GPU 上的 PyTorch 用户体验得到了进一步改进,简化了安装步骤,提供了 Windows 发行版二进制文件,并扩展了受支持的 GPU 型号覆盖范围,包括最新的 Intel® Arc™ B 系列独立显卡。寻求在 英特尔® 酷睿™ Ultra AI PC英特尔® Arc™ 独立显卡上微调、推理和开发 PyTorch 模型的应用程序开发人员和研究人员现在可以直接安装适用于 Windows、Linux 和 Windows Subsystem for Linux 2 的二进制发行版 PyTorch。

  • 简化的 Intel GPU 软件堆栈设置,只需一键安装 torch-xpu PIP wheels 即可开箱即用地运行深度学习工作负载,从而消除了安装和激活 Intel GPU 开发软件包的复杂性。
  • 适用于 Intel GPU 的 torch core、torchvision 和 torchaudio Windows 二进制版本已发布,受支持的 GPU 型号已从配备 Intel® Arc™ 显卡的英特尔® 酷睿™ Ultra 处理器、配备 Intel® Arc™ 显卡的英特尔® 酷睿™ Ultra 系列 2英特尔® Arc™ A 系列显卡扩展到最新的 GPU 硬件 英特尔® Arc™ B 系列显卡
  • 进一步增强了 Intel GPU 上 Aten 算子的覆盖范围,使用 SYCL* 内核实现流畅的 eager 模式执行,并修复了错误并优化了 Intel GPU 上 torch.compile 的性能。

有关 Intel GPU 支持的更多信息,请参阅入门指南

[Prototype] X86 CPU 上对 LLM 的 FlexAttention 支持

FlexAttention 最初在 PyTorch 2.5 中引入,旨在为具有灵活 API 的 Attention 变体提供优化的实现。在 PyTorch 2.6 中,通过 TorchInductor CPP 后端添加了对 X86 CPU 的 FlexAttention 支持。这项新功能利用并扩展了当前的 CPP 模板能力,以支持广泛的 Attention 变体(例如:PageAttention,这对于 LLM 推理至关重要),基于现有的 FlexAttention API,并在 x86 CPU 上带来优化的性能。借助此功能,可以轻松使用 FlexAttention API 在 CPU 平台上组合 Attention 解决方案并实现良好的性能。

[Prototype] Dim.AUTO

Dim.AUTO 允许将自动动态形状与 torch.export 一起使用。用户可以使用 Dim.AUTO 导出并“发现”其模型的动态行为,包括自动推断最小/最大范围、维度之间的关系以及静态/动态行为。

与现有的 named-Dims 方法相比,这是一种更友好的用户体验,后者需要用户在导出时完全了解其模型的动态行为。Dim.AUTO 允许用户编写与模型无关的通用代码,从而提高动态形状导出的易用性。

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

[Prototype] AOTInductor 的 CUTLASS 和 CK GEMM/CONV 后端

CUTLASS 和 CK 后端为 Inductor 中的 GEMM 自动调优添加了内核选择。现在,这在 AOTInductor 中也可用,后者可以在 C++ 运行时环境中运行。这两个后端的一项重大改进是提高了编译时速度,消除了冗余内核二进制编译并支持动态形状。