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.sh/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.sh/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 釋出部落格