原文標題:Amazon opens its AWS Trainium AI chips to every PyTorch developer, no rewrite needed

Amazon 於 10 月 19 日宣布 TorchNeuron,目標是讓使用 PyTorch 的 AI 開發者在 AWS 自研 Trainium 加速器上執行既有程式,而不必為硬體遷移重寫整套模型。官方表示,開發者通常只需修改裝置選擇等設定;這是對 PyTorch 既有硬體後端註冊機制的整合,不代表所有工作負載均不需測試。

TorchNeuron 支援 Trainium2 與 Trainium3,並將於 10 月 20 日的 PyTorch Conference North America 以 Apache 2.0 授權釋出。Amazon 也將部分修改提交 PyTorch 上游,希望降低框架與特定晶片架構之間的耦合。

在工具鏈方面,TorchNeuron 可與 PyTorch 分散式訓練框架 TorchTitan 配合,沿用 Fully Sharded Data Parallelism(FSDP)模式。Hugging Face 的開發工具也加入原生 Trainium 支援;Amazon 另提及對 vLLM 的開源貢獻與效能核心程式公開。

官方舉出 Cartesia 與 Lieber Institute for Brain Development 等早期使用者,應用從語音模型延伸至藥物探索。Amazon 稱每次建置都會對大量模型進行測試,並以示範展示相同訓練程式在不同裝置上可得到一致學習結果;不過新聞未提供可普遍套用的跨晶片效能比較。

Hugging Face 於 10 月 19 日介紹 TorchNeuron,這是讓 AWS Trainium 加速器在 PyTorch 中具備原生裝置支援的新後端。官方表示,從 10 月 20 日起提供公開測試版,研究人員與開發者可安裝 torch-neuronx 套件使用。相較過去透過 Optimum Neuron 與 PyTorch XLA 的工作流程,新方案希望減少特定硬體帶來的程式結構差異。

TorchNeuron 支援 PyTorch 的 eager mode、torch.compile,以及 FSDP、張量平行等分散式運算能力,讓 Trainium 在程式中更接近一般加速器裝置。Hugging Face 表示,Transformers、Diffusers、Accelerate、PEFT 與 TRL 等函式庫可以沿用既有模型類別與訓練程式,不必另行維護 Neuron 專用模型實作或訓練器。

在文字生成推論方面,Hugging Face 統計使用量最高的 100 種文字生成模型架構,TorchNeuron 可直接支援其中 89%,相較 Optimum Neuron 0.4.5 的 46% 提高。開發者可以將模型載入指定為 neuron 裝置;官方範例仍建議使用靜態快取,避免動態形狀造成頻繁重新編譯。這項比例衡量的是模型架構相容性,不是所有模型的執行速度提升百分比。

訓練端則可直接使用 TRL 的 SFTTrainer,透過 DistributedConfig 指定張量平行與 FSDP 分片配置。官方展示以 TP_SIZE=4、FSDP_SIZE=8 的 32 個處理程序執行微調,對應 8 顆 Trainium2 晶片、32 個 NeuronCore,為 trn2.48xlarge 執行個體的一半資源。完成後儲存的仍是標準 Transformers checkpoint,可移往其他硬體環境載入。

實務上仍有相容性限制:官方範例關閉 Dynamo 自動動態形狀功能,將序列填補至固定長度倍數,並在 Trainium 上停用 Transformers 非同步權重載入;這些設定是為了避免後端目前不支援的路徑。

影像生成方面,Diffusers 已有多種 pipeline 可在 eager 或編譯模式運作,官方列出 FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1 與 H3 的張量平行支援。Hugging Face 也鼓勵開發者使用 Neuron Kernel Interface 編寫與分享最佳化 kernel,並舉出視覺 Transformer 核心可達約 2.4 倍加速的特定範例。

發生了什麼事

Amazon 公開 TorchNeuron 方案,讓 PyTorch 程式可透過原生後端使用 AWS Trainium2 與 Trainium3,並宣布開源時程。

Hugging Face 宣布 PyTorch 原生 TorchNeuron 後端將於 10 月 20 日公開測試,整合 Transformers、Diffusers 與多項訓練函式庫,簡化 AWS Trainium 的模型部署。

為什麼重要

既有模型遷移至不同 AI 晶片常涉及程式改寫、除錯與重新驗證;降低相容性門檻有助於擴大硬體選擇。

減少硬體專用模型分支與訓練器,有助讓同一套 PyTorch 程式在不同加速器間移植,降低雲端模型訓練與推論的維護成本;但仍須處理編譯及動態形狀限制。

影響哪些人

主要涉及使用 PyTorch 的研究人員、AI 新創、企業模型平台團隊,以及評估 Trainium 訓練與推論成本的工程師。

在 AWS Trainium 上執行語言模型、影像模型或分散式微調的 AI 平台工程師、研究團隊及維護 Hugging Face 函式庫的開發者。

已確認與未確認

已確認支援晶片世代、Apache 2.0 授權與 10 月 20 日釋出計畫;各模型實際效能、完整相容性與價格效益仍須自行測試。

已確認 10 月 20 日公開測試時程、89% 對 46% 的架構支援比較、8 顆 Trainium2 的訓練範例,以及現有靜態形狀限制。不同模型的效能、成本與完整硬體支援清單仍需實測。