用 torchtitan 在 AWS Trainium 上训练 Qwen3

人工智能

使用 torchtitan 在 AWS Trainium 上训练 Qwen3

在 AWS Trainium 上使用 torchtitan，意味着你可以像往常一样使用它。TorchNeuron 将 Trainium 注册为原生 PyTorch 设备，使 torchtitan 能够通过标准 PyTorch 接口使用它。你可以使用这些配置从随机初始化开始预训练，或微调已发布的权重。你还可以针对自己的工作负载自定义模型实现和内核。在本文中，你将在单台 trn2.48xlarge.

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

前提条件

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

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

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

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

组件	版本
实例	trn2.48xlarge，16 个 Trainium2 加速器，128 个物理 NeuronCore
逻辑 NeuronCore	LNC=2 时为 64 个
torch、torch_neuronx	2.14.0, 2.14.3
transformers	5.16.0
torchtitan-neuron	release-2.33 (0.2.2+neuron.1)

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

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

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 将其与 token 嵌入绑定。这确认了这些文件仍是标准的 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 个 rank 上使用数据并行分片，张量并行度为 1。它不使用专家并行。 num_experts=64 设置模型中的专家数量，而不是专家并行度。

设置 CompileConfig(enable=False) 以 eager 模式运行此示例。TorchNeuron 在 eager 执行和 torch.compile.

中都支持已注册的 NKI 算子。在本次演练使用的 torchtitan 构建中， qwen3_nki_moe_config() 用 NKI 分组专家实现替换 MoE 块。它还会选择 SDPA 和无 gather 的旋转嵌入。当你启用 torch.compile时，将 "grad_clip" 添加到 CompileConfig.components 以编译梯度范数裁剪。

NKI 内核作为 PyTorch 算子进入训练。 wrap_nki 使 NKI 内核可从 PyTorch 调用，而 nki_op registers the operator. register_autograd attaches its backward pass using routing metadata saved during the forward pass.

The example reduces Qwen3-30B-A3B to 4 layers and 64 experts so it fits on one trn2.48xlarge. It routes each token to 4 experts with top_k=4.

Attach your own NKI kernel

In some cases, you might want to optimize an operation for the tensor shapes in your model. A specialized kernel can tailor its tiling and data movement to those shapes. It can also replace a performance-critical sequence of PyTorch operations with one optimized operator. You can write the kernel yourself or start with one from nkilib. This example uses rope_hf from nkilib.core.embeddings, which expects the Hugging Face tensor layout and supports its own backward mode. This example demonstrates kernel registration and autograd integration, not a performance improvement.

Create my_neuron_kernels.py in the repository root, where you launch training. torchtitan stores query and key as [batch, seq, heads, head_dim], whereas rope_hf expects [batch, heads, seq, head_dim]. Convert the layout inside the operator and allocate the output tensors with 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 computes gradients when you call it with backward=True. Register a second operator for that mode. Save the rotary cache during the forward pass so the backward pass can reuse it. Return None for the cache gradient because the cache is not differentiable.

@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)

Because rope_hf requires a supported sequence length, validate that constraint before calling the kernel. Subclass CosSinRoPE and call rope_fwd from apply_rotary_emb:

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)

Register NKIRoPE as the replacement for CosSinRoPE. derive() copies the shared configuration fields into the replacement config:

@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)

To load the custom override, pass its module name to --override.imports. torchtitan imports the registry entry before it builds the model:

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

Before training starts, torchtitan logs each config node that the override replaces. Confirm that the output contains one nki_rope entry for each Qwen3 layer:

[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

Set the sequence length to a multiple of 128 * LNC, which is 256 when LNC is 2. Reshape the rotary cache to [seq, head_dim * 2], with cosine values in the first half and sine values in the second. The kernel validates these constraints before training begins.

For a complete torchtitan example, see torchtitan_neuron/models/gpt_oss/nki_attention.py. It connects nkilib flash-attention kernels through the same operator and override mechanisms. It also handles DTensor inputs for tensor parallelism.

Before you train with a custom kernel, compare its outputs with the implementation it replaces. In the numerical test with float32 inputs, rope_hf and torchtitan’s PyTorch rotary embedding produced a maximum absolute difference of 0.0 for both query and key.

When you register multiple overrides, limit each one to the intended module instances with fqns. For example, fqns=["layers.1[0-9].attention.rope"] targets the rotary embedding in layers 10 through 19. torchtitan stops when two overrides claim the same node and warns when an override matches no nodes.

Review the training output

At startup, confirm that torchtitan reports 64 ranks and the expected model. The first step includes compilation and takes longer than later steps. Each training step reports the loss, gradient norm, high-bandwidth memory (HBM) usage, throughput, and model FLOPs utilization (MFU). For the checkpoint run, confirm that the first-step loss is near 6.54. For the custom kernel run, confirm that the nki_rope override appears once for each Qwen3 layer.

Scale across the instance

The same parallelism settings apply to the Transformers and native model configurations. A trn2.48xlarge exposes 64 logical NeuronCores at LNC=2. Set data_parallel_shard_degree = -1 to shard the model parameters, gradients, and optimizer state across every rank that torchrun starts.

Set config.parallelism.tensor_parallel_degree to shard each layer’s weights across a group of ranks. Each rank computes a slice of every matrix multiplication. With sequence parallelism enabled by default, normalization layers run on activations sharded along the sequence dimension. torchtitan combines tensor parallelism with Fully Sharded Data Parallel (FSDP) on a multidimensional DeviceMesh.

Watch per-rank HBM as you increase the model size, sequence length, or local batch size. These values determine how much memory each rank requires.

Clean up

If you want to keep your checkpoints, copy them to persistent storage before you clean up. If you plan to reuse the instance, stop it. When you no longer need the instance, terminate it.

Conclusion

AWS Trainium works as a supported torchtitan device across all three configurations in this walkthrough. You can train Qwen3 through Hugging Face Transformers v5 and load or save published safetensors directly. You can also train with torchtitan’s native Qwen3 implementation by setting config.model_spec. For mixture-of-experts models, use the torchtitan-neuron NKI recipes. For custom kernels, use the override registry to attach NKI operators through the same launch flow.

For setup guidance, see the AWS Neuron documentation. For current source and training recipes, see the torchtitan repository.

---

About the authors