Qwen3次使用火炬的AWS列车

人工智能

Qwen3次使用火炬的AWS列车

在AWS Trainium上使用火炬剂意味着使用它的方式总是如此。 火炬内龙将Trainium登记为原生的PyTorch设备，它允许火炬通过标准PyTorch接口使用。 您可以使用这些配置从随机初始化或微调公布的权重中预排。 也可以自定义模式执行和内核来完成你的工作量。 在这张帖子中，你用单Trn2.48x大号的三路训练Qwen3.

第一种方法使用变形金刚v5作为火炬的模型后端。 它通过AutoModelForCausalLM运行Hub模型，并直接读写已出版的保温器。 第二个方法通过改变配置。 model spec来切换火炬泰坦的原生Qwen3. 第三个新增 Neuron Kernel 接口(NKI) 混合专家(MoE) 内核从火炬-中微子配方，然后使用火炬Titan的超载注册来附加自己的内核。 所有三种配置都使用相同的火炬run 发布图案和配置结构。

先决条件

要跟踪这个走过，使用一个单一的tn2.48x大号，其中AWS深层学习AMI用于Neuron或Neuron PyTorch训练容器。 这两个环境都提供了Neuron驱动程序，运行时间，以及列车训练所需的编译器。

在安装火炬- 中微子之前，请检查您的环境是否已经包含它。 以下命令打印已安装的版本和软件包路径：

此走行部使用下表中的版本。 为了复制这些例子，将这些版本标注起来。 变形器 v5 集成和 Trainium 特异性覆盖模块可以在火炬发布之间改变。

火炬中微子安装后，火炬神经将Trainium登记为神经元装置，火炬在启动时检测到。 在急切模式下，Neuron会在其执行时编译操作。 用火炬。 compile, Neuron 编译捕获的图表。 这种相同的发布模式在所有三种配置中都有效。

对于此行走中使用的版本，安装火炬- 中微子和火炬的剩余依存点，并使用- no- deps来保存火炬的Neuron 构造：

在开始训练前，从Hugging Face枢纽获得Qwen3型号文物。 保存在您可以传递到培训配置的本地目录中：

每个训练配置都是Python配置功能。 三个实例从一个基础配置开始，您将更新每个步骤。

使用 Hugging Face 变形器 v5 的列车

在第一个配置中，变形金刚v5通过AutoModelForCausalLM构建Qwen3. 枢纽模型定义提供了用于前向和后向通行证的执行。 您通过在配置中设置其Hub ID来选择模型。

用变压器开始建模后端 full(). 它留下映射到 Hugging Face 配置密钥的字段未设置，因此 AutoConfig. from pretrained () 从Hub 中加载模型架构。

在火炬-中微子寄存器根中创建神经元 qwen3.py，在那里运行训练命令。 添加导入、 模型化工艺路径和基础配置：

这个命令发布 64训练过程，每个逻辑NeuronCore一个。 --模块神经元 qwen3导入本地配置模块。 --config v5 qwen3 4b 选择配置函数。 数据并行设定使模型在所有64个级别上都变得僵硬。

装入并保存 Hugging Face 检查站

从已公布的重量到微调，允许检查点加载和点首加载 path to ASSETS. 状态词典适配器调整主型号。 前缀和手柄绑定嵌入。 它将最终检查站保存为Hugging Face安全防护装置。

使用第一个训练步骤损失作为快速检查，火炬侠将公布的重量加载。 对于Qwen/Qwen3-4B-Instruct-2507，经核实的运行产生第1步从公布的权重中损失6.54，从随机初始化中损失12.43.

训练后，将输出的检查站装入AutoModelForCausalLM并装入 state dict. 成功的往返报告除了lm head之外没有丢失或意外的密钥。 重量，Qwen3-4B与符号嵌入关系。 这证实档案仍然是标准Hugging Face检查站。

用火炬泰坦的本地型号进行训练

在第二个配置中，火炬泰坦构建其本土的Qwen3模型。 将本地的 Qwen3 导入和配置函数添加到神经元 qwen3. py:

此示例从qwen3 32b()开始，因为测试过的火炬构造不能提供通用的qwen3 4b()配置构造器。 下一行用4B规格取代32B型号规格。 其余的田地复制了测试的运行。

config.model spec选择模型类，并定义火炬泰坦如何将模型硬化并绘制其检查站密钥。 这里，model registry ("4B", attn backend="sdpa") 选择火炬泰坦的本地 Qwen3 执行，并使用缩放的点产品关注度( SDPA)。

为混合专家添加 NKI 内核

在第三个配置中，火炬泰坦运行着混合专家层，内核用NKI,Neuron内核语言书写。 将所需的导入和 Qwen3 MOE 配置函数添加到神经元 qwen3. py 中：

这种配置在64个排位上使用数据平行硬化，拉伸平行度为1. 它没有使用专家平行主义。 num experts=64设定模型中的专家人数，而非专家平行学位。

设置 CompileConfig( enable=False) 在急切模式下运行此示例。 火炬内龙支持注册的NKI运营商同时急切执行和火炬.compile.

在用于此行走的火炬建筑中，qwen3 nki moe config () 将MOE块替换为NKI分组专家执行。 它还选择了SDPA和无集合旋转嵌入。 当启用 Film.compile 时，在 compileConfig 中加入“ grad clip ”，用于编译渐变-规范剪辑。

NKI内核作为PyTorch操作员进入训练。 graw nki 使 NKI 内核可以从 PyTorch 调用，而 nki op 注册操作员。 使用在前传中保存的路由元数据附加其后传。

该示例将Qwen3-30B-A3B降低为4层，64位专家因此适合1tn2.48x大号。 它给4位专家送去每个符牌，上面为 k=4。

附加您自己的 NKI 内核

在某些情况下，您可能想要优化一个操作，用于您的模型中的收音机形状。 一个专门的内核可以使其板块和数据运动适应这些形状。 它也可以用一个优化的操作员取代PyTorch操作的性能关键序列。 你可以自己写内核，也可以从 nkilib 开始写。 这个例子使用 nkilib.core.embeddings 的绳索 hf，它期望Hugging Face 的拉诺布局，并支持自己的后向模式。 这个例子表明内核登记和自动分级整合，而不是业绩改进。

在寄存器根创建我的 neuron 内核.py，您的发布训练所在。 火炬存储查询和密钥为[批，下，头，头 dim]，而绳 hf预期[批，头，下，头 dim]. 转换运算器内部的布局，用火炬分配输出亮度。 freety like():

绳 hf 计算梯度时使用后向=True表示。 注册该模式的第二个操作员。 前传时保存旋转缓存，使后传能够再利用。 没有返回缓存梯度，因为缓存没有不同。

因为绳索 hf需要支持的序列长度，在调用内核前验证这个约束。 子类 CosSinRoPE 从应用中调用绳索 fwd:

将 NKIROPE 注册为 CosSinRoPE 的替换。 生成( ) 复制共享配置字段到替换配置：

要装入自定义覆盖，请将模块名称传至 -- 覆盖。 进口。 火炬公司在建立该模型之前进口登记条目：

在训练开始前，火炬记录每个配置节点，覆盖将替换。 确认输出包含每个 Qwen3 层的一个 nki rope 条目：

将序列长度设定为128 * LNC的倍数，当LNC为2. 将旋转缓存重排为 [seq, head dim * 2]，上半段为余弦值，第二段为正弦值。 内核在训练开始前验证这些限制。

完整的火炬图例，见火炬图/模型/gpt oss/nki attention.py. 它通过相同的运算符连接nkilib闪存内核并覆盖机制。 它还处理用于拉伸平行主义的Dentensor输入。

在使用自定义内核进行训练之前，先将其输出与它所取代的执行进行比较。 在带有浮点32输入的数值测试中，绳子 hf和火炬泰坦的PyTorch旋转嵌入产生查询和键的最大绝对差为0.0.

当您注册多个覆盖时，请将每个覆盖限制为带有 fqns 的预定模块实例。 例如，fqns=['layers.1[0-9]. attention.rope]针对10至19层的旋转嵌入。 当两个覆盖点声称相同的节点时火炬停止，当一个覆盖点与节点匹配时火炬停止警告。

审查培训产出

启动时确认火炬报64级和预期型号。 第一步包括汇编和采取比以后更长的步骤。 每个培训步骤都报告损失、梯度规范、高频段内存(HBM)使用情况、吞吐量和FLOPs模型利用情况。 对于检查站运行，确认第一步损失接近6.54. 对于自定义内核运行，确认每个 Qwen3 层的 nki rope 覆盖一次。

跨实例缩放

相同的平行设置适用于变形器和原生模型配置。 rn2.48x大号在LNC=2时暴露出64个逻辑的NeuronCores. 设置数据并行硬度= -1 以压缩模式参数，梯度，并优化火炬启动的每个级别上的状态。

设置配置。 paralism. tensor paral 度，使每层的重量在一组排位上被压缩。 每个阶次计算每个矩阵乘法的一块。 由于默认允许序列平行性，激活时的常态层沿序列维度发生硬化。 火炬式在多维设备Mesh上将拉诺并行主义与Fully Sharded Data并行(FSDP)结合。

当您增加模型大小、序列长度或本地批量大小时，注意每个排名的HBM。 这些值决定了每个级别需要多少内存。

清理

如果你想保持你的检查站，复制到持续储存在你清理之前。 如果你打算重用这个实例，就停止。 当您不再需要实例时，终止它。

结论

AWS Trainium作为支持的火炬装置，在这个穿行中贯穿所有三种配置。 可以通过Hugging Face变形器v5对Qwen3进行训练，直接加载或保存已公布的安全器。 您也可以通过设置配置( config. model spec)来使用火炬的本地 Qwen3 执行程序进行训练。 对于混合专家模型，使用火炬-中微子NKI配方。 对于自定义内核，通过相同的发布流，使用覆盖注册来附加NKI运算符。

关于设置指导，见AWS Neuron文档。 关于当前来源和培训食谱，请参见火炬寄存器。

本文由 AI图鉴自动抓取并翻译，版权归原作者所有；如需查看原文请点下方链接。

Artificial Intelligence

Train Qwen3 on AWS Trainium with torchtitan

Using torchtitan on AWS Trainium means using it the way you always have. TorchNeuron registers Trainium as a native PyTorch device, which lets torchtitan use it through standard PyTorch interfaces. You can use these configurations to pre-train from random initialization or fine-tune published weights. You can also customize the model implementation and kernels for your workload. In this post, you train Qwen3 three ways on a single trn2.48xlarge.

The first method uses Transformers v5 as a modeling backend in torchtitan. It runs Hub models through AutoModelForCausalLM and reads and writes published safetensors directly. The second method switches to torchtitan’s native Qwen3 by changing config.model_spec. The third adds Neuron Kernel Interface (NKI) mixture-of-experts (MoE) kernels from the torchtitan-neuron recipes and then uses torchtitan’s override registry to attach your own kernel. All three configurations use the same torchrun launch pattern and config structure.

Prerequisites

To follow this walkthrough, use a single trn2.48xlarge with either the AWS Deep Learning AMI for Neuron or a Neuron PyTorch training container. Both environments provide the Neuron driver, runtime, and compiler required to train on Trainium.

Before you install torchtitan-neuron, check whether your environment already includes it. The following command prints the installed version and package path:

This walkthrough uses the versions in the following table. To reproduce the examples, pin these versions. The Transformers v5 integration and Trainium-specific override modules can change between torchtitan releases.

When torch_neuronx is installed, TorchNeuron registers Trainium as the neuron device and torchtitan detects it at startup. In eager mode, Neuron compiles operations as they execute. With torch.compile, Neuron compiles the captured graph. This same launch pattern works across all three configurations.

For the versions used in this walkthrough, install torchtitan-neuron and torchtitan’s remaining dependencies with --no-deps to preserve the Neuron build of torch:

Before you start training, get the Qwen3 model artifacts from Hugging Face Hub. Store them in a local directory that you can pass to the training config:

Each training configuration is a Python config function. The three examples start from a base config that you update for each step.

Train with Hugging Face Transformers v5

In the first configuration, Transformers v5 builds Qwen3 through AutoModelForCausalLM. The Hub model definition supplies the implementation used for the forward and backward passes. You select the model by setting its Hub ID in the config.

Start with transformers_modeling_backend_full(). It leaves the fields that map to Hugging Face configuration keys unset, so AutoConfig.from_pretrained() loads the model architecture from the Hub.

Create neuron_qwen3.py in the torchtitan-neuron repository root, where you run the training command. Add the imports, model artifact path, and base config:

This command launches 64 training processes, one for each logical NeuronCore. --module neuron_qwen3 imports the local config module. --config v5_qwen3_4b selects the configuration function. The data parallel setting shards the model across all 64 ranks.

Load and save Hugging Face checkpoints

To fine-tune from published weights, enable checkpoint loading and point initial_load_path to ASSETS. The state-dictionary adapter adjusts the leading model. prefix and handles tied embeddings. It saves the final checkpoint as Hugging Face safetensors.

Use the first training-step loss as a quick check that torchtitan loaded the published weights. For Qwen/Qwen3-4B-Instruct-2507, the verified run produced a step 1 loss of 6.54 from published weights and 12.43 from random initialization.

After training, load the exported checkpoint with AutoModelForCausalLM and load_state_dict. A successful round trip reports no missing or unexpected keys other than lm_head.weight, which Qwen3-4B ties to the token embeddings. This confirms that the files remain a standard Hugging Face checkpoint.

Train with torchtitan’s native models

In the second configuration, torchtitan builds its native Qwen3Model. Add the native Qwen3 imports and config function to neuron_qwen3.py:

This example starts from qwen3_32b() because the tested torchtitan build does not provide a generic qwen3_4b() config builder. The next line replaces the 32B model specification with the 4B specification. The remaining fields reproduce the tested run.

config.model_spec selects the model class and defines how torchtitan shards the model and maps its checkpoint keys. Here, model_registry("4B", attn_backend="sdpa") selects torchtitan’s native Qwen3 implementation with scaled dot product attention (SDPA).

Add NKI kernels for mixture of experts

In the third configuration, torchtitan runs mixture-of-experts layers with kernels written in NKI, the Neuron kernel language. Add the required imports and Qwen3 MoE config function to neuron_qwen3.py:

This configuration uses data parallel sharding across 64 ranks with a tensor parallel degree of 1. It does not use expert parallelism. num_experts=64 sets the number of experts in the model, not the expert parallel degree.

Setting CompileConfig(enable=False) runs this example in eager mode. TorchNeuron supports registered NKI operators in both eager execution and torch.compile.

In the torchtitan build used for this walkthrough, qwen3_nki_moe_config() replaces the MoE block with the NKI grouped-experts implementation. It also selects SDPA and the gather-free rotary embedding. When you enable torch.compile, add "grad_clip" to CompileConfig.components to compile gradient-norm clipping.

NKI kernels enter training as PyTorch operators. wrap_nki makes an NKI kernel callable from PyTorch, and 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():

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.

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:

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

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

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:

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.

Auto-fetched and translated by AI Atlas. All rights belong to the original author.