AI-translated from English; not yet reviewed by a fluent editor.

# AWS 的 TorchNeuron 為 Trainium 帶來原生 PyTorch 後端

> AWS 的 Neuron 2.33 版本為 Trainium2 和 Trainium3 推出公開測試版 PyTorch 後端。其支援版本、運算子涵蓋範圍與編譯限制，界定了工程師評估時需要檢查的事項。

By BIG CHANGE Editorial

Published: 2026-10-20T04:20:56.038Z
Updated: 2026-10-20T04:20:56.038Z
Canonical: https://bigchange.ai/blog/aws-torchneuron-native-pytorch-backend-trainium

![Conceptual charcoal illustration of a code-marked monitor beside an open rack drawer containing an unbranded accelerator board.](https://bigchange.ai/api/media/file/torchneuron-native-device-trainium-hero-v1.png)
AI-generated conceptual illustration by BIG CHANGE.

AWS 的 Neuron 2.33 [版本說明](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html)將 TorchNeuron 公開測試版的日期列為 10 月 19 日。同樣標註 10 月 19 日的 [PyTorch 公告](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/)則表示，公開測試版與開放原始碼版本將於 10 月 20 日上線。`torch-neuronx` 套件向 PyTorch 註冊一個 `neuron` 裝置供 Trainium2 和 Trainium3 使用，讓支援的模型與張量運算能使用熟悉的 PyTorch API。對考慮採用這些晶片的工程師而言，實用的改變是從既有 PyTorch 程式進入裝置實測的路徑更短。程式能否正確且有效率地運行，仍取決於其運算、形狀與軟體環境。

過去 Trainium 的 PyTorch 路徑使用以 XLA 為基礎的 `torch-neuronx` 函式庫，以及 `torch_xla` 裝置程式碼。AWS 表示，原生後端使用 PyTorch 的 `PrivateUse1` 擴充機制，並將 `xm.xla_device()` 改為 `neuron` 裝置配置。對已使用 Neuron 的專案而言，即使套件名稱很熟悉，這仍是實質的介面變動。此版本適用於 Trn2 和 Trn3 執行個體上的訓練與推論；支援矩陣未列出 Trn1。

## PyTorch 程式有哪些改變

依據 [AWS 版本說明](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html)，模型可透過 `.to("neuron")` 移動，並以 `device="neuron"` 建立輸入張量。支援的 `torch.nn`、autograd 與最佳化器運算隨後會透過後端執行。依照 PyTorch 的裝置無關 `torch.accelerator` API 編寫的程式碼，無須針對個別裝置重寫即可偵測 Neuron，這是 [TorchNeuron 概觀](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html)所述的能力。

共有兩條執行路徑。即時執行模式會在 Python 呼叫時執行 PyTorch 運算，保留一般控制流程、列印輸出與中斷點。各個運算首次使用時會編譯，後續相符的呼叫則可使用快取。工程師也能透過 `torch.compile(model, backend="neuron", dynamic=False)` 編譯模型或選定區域。原生後端會編譯支援的前向與反向圖；擷取區域以外的程式碼仍以即時執行模式運行。[PyTorch 公告](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/)也介紹了在編譯前設定一次 `torch.compiler.set_default_backend("neuron")` 的做法，如此後續編譯呼叫即可省略後端引數。

此版本包含一個 `neuron` 後端，供 `torch.distributed` 使用，並有文件說明的集合通訊與點對點運算。AWS 列出的支援 API 包括 DistributedDataParallel、FSDP2、DTensor、DeviceMesh 與管線平行處理。這些是分散式訓練專案值得檢視的路徑，但版本說明並未證明既有多節點工作能原封不動運行，或達到特定吞吐量目標。

## 評估從相符的環境開始

AWS 的 [安裝快速入門](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-torch-neuron.html)要求使用 Trn2 或 Trn3 執行個體，並搭配 Neuron Base DLAMI 或支援的 Ubuntu AMI。版本矩陣指定 Ubuntu 24.04 或 26.04 上的 Linux x86\_64。此版本已結束對 Ubuntu 22.04 的支援。它列出 Python 3.12、3.13 與 3.14，並以 `torch-neuronx` 2.13.3.0 搭配 PyTorch 2.13.0，或以 `torch-neuronx` 2.14.3.0 搭配 PyTorch 2.14.0。套件系列必須與 PyTorch 版本相符。

快速入門會安裝 `torch-neuronx==2.14.*`，並一併安裝 `neuronx-cc==2.*` 與 `nki`，套件來源為 Neuron 與 PyTorch 套件索引。接著會檢查裝置是否可用，並將 Neuron 上的小型矩陣運算與 CPU 結果比較。這驗證的是 AWS 範例中已安裝的裝置路徑，並不能驗證讀者自己的模型。合理的下一步評估，是執行實際模型的前向與反向路徑，使用 `torch_neuronx.get_fallback_ops()` 檢視記錄到的 CPU 回退運算，將輸出與現有環境比較，並分別量測編譯與穩態執行。最後這些步驟是 BIG CHANGE 提出的評估問題，並非我們做過的測試。

## 測試版仍可能改變工作方式之處

在即時執行模式下，不支援的運算可能回退至 CPU，並發出警告。在 `torch.compile` 下，若擷取至圖中的運算無法由 Neuron 後端降階轉換，預設會導致編譯失敗。AWS 請開發者參考其 [運算子相容性表](https://github.com/aws-neuron/torch-neuronx/blob/main/docs/torch_api_compatibility.md)，確認各運算的涵蓋範圍。因此，即時執行順利完成，本身並不代表完全在裝置上執行，或具備編譯相容性。

版本說明也表示，`torch.compile` 需要設定 `dynamic=False`；新的輸入形狀可能需要重新編譯。AWS 建議使用填補或分桶來減少形狀變化。若沒有相符的快取產物，第一次執行會包含編譯。為工作負載計時時，這些細節很重要：編譯延遲與輸入形狀的分布，可能改變評估結果。

其他明確限制也縮小了相容性主張的適用範圍。此測試版不支援 `torch.sparse`、`torch.func`、`torch.nested` 或 `torch.distributions`。只有部分 `torch.fft` 轉換可在 Neuron 上運行，而 `float64` 會以 `float32` 執行並發出警告。此版本缺少 TorchAO 整合與複合 MXFP8 資料型別支援。使用這些功能介面的團隊，在把裝置配置視為遷移計畫之前，需要先檢查自己的程式路徑。

### 重大改變

TorchNeuron 在公開測試版中，將 Trainium 接到 PyTorch 原生裝置與編譯介面之後。這讓支援的 PyTorch 程式能更直接地進行初次試用，也讓既有 Neuron 使用者有一條不同於 XLA 程式碼的路徑。採用與否的決定性證據，將來自各個工作負載：運算子是否留在預定裝置上、數值結果是否仍可接受，以及量測到的效能與成本是否足以支持執行個體與遷移工作。BIG CHANGE 為這篇分析查閱了文件與公告；我們沒有運行 TorchNeuron，也沒有獨立驗證 AWS 的效能主張。

## 資料來源與延伸閱讀

- [AWS Neuron 2.33 TorchNeuron 版本說明](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html)：發布日期、測試版狀態、版本與作業系統矩陣、API，以及已知限制。
- [AWS 的 TorchNeuron 概觀](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html)：裝置整合、即時與編譯執行、CPU 回退行為，以及經由 Neuron 的軟體路徑。
- [AWS 安裝與驗證快速入門](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-torch-neuron.html)：執行個體與映像檔先決條件、套件安裝，以及文件中的裝置檢查。
- [PyTorch 部落格：介紹 TorchNeuron](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/)：從 PyTorch 角度說明裝置切換、編譯與分散式 API。此文是 AWS 貢獻者在 PyTorch 部落格發布的公告，並非獨立效能評測。
- [Amazon 公告](https://www.aboutamazon.com/news/innovation-at-amazon/torchneuron-pytorch-aws-trainium-chips-ai)：AWS 對其目標與早期採用者的說明。其遷移、效能與經濟效益主張，需要針對個別工作負載驗證。

## Sources

- [AWS Neuron 版本說明：TorchNeuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html) — 第一手資料；將公開測試版日期列為 2026-10-19，比 PyTorch 公告所述的公開測試版與開放原始碼上線日期早一天；提供版本矩陣與明列限制。
- [什麼是 TorchNeuron？](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html) — 裝置與後端使用方式的第一手文件。
- [PyTorch 部落格：TorchNeuron，AWS Trainium 的原生 PyTorch 後端](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/) — AWS 貢獻者在 PyTorch 端發布的公告；標註 10 月 19 日，並表示公開測試版與開放原始碼版本於 10 月 20 日上線。有效網址由 Amazon 文章連入；並非獨立效能驗證。
- [Amazon 公告](https://www.aboutamazon.com/news/innovation-at-amazon/torchneuron-pytorch-aws-trainium-chips-ai) — 廠商公告；遷移與經濟效益主張是 AWS 的說法，未經獨立驗證。
- [安裝並驗證 torch-neuronx](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-torch-neuron.html) — AWS 快速入門，涵蓋先決條件、相符套件安裝，以及文件中的 CPU／Neuron 裝置檢查。
- [TorchNeuron API 相容性表](https://github.com/aws-neuron/torch-neuronx/blob/main/docs/torch_api_compatibility.md) — 維護者針對各個 API 列出的原生支援、CPU 回退與不支援分類。