用 torchtitan 在 AWS Trainium 上训练 Qwen3

在 AWS Trainium 上使用 torchtitan，意味着按照你一直以来的方式使用它。TorchNeuron 将 Trainium 注册为原生 PyTorch 设备，这让 torchtitan 能够通过标准 PyTorch 接口使用它。你可以使用这些配置从随机初始化进行预训练，或对已发布的权重进行微调。你也可以针对自己的工作负载自定义模型实现和内核。在这篇文章中，你将在单个 trn2.48xlarge 上以三种方式训练 Qwen3。

第一种方法将 Transformers v5 作为 torchtitan 的建模后端。它通过 AutoModelForCausalLM 运行 Hub 模型，并直接读写已发布的 safetensors。第二种方法通过修改 config.model_spec 切换到 torchtitan 原生的 Qwen3。第三种方法添加 torchtitan-neuron 配方中的 Neuron Kernel Interface（NKI）混合专家（MoE）内核，然后使用 torchtitan 的覆盖注册表挂载你自己的内核。所有三种配置都使用相同的 torchrun 启动模式和配置结构。

前提条件

要按照本演练操作，请使用单个 trn2.48xlarge，并搭配适用于 Neuron 的 AWS Deep Learning AMI 或 Neuron PyTorch 训练容器。这两种环境都提供了在 Trainium 上训练所需的 Neuron 驱动程序、运行时和编译器。

在安装 torchtitan-neuron 之前，请检查你的环境中是否已经包含它。以下命令会打印已安装的版本和包路径：

python3 -c "import torchtitan, torchtitan_neuron; print(torchtitan.__version__, torchtitan_neuron.__file__)"

本演练使用下表中的版本。要复现这些示例，请固定这些版本。Transformers v5 集成和 Trainium 专属的覆盖模块在不同 torchtitan 版本之间可能会有变化。

当 torch_neuronx 安装后，TorchNeuron 会将 Trainium 注册为 neuron 设备，torchtitan 会在启动时检测到它。在 eager 模式下，Neuron 会在操作执行时进行编译。使用 torch.compile 时，Neuron 会编译捕获的计算图。这种相同的启动模式适用于所有三种配置。

对于本演练使用的版本，请使用 --no-deps 安装 torchtitan-neuron 和 torchtitan 的其余依赖，以保留 torch 的 Neuron 构建：

git clone https://github.com/aws-neuron/torchtitan-neuron
cd torchtitan-neuron
pip install -r requirements.txt

在开始训练之前，请从 Hugging Face Hub 获取 Qwen3 模型工件。将它们存储在一个本地目录中，并在训练配置中传入该目录：

hf download Qwen/Qwen3-4B-Instruct-2507 \
  --local-dir "$PWD/assets/hf/Qwen3-4B-Instruct-2507"

每个训练配置都是一个 Python 配置函数。这三个示例都从一个基础配置开始，然后针对每一步进行更新。

使用 Hugging Face Transformers v5 训练

在第一种配置中，Transformers v5 通过 AutoModelForCausalLM 构建 Qwen3。Hub 模型定义提供用于前向和反向传播的实现。你可以通过在配置中设置其 Hub ID 来选择模型。

从 transformers_modeling_backend_full() 开始。它将映射到 Hugging Face 配置键的字段留空，这样 AutoConfig.from_pretrained() 就会从 Hub 加载模型架构。

在 torchtitan-neuron 仓库根目录（也就是你运行训练命令的地方）创建 neuron_qwen3.py。添加导入、模型工件路径和基础配置：

import os
from torchtitan.experiments.transformers_modeling_backend.config_registry import (
    transformers_modeling_backend_full,
)
from torchtitan.experiments.transformers_modeling_backend.configs import (
    TransformersBackendConfig,
)

ASSETS = os.path.abspath("./assets/hf/Qwen3-4B-Instruct-2507")

def v5_qwen3_4b() -> TransformersBackendConfig:
    config = transformers_modeling_backend_full()
    config.hf_model = "Qwen/Qwen3-4B-Instruct-2507"
    config.hf_assets_path = ASSETS
    config.dataloader.dataset = "c4_test"
    config.training.steps = 10
    config.training.local_batch_size = 1
    config.parallelism.data_parallel_shard_degree = -1
    return config

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 v5_qwen3_4b

此命令启动 64 个训练进程，每个逻辑 NeuronCore 一个进程。--module neuron_qwen3 导入本地配置模块，--config v5_qwen3_4b 选择配置函数。数据并行设置将模型分片到所有 64 个 rank 上。

加载和保存 Hugging Face 检查点

要从已发布的权重进行微调，请启用检查点加载并将 initial_load_path 指向 ASSETS。状态字典适配器会调整开头的 model. 前缀并处理共享的嵌入。它将最终检查点保存为 Hugging Face safetensors 格式。

def v5_qwen3_4b_finetune() -> TransformersBackendConfig:
    config = v5_qwen3_4b()
    config.checkpoint.enable = True
    config.checkpoint.folder = "checkpoint_qwen3_4b_ft"
    config.checkpoint.initial_load_path = ASSETS # absolute path required
    config.checkpoint.initial_load_in_hf = True
    config.checkpoint.initial_load_model_only = True
    config.checkpoint.last_save_in_hf = True
    config.checkpoint.last_save_model_only = True
    return config

使用第一个训练步骤的损失作为快速检查，验证 torchtitan 是否成功加载了已发布的权重。对于 Qwen/Qwen3-4B-Instruct-2507，经核实的运行中，已发布权重产生的第 1 步损失为 6.54，而随机初始化为 12.43。

训练完成后，使用 AutoModelForCausalLM 和 load_state_dict 加载导出的检查点。如果往返转换成功，则除了 lm_head.weight 之外不会报告任何缺失或多余的键，Qwen3-4B 将该键与词元嵌入绑定在一起。这证实了这些文件仍然是标准的 Hugging Face 检查点。

图 1：已发布的检查点起始损失为 6.54，且导出的权重加载时没有缺失或意外的键

使用 torchtitan 的原生模型进行训练

在第二种配置中，torchtitan 构建其原生的 Qwen3Model。将原生的 Qwen3 导入和配置函数添加到 neuron_qwen3.py 中：

此示例从 qwen3_32b() 开始，因为所测试的 torchtitan 构建没有提供通用的 qwen3_4b() 配置构建器。下一行将 32B 模型规格替换为 4B 规格。其余字段复现了所测试的运行。

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

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 native_qwen3_4b

config.model_spec 选择模型类，并定义 torchtitan 如何对模型进行分片以及如何映射其检查点键。此处，model_registry("4B", attn_backend="sdpa") 选择 torchtitan 原生的 Qwen3 实现，并使用缩放点积注意力（SDPA）。

为专家混合添加 NKI 内核

在第三种配置中，torchtitan 使用 NKI（Neuron 内核语言）编写的内核运行专家混合层。将所需的导入和 Qwen3 MoE 配置函数添加到 neuron_qwen3.py 中：

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

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 nki_qwen3_30b_a3b

此配置在 64 个排名上使用数据并行分片，张量并行度为 1。它不使用专家并行。num_experts=64 设置模型中专家的数量，而不是专家并行度。

设置 CompileConfig(enable=False) 会以即时模式运行此示例。TorchNeuron 在即时执行和 torch.compile 两种模式下都支持已注册的 NKI 算子。

在本演练使用的 torchtitan 构建中，qwen3_nki_moe_config() 将 MoE 块替换为 NKI 分组专家实现。它还会选择 SDPA 和无需 gather 的旋转位置嵌入。启用 torch.compile 时，请将 "grad_clip" 添加到 CompileConfig.components 以编译梯度范数裁剪。

NKI 内核作为 PyTorch 算子进入训练。wrap_nki 使 NKI 内核可从 PyTorch 调用，nki_op 注册该算子。register_autograd 使用前向传播期间保存的路由元数据来附加其反向传播。

该示例将 Qwen3-30B-A3B 缩减为 4 层和 64 个专家，使其能够装入单个 trn2.48xlarge。它通过 top_k=4 将每个词元路由到 4 个专家。

附加你自己的 NKI 内核

在某些情况下，你可能希望针对模型中的张量形状优化某个操作。专门的内核可以针对这些形状定制其分块和数据移动方式。它还可以用一个优化的算子替换性能关键的 PyTorch 操作序列。你可以自己编写内核，也可以从 nkilib 中选择一个作为起点。此示例使用 nkilib.core.embeddings 中的 rope_hf，它期望 Hugging Face 的张量布局并支持自己的反向模式。此示例演示的是内核注册和自动求导集成，而非性能改进。

在仓库根目录（即你启动训练的地方）创建 my_neuron_kernels.py。torchtitan 将 query 和 key 存储为 [batch, seq, heads, head_dim]，而 rope_hf 期望 [batch, heads, seq, head_dim]。在算子内部转换布局，并使用 torch.empty_like() 分配输出张量：

import torch
from nkilib.core.embeddings.rope_hf import rope_hf
from torch_neuronx import nki_op
from torch_neuronx.nki_hop import wrap_nki
from torch_neuronx.utils import get_logical_neuron_cores

LNC = (int(get_logical_neuron_cores()),)
wrapped_rope = wrap_nki(rope_hf)

@nki_op("my_kernels::rope_fwd")
def rope_fwd(query, key, rope_cache):
    q = query.permute(0, 2, 1, 3).contiguous()
    k = key.permute(0, 2, 1, 3).contiguous()
    cache = rope_cache.reshape(-1, rope_cache.shape[-1])
    q_out, k_out = wrapped_rope[LNC](
        q,
        k,
        torch.empty_like(q),
        torch.empty_like(k),
        rope_cache=cache,
    )
    return (
        q_out.permute(0, 2, 1, 3).contiguous(),
        k_out.permute(0, 2, 1, 3).contiguous(),
    )

rope_hf 在你以 backward=True 调用时会计算梯度。为该模式注册第二个算子。在前向传播过程中保存 rotary 缓存，以便反向传播可以复用它。由于缓存不可微分，返回 None 作为其梯度。

@nki_op("my_kernels::rope_bwd")
def rope_bwd(grad_query, grad_key, rope_cache):
    gq = grad_query.permute(0, 2, 1, 3).contiguous()
    gk = grad_key.permute(0, 2, 1, 3).contiguous()
    cache = rope_cache.reshape(-1, rope_cache.shape[-1])
    out_q, out_k = wrapped_rope[LNC](
        gq,
        gk,
        torch.empty_like(gq),
        torch.empty_like(gk),
        rope_cache=cache,
        backward=True,
    )
    return (
        out_q.permute(0, 2, 1, 3).contiguous(),
        out_k.permute(0, 2, 1, 3).contiguous(),
    )

def _setup_context(ctx, inputs, output):
    ctx.save_for_backward(inputs[2])

def _backward(ctx, grad_query, grad_key):
    (rope_cache,) = ctx.saved_tensors
    grad_q, grad_k = rope_bwd(grad_query, grad_key, rope_cache)
    return grad_q, grad_k, None

rope_fwd.register_autograd(_backward, setup_context=_setup_context)

由于 rope_hf 要求特定的序列长度支持，请在调用内核之前校验该约束。继承 CosSinRoPE 并在 apply_rotary_emb 中调用 rope_fwd：

from dataclasses import dataclass
from torchtitan.config import derive, override
from torchtitan.models.common.rope import CosSinRoPE

class NKIRoPE(CosSinRoPE):
    @dataclass(kw_only=True, slots=True)
    class Config(CosSinRoPE.Config):
        pass

    @staticmethod
    def apply_rotary_emb(query, key, rope_cache):
        stride = 128 * LNC[0]
        if query.shape[1] % stride:
            raise ValueError(
                f"rope_hf requires a sequence length that is a multiple of "
                f"128 * LNC={stride}, got {query.shape[1]}"
            )
        q_out, k_out = rope_fwd(query, key, rope_cache)
        return q_out.type_as(query), k_out.type_as(key)

将 NKIRoPE 注册为 CosSinRoPE 的替换项。derive() 会将共享的配置字段复制到替换配置中：

@override(
    "nki_rope",
    target=CosSinRoPE.Config,
    exact=True,
    description="nkilib rope_hf kernel in place of the PyTorch rotation",
)
def _nki_rope(cfg: CosSinRoPE.Config) -> NKIRoPE.Config:
    return derive(cfg, NKIRoPE.Config)

要加载自定义覆盖项，请将其模块名传递给 --override.imports。torchtitan 会在构建模型之前导入注册表条目：

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 native_qwen3_4b \
  --override.imports my_neuron_kernels

在训练开始之前，torchtitan 会记录覆盖项所替换的每个配置节点。确认输出中每个 Qwen3 层都包含一个 nki_rope 条目：

[Override] nki_rope: model_spec.model.layers.0.attention.rope CosSinRoPEConfig -> NKIRoPE.Config
[Override] nki_rope: model_spec.model.layers.1.attention.rope CosSinRoPEConfig -> NKIRoPE.Config

将序列长度设置为 128 * LNC 的倍数，在 LNC 为 2 时该值为 256。将 rotary 缓存重塑为 [seq, head_dim * 2]，前半部分为余弦值，后半部分为正弦值。内核会在训练开始前校验这些约束。

完整的 torchtitan 示例请参见 torchtitan_neuron/models/gpt_oss/nki_attention.py。它通过相同的算子和覆盖机制接入 nkilib flash-attention 内核。它还处理用于张量并行的 DTensor 输入。

在使用自定义内核训练之前，请将其输出与被替换的实现进行比较。在使用 float32 输入的数值测试中，rope_hf 与 torchtitan 的 PyTorch rotary embedding 在 query 和 key 上的最大绝对差为 0.0。

注册多个覆盖项时，请使用 fqns 将每个覆盖项限制在其目标模块实例上。例如，fqns=["layers.1[0-9].attention.rope"] 针对第 10 到 19 层的 rotary embedding。当两个覆盖项声明同一节点时，torchtitan 会停止；当某个覆盖项未匹配到任何节点时，它会发出警告。

查看训练输出

在启动时，确认 torchtitan 报告了 64 个 rank 以及预期的模型。第一步包含编译，耗时长于后续步骤。每个训练步骤都会报告损失、梯度范数、高带宽内存（HBM）使用量、吞吐量以及模型 FLOPs 利用率（MFU）。对于检查点运行，确认第一步损失接近 6.54。对于自定义内核运行，确认 nki_rope 覆盖项在每个 Qwen3 层出现一次。

跨实例扩展

相同的并行度设置适用于 Transformers 和原生模型配置。trn2.48xlarge 在 LNC=2 时暴露 64 个逻辑 NeuronCore。设置 data_parallel_shard_degree = -1 以在 torchrun 启动的所有 rank 之间分片模型参数、梯度和优化器状态。

设置 config.parallelism.tensor_parallel_degree 以在一组 rank 之间分片每层的权重。每个 rank 计算每个矩阵乘法的一部分。在默认启用序列并行的情况下，归一化层在沿序列维度分片的激活上运行。torchtitan 在多维 DeviceMesh 上将张量并行与全分片数据并行（FSDP）相结合。

在增大模型规模、序列长度或本地批大小时，观察每个 rank 的 HBM。这些值决定了每个 rank 所需的内存量。

清理

如果你想保留检查点，请在清理之前将其复制到持久存储中。如果你打算复用该实例，请停止它。当你不再需要该实例时，请终止它。

结论

AWS Trainium 作为受支持的 torchtitan 设备，可用于本指南中的全部三种配置。你可以通过 Hugging Face Transformers v5 训练 Qwen3，并直接加载或保存已发布的 safetensors。你也可以通过设置 config.model_spec，使用 torchtitan 原生的 Qwen3 实现进行训练。对于混合专家模型，请使用 torchtitan-neuron NKI 配方。对于自定义内核，可以使用 override registry 通过相同的启动流程挂载 NKI 算子。

有关设置指南，请参阅 AWS Neuron 文档。有关最新源代码和训练配方，请参阅 torchtitan 仓库。

关于作者