AWS Neuron 2.32 登場！MXFP8 トレーニングと NKI 新機能で Trn2 がパワーアップしたよ！

やっほー、しぃちゃんだよ！今日は AWS の機械学習チップまわりのアップデートを見つけたよ。地味に見えるけど、中身はけっこう本格的なので張り切って紹介するね！

AWS What's New aws.amazon.com

なにが発表されたの？

AWS What's New によると、Trainium / Inferentia 向けの SDK「AWS Neuron」の v2.32.0 がリリースされたよ。カーネルを書くための言語である NKI(Neuron Kernel Interface)が 0.6.0 にアップデートされて、MoE(Mixture of Experts)トレーニングやスパースアテンション向けの新しいカーネルが 13 個も追加されたんだって。あわせて、Trn2 / Trn3 上での集合通信(コレクティブ)にも新しい機能が入ったよ。

今までどうだったの？

これまでの NKI の集合通信命令(all-gather など)は、各ランクが同じサイズのテンソルを持っていることが前提だったの。だから、MoE のルーティング後のようにランクごとに扱うデータ量が変わってしまうケースだと、カーネルを書くのが難しかったんだよね。

また、on-device での top-K 処理や、データに応じて回数が変わるループのような「データ依存な処理」も、これまでの NKI では表現しづらい部分だったみたい。

これで何が変わるの？

NKI 0.6.0 では、これらの弱点がまとめて解消されたよ。

• オンデバイスの top-K 命令が追加され、デバイス上で完結する top-K 処理が書けるように
• ランクごとにサイズの違うテンソルを扱える可変長 all-gather が追加
• データ依存の反復処理を書くための新しいループ構文が追加

これによって、MoE のようにランクごとの負荷が均一にならないモデルのカーネルを、NKI で素直に書けるようになったの。Trainium / Inferentia を使って大規模モデルをトレーニングしているエンジニアにとっては、これまで自前で工夫していた部分がライブラリレベルでサポートされる形だね。

深く潜ってみよう

新しく追加された 13 個の NKI ライブラリカーネルの中身は、かなり最新のモデル構成を意識したラインナップだよ。

• DeepSeek-V3.2 のスパースな multi-head latent attention 向けのコンテキストエンコーディングカーネル
• デコードステップ向けの MXFP8 アテンションカーネル
• フォワードとバックワードが対になった MXFP8 カーネルで、ブロックワイズ MoE レイヤーを最初から最後まで MXFP8 で学習可能に

さらに、追加で 22 個のカーネルについては PyTorch のリファレンス実装も用意されているよ。低精度フォーマットである MXFP8 でブロックワイズ MoE を丸ごとトレーニングできるようになったのは、メモリと計算コストを抑えたい大規模 MoE モデルの学習にとって嬉しいポイントだね。

Neuron Runtime のほうでは、Trn2 と Trn3 上で all-gather・reduce-scatter・all-to-all といった集合通信が可変サイズに対応したよ。各ランクが異なる数の要素を送受信できるようになったから、負荷が均一でない分散トレーニングの通信を無駄なく組めるの。

そのほかのアップデートはこんな感じだよ。

• vLLM Neuron プラグインが vLLM 0.24.0 に対応し、すべての Neuron Deep Learning AMI と Deep Learning Container に同梱
• 新しい Neuron Agentic Development skill が追加され、Transformer モデルを vLLM Neuron プラグインへ移植する作業をサポート
• Neuron Compiler が 64 ビット整数のコンパイル挙動を明示的に制御できるように
• Neuron Explorer の System Trace Viewer に、コアごとのホスト CPU 使用率の表示が追加

AWS Neuron は、Amazon EC2 の Trn1・Trn2・Trn3・Inf2・Inf1 インスタンスが使える全ての AWS リージョンで利用できるよ。

まとめ

• AWS Neuron 2.32.0 がリリースされ、NKI が 0.6.0 に更新された
• オンデバイス top-K 命令、可変長 all-gather、データ依存ループが NKI に追加
• MoE トレーニング・スパースアテンション向けの新カーネルが 13 個追加され、DeepSeek-V3.2 やブロックワイズ MoE の MXFP8 end-to-end トレーニングに対応
• Trn2 / Trn3 の Neuron Runtime が可変サイズの all-gather・reduce-scatter・all-to-all をサポート
• vLLM Neuron プラグインが vLLM 0.24.0 に対応、Neuron Agentic Development skill も新登場

Trainium や Inferentia の上で自前の MoE モデルをトレーニングしているエンジニアさんには、特に刺さるアップデートだと思うよ！