詳しい解説 · 読了目安 9 分
TrainiumでQwen3を学習する：重みとカーネルを検証するtorchtitanの3つの経路
AWSの10月21日の解説は、1台のTrainium2インスタンスでTransformers、ネイティブQwen3、NKIカーネルを接続する方法を示す。速度や費用の優位性を実証したものではない。
- 当サイトの公開日
- 原典の公開日
- 出典確認日
AWSが10月21日に公開した技術解説は、Trainium上のtorchtitanでQwen3を学習する3つの構成を紹介する。Transformers v5バックエンド、torchtitanのネイティブモデル、NKIカーネルを使う縮小版の混合エキスパート例だ。具体的な統合手順に加え、チェックポイントと数値結果の確認方法が示されている点が重要である。Qwenの新モデル発表でも、学習費用の安さを証明するベンチマークでもない。
ニュースを聴く
この記事を読み上げるか、ニュースを続けて聴けます。
読み上げを準備中…
端末の音声
読み上げはブラウザーまたは端末が処理します。AISubHub は文章を音声 API に送りません。利用できる音声やロック画面での再生は端末によって異なります。
再開時に短い部分を繰り返す場合があります。
要点
- 使用するのはtrn2.48xlarge 1台。Trainium2を16基、物理NeuronCoreを128個備え、LNC=2では論理NeuronCoreが64個となる。起動する学習プロセスも64個である。
- 密なモデルの例はQwen3-4B-Instruct-2507のファイルを使い、c4_testで10ステップ学習する。MoE例は4層、64エキスパートに縮小し、各トークンで4エキスパートを選択する。元の30B-A3B全体を変更せず学習する例ではない。
- AWSは公開済み重みでの初回ステップの損失を6.54、ランダム初期化では12.43と報告する。独自の回転位置埋め込みのfloat32テストでは最大絶対差が0.0だった。これらは別の問題を確認する数値である。
同じデバイス統合の上でモデル実装を3通り選ぶ
TorchNeuronはTrainiumをPyTorchのネイティブなneuronデバイスとして登録し、torchtitanが通常のインターフェースを使えるようにする。第1の経路ではTransformers v5のAutoModelForCausalLMでモデルを構築する。第2ではconfig.model_specを通じてtorchtitan独自のQwen3実装に切り替える。第3ではMoE層にTrainium向けNKI演算子を組み込み、独自の置換も実演する。共通の学習入口の背後にある実装の違いであり、3つの新しいQwenモデルではない。
ガイドはtorch 2.14.0、torch_neuronx 2.14.3、Transformers 5.16.0、torchtitan-neuron release-2.33のバージョン0.2.2+neuron.1を固定する。ドライバー、ランタイム、コンパイラを含むNeuron向けDeep Learning AMIか学習コンテナが必要だ。使い慣れたPyTorchコードで統合箇所を絞れる一方、試験済み依存関係の組み合わせを省略できるわけではない。AWSはTransformersとの統合や置換モジュールがtorchtitanのリリース間で変わり得ると明記している。
4B例で確認するのは初期化とチェックポイントの往復
最初の構成はQwen/Qwen3-4B-Instruct-2507を指定し、c4_test、10ステップ、ローカルバッチサイズ1を設定して、データ並列の状態を起動済み64ランクに分割する。アーキテクチャを読み込むことと、公開済みの重みを読み込むことは別だ。微調整の構成ではダウンロードしたファイルからのチェックポイント読み込みを明示的に有効にし、最後にHugging Face safetensorsとしてモデルを出力する。報告された初回損失は公開済み重みで6.54、ランダム初期化で12.43だった。
この差は初期化を確認する手掛かりであり、10ステップで優れたアシスタントができることや最終収束の証明ではない。AWSは出力したチェックポイントをTransformersで再読み込みし、埋め込みと共有されるlm_head.weightを考慮しながらキーも検証する。ネイティブ経路がqwen3_32bビルダーから始まるのは、試験した版に汎用4Bビルダーがないためだ。その後でモデル仕様を4Bに置き換え、グローバルバッチサイズを64にする。関数名だけで32Bを学習したと判断してはいけない。
MoE例の規模はモデルファミリー名より小さい
Qwen3-30B-A3Bの構成は、1台に収めるため4層と64エキスパートに縮小される。top_k=4は各トークンを4つのエキスパートに送る設定だ。資料では64ランクでデータ並列の分割を行い、テンソル並列度は1、エキスパート並列は使用しない。エキスパート数はモデル内部の設定であり、エキスパート並列に使うデバイス数ではない。この区別により、小さな統合例を元のフルモデルの学習成果と取り違えずに済む。
この例はコンパイルを無効にしたeagerモードで動く。登録済みNKI演算子はtorch.compileにも対応するが、実行方式の切り替えは別に判断すべき事項だ。試験した版ではMoEブロックをグループ化エキスパート用カーネルに置換し、スケールド内積アテンションとgather不要の回転位置埋め込みを選ぶ。AWSは順伝播で保存した情報を使って演算子にautogradを接続する方法を説明している。この構成の高速性を証明するスループット比較は載せていない。
独自カーネルは配置、勾配、適用範囲まで確認する
回転位置埋め込みの例は、登録が単なる関数呼び出しではないことを示す。torchtitanのqueryとkeyの配置はバッチ、系列、ヘッド、ヘッド次元だが、rope_hfが期待するのはバッチ、ヘッド、系列、ヘッド次元である。ラッパーは配置を変換し、出力用領域を確保してから元の配置に戻す。別の逆伝播演算子が保存済み回転キャッシュを再利用し、キャッシュ自体には勾配を与えない。系列長は128×LNCの倍数である必要があり、LNC=2なら256の倍数になる。
次に置換レジストリで目的のモジュール構成を差し替える。起動ログにはQwen3の各層につき1つのnki_rope置換が必要で、より限定した名前パターンで特定の層だけを指定することもできる。AWSはPyTorch実装とのfloat32比較で、queryとkeyの最大絶対差がともに0.0だったと報告する。検証できたのはその数値テストであり、全データ型、形状、学習実行の同等性ではない。著者も性能向上ではなく、カーネルと自動微分の統合例だと明示している。
解説を測定可能な試験に変える
初回ステップにはコンパイルが含まれ、後のステップより時間がかかる。ログには損失、勾配ノルム、高帯域幅メモリの使用量、スループット、モデルFLOPs利用率が出るが、記事は学習全体の費用を管理条件下で比較していない。モデル規模、系列長、ローカルバッチを増やす際は各ランクのメモリが制約になる。多次元DeviceMeshでテンソル並列と完全分割データ並列を組み合わせることは可能だが、掲載されたMoE設定をその大規模実験と同一視してはいけない。
AISubHubが提案する順序は、試験済みバージョンを維持し、選択モデルと64ランクを確認し、チェックポイントの読み込みと出力を検証し、長時間実行の前に独自カーネルの出力を比較することだ。その後、ウォームアップと継続実行を分けて測定し、対象データでの収束を評価する。借りたハードウェアを停止・終了する前には必要なチェックポイントを保存する。これらはガイドから導いた評価手順であり、本站の再現結果ではない。サブスクリプション特典、無料学習枠、Trainiumの普遍的優位性も示すものではない。
注目する理由
AISubHubの分析：本格的な学習に進む前に、モデル互換性、チェックポイントの移植性、カーネルの正しさを小さな検証に分けられる。共通の起動構造は手順の差を減らすが、バージョン、メモリ、コンパイル、数値計算の要件は残る。
制約と提供状況
AWSのJosh Longenecker氏とAlex Hoffman氏による技術ガイドであり、AISubHubによるハードウェア試験ではない。比較可能な学習全体の費用、最終的な収束結果、独自カーネルによる速度向上は示されていない。インスタンスの空きと料金、正確なソフトウェア構成、対象処理の検証は運用者が確認する必要がある。個人向けサブスクリプションの特典も発表されていない。
出典と根拠
AI を活用した編集記事を、記載した一次情報で確認しています。企業の主張は出典を明示し、掲載は推奨を意味しません。