AWS Machine Learning Blog

Train Qwen3 on AWS Trainium with torchtitan

6.9内容质量
Train Qwen3 on AWS Trainium with torchtitan

TL;DR · AI 摘要

Train Qwen3 on AWS Trainium with torchtitan Artificial Intelligence Train Qwen3 on AWS Trainium with torchtitan Using to...

核心要点

  • 主题聚焦：Train Qwen3 on AWS Trainium with torchtitan
  • 来源：AWS Machine Learning Blog，建议结合原文判断细节。
  • AI 分析暂不可用，本条为保底评分与摘要。
#AI#后端#安全#产品
打开原文

在 AWS Trainium 上使用 torchtitan 训练 Qwen3 | 人工智能

在 AWS Trainium 上使用 torchtitan 训练 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 内核接口（NKI）专家混合（MoE）内核，然后使用 torchtitan 的覆盖注册表附加你自己的内核。所有三种配置都使用相同的 torchrun 启动模式和配置结构。

先决条件

要跟随本指南，使用带有 AWS Neuron 的深度学习 AMI 或 Neuron PyTorch 训练容器的单个 trn2.48xlarge 实例。这两个环境都提供了在 Trainium 上训练所需的 Neuron 驱动程序、运行时和编译器。

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

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

本指南使用下表中的版本。要重现示例，请固定这些版本。Transformers v5 集成和 Trainium 特定的覆盖模块可能在 torchtitan 版本之间发生变化。

| 组件 | 版本 | |------|------| | 实例 | trn2.48xlarge，16 个 Trainium2 加速器，128 个物理 NeuronCores | | 逻辑 NeuronCores | 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 注册为神经元设备，torchtitan 在启动时会检测到它。在急切模式下，Neuron 会在操作执行时进行编译。使用 torch.compile 时，Neuron 会编译捕获的图。这种相同的启动模式适用于所有三种配置。

在本指南中使用的版本中，通过 --no-deps 参数安装 torchtitan-neuron 和 torchtitan 的其余依赖项，以保留 torch 的 Neuron 构建：

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

在开始训练之前，从 Hugging Face Hub 获取 Qwen3 模型资源。将其存储在可以传递给训练配置的本地目录中：

code
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，这是你运行训练命令的位置。添加导入、模型资源路径和基础配置：

code
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
code
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。状态字典适配器会调整主模型的前缀并处理绑定的嵌入层。它会将最终检查点保存为 Hugging Face safetensors 格式。

code
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 # 需要绝对路径
    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。向 neuron_qwen3.py 添加原生 Qwen3 的导入和配置函数：

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

code
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
code
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 内核语言）编写的内核运行专家混合层。向 neuron_qwen3.py 添加所需的导入项和 Qwen3 MoE 配置函数：

code
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
code
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) 会以急切模式运行此示例。TorchNeuron 支持在急切执行和 torch.compile 中使用已注册的 NKI 操作符。

在本教程使用的 torchtitan 构建中，qwen3_nki_moe_config() 会将 MoE 模块替换为 NKI 分组专家实现，同时选择 SDPA 和无收集旋转嵌入。启用 torch.compile 时，向 CompileConfig.components 添加 "grad_clip" 以编译梯度范数裁剪。

NKI 内核以 PyTorch 操作符形式参与训练。wrap_nki 使 NKI 内核可以从 PyTorch 调用，nki_op 注册操作符。register_autograd 通过前向传播时保存的路由元数据附加其反向传播。

该示例将 Qwen3-30B-A3B 缩减为 4 层和 64 个专家，以便适配单个 trn2.48xlarge 实例。它通过 top_k=4 将每个 token 路由到 4 个专家。

添加你自己的 NKI 内核

在某些情况下，你可能需要针对模型中张量的形状优化操作。专用内核可以针对这些形状定制分块策略和数据移动方式。它还可以将性能关键的PyTorch操作序列替换为经过优化的操作符。你可以自己编写内核，也可以从nkilib开始。这个示例使用了nkilib.core.embeddings中的rope_hf，它期望Hugging Face张量布局并支持自己的反向模式。这个示例演示了内核注册和autograd集成，而不是性能改进。

在仓库根目录创建my_neuron_kernels.py，这是你启动训练的地方。torchtitan将查询和键存储为[batch, seq, heads, head_dim]，而rope_hf期望[batch, heads, seq, head_dim]。在操作符内部转换布局，并使用torch.empty_like()分配输出张量：

code
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计算梯度。为该模式注册第二个操作符。在正向传递期间保存旋转缓存，以便反向传递可以重复使用它。由于缓存不可微分，返回None作为缓存梯度。

code
@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(),
    )
code
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：

code
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
code
@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() 会将共享的配置字段复制到替代配置中：

code
@override(
    "nki_rope",
    target=CosSinRoPE.Config,
    exact=True,
    description="nkilib rope_hf 内核替代 PyTorch 旋转",
)
def _nki_rope(cfg: CosSinRoPE.Config) -> NKIRoPE.Config:
    return derive(cfg, NKIRoPE.Config)

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

code
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 条目：

code
[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。将旋转缓存重塑为 [seq, head_dim * 2]，前半部分为余弦值，后半部分为正弦值。内核会在训练开始前验证这些约束条件。

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

在使用自定义内核训练前，将其输出与被替换的实现进行比较。在使用 float32 输入的数值测试中，rope_hf 和 torchtitan 的 PyTorch 旋转嵌入在查询和键上产生的最大绝对差异均为 0.0。

当注册多个覆盖时，应通过 fqns 限制每个覆盖到目标模块实例。例如，fqns=["layers.1[0-9].attention.rope"] 会定位第 10 到 19 层的旋转嵌入。当两个覆盖声称同一个节点时，torchtitan 会停止执行，并在覆盖未匹配任何节点时发出警告。

查看训练输出

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

在实例内扩展

相同的并行设置适用于 Transformers 和原生模型配置。trn2.48xlarge 在 LNC=2 时暴露 64 个逻辑 NeuronCores。将 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 配方。对于自定义内核，请使用覆盖注册表通过相同的启动流程附加 NKI 操作符。

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

作者信息

'“