把一段跑在 GPU 上的 PyTorch 训练代码原样搬到 Trainium，通常不会一次成功：多数人第一次遇到的不是算错，而是“卡在编译”“每步都在重新编译”“吞吐只有预期的一小部分”。这篇教程给出一条可执行的路径：先确认环境能看见设备，再让一个最小模型在 Trainium 上完成前向、反向和参数更新，然后讲清楚图捕获这一步到底在做什么，最后用一张自己算出来的对照表，把 Trainium 和 GPU 的吞吐与单位成本放在一起比较。
读完并照着做完，你会拿到四样东西：
- 一台
neuron-ls能列出设备、neuron-top能看到利用率的 Trainium 环境； - 一个能在 Trainium 上完成若干步训练的小模型脚本；
- 一次明确的图捕获（trace / 惰性编译）体验，知道什么代码会把图打碎；
- 一张对照表：GPU 与 Trainium 的稳态吞吐、首次编译耗时，以及换算出来的“每千步成本”。
前置条件清单
动手前把这些准备好，能省掉一半的排错时间。
- 一个 AWS 账号，以及一台 Trainium 实例（trn 系列）。具体哪些实例类型、哪些区域可用，以 AWS 官方文档当前列出的为准。
- 确认账号的 Trainium 按需配额够用。新账号常见的情况是配额为 0，需要单独提交申请，审批不是即时的，建议提前一天处理。
- 一个对照用的 GPU 环境。可以是本地卡，也可以是云上的 GPU 实例。没有对照就没法算成本，这一步别省。
- 熟悉 PyTorch 的基础训练循环：
forward/backward/optimizer.step()。 - 实例系统盘留出足够空间（百 GB 量级比较稳妥）。Neuron 的编译产物很占地方，缓存涨得比想象中快。
- 时间预算：第一次跑通要预留一到两小时，其中相当一部分会花在编译上。
- 一个沙盒账号或独立项目。迁移过程会反复开关实例，别在生产账号上直接试。
步骤一：起实例，先确认“看得见设备”
推荐从 AWS 提供的、已集成 Neuron 软件栈的深度学习 AMI 起步，能跳过大部分驱动安装的麻烦。如果用的是干净系统，则按官方安装文档装 Neuron 驱动、运行时和工具链。包名会随版本演进，命令以官方安装页当前给出的为准，大致形态是：
```bash
sudo apt-get update
sudo apt-get install -y \
aws-neuronx-dkms \
aws-neuronx-collectives \
aws-neuronx-runtime-lib \
aws-neuronx-tools
```
装完重启一次，然后用工具确认设备在位：
```bash
neuron-ls
```
正常情况下会列出若干 Neuron 设备以及每个设备的核心数。看不到设备时不要往下走，先解决驱动问题（见排错部分）。
再开一个终端跑：
```bash
neuron-top
```
这是一个类似 nvidia-smi + top 的工具，后面判断“模型到底有没有真的在加速器上跑”全靠它。现在它应该基本空闲。
步骤二：建一个干净的 Python 环境
不要在系统 Python 里装框架。Neuron 对 PyTorch 版本和扩展包版本有对应关系，混装是排错噩梦。
```bash
sudo apt-get install -y python3-venv
python3 -m venv ~/neuron-env
source ~/neuron-env/bin/activate
pip install -U pip wheel
```
然后安装 Neuron 的 PyTorch 扩展。软件源地址与包名以官方安装文档当前版本为准：
```bash
pip install torch-neuronx neuronx-cc \
--extra-index-url <官方文档给出的软件源地址>
```
验证一下两边版本是否对得上：
```bash
python -c "import torch, torch_neuronx; print(torch.__version__, torch_neuronx.__version__)"
```
能打印出版本号就说明导入链路通了。注意：torch 与 torch-neuronx 必须匹配官方兼容矩阵，装错组合后面会以各种奇怪的方式失败。
步骤三：写一个两边都能跑的小模型
对照实验的第一原则是“同一份代码、同一份数据、同一个 batch”。所以先写一个足够小、GPU 和 Trainium 都能跑的模型，别一上来就搬大语言模型。
```python
model.py
import torch
import torch.nn as nn
class TinyModel(nn.Module):
def __init__(self, vocab=32000, d_model=512, nhead=8, layers=4):
super().__init__()
self.embed = nn.AI 词典：Embedding">Embedding(vocab, d_model)
layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=2048,
batch_first=True,
)
self.encoder = nn.TransformerEncoder(layer, num_layers=layers)
self.head = nn.Linear(d_model, vocab)
def forward(self, input_ids):
h = self.embed(input_ids)
h = self.encoder(h)
return self.head(h)
```
再造一批假数据，保证两边喂进去的张量完全一致：
```python
data.py
import torch
def make_batch(batch_size=8, seq_len=512, vocab=32000):
input_ids = torch.randint(0, vocab, (batch_size, seq_len))
labels = torch.randint(0, vocab, (batch_size, seq_len))
return input_ids, labels
```
先在有 GPU 的机器上跑通，记下基线数字。这一步的意义是：如果 Trainium 那侧结果对不上，你能确定问题出在迁移，而不是模型本身。
步骤四：理解图捕获，这是迁移的分水岭
GPU 是“来一条算子执行一条”，Trainium 不是。它需要一个完整的静态计算图，交给编译器统一编排，然后在设备上执行。所以“图捕获”是整个迁移里最需要理解的一步。
有两种用法，别混：
推理/静态导出，用 torch_neuronx.trace：
```python
import torch
import torch_neuronx
from model import TinyModel
model = TinyModel().eval()
example = torch.randint(0, 32000, (1, 128))
traced = torch_neuronx.trace(model, example)
torch.jit.save(traced, "tiny_neuron.pt")
```
这里 example 的形状会参与编译。把 batch 固定成 1，后面喂 batch 8 就会触发重新编译。需要可变 batch 时，用动态 batch 的开关（具体参数名以官方 API 文档为准）：
```python
traced = torch_neuronx.trace(model, example, dynamic_batch_size=True)
```
训练走的是惰性执行路线：代码看起来和普通 PyTorch 一模一样，但第一次 forward 不会立刻算出结果，而是先记录图、编译、再执行。所以第一次迭代会非常慢，几秒到几十分钟都正常，取决于模型大小。
无论走哪条路，都要记住一句话：任何依赖张量具体数值的 Python 控制流，都会把图打碎。 典型杀手是 .item()、if loss > 0.5:、print(loss)、以及每步都在变的 seq_len。打碎之后不是报错，而是悄悄重新编译，你会看到 GPU 上十分钟的活儿在 Trainium 上跑了三小时。
步骤五：把训练循环搬过去
训练侧用 PyTorch/XLA 的接口接 Neuron 设备：
```python
train_neuron.py
import torch
import torch.nn as nn
import torch_xla.core.xla_model as xm
import torch_neuronx.xla_impl # 让 XLA 运行时识别 Neuron 设备
from model import TinyModel
from data import make_batch
device = xm.xla_device()
model = TinyModel().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for step in range(50):
input_ids, labels = make_batch()
input_ids = input_ids.to(device)
labels = labels.to(device)
logits = model(input_ids)
loss = nn.functional.cross_entropy(
logits.reshape(-1, logits.size(-1)),
labels.reshape(-1),
)
optimizer.zero_grad()
loss.backward()
xm.optimizer_step(optimizer)
if step % 10 == 0:
print(step, float(loss.detach().cpu()))
```
几点说明：
import torch_neuronx.xla_impl的具体导入路径可能随版本调整，以官方教程当前写法为准。漏了这一句，xm.xla_device()找不到 Neuron 设备。- 打印 loss 用了
.detach().cpu()并且只在每 10 步一次，就是为了减少设备同步。真在意吞吐时，把它改成每 100 步一次，或者干脆写到文件里事后看。 - 想要 bf16 训练，通常通过环境变量开启（变量名以官方文档为准），例如
XLA_USE_BF16=1。开启后数值会和 fp32 有差异，这是正常的，对照时要保证两边精度设置一致。
先只跑 20 步，目的是把那面“编译墙”撞过去。跑通之后再谈性能。
步骤六：用并行编译缩短等待
模型稍大一点，串行编译会让人失去耐心。Neuron 提供了并行编译的开关，先只编译、不训练：
```bash
NEURON_PARALLEL_COMPILE=1 python train_neuron.py
```
它会把这批图分发到多个核心上并行编译并落盘缓存。编译完成后再正常执行，就会直接命中缓存。
多台机器、多人协作时，可以指定共享的编译缓存目录（变量名以官方文档为准）：
```bash
export NEURON_COMPILE_CACHE_URL=file:///mnt/shared/neuron-cache
```
缓存是双刃剑：命中时快得感人，不命中时磁盘涨得也感人。定期清理，或者挂一块独立的大盘专门放它。
步骤七：测吞吐，再折算成本
这一步最容易被做错。正确的做法是：丢掉编译时间，只算稳态。
在脚本里加一段计时代码：
```python
import time
前 5 步预热，触发编译
for _ in range(5):
step_once()
torch.cuda.synchronize() if torch.cuda.is_available() else None
start = time.time()
for _ in range(20):
step_once()
elapsed = time.time() - start
steps_per_sec = 20 / elapsed
samples_per_sec = steps_per_sec * BATCH_SIZE
print(f"steps/s={steps_per_sec:.3f} samples/s={samples_per_sec:.1f}")
```
Trainium 侧的同步语义和 GPU 不完全一样，稳妥起见在计时前显式做一次设备同步（例如 xm.mark_step() 或框架提供的同步调用，以官方文档为准），否则你会测到一个“很快但没算完”的数字。
拿到数字后填这张表，两边各一份。价格去 AWS 官方定价页现查，不要用记忆里的数字：
```text
项目 GPU 实例 Trainium 实例
单步耗时（稳态） ? ?
samples/s ? ?
首次编译耗时 ? ?
跑到目标 loss 需要步数 ? ?
实例按需小时价 ? ?
```
成本换算公式只有一行：
```text
每千步成本 = 小时价 × (总耗时秒数 / 3600) / (总步数 / 1000)
```
更贴近业务的版本是按“有效工作量”算：
```text
每百万 token 成本 = 小时价 × 训练小时数 / 训练出的 token 数（百万）
```
三个必须遵守的对照纪律：
1. 同等收敛。要比“跑到同一个验证集 loss 需要多少成本”，而不是“跑完固定步数谁快”。两边学习率、batch、数据顺序都要对齐。
2. 同等精度。一边 fp32 一边 bf16 的对比没有意义。
3. 算上编译。如果任务只跑一次、规模很小，首次编译时间会显著拉高总成本；如果是长期反复训练，编译缓存可以摊薄它。这个差别要写进结论里。
另外一个现实结论：小模型、小 batch 的场景，GPU 往往更省事也更便宜；Trainium 的甜点区在大 batch、大模型、长时间训练。别为了迁移而迁移。
常见坑与排错
neuron-ls 什么都看不到。
驱动没装好或没重启。回到步骤一，确认 aws-neuronx-dkms 装上了，lsmod 里能看到 Neuron 相关模块，然后重启实例。
导入就报错，提示版本不兼容。
torch 和 torch-neuronx 版本没对上。不要自己猜组合，直接查官方兼容矩阵，按组合重装，必要时重建虚拟环境。
跑起来极慢，且 neuron-top 显示反复编译。
八成是图被打碎了。逐个排查：.item()、if 依赖张量值、print(tensor)、每步变化的 shape、DataLoader 里返回了非张量结构。把控制流挪到图外，把 shape 固定下来。
设备内存不足。
减小 batch 或序列长度，用梯度累积凑等效 batch。不要指望像 GPU 那样靠“碎片整理”蒙混过去，静态图对内存的规划是编译期定死的。
报“算子不支持”。
错误信息里会给出算子名。找等价实现替换，比如用一个更基础的组合代替冷门算子。这也是迁移里最耗时的一类活儿，模型越“花”，改得越多。
磁盘突然满了。
编译缓存。清理缓存目录，或者把缓存指到大盘上。
设备利用率上不去。
先看 neuron-top。设备闲、CPU 忙，说明卡在数据加载：加 num_workers、做预取、把预处理提前离线做好。设备忙但吞吐仍低，说明 batch 太小或模型本身不适合这个硬件。
结果和 GPU 对不上。
先在完全相同精度（都 fp32）、相同 seed、相同数据顺序下比。确认能对上之后再开 bf16，此时出现小幅数值差异属于正常范围。
账单失血。
Trainium 实例按小时计费，忘了关就是纯亏损。养成“跑完即停”的习惯，或者用脚本定时关机。
下一步建议
跑通这条最小路径之后，按这个顺序往下走：
1. 换真模型。把 TinyModel 换成你实际业务里的模型，这一步会暴露算子兼容性和 shape 动态性的真实问题，通常需要几天而不是几小时。
2. 上分布式。单卡装不下时，用 Neuron 的分布式训练库做数据并行或张量并行。并行策略会改变编译行为，先在小模型上验证配置，再上大模型。
3. 接生态库。Hugging Face 模型有对应的 Neuron 适配路径，能省掉大量算子改写工作，具体支持范围以官方仓库当前说明为准。
4. 固化冒烟测试。在 CI 里放一个“能编译、能跑三步、loss 不炸”的最小用例，防止升级 SDK 后悄悄坏掉。
5. 建成本看板。把每千步成本、每百万 token 成本做成周度指标。迁移的最终说服力不是跑分，而是这张表。
迁移到新硬件从来不是“换个设备跑一下”，而是一次把训练代码里所有隐式假设显式化的过程。图捕获这一步逼你把动态的东西固定下来，短期是负担，长期反而是收益——你会发现自己的训练循环变得更快、更可预测，也更容易排查问题。