AWS、TrainiumでQwen3を学習する手順を公開

この投稿は、ジョブの起動方法ではなくモデルの構築方法が異なる3つの設定を中心に構成されている。1つ目ではTransformers v5がAutoModelForCausalLMを通じてQwen3の実装を提供し、設定はHub IDを指定するだけだ。2つ目はconfig.model_specを変更してtorchtitanのネイティブQwen3Modelに差し替え、qwen3_32b()から始めてモデル仕様を4B版に置き換える。テストしたビルドには汎用の4Bビルダーがないためだ。3つ目はさらに進み、mixture-of-expertsブロックをqwen3_nki_moe_config()によるNKIのgrouped-expertsカーネルに置き換え、1台のtrn2.48xlargeに収まるよう4層・64エキスパートに縮小している。

カスタムカーネルの節が最も手順的だ。nkilibのカーネルをwrap_nkiでラップしてPyTorchオペレータとして登録し、2つ目のオペレータが保存されたルーティングのメタデータを使って逆伝播を処理する。数値チェックではrope_hfとtorchtitanのPyTorch rotary embeddingをfloat32入力で比較し、クエリとキーの最大絶対差は0.0だったと報告している。著者らはこの部分を、性能向上ではなく登録と自動微分の統合を示すものと位置づけている。

3つの経路を結びつけるのはランチャーとオーバーライドの仕組みだ。データ並列シャーディングは64ランクにわたりdata_parallel_shard_degree = -1で設定され、テンソル並列は同じ多次元DeviceMesh上に重ねられる。オーバーライドは完全修飾名で範囲を限定するため、例えばfqns=["layers.1[0-9].attention.rope"]は10層から19層のrotary embeddingを対象にする。2つのオーバーライドが同じノードを指定するとtorchtitanは停止する。これは、Hub定義の実装とネイティブ実装の切り替えやカスタムカーネルの組み込みが、別の学習スタックではなく設定変更で済むワークフローを示している。