AWS发布教程：在单台trn2.48xlarge上用torchtitan训练Qwen3
AWS机器学习博客给出三种在Trainium上训练Qwen3的配置：Hugging Face Transformers v5后端、torchtitan原生Qwen3，以及NKI混合专家内核；三种方式共用同一套torchrun启动方式和配置结构。文中数字均来自这次验证运行。
AI解读：第一种是Transformers v5建模后端，通过AutoModelForCausalLM加载Hub模型，直接读写已发布的safetensors；第二种将config.model_spec切换为torchtitan原生Qwen3，示例用 4B规格替换 32B；第三种在torchtitan-neuron配方中启用NKI混合专家内核，并通过override registry挂载自定义内核。
作者称三种配置使用相同的torchrun启动模式和配置结构。要复现示例需固定版本：torch/torch_neuronx 2.14.0/2.14.3、transformers 5.16.0、torchtitan-neuron release-2.33。
作者报告的可核对数值只有两项：Qwen3-4B-Instruct-2507从已发布权重加载时第 1 步loss为 6.54，随机初始化为 12.43；自定义rope_hf内核在float32下与torchtitan PyTorch旋转位置编码的最大绝对差为 0.0。后一示例用于演示内核注册与自动求导，并非性能提升。
使用范围也有限定：MoE示例将Qwen3-30B-A3B缩减为 4 层、64 个专家、top_k=4 以塞进单台实例；该示例关闭torch.compile以eager模式运行，使用数据并行分片、张量并行度为 1，不使用专家并行。自定义rope内核要求序列长度为 128×LNC的倍数。
AWS机器学习博客发布一篇教程，介绍在单台trn2.48xlarge上使用torchtitan训练Qwen3的三种配置。作者Josh Longenecker称，三种方式共用同一套torchrun启动方式和配置结构。
三种路径分别是：以Hugging Face Transformers v5作为建模后端；切换到torchtitan原生Qwen3实现；以及在混合专家层使用Neuron Kernel Interface（NKI）内核。
运行环境与版本
教程要求使用一台trn2.48xlarge，配置 16 个Trainium2加速器；当LNC=2 时，暴露 64 个逻辑NeuronCore。软件环境为AWS Deep Learning AMI for Neuron或Neuron PyTorch训练容器，两者都提供Neuron驱动、运行时和编译器。
作者称，torchtitan-neuron release-2.33（0.2.2+neuron.1）已在环境中安装或需要先安装；示例固定使用torch、torch_neuronx 2.14.0、2.14.3 及transformers 5.16.0。文中称Transformers v5集成和Trainium专用覆盖模块可能随torchtitan版本变化。
Qwen3模型文件需从Hugging Face Hub下载到本地目录，例如Qwen/Qwen3-4B-Instruct-2507。
三种配置的具体做法
第一种配置从transformers_modeling_backend_full() 开始，将hf_model设为Hub ID，由AutoConfig.from_pretrained() 从Hub加载模型架构。启动命令为torchrun --nproc_per_node=64，每个逻辑NeuronCore对应一个训练进程，数据并行分片设为 -1，即把模型切分到全部 64 个rank。
第二种配置从qwen3_32b() 开始，再用config.model_spec = model_registry("4B", attn_backend="sdpa") 替换为 4B规格。作者说明，所用torchtitan构建没有通用的qwen3_4b() 配置函数，所以从 32B配置改起。
第三种配置调用qwen3_nki_moe_config("30B-A3B", ...)，将MoE块替换为NKI分组专家实现，并选用SDPA和不带gather的旋转嵌入。该示例设置n_layers=4、num_experts=64、top_k=4 以适配单台实例，并用CompileConfig(enable=False) 在eager模式下运行。
加载检查点与自定义内核
要基于已发布权重微调，需打开checkpoint.enable，把initial_load_path指向本地模型目录，并开启initial_load_in_hf与last_save_in_hf。作者称状态字典适配器会处理model.前缀和共享嵌入，最终以Hugging Face safetensors保存。
教程给出的加载校验是第 1 步loss：从已发布权重初始化为 6.54，随机初始化为 12.43。导出检查点后可用AutoModelForCausalLM和load_state_dict加载；作者称往返加载没有缺失或意外键，只有lm_head.weight例外，因为Qwen3-4B将其与词嵌入绑定。
自定义内核示例用nkilib的rope_hf替换旋转位置编码，并通过wrap_nki、nki_op和register_autograd接入训练。作者明确该示例演示内核注册与自动求导集成，并不代表性能提升。数值测试中，float32输入下rope_hf与原PyTorch实现的查询、键最大绝对差均为 0.0。
rope_hf要求序列长度为 128×LNC的倍数，LNC=2 时为 256；示例在NKIRoPE中先检查这一条件再调用内核。多个覆盖同时启用时，作者建议用fqns限定作用模块，例如fqns=["layers.1[0-9].attention.rope"] 只作用于第 10 至 19 层。