PyTorch 2.9 版本发布博客 – PyTorch

我们很高兴宣布发布 PyTorch® 2.9 (发行说明)！此版本包含：

libtorch 稳定 ABI 的更新，适用于第三方 C++/CUDA 扩展
对称内存，可轻松编程多 GPU 内核
在 torch.compile 中，能够任意切换图断裂时的错误或恢复行为
扩展的 wheel 变体支持，包括 ROCm、XPU 和 CUDA 13
在 Intel GPU 上启用 FlexAttention
基于 FlexAttention 的 X86 CPU Flash 解码优化
Arm 平台改进和优化
在所有支持的 CUDA 版本上启用 Linux aarch64 二进制 wheel 构建

自 PyTorch 2.8 以来，此版本包含来自 452 名贡献者的 3216 次提交。我们衷心感谢我们敬业的社区所做的贡献。一如既往，我们鼓励您尝试这些新功能并报告任何问题，以便我们改进 2.9。有关 PyTorch 2 系列入门的更多信息，请参阅我们的入门页面。

API 不稳定功能

[API 不稳定] torch::stable::Tensor

如果您使用 PyTorch 维护和构建自己的自定义 C++/CUDA 扩展，此更新适合您！我们一直在构建一个带有 C++ 便利包装器的稳定 ABI，使您能够用一个 torch 版本构建扩展，并用另一个版本运行。自上次发布以来，我们添加了以下 API：

引入设备工具（如 Device Guard 和 Stream），可在torch/csrc/stable/accelerator.h中访问。
更多torch::stable::Tensor API：默认构造函数、is_cpu、scalar_type 和 get_device_index
更多稳定的 ATen 操作，可在torch/csrc/stable/ops.h中访问：amax、narrow、new_empty + new_zeros dtype 变体、pad

通过这些 API，我们已经能够为 Flash-Attention 3 启用 libtorch-ABI wheel：请参阅此处的 PR。虽然我们在 API 设计时有意确保最大的稳定性，但请注意，高级 C++ API 仍处于预览阶段！我们正在进行许多下一步工作：构建 ABI 接口、建立版本控制、编写更多文档以及使更多自定义内核实现 ABI 稳定。

[API 不稳定] 对称内存编程

我们引入 PyTorch 对称内存，以方便编程多 GPU 内核，这些内核可在 NVLinks 和 RDMA 网络上工作。对称内存解锁了三个新的编程机会：

内核内通信：GPU 内核现在能够发出通信原语，如 puts 和 gets，并与计算指令交错，从而实现最小粒度的融合。
超低延迟远程访问：远程访问可以单向进行，无需等待远程 GPU 发出相应的命令。RDMA 和 IB-GDA 等协议支持通过网络直接进行内存访问。
定制通信模式：增加的灵活性允许编写针对应用程序通信需求量身定制的内核。即使标准库中尚未提供，添加对新数据布局的支持也将变得简单。

2.9 版本将包括：

允许远程直接访问的对称张量分配（目前支持的后端：CUDA和NVSHMEM）。
利用直接访问加速集合操作，例如：one_shot_all_reduce, two_shot_all_reduce_, multimem_all_gather_out 等。
用于 MoE 模型的非传统 all_to_all_v （all_to_all_vdev, all_to_all_vdev_2d用于令牌分发，all_to_all_vdev_2d_offset 用于令牌组合）。
对定制多 GPU 内核的编程支持：Triton 的 NVSHMEM 插件。持续支持 Async TP 并推广到其他模式。

对称内存操作可在torch.ops.symm_mem下使用。

[API 不稳定] 在 torch.compile 中，能够切换图断裂时的错误或恢复行为

此功能通过引入 torch._dynamo.error_on_graph_break()（一个上下文管理器/装饰器）扩展了 torch.compile 的图断裂选项，该管理器/装饰器允许用户标记 torch.compiled 代码中在遇到图断裂时 torch.compile 应该报错或恢复的区域。与现有的 fullgraph 功能不同，error_on_graph_break 可以任意切换（一旦 fullgraph 设置为 true，就不能再设置为 false）。更多详细信息请参阅教程。

[API 不稳定] 扩展的 wheel 变体支持

PyTorch 继续通过参与 WheelNext 计划来改进 Python 打包生态系统。Jonathan Dekhtiar (NVIDIA) 和 Eli Uriegas (Meta) 将在 2025 年 PyTorch 大会上介绍该领域的持续工作：演讲。

PyTorch 2.9.0 版本通过添加 AMD (ROCm)、Intel (XPU) 和 NVIDIA CUDA 13 扩展了其 wheel 变体支持矩阵的范围。这项工作是PyTorch 2.8 中此功能的初步启用的后续。这包括所有上述硬件平台的提供商插件，它们可以自动检测平台属性，包括支持的软件类型和安装的硬件。

虽然 NVIDIA CUDA wheels 支持 Windows 和 Linux，但 ROCm （完整博客在此）和 XPU 平台目前仅支持 Linux。

注意：此特定功能是实验性的，基于 wheel 变体提案。（正在进行中）

安装说明

Linux x86 和 aarch64，MacOS

curl -LsSf https://astral.ac.cn/uv/install.sh | 
INSTALLER_DOWNLOAD_URL=https://wheelnext.astral.sh/v0.0.2 sh
uv venv
uv pip install torch

Windows x86

powershell -c { $env:INSTALLER_DOWNLOAD_URL = 'https://wheelnext.astral.sh/v0.0.2'; irm 
https://astral.ac.cn/uv/install.ps1 | iex }
uv venv
uv pip install torch

[API 不稳定] 在 Intel GPU 上启用 FlexAttention

Intel GPU 上新的 FlexAttention 前向和后向支持，与 PyTorch 的 GPU 常见行为保持一致，为用户提供了在不同 GPU 上更一致、更便携的性能。这意味着开发人员可以一次编写代码，或者运行 Hugging Face/Transformers，并广泛支持 FlexAttention，在不修改任何代码的情况下实现显著的注意力机制效率。

[API 不稳定] 基于 FlexAttention 的 X86 CPU Flash 解码优化

Flash 解码是一种在 LLM 推理中常用的技术，用于加速注意力并加快长序列的生成，目前仅适用于 CUDA 路径。在此版本中，我们引入了基于 FlexAttention 的 X86 CPU inductor 后端 Flash 解码优化支持，它通过分区和归约实现 KV 序列的并行化。当原始并行度不足时（例如小批量大小/头数/Q 序列长度和长 KV 序列长度），此优化可以大大提高 CPU 利用率。预计这将有助于 PyTorch 用户提高 LLM 解码阶段的性能，尤其是在长上下文长度的情况下。

Arm 平台启用和优化

PyTorch 2.9 在 Arm 上提供了关键的后端改进，具有更好的性能和测试覆盖率。

torch.compile 改进： TorchBench、HuggingFace 和 TIMM 测试套件在 torch.compile 模式下比 Eager 模式更快。请参阅PyTorch HUD Dashboard上的结果。
运算符改进： 在 AArch64 上扩展和优化了卷积、激活和量化操作，以实现更快的模型执行。
拓宽 Arm CI 覆盖范围：通过添加基于 Arm Neoverse V2 的 AWS Graviton 4 实例。

PyTorch 2.9 发布博客