AWS 發布 torchtitan-neuron：在 Trainium 上訓練 Qwen3

背景

TorchTitan 是 PyTorch 生態系中用於生產級 LLM 預訓練的解決方案。AWS Neuron 團隊發布了名為 TorchTitan Neuron 的分支專案，透過 TorchNeuron（Neuron 的原生 PyTorch 後端）為 AWS Trainium 硬體增加支援。此專案旨在讓開發者能使用標準 PyTorch 介面在 Trainium 上訓練大型語言模型，並整合了 Neuron 專屬的配方與 NKI（Neuron Kernel Interface，Neuron 核心介面）核心整合。 官方原文・github.com AWS ML Blog

發生什麼事

AWS Neuron 於 2026 年 3 月 3 日在 GitHub 發布了 torchtitan-neuron 專案。該專案是 pytorch/torchtitan 的分支，主要功能是透過 TorchNeuron 後端支援 AWS Trainium。官方文件指出，此專案包含了 Neuron 專屬的訓練配方、NKI 核心整合以及相關文件，這些內容位於 torchtitan_neuron/ 目錄下。 官方原文・github.com

重點內容

根據官方文件與 AWS ML Blog 的說明，此專案支援的模型包括 Llama3、Qwen3、DeepSeek-V3 和 GPT-OSS。在硬體方面，支援 Trn2 和 Trn3 實例。執行模式支援 Eager 以及使用 --compile.backend neuron 參數的 torch.compile（僅限靜態形狀）。並行策略支援 FSDP2、SimpleFSDP、TP（Tensor Parallelism，張量並行）、CP（Context Parallelism，上下文並行）、EP（Expert Parallelism，專家並行）、PP（Pipeline Parallelism，流水線並行）、DDP（Distributed Data Parallel，分散式資料並行）以及 activation checkpointing。 官方原文・github.com AWS ML Blog

AWS ML Blog 提供了一個在單一 trn2.48xlarge 實例上訓練 Qwen3 的範例。該實例配備 16 個 Trainium2 加速器，在 LNC=2 設定下有 64 個邏輯 NeuronCores。使用的版本包括 torch 2.14.0、torch_neuronx 2.14.3、transformers 5.16.0 以及 torchtitan-neuron release-2.33 (0.2.2+neuron.1)。文章展示了三種訓練方式：第一種使用 Transformers v5 作為建模後端，透過 AutoModelForCausalLM 執行 Hub 模型；第二種切換至 torchtitan 的原生 Qwen3 實現；第三種加入 NKI MoE（Mixture-of-Experts，混合專家）核心。 AWS ML Blog

在驗證測試中，使用 Qwen/Qwen3-4B-Instruct-2507 模型，從公開權重載入時第一步損失值為 6.54，而從隨機初始化開始訓練時第一步損失值為 12.43。訓練完成後導出的檢查點可透過 AutoModelForCausalLM 和 load_state_dict 載入，確認除 lm_head.weight（Qwen3-4B 將其與 token embeddings 綁定）外無缺失或意外鍵值，證明檔案仍為標準 Hugging Face 檢查點格式。 AWS ML Blog

對使用者的意義

開發者可以使用此專案在 AWS Trainium 硬體上進行 LLM 預訓練或微調。使用者可以透過標準 PyTorch 介面操作，無需改變原有的程式碼結構。根據官方文件，使用者可以自訂模型實現和核心以適應特定工作負載。開始使用前，需在 Trn2 或 Trn3 實例上安裝並驗證 torch-neuronx，然後克隆並安裝 torchtitan-neuron 分支專案。 官方原文・github.com AWS ML Blog

限制與未定事項

官方文件指出，torch.compile 模式僅支援靜態形狀（static shapes only）。在 AWS ML Blog 的範例中，為了讓 Qwen3-30B-A3B 模型能在單一 trn2.48xlarge 實例上運行，將其層數減少至 4 層並設定 64 個專家。此外，該專案目前處於 Neuron SDK 2.33.0 的 TorchNeuron 公開測試階段（public beta）。 官方原文・github.com AWS ML Blog