torchtitan を使って AWS Trainium 上で Qwen3 をトレーニングする

torchtitanを以下で使用する AWS Trainium これまで通りの使い方をするだけで済む、ということです。TorchNeuronはTrainiumをネイティブのPyTorchデバイスとして登録するため、torchtitanは標準的なPyTorchインターフェースを通じてそれを利用できます。これらの構成を使って、ランダム初期化からの事前学習や、公開済み重みのファインチューニングを行うことができます。また、自分のワークロードに合わせてモデル実装やカーネルをカスタマイズすることも可能です。この投稿では、Qwen3を1つの trn2.48xlarge.

最初の方法は、torchtitanでTransformers v5をモデリングバックエンドとして使用するものです。Hubのモデルを次を通じて実行します AutoModelForCausalLM safetensors形式で公開されたファイルを直接読み書きします。2つ目の方法では、変更することでtorchtitanのネイティブなQwen3に切り替えます。 config.model_spec。3番目では、torchtitan-neuronレシピからNeuron Kernel Interface（NKI）のmixture-of-experts（MoE）カーネルを追加し、その後torchtitanのオーバーライドレジストリを使って独自のカーネルをアタッチします。3つの構成すべてで同じ torchrun 起動パターンと設定構造。

前提条件

このウォークスルーに従うには、単一の trn2.48xlarge 「AWS Deep Learning AMI for Neuron」または「Neuron PyTorch training container」のいずれかを使用します。どちらの環境でも、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.compileNeuron はキャプチャされたグラフをコンパイルします。この同じ起動パターンは、3 つの構成すべてで機能します。

このウォークスルーで使用するバージョンでは、torchtitan-neuron と torchtitan のその他の依存関係を次のコマンドでインストールします --no-deps 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のconfig関数です。3つの例は、基本となるconfigを出発点として、各ステップごとに更新していく形になっています。

Hugging Face Transformers v5 でトレーニングする

最初の構成では、Transformers v5 が Qwen3 を以下のように構築します。 AutoModelForCausalLM. Hubモデル定義がフォワードパスおよびバックワードパスの実装を提供します。設定でHub IDを指定することでモデルを選択します。

まずはここから始めます transformers_modeling_backend_full()「. Hugging Face の設定キーに対応するフィールドは未設定のままになるため、」 AutoConfig.from_pretrained() Hubからモデルアーキテクチャを読み込みます。

作成 neuron_qwen3.py torchtitan-neuronリポジトリのルートで、トレーニングコマンドを実行する場所です。インポート、モデルアーティファクトのパス、およびベース設定を追加します:

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

このコマンドは、各論理 NeuronCore に1つずつ、合計 64 個のトレーニングプロセスを起動します。 --module neuron_qwen3 ローカルのconfigモジュールをインポートします。 --config v5_qwen3_4b 構成関数を選択します。データ並列設定では、モデルが全64ランクにまたいでシャーディングされます。

Hugging Faceのチェックポイントの読み込みと保存

公開済みの重みからファインチューニングするには、チェックポイントの読み込みを有効にして、 initial_load_path に ASSETS。state-dictionaryアダプタは先頭の 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

最初のtraining-stepのlossを使って、torchtitanが公開済みの重みを正しくロードできたかの簡易確認を行えます。For Qwen/Qwen3-4B-Instruct-2507検証済みの実行では、公開済みの重みからはステップ 1 の損失が 6.54、ランダム初期化からは 12.43 が得られました。

学習後、エクスポートされたチェックポイントを以下で読み込みます。 AutoModelForCausalLM および load_state_dict。ラウンドトリップが成功すると、以下のキーを除いて欠落や予期しないキーは報告されません。 lm_head.weight,これはQwen3-4Bがトークン埋め込みと結び付けているものである。このことから、これらのファイルが標準的なHugging Faceチェックポイントのままであることが確認される。

Figure 1: 公開済みチェックポイントは損失6.54から開始し、エクスポートされた重みは欠落キーや予期しないキーなしで読み込まれます

torchtitanのネイティブモデルでトレーニングする

2番目の構成では、torchtitanがネイティブの Qwen3Model。ネイティブの Qwen3 のインポートと config 関数を次に追加します neuron_qwen3.py:

この例は次から始まります。 qwen3_32b() テストされたtorchtitanビルドが汎用的なものを提供していないため qwen3_4b() config builder です。次の行で、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)付きで選択します。

专家混合（Mixture of Experts）用の NKI カーネルを追加

3番目の構成では、torchtitanが、Neuronカーネル言語であるNKIで書かれたカーネルを使ってmixture-of-experts層を実行します。必要なimportと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) この例をeagerモードで実行します。TorchNeuronは、登録されたNKIオペレーターをeager実行と torch.compile.

このウォークスルーで使用したtorchtitanビルドでは、 qwen3_nki_moe_config() MoEブロックをNKIのgrouped-experts実装に置き換えます。また、SDPAとgather-freeのrotary embeddingを選択します。有効にすると torch.compile、追加 "grad_clip" に CompileConfig.components 勾配ノルムクリッピングを組み込むため。

NKIカーネルがPyTorchオペレーターとして学習に組み込まれます。 wrap_nki PyTorchからNKIカーネルを呼び出せるようにし、 nki_op オペレーターを登録します。 register_autograd フォワードパス中に保存されたルーティングメタデータを使用して、そのバックワードパスを接続します。

この例では、Qwen3-30B-A3Bを1つのデバイスに収まるように4層・64エキスパートまで削減しています。 trn2.48xlarge。各トークンを4人のエキスパートにルーティングし、 top_k=4.

独自のNKIカーネルを取り付ける

場合によっては、モデル内のテンソル形状に合わせて演算を最適化したいことがあります。特殊化されたカーネルは、それらの形状に応じてタイリングやデータ移動を調整できます。また、パフォーマンス上重要な一連のPyTorch演算を、1つの最適化されたオペレーターに置き換えることもできます。カーネルは自身で作成してもよいですし、次から提供されているものをベースにしても構いません。 nkilib。この例では以下を使用します rope_hf から nkilib.core.embeddingsHugging Faceのテンソルレイアウトを前提とし、独自のbackwardモードを備えています。この例は、カーネル登録とautogradの統合を実演するものであり、パフォーマンス向上を目的としたものではありません。

作成 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。そのモードに対して2番目のオペレーターを登録してください。順伝播の際にロータリーキャッシュを保存し、逆伝播で再利用できるようにしてください。返します 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 そして呼び出し rope_fwd から 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)

登録 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はオーバーライドが置き換える各設定ノードをログに出力します。出力に1つの nki_rope 各Qwen3レイヤーのエントリー:

[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 * LNCLNC が 2 のときは 256 になります。ロータリーキャッシュを以下のように再形成します。 [seq, head_dim * 2]、前半にコサイン値、後半にサイン値を配置します。カーネルは訓練開始前にこれらの制約を検証します。

torchitannの完全な例については、次を参照してください: torchtitan torchtitan_neuron/models/gpt_oss/nki_attention.py「. It connects」 nkilib flash-attention カーネルに対しても、同じ operator および override のメカニズムを通じて処理します。さらに対応するのは DTensor テンソル並列処理の入力。

カスタムカーネルで学習する前に、その出力を置き換え対象の実装と比較してください。float32入力を用いた数値テストでは、 rope_hf また、torchtitan の PyTorch rotary embedding では、query と key の両方で最大絶対差が 0.0 でした。

複数のオーバーライドを登録する場合は、それぞれを対象のモジュールインスタンスに限定してください。その際に使用するのが次のものです: fqns。例えば、 fqns=["layers.1[0-9].attention.rope"] レイヤー10から19までの回転埋め込み（rotary embedding）を対象とします。2つのオーバーライドが同じノードを要求するとtorchtitanは停止し、オーバーライドがどのノードにも一致しない場合は警告を出します。

トレーニング出力を確認する

起動時に、torchtitanが64ランクと想定されるモデルを報告することを確認します。最初のステップにはコンパイルが含まれるため、後のステップよりも時間がかかります。各トレーニングステップでは、損失、勾配ノルム、高帯域幅メモリー(HBM)使用量、スループット、モデルFLOPs利用率(MFU)が報告されます。チェックポイント実行では、最初のステップの損失が6.54に近いことを確認します。カスタムカーネル実行では、 nki_rope オーバーライドが各Qwen3レイヤーについて1回出現することを確認します。

インスタンス間のスケーリング

同じ並列化設定が、Transformers構成とネイティブモデル構成の両方に適用されます。 trn2.48xlarge はLNC=2で64個の論理NeuronCoreを公開します。 data_parallel_shard_degree = -1 を設定して、 torchrun が起動するすべてのランクにわたって、モデルのパラメーター、勾配、オプティマイザーの状態をシャーディングします。

を設定して、各レイヤーの重みをランクのグループ全体にシャーディングします。各ランクは、すべての行列乗算のスライスを計算します。シーケンス並列がデフォルトで有効になっているため、正規化レイヤーはシーケンス次元に沿ってシャーディングされたアクティベーションに対して実行されます。torchtitanは、テンソル並列を多次元の config.parallelism.tensor_parallel_degree 各ランクのグループに対して各層の重みをシャーディング(分割保持)します。各ランクはすべての行列乗算のスライスを計算します。シーケンス並列がデフォルトで有効になっている場合、正規化層はシーケンス次元に沿ってシャーディングされた活性化値に対して実行されます。torchtitan は、多次元の構成においてテンソル並列と Fully Sharded Data Parallel (FSDP) を組み合わせます。 DeviceMesh.

モデルサイズ、シーケンス長、またはローカルバッチサイズを増やす際には、ランクごとのHBMを監視してください。これらの値が、各ランクに必要なメモリー量を決定します。

クリーンアップ

チェックポイントを保持したい場合は、クリーンアップの前に永続ストレージにコピーしてください。インスタンスを再利用する予定がある場合は、 停止します。インスタンスが不要になったら、 終了します.

結論

AWS Trainium は、このウォークスルーの3つの構成すべてにおいて、サポート対象のtorchtitanデバイスとして機能します。Hugging Face Transformers v5を通じてQwen3をトレーニングし、公開されているsafetensorsを直接読み込んだり保存したりできます。また、 config.model_specを設定することで、torchtitanのネイティブQwen3実装でトレーニングすることもできます。mixture-of-expertsモデルには、torchtitan-neuron NKIレシピを使用します。カスタムカーネルには、同じ起動フローを通じてNKIオペレーターをアタッチするために、オーバーライドレジストリを使用します。

セットアップのガイダンスについては、 AWS Neuronドキュメントを参照してください。最新のソースとトレーニングレシピについては、torchtitanリポジトリを参照してください。

著者について