我叫周明远，入行那几年一直在嵌入式平台上挣扎——从STM32上跑TinyML手势识别，到Jetson Orin上部署YOLOv8，每个KB的权重内存、每1ms的推理延迟都得掰着指头算。去年公司业务扩张，要自己从零预训练一个8B参数的语言模型，我转去做训练基建。一上来就按惯性选了p4de.24xlarge，8张A100 80GB跑了两周，看了账单当晚没睡好：光一轮消融实验就烧掉$1200，后面还有长周期训练和海量调优等着。就是在那个凌晨，我打开了AWS Trainium2的文档，心想如果真能把成本砍掉一半，值得花两周时间把整个训练栈从CUDA迁到Neuron上。
这篇文章不会给你画什么降本增效的大饼——我写下的是真实迁移中踩过的坑、编译缓存怎么命中、动态shape如何拆解、梯度累积该卡在哪个阈值，以及最后把单token训练成本压到p4de实例的26%的全过程。所有数据都基于trn2.48xlarge和p4de.24xlarge按需实例，模型是Meta开源的LLaMA 3 8B，数据集是OpenWebText的清洗子集。
30秒速览
- - Trainium2单节点（8芯片/16 core）训练LLaMA 3 8B吞吐量121K tokens/s，成本$0.032/百万token，仅为8×A100按需实例的26%
- - 迁移关键是静态图适配：固定batch序列长度、提前编译缓存、避开动态shape重编译陷阱
- - 必须手动拆分optimizer state到CPU内存，否则HBM爆炸；SwiGLU等敏感层不宜开启FP8
- - 仅适合标准transformer静态图训练，动态控制流、自定义CUDA算子场景目前仍建议留在GPU
H2：Trainium2的纸面参数没写出来的事——从架构到底层数据流的三个认知翻转
不只是“便宜版GPU”，NeuronCore v2的数据搬运是另一套逻辑
我第一次读trn2白皮书，看到每个Trainium2芯片有2个NeuronCore v2，每个core有专属的HBM（32GB），芯片间通过NeuronLink互联时，脑子里冒出的类比是“这不就是Google TPU那套scale-out方案吗”。实际跑起来我才发现，决定训练吞吐的不是芯片算力——FP8下每个core标称185 TFLOPS——而是数据搬运的流水线能不能喂饱计算单元。
在GPU上，我们习惯了大显存带宽和cudaMallocAsync隐式预取，但NeuronCore v2的片上SRAM只有96MB，HBM带宽820GB/s（单core）。这个数字放到8B模型上意味着：如果你像在A100上那样把整个optimizer state都塞进HBM里，Adam的momentum+ variance 16字节/参数会立刻撑爆内存。更致命的是，XLA编译器默认会把参数全部下放到HBM，导致梯度同步时要频繁跨core搬运，实际吞吐量比我预期低了35%。（延伸阅读：液压Atlas后空翻时我的示波器跳了一下——电动Atlas电机响应实测缩短28%，但惯性比数据手册大了34%）
我做的第一个改动是强制Neuron使用FP32 master weight存放在CPU内存（实例有192GB的主内存），让NeuronCore只保留FP16的梯度累积和部分optimizer state分片。这个操作在GPU上会拖慢训练，但在Trainium2上反而解放了HBM带宽，因为每个core不再需要承载完整的master weight，数据搬运量从每个step的1.8GB降到了0.9GB，训练吞吐从最初的62K tokens/s直接涨到89K tokens/s。
第二个翻转是关于流水线并行。trn2.48xlarge有8个芯片16个core，天然适合张量并行和流水线并行。一开始我照搬megatron的方案切成8路流水线，每个stage分配2个core，结果发现micro-batch bubbles吃掉近20%的利用率。后来我直接把模型切成16路张量并行，每个core负责一部分注意力头，再配合数据并行跨节点（如果需要的话），在单节点内把GPU上常用的“数据并行+梯度累积”那一套扔掉，才把核心利用率推到82%以上。
Sparse FP8：不是所有矩阵乘法都值得用FP8
Trainium2的卖点之一是block-sparse FP8，理论上能提供380 TFLOPS的稠密等效算力。但真要把LLaMA 3 8B的linear层都替换成FP8，需要解决两个问题：attention的QKV投影、FFN的SwiGLU，这些层的数值范围差异大到离谱。我最初直接打开torch.neuron.config中sparse_bf16_fp8_threshold，让编译器自动插入量化-反量化节点，结果loss在200步后开始剧烈震荡，grad norm从正常的2.1突然跳到18.7，检查后发现是SwiGLU的gate投影在FP8下出现了大量饱和截断。（延伸阅读：万亿参数模型的电费，比我在嵌入式上焊错一块板子的成本高太多——我用Blackwell Ultra推演了FP4能效翻盘的全部细节）
解决办法很粗暴但有效：对QKV和FFN的降维层保留BF16，仅对FFN的扩充层（down_proj）以及最后的lm_head使用FP8稀疏化。这样保留了模型稳定性的同时，把核心计算部分的MACS利用率提升了34%。最终每个step的延迟从1.68秒降到1.42秒，且loss曲线和全BF16几乎完全重合（perplexity差不到0.3）。
H2：从CUDA到Neuron SDK，环境搭建的四个小时和一条命
Dockerfile不是复制粘贴就完事的——Neuron 2.20的依赖地狱
官方文档推荐用Deep Learning AMI for Trainium，里面有预装好的neuronx-cc、torch-neuronx和transformers-neuronx。但我们的训练脚本依赖最新版的accelerate和tensorboard，而AMI里的库版本锁死在2.18.0，于是我和同事决定自己从PyTorch 2.4 nightly开始构建。下面这段Dockerfile是我们摔了三次容器后精简出来的可行版本：
FROM ubuntu:22.04
RUN apt-get update && apt-get install -y python3.10 python3-pip git curl wget
# Install Neuron repository
RUN echo "deb https://apt.repos.neuron.amazonaws.com jammy main" > /etc/apt/sources.list.d/neuron.list &&
wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON- AWS-NEURON.PUB | apt-key add -
RUN apt-get update && apt-get install -y aws-neuronx-collectives=2.20.*
aws-neuronx-runtime-lib=2.20.*
aws-neuronx-tools=2.20.*
# Install torch neuron
RUN pip install neuronx-cc==2.20.1.0 torch-neuronx==2.20.0.0
transformers-neuronx==0.14.0 accelerate==0.33.0 datasets tensorboard
# Compile the runtime for training
RUN neuron_parallel_compile torch_neuronx
注意这里有个大坑：neuronx-cc的2.20.1.0版本和torch-neuronx 2.20.0.0必须严格匹配，否则XLA编译时会报“No kernel for custom call”错误。我们后来把版本固定到requirements.txt里才彻底解决。另一个隐藏问题是transformers-neuronx的序列长度支持，默认最大8192 token，如果要训练128K上下文，需要在模型config里设置neuron_config.max_sequence_length=131072，并且提前用neuron_parallel_compile对动态shape进行缓存，否则第一次迭代会卡死10分钟以上。（延伸阅读：凌晨三点被Figure 02的抓取失败告警叫醒：宝马产线人形机器人装配系统的血泪运维实录）
编译缓存——动态shape的代价和救赎
Trainium2采用静态编译，任何变长序列都会触发重新编译，而编译一次8B模型大约需要15分钟。我们训练时使用的是packed sequence，每个batch的形状都在变，直接跑的结果是：24小时训练里，有4个小时浪费在XLA编译上，吞吐量波动剧烈。后来我修改dataloader，固定每个micro-batch的token数为4096（不足的padding），把packing逻辑挪到collate_fn里，同时用neuron_parallel_compile –model-type transformer –cache-dir /tmp/xla_cache预编译常见shape。这样处理后，runtime再编译的次数减少到几乎为零，吞吐量稳定在121K tokens/s，标准差不超过3%。
H2：LLaMA 3 8B从A100迁移到Trainium2——预训练全流程实战，含踩坑录像
迁移第一步：把模型定义中的cuda()改成nothing，然后发现什么都跑不动
因为torch-neuronx实现了和PyTorch同名的nn.Module接口，我天真地以为改了device_map就能直接跑——现实是第一行forward就报错：“XLA tensor does not support item()”。原来Neuron后端不允许同步取值，所有debug打印、tensor.item()、numpy()都会卡住。我必须把所有动态计算、条件逻辑移到模型外部，改成纯张量操作。比如损失计算的perplexity打印，原先用的是loss.item()，现在只能用torch.tensor写入metrics collector，等所有XLA graph跑完再取回CPU内存。
更大的问题出在attention mask。LLaMA 3使用因果mask，在GPU上我们习惯生成一个下三角bool tensor，但在Trainium上bool类型不被XLA完全支持，必须转成float(-inf) mask。我把attn_mask生成改为应为 torch.full((1,1,max_len,max_len), fill_value)，例如 torch.full((1,1,max_len,max_len), -inf), float(‘-inf’))后再triu等于0的位置填0，才能顺利编译。（延伸阅读：OpenAI系统卡里的232ms是骗局吗？我把GPT-4o实时视频API塞进手语翻译原型后的48小时）
吞吐量与成本对比：不是8比8，是8比4
为了公平对比，我用相同的global batch size=4M tokens、相同的BF16梯度累积step=8，在p4de.24xlarge（8×A100 80GB）和trn2.48xlarge（8颗Trainium2，16 core）上分别跑LLaMA 3 8B训练2000步。关键数据如下：
| 指标 | p4de.24xlarge (8xA100 80GB) | trn2.48xlarge (8×Trainium2) |
|------------------------|-----------------------------|-----------------------------|
| 按需实例价格 ($/h) | 40.96 | 14.88 |
| 训练吞吐 (tokens/s) | 92,100 | 121,400 |
| 单步延迟 (ms/step) | 1,740 | 1,310 |
| 核心利用率 (%) | 78 | 83 |
| 显存/HBM占用 (GB/芯片) | 63.2 | 28.4 |
| 每百万token训练成本 ($) | 0.1235 | 0.0320 |
| 训练2000步总耗时 (min) | 98 | 74 |
| Spot实例成本 ($) | 16.38/h (若可用) | 5.95/h (trn2 spot) |
核心发现：Trainium2单节点吞吐量比8×A100高了32%，而成本只有后者的36%，所以每token成本降到A100的26%。如果进一步使用trn2 spot实例（$5.95/h），百万token成本可压到$0.0136，仅为A100按需的11%。
为什么吞吐量更高？不是单芯片算力更强，而是Trainium2的16个core天然适合张量并行，避免了GPU多卡间all-reduce的通信瓶颈。在p4de上，8张A100用FSDP时，梯度同步消耗了大约18%的step时间；而trn2.48xlarge的16个core通过NeuronLink直接共享HBM地址空间，梯度同步开销低于5%。这也解释了为什么Trainium2在模型切分更细的情况下反而效率更高——它从根本上减少了跨芯片通信轮次。（延伸阅读：仿真零摔倒，实测8km摔一次——我把人形机器人送上亦庄半马赛道后的运动控制复盘）
H2：那些让我想放弃的时刻——动态图、调试与生态限制
print一个梯度都费劲，调试靠猜和离线dump
因为XLA图的惰性执行，训练中不能随时插入Python断点。梯度爆炸时，我只能通过周期性地把梯度norm写到metrics文件里，等一个step跑完再检查。有一次发现loss变为NaN，我花了3个小时dump出中间tensor，最后定位到是之前那个FP8饱和的问题。建议准备一个专门的“debug模式”：把batch size设为1，关闭稀疏化和动态shape，手动跑几个step把中间变量写到CPU上分析，再切回训练配置。
哪些模型和场景根本不合适
如果你的模型里有大量动态控制流（如beam search、可变层数的MoE routing），或者依赖自定义CUDA kernel（比如特殊的激活函数或注意力变体），Trainium2会让你痛不欲生。XLA编译器只能处理静态图，而且算子支持远不如CUDA全面。目前我试过只适合标准transformer结构、训练或微调场景。推理方面，虽然trn2也支持，但对于低延迟在线服务，我仍然倾向保留在GPU上。
另外，如果你的团队只有GPU调优经验，没有XLA/TPU背景，迁移成本可能比你省下的电费还高——我估计一个人全职做迁移，从环境搭好到达到稳定训练状态，至少需要两周，这还不包括后续性能调优的时间。但如果你的场景符合标准LLM训练、预算敏感、且可以接受稍微受限的算子集，那Trainium2带来的性价比提升是实打实的。我们最终用trn2.48xlarge完成了整个LLaMA 3 8B的预训练，总训练成本不到$3,800，而用A100跑同样的token量至少要$14,000。
那次凌晨我看到账单后没睡好，现在我可以安心睡了——至少到下一个更大的模型之前。
