Trainium で素の PyTorch が動く！AWS Neuron 2.33 に TorchNeuron が登場

やっほー、しぃちゃんだよ！今日は AWS の AI チップ Trainium まわりの大きなアップデート、AWS Neuron 2.33 を紹介するね！

AWS What's New aws.amazon.com

なにが発表されたの？

AWS の What's New で、AWS Neuron 2.33 が公開されたよ。目玉は 2 つで、どちらもパブリックベータなの。

• TorchNeuron: AWS Trainium 向けのネイティブ PyTorch バックエンド
• Neuron Observer: Trainium の不調を検知する仕組み

今までどうだったの？

原文には従来の仕組みの詳しい説明はなかったので、ここは今回できるようになったことを中心にお話しするね。

これで何が変わるの？

TorchNeuron のおかげで、標準の PyTorch API で書いたモデルを、Amazon EC2 の Trn2 と Trn3 インスタンスで動かせるよ。eager モードでも torch.compile でも使えるの。PyTorch で書いたコードを Trainium に持ってきたい人には、うれしい変化だよね！

Neuron Observer は、Trainium のヘルス問題を 1 秒未満で検知して、影響を受けたノードやデバイスを特定してくれるよ。

深く潜ってみよう

推論まわり

• vLLM Neuron プラグインが vLLM 0.29.0 にアップグレード
• Wan2.2 の動画生成向けの vLLM Omni Neuron プラグインが Neuron リリースの一部になった
• 両方とも TorchNeuron 経由のネイティブ PyTorch コンパイルに対応

学習環境

新しい Neuron Deep Learning AMI と、PyTorch 学習用の Deep Learning Container が登場したよ。どちらも次を含んでいるの。

• TorchNeuron
• PyTorch 2.14
• TorchTitan

プロファイリングとデバッグ

• Neuron Explorer が torch.profiler で取得したプロファイルを開ける
• PyTorch のメモリスナップショットから、メモリ割り当てを、それを行った Python のソース行に紐づけて表示できる。eager モードでの OOM（メモリ不足）のデバッグに役立つよ

NKI

• Neuron Kernel Interface（NKI）で書くカーネルで、ジェネレータやクラス継承を含む Python 言語のフル機能がメタプログラミングに使えるようになった
• NKI Library に、実験的なカーネルが 25 個追加された。DeepSeek-V3.2、Qwen3.5、Kimi Delta Attention 向けのものも含まれているよ

詳しくは AWS Neuron 2.33.0 のリリースノートや、TorchNeuron・Neuron Observer・NKI のドキュメントを見てね！

まとめ

• TorchNeuron（パブリックベータ）で、標準の PyTorch API のモデルを Trn2 / Trn3 で実行できる
• Neuron Observer（パブリックベータ）が Trainium の不調を 1 秒未満で検知してノードやデバイスを特定
• vLLM 0.29.0 対応と、Wan2.2 向けの vLLM Omni プラグインを同梱
• PyTorch 2.14 と TorchTitan 入りの AMI と Container が登場
• Neuron Explorer のメモリ帰属と、NKI の Python フル機能、実験的カーネル 25 個

Trainium で PyTorch の学習や推論を回したい、機械学習エンジニアのみんなに刺さる更新だよ！