2026 年 10 月 20 日，Amazon 在 PyTorch Conference North America 以 Apache 2.0 授權開源 TorchNeuron，並同步進入公開測試。這是一個原生 PyTorch 後端，讓現有模型程式、訓練腳本與周邊工具直接跑在 AWS Trainium2 與 Trainium3 晶片上；開發者只要把模型改為 .to(neuron)，即完成遷移切換。[1][2]
Amazon 指出，目前約 85% 的 AI 研究人員與開發者以 PyTorch 搭配程式，但可執行的硬體選擇遠比程式碼本身狹窄，每次遷移到新晶片往往要重寫並重新驗證花費數月才調好的工作。[2] TorchNeuron 的目標正是移除這道門檻。
以標準寫法接入既有生態
TorchNeuron 建基於 PyTorch 為自訂硬體而設的 PrivateUse1 寫法後端，以及 torch.accelerator 這組寫法中立 API。Trainium 證為標準 PyTorch 寫法之後，torch.nn、torch.tensor、torch.autocast、torch.distributed 與 torch.profiler 均以原有寫法運作。[1]
已按 torch.accelerator 抽象的現代程式碼可原封不動執行；需要編譯的模型只需以 torch.compiler.set_default_backend("neuron") 設置一次，之後直接呼叫 torch.compile，或於每次呼叫時傳入 backend="neuron"。這個做法借用社區建立的擴充鑰匙，而非另建一套專有介面，因此上游新版本與新函式庫可以較快獲得支持。[1]
訓練與推理共用同一套工具
TorchNeuron 支援兩種執行模式。Eager 模式逐步執行運算，保留 Python 控制流、print 與斷點，適合研究與除錯；官方另提供實驗性的 adaptive eager 模式，以減少每次喚醒的開銷並尋找算子融合機會。torch.compile 模式把模型掃描成圖，經 Torch-MLIR 降至 StableHLO，再由 Neuron 編譯器產生可於硬體執行的針對性、高效能的生產模型。[1]
分散式訓練沿用開發者熟悉的介面，包含 DDP、FSDP2、SimpleFSDP、張量平行與流水線平行，並支持 DTensor 與 DeviceMesh；集群內通訊在節點內經 Trainium 專用的 collective cores，跨節點經 Elastic Fabric Adapter，PyTorch API 無須改動。資料類型方面支援 BF16、FP32 與 FP16，TorchAO 與 FP8 仍在私人測試，將於下一版本全面開放。[1]
推理一樣建構於 torch.compile 支持之上，vLLM Neuron 與 vLLM Omni Neuron 分別服務語言模型與影像、影片生成模型。開發者還可透過 Neuron Kernel Interface 擴充自訂 kernel，並以 Neuron Explorer 配合 torch.profiler，把每個事件追蹤至原始碼行以定位瓶頸。[1]
對晶片競爭格局的意味
Amazon 在示範中於傳統晶片與 Trainium 上執行同一個訓練腳本，唯一區別是指定晶片的設置，結果顯示模型以相同速度學習相同的內容。[2] 這種「無需重寫」的遷移能力，直接針對 Nvidia 多年來把 CUDA 深度嵌入生態所形成的黏著度。
對企業而言，實際價值在於同一塊晶片可兼做訓練與推理：離峰時段可把推理容量轉為訓練工作，減少專用硬體閒置。[2] Amazon 期早期採用者跨越機構研人、影片生成、世界模型、自動駕駛、語音音訊、藥物研發、大型語言模型與金融等領域。[2]
當開發者能在晶片之間低成本遷移，硬體廠商便需以性價比與工具鏈質素搶取客戶，而非仰賴既有程式碼的慣性。TorchNeuron 能否兌現這個承諾，仍取決於實際負載下的效能與算子覆蓋程度，但至少把「換晶片等於重寫」這個前提移除了。[2]
參考資料
[1] Introducing TorchNeuron: Native PyTorch backend for AWS Trainium(2026年10月19日)
[2] Amazon opens its AWS Trainium AI chips to every PyTorch developer, no rewrite needed(2026年10月19日)