在 AWS Trainium 上用 torchtitan 训练 Qwen3
原标题:Train Qwen3 on AWS Trainium with torchtitan
AI 摘要
AWS 机器学习博客发布教程,介绍如何在单台 trn2.48xlarge 实例上使用 torchtitan 在 AWS Trainium 上训练 Qwen3。教程给出三种配置:以 Transformers v5 作为建模后端、使用 torchtitan 原生 Qwen3 模型,以及接入 NKI 的 MoE 内核。TorchNeuron 将 Trainium 注册为原生 PyTorch 设备,三种配置共用相同的 torchrun 启动方式。
正文节选
Train Qwen3 on AWS Trainium with torchtitan Using torchtitan on AWS Trainium means using it the way you always have. TorchNeuron registers Trainium as a native PyTorch device, which lets torchtitan use it through standard PyTorch interfaces. You can use these configurations to pre-train from random initialization or fine-tune published weights. You can also customize the model implementation and kernels for your workload. In this post, you train Qwen3 three ways on a single trn2.48xlarge. The fi