AWS 实战教程：用 torchtitan 在 Trainium 上以三种方式训练 Qwen3

AWS ML Blog · rss · 2026-10-22

AWS ML Blog 发布完整教程，演示在单台 trn2.48xlarge（16 颗 Trainium2、128 物理核）上用 torchtitan-neuron 训练 Qwen3-4B 的三种配置：

Transformers v5 后端：通过 AutoModelForCausalLM 加载 Hub 模型，直接读写 safetensors，可从公开权重微调（验证 run 首步 loss 6.54，随机初始化为 12.43），导出 checkpoint 可被 HF 标准加载。
torchtitan 原生 Qwen3：改 config.modelspec 切到原生实现（SDPA 注意力）。
NKI MoE 内核：用 Neuron 内核语言编写的混合专家内核（torchtitan-neuron recipes），并通过 override registry 挂载自定义 kernel。

三种配置共用同一 torchrun 启动模式和 config 结构；文内给出完整版本锁定表（torch 2.14.x、transformers 5.16.0、torchtitan-neuron release-2.33）、安装命令与配置代码，可照做复现。

「Infra」频道最新