AWS 博客：用 torchtitan 在 Trainium 上训练 Qwen3

核心事实

AWS 机器学习博客发布了一篇技术教程，介绍如何在单台 trn2.48xlarge 实例上使用 torchtitan 训练 Qwen3。文章称，TorchNeuron 会把 Trainium 注册为原生 PyTorch 设备，因此 torchtitan 可以通过标准 PyTorch 接口调用它。教程给出三种配置：一是以 Transformers v5 作为建模后端，通过 AutoModelForCausalLM 读写 Hub 上的 safetensors；二是改用 torchtitan 原生 Qwen3，只需修改 config.model_spec；三是在此基础上接入 torchtitan-neuron 配方中的 Neuron Kernel Interface（NKI）混合专家（MoE）内核，并通过 override registry 挂载自定义内核。三种配置共用同一套 torchrun 启动方式和配置结构。

背景与影响

文章给出的环境为 trn2.48xlarge（16 个 Trainium2 加速器、128 个物理 NeuronCore，LNC=2 时 64 个逻辑 NeuronCore），并列出 torch、torch_neuronx 2.14.0/2.14.3、transformers 5.16.0、torchtitan-neuron release-2.33 等版本。教程以 Qwen/Qwen3-4B-Instruct-2507 为例，启动 64 个训练进程，每个逻辑 NeuronCore 一个。文中提到，验证运行中从已发布权重加载时第 1 步 loss 为 6.54，随机初始化时为 12.43；导出检查点用 AutoModelForCausalLM 加载后，除 Qwen3-4B 与词嵌入绑定的 lm_head.weight 外没有缺失或意外键。这些数字来自该博客的示例运行，并非通用性能结论。

限制与来源

文章提醒，Transformers v5 集成和 Trainium 专用 override 模块可能随 torchtitan 版本变化，复现时需固定上述版本；示例使用 c4_test 数据集、10 个训练步和本地批大小 1，属于演示配置。由于测试所用构建没有通用的 qwen3_4b() 配置构造器，原生模型示例从 qwen3_32b() 出发再替换为 4B 规格。具体环境、版本与可用性请以 AWS 官网及项目仓库当前信息为准。来源：aws.amazon.com 博客文章《Train Qwen3 on AWS Trainium with torchtitan》。