他已經有一支在 GPU 上跑得好好的 Hugging Face 文字生成腳本，打算只改幾行就搬過去。
Trainium 和 Inferentia 是 AWS 設計的 AI 晶片專門用來做人工智慧運算的處理器，功能跟顯卡類似，但是由雲端公司自己設計。Trainium 主打用較低成本做訓練，Inferentia 主打用較低成本做大規模推論。過去在這些晶片上用 Hugging Face，靠的是一個叫 optimum-neuron 的函式庫。它底層用 torch-neuronx，而 torch-neuronx 又建立在 PyTorch XLA 上。
舊做法的代價，是你得「像編譯器一樣思考」。第一，它用 lazy tensor惰性張量，意思是你寫的運算不會馬上執行，而是先記下來，等到真的需要結果時才一起算，不是一行一行立刻跑。第二，它只接受 靜態形狀資料的大小（例如一次處理幾個字）必須事先固定，不能臨時變。第三，每次出現 graph break（計算流程被打斷），就要重新編譯一次。optimum-neuron 的任務就是把這些麻煩藏起來，所以得自己維護大量客製程式碼，例如模型類別和訓練器。
TorchNeuron 是一個 PyTorch 外掛，它改走另一條路：直接把 neuron 做成 PyTorch 認得的原生裝置。它支援 eager 模式PyTorch 的預設跑法，寫一行就立刻執行一行，方便除錯，支援 torch.compile，也支援 PyTorch 自己的分散式功能，例如 FSDP 和 TP。這樣一來，Transformers、Diffusers、Accelerate、PEFT、TRL 都能在 Trainium 上直接跑，用的類別和腳本跟在其他加速器上一樣。
時間點是這樣：TorchNeuron 預計從 2026 年 10 月 20 日起開放公開 beta，研究者和開發者可用 pip install torch-neuronx 安裝。Hugging Face 也明說，之後建議 Neuron 使用者直接用這些函式庫，不必再繞經 Optimum Neuron。
| 項目 | Optimum Neuron（基於 PyTorch XLA） | TorchNeuron |
|---|---|---|
| 執行方式 | 惰性張量 | eager 執行，另支援 torch.compile |
| 張量形狀 | 只支援靜態形狀 | 文章建議生成時用靜態快取避免重編譯 |
| Hugging Face 整合 | 靠 optimum-neuron 的客製類別與訓練器 | Transformers、Diffusers、Accelerate、PEFT、TRL 直接執行 |
| 分散式 | 由函式庫包裝處理 | PyTorch 原生的 DTensor、FSDP2、TP |
| 前 100 名文字生成架構支援度 | 46%（v0.4.5） | 89% |
這是文章範例的前半段：在 Trainium 上用張量平行載入一個大型語言模型。原文在 tokenizer 那一行之後被截斷，所以這裡只講到載入模型為止。
import torch_neuronx # noqa: F401from transformers.distributed import DistributedConfigfrom transformers.generation.configuration_utils import CompileConfigMODEL_ID = "Qwen/Qwen3.8-27B"MAX_SEQ_LEN = 1024PREFILL_CHUNK_SIZE = 512dist.init_process_group(backend="neuron")world_size = dist.get_world_size()distributed_config=DistributedConfig(tp_size=world_size),device_mesh=DeviceMesh("neuron", list(range(world_size))),這是一個真實小專案：下載（或複製）檔案，照步驟在你電腦上跑起來。
import importlib.util
import sys
def has(name):
return importlib.util.find_spec(name) is not None
for pkg in ('torch', 'transformers', 'torch_neuronx'):
print(pkg, 'OK' if has(pkg) else 'MISSING')
if not has('torch_neuronx'):
print('尚未安裝 torch-neuronx：請在 Trainium 機器上執行 pip install torch-neuronx')
sys.exit(1)
import torch
import torch_neuronx # noqa: F401
print('torch 版本：', torch.__version__)
print('環境檢查完成，可以執行 run_generate.py')
import sys
import torch
import torch_neuronx # noqa: F401 # 匯入後 neuron 才會成為 PyTorch 認得的裝置
from transformers import AutoModelForCausalLM, AutoTokenizer
# 預設沿用文章範例的模型；若單顆晶片放不下，請在指令後面帶入較小的模型名稱
MODEL_ID = sys.argv[1] if len(sys.argv) > 1 else 'Qwen/Qwen3.8-27B'
PROMPT = '用一句話解釋什麼是 Trainium。'
# 文章重點：只要 device_map='neuron'，模型就會被搬到 Neuron 裝置上
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
dtype=torch.bfloat16,
device_map='neuron',
).eval()
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
inputs = tokenizer(PROMPT, return_tensors='pt').to('neuron')
# 文章強烈建議在 Neuron 上用靜態快取，避免頻繁重新編譯
with torch.no_grad():
out = model.generate(**inputs, max_new_tokens=64, cache_implementation='static')
print(tokenizer.decode(out[0], skip_special_tokens=True))