AWS 公測 TorchNeuron，讓 Trainium 原生執行 PyTorch，編譯仍限靜態形狀

AWS 於 10 月 20 日發布 TorchNeuron 公開測試版，為 Trainium2、Trainium3 提供原生 PyTorch 後端。開發者可透過 .to('neuron') 搬移模型與輸入，沿用熟悉的張量操作與訓練介面，並以 Neuron Kernel Interface（NKI）加入自訂核心。官方公告

技術上的關鍵是透過 PyTorch 的 PrivateUse1 機制註冊 neuron 裝置，預設逐一執行運算，也可用 torch.compile 的 neuron 後端編譯模型。專案另支援 DTensor、分散式集合通訊及分片訓練；NKI 核心可放入編譯圖，讓硬體最佳化保留在 PyTorch 工作流程內。原始碼與使用說明

不過，裝置切換只是遷移起點。文件列明編譯模式僅支援靜態形狀，新輸入形狀會重新編譯；非連續張量每次呼叫都需複製。float64 會以 float32 執行，TorchAO 整合也尚未提供。API 相容性表還區分原生支援、回退 CPU 與未支援項目，工程團隊需要檢查精度與資料搬移成本。限制說明

同日公布的 Laplace 3B 案例呈現了實際移植工作：團隊把 GPU 專用 Triton 注意力核心改寫為標準 PyTorch 運算，重設 Muon 最佳化器的通訊方式，並調整反向傳播的展開。AWS 與合作團隊回報，單台 trn2.48xlarge 訓練吞吐量約每秒十二萬 token；這是特定架構與配置的結果，不能直接推廣至其他模型。技術案例

對評估替代加速器的團隊而言，新後端降低了試跑門檻，也提供檢查實作的開源入口。接下來應觀察算子覆蓋率、首次編譯時間、形狀分桶成本與端到端數值一致性。安裝時也須留意版本：儲存庫提醒，較早的同名 torch-neuronx 套件屬於 XLA 路線，應依新版指定的版本組合部署。安裝文件