Qwen3 在 Trainium 上的三条训练路径

在 AWS Trainium 上训练 Qwen3，可以继续沿用 torchtitan 的配置方式和 PyTorch 接口。作者在一台 trn2.48xlarge 上演示了三条路径：使用 Transformers v5 的模型实现、切换到 torchtitan 原生实现，以及接入 NKI 专家内核和自定义算子。

这篇值得看的地方，是模型、权重格式和硬件优化如何衔接。但它证明的是训练链路能够跑通，没有证明 Trainium 更快、更便宜，也没有证明示例能直接用于完整模型训练。

三条训练路径，为什么能共用一个入口

连接 torchtitan 与 Trainium 的是 TorchNeuron。安装 torch_neuronx 后，Trainium 会注册为 PyTorch 的 neuron 设备，torchtitan 在启动时识别它。上层训练流程因此可以继续使用标准 PyTorch 接口。

执行方式也有两种：eager 模式在算子执行时编译；使用 torch.compile 时，Neuron 编译捕获到的计算图。选择哪种模型实现，并不要求重新设计启动流程。

第一条路径以 Transformers v5 为建模后端，通过 AutoModelForCausalLM 构建 Qwen3。配置从 transformers_modeling_backend_full() 开始，将 config.hf_model 设为 Qwen/Qwen3-4B-Instruct-2507。与 Hugging Face 架构配置对应的字段保持未设置，让 AutoConfig.from_pretrained() 读取模型结构。

第二条路径换成 torchtitan 原生 Qwen3Model。下面这段配置展示了真正控制模型实现的位置：

from torchtitan.models.qwen3 import model_registry
from torchtitan.models.qwen3.config_registry import qwen3_32b
from torchtitan.trainer import Trainer

def native_qwen3_4b() -> Trainer.Config:
    config = qwen3_32b()
    config.model_spec = model_registry("4B", attn_backend="sdpa")
    config.hf_assets_path = ASSETS
    config.dataloader.dataset = "c4_test"
    config.training.steps = 10
    config.training.local_batch_size = 1
    config.training.global_batch_size = 64
    config.parallelism.data_parallel_shard_degree = -1
    return config

要看的是 config.model_spec：作者从 qwen3_32b() 起步，是因为测试版本没有通用的 qwen3_4b() 配置构造器，下一行已经将模型规格替换为 4B，并选择缩放点积注意力 SDPA。这里没有训练 32B 模型。

两条路径都使用 torchrun，核心参数为 --nproc_per_node=64 --rdzv_backend c10d --rdzv_endpoint localhost:0，同时带上 --local-ranks-filter 0 --role rank --tee 3。入口是 -m torchtitan.train --module neuron_qwen3，再通过 --config 选择配置函数。

在 LNC=2 时，一台 trn2.48xlarge 提供 64 个逻辑 NeuronCore，对应启动的 64 个进程。data_parallel_shard_degree=-1 将参数、梯度和优化器状态分片到全部 rank。

权重能读写，才算接上现有工作流

能启动模型，不代表能正确加载已有权重。第一条路径的实用价值，是直接读写 Hugging Face 的 safetensors，减少训练前后转换格式的工作。

作者先用 hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir "$PWD/assets/hf/Qwen3-4B-Instruct-2507" 下载模型文件，再将本地目录的绝对路径赋给 ASSETS。

微调配置需要开启 checkpoint.enable，把 initial_load_path 指向 ASSETS，并设置 initial_load_in_hf=True、initial_load_model_only=True。导出时使用 last_save_in_hf=True 和 last_save_model_only=True。状态字典适配器负责处理键名前面的 model. 前缀，以及共享嵌入权重。

作者给出了一个实测信号：加载已发布权重时，第一步 loss 为 6.54；随机初始化时为 12.43。训练结束后，再通过 AutoModelForCausalLM 和 load_state_dict 回读导出文件，除共享词嵌入涉及的 lm_head.weight 外，没有其他缺失或多余的键。

这些检查支持「权重加载和导出链路有效」，不能直接支持「微调效果更好」。第一步 loss 适合快速排错，回读检查适合确认结构兼容，两者都不能替代任务评测。

示例只在 c4_test 上跑 10 步，local batch size 为 1。虽然模型名称含有 Instruct，这仍不是一套完整的指令微调方案。

接入 NKI 内核，难点在接口契约

第三条路径面向混合专家模型（MoE），通过 qwen3_nki_moe_config() 接入 NKI 内核：

from torchtitan.config import CompileConfig
from torchtitan_neuron.models.qwen3.config_registry import qwen3_nki_moe_config

def nki_qwen3_30b_a3b() -> Trainer.Config:
    config = qwen3_nki_moe_config(
        "30B-A3B",
        compile_config=CompileConfig(enable=False, backend="neuron"),
        n_layers=4,
        num_experts=64,
        top_k=4,
    )
    config.hf_assets_path = ASSETS
    return config

这里要同时看 n_layers=4 和 num_experts=64。作者为了让示例装进单机，缩减了模型；不能把它理解为完整 Qwen3-30B-A3B 的训练结果。每个 token 路由到 4 个专家，但专家数量 64 不代表专家并行度。该配置使用 64 rank 的数据并行分片，tensor parallel degree 为 1，没有启用专家并行。

CompileConfig(enable=False, backend="neuron") 选择 eager 模式。该配方还替换了 grouped-experts 实现，选择 SDPA 和 gather-free rotary embedding。若启用 torch.compile，原文要求向 CompileConfig.components 加入 "grad_clip"，以编译梯度范数裁剪。

自定义内核则展示了更底层的接法：wrap_nki 将内核包装成 PyTorch 可调用对象，nki_op 注册算子，register_autograd 接上反向传播。

作者用 nkilib.core.embeddings 中的 rope_hf 替换旋转位置编码。这里最大的坑是张量布局：torchtitan 使用 [batch, seq, heads, head_dim]，内核要求 [batch, heads, seq, head_dim]，前后都要转换，并用 torch.empty_like() 分配输出。反向调用还要传入 backward=True，复用前向保存的缓存，缓存梯度返回 None。

替换通过 @override 注册，derive() 复制共有配置字段，启动时增加 --override.imports my_neuron_kernels。日志应当为每个 Qwen3 层打印一条 nki_rope 替换记录。

序列长度必须是 128 * LNC 的倍数，LNC=2 时就是 256；缓存形状要求为 [seq, head_dim * 2]，前半部分存 cosine，后半部分存 sine。这些约束不满足，内核就不能直接接入。

作者在 float32 输入下比较了新旧实现，query 和 key 的最大绝对误差都为 0.0。这验证了该数值测试中的一致性，没有验证性能收益。

我的判断

我认可这篇的接入顺序：先打通标准模型与权重，再切换实现，最后替换局部内核。每次只改变一层，出现问题时比较容易定位。

但复现材料存在缺口。原文说要固定版本，并提到版本表，提供的正文却没有表格。它还强调使用 --no-deps 保留 Neuron 版 torch，展示的安装命令却是 pip install -r requirements.txt，没有解释参数放在哪里。考虑到 Transformers v5 集成和 override 模块会随 torchtitan 版本变化，这两处不能靠读者猜。

性能判断也要收住。原文没有吞吐、成本或替代硬件的对照数据，自定义内核示例更明确说明不用于展示加速。对普通前端、Node 团队，它提供的是训练基础设施接入方法，不能据此决定将线上模型服务迁到 Trainium。

此外，自定义 rope_hf 包装并没有覆盖所有并行情况。原文另指向 torchtitan_neuron/models/gpt_oss/nki_attention.py，说明其中还处理了张量并行下的 DTensor 输入。单机示例跑通后，扩大并行规模仍需继续验证。

能拿来做什么

• 如果团队准备在 Trainium 上微调 Qwen3，先在 AWS Deep Learning AMI for Neuron 或 Neuron PyTorch training container 中确认依赖，再跑 10 步配置，检查 64 ranks、模型规格和权重回读结果。
• 如果要定制训练算子，先核对布局、序列长度、缓存格式和反向传播，再比较输出。性能优化要另测吞吐与 HBM，首步包含编译，不能直接拿它代表稳定训练速度。
• 如果要扩展模型、序列长度或 local batch size，持续观察每个 rank 的 HBM。设置 config.parallelism.tensor_parallel_degree 前，先弄清它与 FSDP、默认开启的 sequence parallelism 如何在 DeviceMesh 上组合。
• 如果要叠加多个 override，用 fqns=["layers.1[0-9].attention.rope"] 这类范围限制目标层。该表达式匹配第 10 至 19 层；两个 override 抢占同一节点会报错，没有匹配节点则会警告。