跳转到主要内容
博客

PyTorch 2.6 发布博客

作者: 2025 年 1 月 29 日2025 年 5 月 3 日暂无评论

我们很高兴宣布发布 PyTorch® 2.6(发行说明)!此版本为 PT2 带来了多项改进:torch.compile 现在可以与 Python 3.13 一起使用;新的性能相关旋钮 torch.compiler.set_stance;多项 AOTInductor 增强功能。除了 PT2 改进之外,另一个亮点是 X86 CPU 上的 FP16 支持。

注意:从本版本开始,我们将不再在 Conda 上发布,详情请参阅[公告] 弃用 PyTorch 官方 Anaconda 频道

对于此版本,随 CUDA 12.6.3 发布的实验性 Linux 二进制文件(以及 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 系列的更多信息,请访问我们的入门页面。

测试版 原型
torch.compiler.set_stance 改进了英特尔 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 支持  

*要查看完整的公开功能提交列表,请点击此处

测试版功能

[测试版] torch.compiler.set_stance

此功能允许用户指定 torch.compile 在编译函数的不同调用之间可以采取的不同行为(“立场”)。例如,其中一个立场是

“eager_on_recompile”,它指示 PyTorch 在需要重新编译时急切地生成代码,并在可能的情况下重用缓存的已编译代码。

有关更多信息,请参阅set_stance 文档使用 torch.compiler.set_stance 进行动态编译控制教程。

[测试版] torch.library.triton_op

torch.library.triton_op 提供了一种创建由用户定义的 Triton 内核支持的自定义运算符的标准方法。

当用户将用户定义的 Triton 内核转换为自定义运算符时,torch.library.triton_op 允许 torch.compile 深入了解实现,从而使 torch.compile 能够优化其中的 Triton 内核。

有关更多信息,请参阅triton_op 文档将用户定义的 Triton 内核与 torch.compile 结合使用教程。

[测试版] torch.compile 对 Python 3.13 的支持

torch.compile 以前只支持到 Python 3.12 版本。现在用户可以使用 torch.compile 在 Python 3.13 中优化模型。

[测试版] AOTInductor 的新打包 API

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

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

[测试版] AOTInductor:迷你化器

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

有关更多信息,请参阅AOTInductor Minifier 文档

[测试版] AOTInductor:ABI 兼容模式代码生成

AOTInductor 生成的模型代码依赖于 PyTorch C++ 库。由于 PyTorch 发展迅速,确保之前 AOTInductor 编译的模型可以在更新的 PyTorch 版本上继续运行至关重要,即 AOTInductor 向后兼容。

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

[测试版] X86 CPU 的 FP16 支持(急切模式和 Inductor 模式)

Float16 数据类型通常用于 AI 推理和训练中,以减少内存使用并加快计算速度。最近推出的Intel® Xeon® 6 with P-Cores 等 CPU 支持带原生加速器AMX 的 Float16 数据类型。X86 CPU 上的 Float16 支持在 PyTorch 2.5 中作为原型功能引入,现在它在急切模式和 Torch.compile + Inductor 模式下都得到了进一步改进,使其成为 Beta 级别功能,并在广泛的工作负载范围内验证了功能和性能。

原型功能

[原型] 改进了英特尔 GPU 上的 PyTorch 用户体验

PyTorch 在英特尔 GPU 上的用户体验通过简化的安装步骤、Windows 发布二进制分发和扩展的支持 GPU 模型(包括最新的 Intel® Arc™ B 系列独立显卡)得到了进一步改进。寻求在Intel® Core™ Ultra AI PCIntel® Arc™ 独立显卡 上微调、推理和开发 PyTorch 模型的应用程序开发人员和研究人员现在可以直接通过 Windows、Linux 和适用于 Linux 2 的 Windows 子系统 的二进制版本安装 PyTorch。

  • 简化的英特尔 GPU 软件堆栈设置,支持一键安装 torch-xpu PIP wheel,以开箱即用的方式运行深度学习工作负载,消除了安装和激活英特尔 GPU 开发软件包的复杂性。
  • 适用于英特尔 GPU 的 torch 核心、torchvision 和 torchaudio 的 Windows 二进制版本已提供,支持的 GPU 模型已从带 Intel® Arc™ Graphics 的 Intel® Core™ Ultra 处理器、带 Intel® Arc™ Graphics 的 Intel® Core™ Ultra 系列 2Intel® Arc™ A 系列显卡 扩展到最新的 GPU 硬件Intel® Arc™ B 系列显卡
  • 通过 SYCL* 内核进一步增强了英特尔 GPU 上 Aten 运算符的覆盖范围,以实现流畅的急切模式执行,并修复了错误并优化了英特尔 GPU 上 torch.compile 的性能。

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

[原型] X86 CPU 上对 LLM 的 FlexAttention 支持

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

[原型] Dim.AUTO

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

与现有用于指定动态形状的命名维度方法相比,这是一种更用户友好的体验,后者要求用户在导出时完全理解其模型的动态行为。Dim.AUTO 允许用户编写不依赖于模型的通用代码,从而提高了使用动态形状进行导出的易用性。

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

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

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