Qwen3をTrainiumで学習するには？torchtitanの実装選択と検証の進め方

AWSが、torchtitanでQwen3を学習する3つの構成を紹介しました。Transformers v5、ネイティブ実装、NKIカーネルの役割を整理し、並列化と性能評価へ進む前に押さえたい検証条件を解説します。

こはるの読みどころ

モデル実装を選ぶ話と、計算を分散・最適化する話を分けると読みやすいよ。MoEの例がどこまで検証しているかにも目を向けてみよう！

学習用アクセラレータを選ぶとき、気になるのは計算性能だけではありません。使いたいモデルの実装や、学習を起動・検証する手順をどこまで引き継げるかも、導入の判断材料になります。

2026年10月21日に公開されたAWSのQwen3学習ガイドは、その接点を考える材料です。torchtitanで何を共通化でき、どこから個別の検証が必要になるのか。モデル実装、並列化、カーネル最適化の順に見ると、試し始める場所が分かります。

TorchNeuronがモデルとTrainiumの間をつなぐ

TorchNeuronは、TrainiumをPyTorchのneuronデバイスとして扱うためのバックエンドです。torchtitanは、その上でモデルや学習処理を動かします。従来のXLAベースのtorch-neuronxと、PyTorch 2.14向けのネイティブバックエンドは別の系統なので、同じパッケージ名だけで環境を判断しないようにしたいですね。TorchNeuronの公式説明に、この区別が示されています。

この構成では、ハードウェアへの接続とモデル実装の選択を分けて考えられます。ただし、Neuron 2.33のTorchTitan Neuronはベータ版です。導入時は、使いたいモデルと実行方式の組み合わせを具体的に絞るのが出発点になります。

Qwen3の実装は、変えたい部分から選ぶ

AWSの例では、同じtorchrunの起動パターンと設定構造で、次の3構成を扱います。右列は、それぞれの役割から考えた選び方です。

構成	モデル・演算を担うもの	選ぶときの着眼点
Transformers v5	Hugging Faceのモデル実装。safetensorsの入出力にも対応	公開済みの重みを使う流れを軸にしたい
torchtitanネイティブ	torchtitan自身のQwen3実装	モデル内部を追いながら変更したい
NKIを使うMoE構成	Neuron向けのexpert演算カーネル	演算実装を調整する足場が欲しい

この選択を、単純な速さの順位と捉える必要はありません。まず「どの実装を保守し、何を変更したいか」を決めると、比較すべき対象が絞れます。

また、Qwen3-30B-A3Bを使うMoE例は4層・64 experts・各トークンの選択数4に縮小され、expert parallelismは使いません。フル構成の学習や収束を示す例として扱わず、カーネルを学習へ組み込むための検証として読むのが適切です。

64プロセスをどう使うかはTPとFSDPで決まる

モデル実装が決まっても、並列化の設計は残ります。Tensor Parallelism（TP）は層内の計算を分担し、Fully Sharded Data Parallel（FSDP）はパラメータなどの学習状態を分割します。torchtitanはDeviceMeshやDTensorを使ってこれらを組み合わせます。TorchTitan Neuronの構成説明では、この仕組みと対応範囲を確認できます。

背景を具体化する例として、公式のQwen3-8B基本チュートリアルには、1台のtrn2.48xlargeでTP4 × FSDP16を使う構成があります。64の論理NeuronCoreに対して、4 × 16 = 64という並列化の配置です。これは先ほどのAWSブログの3構成とは別の、分散配置を理解するための例です。

同じ64プロセスでも、どの軸へ割り当てるかで通信やメモリの使い方が変わります。モデル名とプロセス数だけをそろえて比較するのではなく、TPとFSDPの次数まで記録すると、差を説明しやすくなります。

NKIへ進む前に、逆伝播と入力形状を検証する

NKIは、Trainium向けの演算を細かく制御したいときの選択肢です。学習に組み込むには、順伝播で値を返すだけでなく、逆伝播も用意する必要があります。NKI連携の説明では、カスタム演算にregister_autograd()で逆伝播を登録する仕組みが案内されています。

ここから導ける実務上の判断は、まず既存実装を比較基準にすることです。カーネルを差し替えたら、出力だけでなく勾配も比較する。そのうえで学習全体の時間を測れば、演算の変更と学習結果の変化を切り分けやすくなります。

コンパイルを使う際は、入力形状も比較条件です。Hugging FaceのTrainiumガイドは、新しい入力形状が再コンパイルを起こすことを説明しています。長さを固定したpaddingや、サイズの異なる最終バッチを避ける設定が紹介されており、性能測定ではデータの渡し方までそろえる必要があると分かります。

短い学習を基準にしてから最適化を重ねる

最初の目標は、使う構成で学習が最後まで進むことです。公式セットアップ手順は、Neuron SDK 2.33以降とPyTorch 2.14、動作確認済みのtorch-neuronxを前提に、少数ステップの実行を案内しています。デバイス認識、パッケージの読み込み、有限なlossの出力という順序なら、環境の問題を切り分けやすくなります。

次に、eagerでの短い実行とチェックポイント保存を基準にし、並列化の軸を一つずつ増やします。TorchTitan Neuronのベストプラクティスは、このときステップ当たりのトークン数を固定するよう勧めています。検証できた構成へコンパイルやカーネル変更を重ねると、結果が変わった理由を追えます。

Qwen3をTrainiumで試す際の判断軸は、モデル実装をどこまで変えたいかと、その変更を比較できる基準があるかです。共通の学習基盤を活かしながら、まず一つの構成を再現可能にする。その基準があれば、次に手を入れる場所を実測から選べます。