在 AWS Trainium 上用 torchtitan 训练 Qwen3

AWS Machine Learning Blog 发布了一篇操作教程，介绍如何在单台 trn2.48xlarge 实例上使用 torchtitan 训练 Qwen3。核心机制是 TorchNeuron 把 Trainium 注册为原生 PyTorch 设备，torchtitan 通过标准 PyTorch 接口即可识别，因此启动方式和配置结构与常规用法保持一致。

三种训练配置

教程以 Qwen/Qwen3-4B-Instruct-2507 为例，给出三条路径：

• 以 Transformers v5 作为建模后端，通过 AutoModelForCausalLM 加载 Hub 模型，直接读写已发布的 safetensors；
• 切换到 torchtitan 原生 Qwen3 实现，只需修改 config.model_spec；
• 在 torchtitan-neuron recipes 基础上加入 Neuron Kernel Interface（NKI）的 MoE 内核，并通过 torchtitan 的 override registry 挂载自定义内核。

三种配置都使用相同的 torchrun 启动模式和配置结构。示例命令以 64 个进程启动，对应 64 个逻辑 NeuronCore（LNC=2）。

环境与版本

教程要求使用 AWS Deep Learning AMI for Neuron 或 Neuron PyTorch 训练容器，二者均提供 Neuron 驱动、运行时和编译器。文中固定了一组版本以便复现：torch 与 torch_neuronx 为 2.14.0 和 2.14.3，transformers 5.16.0，torchtitan-neuron 为 release-2.33（0.2.2+neuron.1）。安装 torchtitan-neuron 及其余依赖时使用 --no-deps，以保留 Neuron 构建的 torch。

微调与检查点

从已发布权重微调时，需开启 checkpoint 并设置 initial_load_path，状态字典适配器会处理模型前缀和 tied embeddings，最终以 Hugging Face safetensors 格式保存。教程给出的验证数据是：加载已发布权重时首步 loss 为 6.54，随机初始化则为 12.43。导出的检查点可用 AutoModelForCausalLM 和 load_state_dict 重新加载，除 Qwen3-4B 与 token embedding 绑定的 lm_head.weight 外，没有缺失或意外键。

编辑部观察

这篇教程的价值在于把 Trainium 上的训练入口收敛到标准 PyTorch 接口，降低迁移时的改动面。需要注意的是，文中版本组合被明确要求固定，Transformers v5 集成和 Trainium 专用 override 模块可能随 torchtitan 版本变化；实际落地前应先核对自身环境的 torchtitan-neuron 版本，并用首步 loss 做一次加载正确性检查。教程只覆盖单机 64 逻辑核的配置，未涉及多机扩展或性能对比数据。