torchtitan とカスタムカーネルを用いた AWS Trainium での Qwen3 のトレーニング

AWS は、Amazon Web Services インフラストラクチャ上で事前学習およびファインチューニングワークフローを最適化しようとする機械学習チーム向けの実際的な参考資料として、torchtitan フレームワークを使用して自社の Trainium2 アクセラレータで Qwen3 大規模言語モデルをトレーニングする方法を示す詳細な技術的ウォークスルーを公開しました。2026年10月に公開されたこのガイドは、専用ハードウェア上のネイティブ PyTorch インターフェースを活用したいエンジニア向けに、3つの異なる構成パスを概説しています。

何が起きたか

この記事では、16個の Trainium2 アクセラレータと128個の物理的な NeuronCore を備えた単一の trn2.48xlarge インスタンスで Qwen3 モデルを実行するためのステップバイステップの実装ガイドを提供しています。著者らは、TorchNeuron が Trainium をネイティブ PyTorch デバイスとして登録しているため、開発者が標準的な GPU クラスターで使用する場合と同様に、Trainium 上で torchtitan を使用できることを実証しています。これにより、フレームワークは起動時にハードウェアを検出し、操作をイージモードまたは torch.compile 経由でコンパイルできます。

ウォークスルーでは、3つの具体的なトレーニング構成を取り上げます。最初の方法は、モデリングバックエンドとして Hugging Face Transformers v5 を利用し、ユーザーが AutoModelForCausalLM を使用して Hub から直接モデルを読み込めるようにします。2番目のアプローチは、torchtitan のネイティブ Qwen3 実装に切り替え、モデルアーキテクチャとシャーディング戦略に対するより細かい制御を提供します。3番目の方法は、Neuron Kernel Interface (NKI) の混合専門家（Mixture-of-Experts）カーネルを統合することで高度なカスタマイズを導入し、開発者がパフォーマンス上重要なレイヤーにカスタムオペレーターを追加できるようにします。

各構成では、異なる実験セットアップ間の一貫性を確保するために、同じ torchrun 起動パターンを使用します。ガイドでは、PyTorch 2.14.0、Transformers 5.16.0、torchtitan-neuron release-2.33 など、正確なソフトウェアバージョンを指定することにより再現性を強調しています。また、チェックポイントの管理方法についても詳しく説明しており、Hugging Face safetensors 形式で保存された重みが欠落キーなしで再読み込みできることを確認し、トレーニングパイプラインの整合性を検証しています。

仕組み

コアメカニズムは、torchtitan と AWS Neuron スタック間の統合に基づいています。torch_neuronx がインストールされると、Trainium が neuron デバイスとして公開されます。その後、torchtitan は完全シャードデータ並列（FSDP）やテンソル並列などの標準的な PyTorch 分散トレーニングプリミティブを使用し、インスタンス上で利用可能な64個の論理 NeuronCore にわたってワークロードを分散させます。フレームワークはモデルパラメータ、勾配、オプティマイザ状態のシャーディングを管理し、シーケンス次元に沿ったアクティベーションシャーディングはシーケンス並列処理によって扱われます。

カスタムカーネルの統合については、ガイドでは torchtitan のオーバーライドレジストリの使用方法を示しています。開発者は Neuron カーネル言語である NKI でカーネルを作成し、wrap_nki および nki_op を使用して PyTorch オペレーターとしてラップできます。これらのオペレーターは、モデルグラフ内の特定のモジュールを置き換えるために登録されます。例えば、カスタムのロータリー埋め込みカーネルはデフォルトの PyTorch 実装を置き換えることができます。システムは起動時にこれらのオーバーライドを検証し、各置換がモデル構造内の正しい完全修飾名を対象としていることを確認します。

トレーニングループは、損失、勾配ノルム、高帯域幅メモリ使用量、モデル FLOPs 利用率などの主要な指標を報告します。この可視性により、エンジニアはリソース消費を監視し、並列設定を動的に調整できます。コンパイルプロセスは最初のトレーニングステップ中に発生し、その後のステップよりも時間がかかりますが、実行の残りの部分における最適化された実行を可能にします。

主要な詳細

• セットアップには、LNC=2 で構成された64個の論理 NeuronCore を備えた trn2.48xlarge インスタンスが必要です。
• ソフトウェア依存関係には、PyTorch 2.14.0、Transformers 5.16.0、torchtitan-neuron 0.2.2+neuron.1 が含まれます。
• Hugging Face チェックポイントの読み込みは、Qwen3-4B-Instruct-2507 モデルにおいてステップ1の損失が約6.54であることを確認することで正しさを検証します。
• カスタム NKI カーネルは、128 * LNC の倍数であるなど、特定のシーケンス長制約に従う必要があります。
• テンソル並列はレイヤー重みをランク間でシャーディングし、データ並列はパラメータとオプティマイザ状態をシャーディングします。
• オーバーライドレジストリは、2つのオーバーライドが同じノードを主張した場合に停止し、一致するノードがない場合に警告を出すことで競合を防ぎます。

なぜ重要なのか

ソフトウェアエンジニアや ML プラクティショナーにとって、このガイドは NVIDIA 以外のハードウェアで大規模モデルをトレーニングするための参入障壁を下げるものです。標準的な PyTorch インターフェースおよび Hugging Face エコシステムとの互換性を維持することで、広範なコードリファクタリングの必要性を減らします。チームは既存のモデルリポジトリやチェックポイント形式を活用でき、新しいアクセラレータアーキテクチャを採用することに伴う運用オーバーヘッドを最小限に抑えます。

カスタム NKI カーネルを注入できる機能は、特にパフォーマンス最適化において重要です。モデルの複雑さが増すにつれ、汎用的な実装ではハードウェア機能を十分に活用できない場合があります。アテンションメカニズムやロータリー埋め込みなどの特定の操作を手動でチューニングしたカーネルに置き換えることを開発者に許可することで、計算効率に対するより細かい制御が可能になります。この柔軟性は、大規模なトレーニングジョブにおいて競争力のあるスループットとコスト効率を実現するために不可欠です。

さらに、詳細な検証手順は、トレーニングプロセスの信頼性に対する確信を提供します。エクスポートされたチェックポイントが正しく読み込まれることを確認することは、Trainium 上でトレーニングされたモデルが、標準的な推論スタックに依存する本番環境でシームレスに展開できることを保証します。このエンドツーエンドの検証は、堅牢な MLOps 実践をサポートし、異なるハードウェアバックエンドを使用するチーム間のコラボレーションを促進します。

できること

• Neuron 用の AWS Deep Learning AMI または Neuron PyTorch コンテナを使用して trn2.48xlarge インスタンスを設定します。
• 提供された requirements.txt を使用して、指定されたバージョンの torchtitan-neuron と依存関係をインストールします。
• トレーニングのために、Hugging Face Hub から Qwen3 モデルアーティファクトをローカルディレクトリにダウンロードします。
• data_parallel_shard_degree = -1 を設定して、すべての64ランクにわたるデータ並列シャーディングを使用するようにトレーニングスクリプトを構成します。
• float32 入力を使用してベースライン実装に対してカスタムカーネル出力を検証し、数値精度を確保します。
• バッチサイズとシーケンス長を適切に調整するために、ランクごとの高帯域幅メモリ使用量を監視します。