AWS Trainium 三种方式训练 Qwen3

# AWS Trainium 三种方式训练 Qwen3

摘要： 在 AWS Trainium 上用 torchtitan 以三种方式训练 Qwen3：Hugging Face Transformers v5 建模后端、torchtitan 原生 Qwen3 实现，以及 NKI 混合专家内核。全部运行于单个 trn2.48xlarge（64 个逻辑 NeuronCore）。本文给出可复现的配置示例、启动命令和验证方法，适合已有 AWS Neuron 环境的开发者。

分类：算力

品牌量身

GrokCode 专注 API 中转、模型天梯与本地部署实验室。在算力领域，我们帮助开发者快速验证模型工程边界。AWS Trainium + torchtitan 的组合正是本地部署实验室的经典案例——它用硬件原生内核优化训练路径，而非依赖云厂商全托管服务。

适用场景与决策框架

如果你已在 AWS 上部署了 Neuron 环境，并想用 torchtitan 针对 Qwen3 做预训练或微调，单个 trn2.48xlarge 就能跑通三种方式。这三种方式的差异主要在模型构建、内核优化和兼容性上。

• Hugging Face Transformers v5：用官方库构建模型，最接近标准 HF 训练流程，适合已有 HF 权重或快速切换模型。
• torchtitan 原生 Qwen3：切换模型规格后直接用 torchtitan 内置实现，关注 sharding 和内存映射。
• NKI 混合专家：添加 Neuron 内核语言（NKI）优化的 MoE 层，可针对特定张量形状微调，但需手动注册内核。

决策 Checklist（仅供参考，非绝对指标）：

• 是否已有 HF 模型权重？选 v5。
• 是否需要精确的 sharding 控制？选原生 Qwen3。
• 是否希望用 Neuron 内核加速 MoE 特定算子？选 NKI 方式。
• 硬件限制：单个 trn2.48xlarge 64 逻辑 NeuronCore，需用数据并行 sharding 或减少层数。

不同配置的 Token 计算量与成本在 AWS 官网定价页面可实时查阅。

AWS Trainium 基础介绍

AWS Trainium 是专为 AI 训练设计的芯片，集成 Neuron 驱动和编译器。TorchNeuron 将 Trainium 注册为 PyTorch 原生设备，让 torchtitan 直接通过标准接口使用。它支持 eager 模式（即时编译）和 torch.compile（图编译），并提供与 Neuron 内核的直接绑定。

Key Points：

• 逻辑 NeuronCore：配置 LNC=2 可获得 64 个核心。
• 实例规格：trn2.48xlarge，128 个物理 NeuronCore。
• 容器/AMI：推荐 AWS Deep Learning AMI for Neuron 或 Neuron PyTorch 训练容器。

安装 torchtitan-neuron 后，执行 python3 -c "import torchtitan, torchtitan_neuron; print(torchtitan.__version__, torchtitan_neuron.__file__)" 检查版本。推荐 pinned 到 release-2.33（0.2.2+neuron.1），并用 --no-deps 保留 Neuron Torch 构建。

三种训练方式对比

下面是三种配置的横向对比表格（移动端横向滚动友好）：

所有方式使用同一 torchrun 启动模式，区别仅在 Python 配置函数和 model_spec 上。

1. Hugging Face Transformers v5 方式

此方式用 AutoModelForCausalLM 通过 Hub 加载模型。创建 neuron_qwen3.py：

```python 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 ```

启动命令： ``bash 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 ``

微调示例：启用 checkpoint.enable 并设置 initial_load_path 指向 HF 资产路径，即可从已发布权重开始训练。验证 round-trip 加载时，检查 missing keys（通常仅 lm_head.weight 为 tied embedding）。

2. torchtitan 原生 Qwen3 方式

切换到内置实现，核心是 config.model_spec。示例（基于 4B 规格）：

```python 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 ```

启动命令与上文完全相同，仅替换配置函数。

3. NKI 混合专家方式

添加 NKI 内核后，构建 MoE 模型。示例（30B-A3B 简化版）：

```python 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 ```

启动命令相同。启用 torch.compile 时可添加 "grad_clip" 到 CompileConfig.components。

自定义内核示例：可通过 nki_op 注册自定义 NKI 算子（如 ROPE 布局转换），并用 wrap_nki 绑定到 PyTorch 接口。

实用检查清单

• [ ] 已安装 torchtitan-neuron 并 pin 版本
• [ ] HF 模型资产已下载到本地路径
• [ ] 配置函数正确设置 hf_assets_path、dataloader.dataset 和并行度
• [ ] 使用 --local-ranks-filter 0 过滤（调试用）
• [ ] 验证 loss 曲线平滑（示例中 v5 方式从随机初始化到 HF 权重 loss 下降明显）

延伸阅读

• GrokCode API 中转页面：了解如何在模型训练后接入 Grok API 中转服务
• GrokCode 模型天梯页面：查看最新开源模型训练性能对比
• GrokCode 本地部署实验室页面：探索更多 Neuron 与 torchtitan 结合的实战案例
• GrokCode 工具箱：查看 vLLM 等本地部署工具集成说明
• GrokCode 官方 API 页面：获取 API 中转服务定价与接入指南

风险与边界

由于涉及 GPU/Neuron 硬件资源分配和模型训练过程，需注意账单变化、实验失败重试和长时间占用资源。实际效果取决于硬件规格与配置细节，以 AWS 官网最新硬件与定价数据为准。非法律意见，建议咨询专业服务。

English summary Train Qwen3 on AWS Trainium using torchtitan in three ways on a single trn2.48xlarge: Hugging Face Transformers v5 modeling backend, torchtitan’s native Qwen3 implementation, and NKI mixture-of-experts kernels. All use the same torchrun launch pattern with TorchNeuron registering Trainium as a native PyTorch device. The first method loads Hub models via AutoModelForCausalLM and reads/writes safetensors. The second switches config.model_spec to the built-in registry. The third registers NKI operators for MoE layers and attaches custom kernels with wrap_nki. Prerequisites include a Neuron-enabled trn2.48xlarge, torchtitan-neuron (pinned to release-2.33), and Qwen3 assets from Hugging Face. Each configuration is a Python config function with data parallel sharding across 64 logical NeuronCores. Examples use 10 training steps, C4 dataset, and local batch size 1. Checkpoints load/save as standard HF format when enabled. This setup demonstrates native PyTorch compatibility on Trainium for pre-training or fine-tuning Qwen3.