AWS Neuron 2.32がNKI拡張、MXFP8トレーニングカーネル、Trn2/3向け可変サイズ集団通信を追加

AWS Neuron 2.32.0 が正式リリースされ、以下の主要機能が追加された。

1. Neuron Kernel Interface (NKI) 0.6.0

   • オンデバイスでの top-K命令 追加
   • テンソルサイズが異なるランク向けの 可変長 all-gather 対応
   • データ依存型の反復処理を可能にする ループ構造 の新設
   • Mixture of Experts (MoE) トレーニング および スパースアテンション 向けに 13種類の新NKIライブラリカーネル を提供（DeepSeek-V3.2のスパース多頭潜在アテンション用コンテキストエンコーディング、MXFP8デコードステップ用アテンション、ブロックワイズMoE層のMXFP8エンドツーエンドトレーニング対応など）
   • PyTorch参照実装を持つ 22種類の追加カーネル をサポート

2. Neuron Agentic Development

   • トランスフォーマーモデルを vLLM Neuronプラグイン に移植する新スキルを追加（vLLM 0.24.0対応）

3. Neuron Runtime

   • Trn2/Trn3インスタンス で 可変サイズの all-gather、reduce-scatter、all-to-all 集団通信 をサポート（各ランクが異なる要素数を送受信可能）

4. Neuron Compiler

   • 64ビット整数コンパイル の明示的制御を追加

5. Neuron Explorer

   • System Trace Viewer に ホストCPUのコア単位利用率 を表示する機能を追加

リージョン対応状況: AWS Neuronは Trn1、Trn2、Trn3、Inf2、Inf1インスタンス が利用可能な全リージョンで提供。ただし、Trn2/Trn3インスタンスの東京（ap-northeast-1）および大阪（ap-northeast-3）での利用可否は公式ドキュメントで明示されていないため、AWS Management ConsoleまたはAWSサポートで確認が必要。

エンジニア視点での考察・アドバイス

解決された技術的課題

• 可変サイズ集団通信 のサポートにより、異なるテンソルサイズを持つランク間での効率的なデータ交換が可能に。特に MoEモデル や スパースアテンション を用いた大規模トレーニングで、メモリ効率と計算パフォーマンスが向上する。
• MXFP8（8ビット浮動小数点） のエンドツーエンドトレーニング対応は、精度を維持しつつメモリ使用量と計算コストを削減するブレークスルー。特に ブロックワイズMoE層 での活用が期待される。
• NKIのループ構造 と top-K命令 により、データ依存型の反復処理が柔軟に記述可能に。カスタムカーネル開発のハードルが下がり、研究開発のスピードアップに寄与する。

日本リージョンでの注意点

• 東京（ap-northeast-1）:
  • Trn1、Inf1、Inf2インスタンス は利用可能だが、Trn2/Trn3インスタンス の利用可否は公式ドキュメントで明示されていない。
  • 可変サイズ集団通信 などの新機能を利用する場合は、AWS Management ConsoleまたはAWSサポートで インスタンスの利用可否を事前に確認 すること。
• 大阪（ap-northeast-3）:
  • Inf1インスタンスのみ利用可能。Trn1、Trn2、Trn3、Inf2インスタンスは非対応。

運用上のインパクトとアドバイス

• コスト最適化: MXFP8トレーニングはメモリ使用量を削減し、同一インスタンスでより大規模なモデルを扱える可能性がある。ただし、精度への影響を検証するため、まずは 小規模なテストジョブ で評価することを推奨。
• セキュリティ: 新しいNKIカーネルやvLLMプラグインを利用する際は、依存ライブラリのバージョン管理と脆弱性チェックを徹底。特に Deep Learning AMI/Container を使用する場合は、最新のセキュリティパッチが適用されていることを確認する。
• 移行のポイント:
  • vLLM Neuronプラグイン を使用している場合は、vLLM 0.24.0へのアップグレードが必要。互換性の問題がないか事前にテスト環境で検証する。
  • 64ビット整数コンパイル の明示的制御は、既存のモデルで意図しない動作を引き起こす可能性があるため、コンパイルオプションの設定を慎重に行う。

次のアクション

1. リージョン確認: Trn2/Trn3インスタンスを利用予定の場合は、AWS Management ConsoleまたはAWSサポートで 利用可否を確認 する。
2. MXFP8トレーニングの検証: 既存のモデルをMXFP8でトレーニングし、精度とパフォーマンスを比較。特に MoEモデル での効果を重点的に評価する。
3. NKIカーネルの活用: 新しいNKIライブラリカーネルを試し、カスタムカーネル開発の効率化を図る。
4. vLLMプラグインのアップグレード: vLLM 0.24.0への移行を計画し、テスト環境で動作確認を行う。

参考ドキュメント

• AWS Neuron 2.32.0 リリースノート
• Neuron Kernel Interface (NKI) ドキュメント
• AWSリージョン別サービス対応表