TorchNeuronがパブリックベータに、AWS TrainiumでネイティブPyTorchを実現

AmazonのTrainiumチップでPyTorchのコードを動かす新しい方法、TorchNeuronが2026年10月20日にパブリックベータになりました。人気のあるオープンモデルの大半が、変更なしでTrainium上で動くようになった点が重要です。Hugging Faceによると、最もよく使われているテキスト生成モデルの設計100種のうち89%が、そのまま動きます。古いツールでは46%でした。

TorchNeuronとは

Trainiumは、Amazonが設計し、クラウド事業のAWSを通じて貸し出しているAIチップです。PyTorchは、AIモデルの構築と学習に最も広く使われているライブラリです。これまでPyTorchのコードは変換レイヤーを通じてTrainiumに届いており、多くのモデルは動かすために追加の作業が必要でした。

TorchNeuronは、Trainiumをネイティブなデバイスとして追加するPyTorchプラグインです。AWSのドキュメントには、「Trainiumを標準のPyTorchデバイスとして登録する」とあります。実際には、開発者はモデルをneuronという名前のデバイスに移します。GPUにcudaで移すのと同じやり方です。

GitHubページによると、このプロジェクトはApache-2.0ライセンスのオープンソースです。AWSは、チップ向けのソフトウェアキットであるNeuronの2.33リリースで、これをパブリックベータとして掲載しています。

古い方法からの変更点

古いルートは、Hugging FaceのライブラリであるOptimum Neuronを経由していました。このライブラリはPyTorch XLAに依存していました。XLAは、実行前にモデルを固定のグラフに変換するコンパイラーレイヤーです。この方式では、モデルごとに個別の調整が必要だったため、多くのモデルのサポートが難しくなっていました。

TorchNeuronは、このレイヤーをなくします。Hugging Face自身のライブラリが、Trainium上で直接動くようになりました。言語モデル向けのTransformers、画像モデル向けのDiffusers、さらに学習とファインチューニング向けのAccelerate、PEFT、TRLです。以下の表は、AWSのドキュメントとHugging Faceの投稿をもとに機能をまとめたものです。

機能	意味
Eagerモード	通常のPyTorchと同じく、コードが1行ずつ実行される
torch.compile	neuronバックエンドを使った、速度向上のための任意のコンパイル
分散学習	FSDP、DDP、DTensor、テンソル並列とパイプライン並列
カスタムカーネル	AWSのNKI言語による、手書きのチップ向けコード
チップ	Trainium 2とTrainium 3。FP8とMXFP8にはTrainium 3が必要

FSDPとDDPは、学習を多数のチップに分散させるPyTorchの標準的な方法です。FP8は、メモリと時間を節約できる、より小さな数値形式です。

モデルと速度

Hugging Faceの投稿では、TRLを使ったQwen3-8Bのファインチューニングを、テンソル並列とFSDPで複数のチップに分散して実行しています。また、FLUX.1-dev、FLUX.2、Qwen-Imageといった画像パイプラインも、テンソル並列で実行しています。

AWSのドキュメントには、混合エキスパート（mixture-of-experts）モデルが掲載されています。これは、各入力を少数の専門のサブネットワークに振り分けるモデルです。DeepSeek、GPT-OSS、Qwen3がサポートされています。これらのモデルの学習には、別のツールであるTorchTitan Neuronを使います。

速度に関する数字は、範囲の狭いものが1つだけです。Hugging Faceの投稿によると、AWSのNeuronチームによるビジョントランスフォーマーのカーネルは、「Trainium上で最大2.4倍の高速化」をもたらします。何と比べた数字なのかは書かれていません。

インストール方法

GitHubページには、バージョンを固定したインストール手順が載っています。Neuronのドライバーとランタイムのパッケージを導入した後に実行します。

pip install torch==2.14.0 --index-url https://download.pytorch.org/whl/cpu
pip install "torch-neuronx==2.14.*" "neuronx-cc==2.*" "nki>=0.6" --extra-index-url https://pip.repos.neuron.amazonaws.com

バージョンの固定が重要です。TorchNeuronは、古いXLAベースの製品と同じtorch-neuronxというパッケージ名で配布されています。それより前のバージョンでは古い製品がインストールされると、リポジトリは警告しています。python -m torch_neuronx.collect_envを実行すると、PyTorchからデバイスを使えることを確認できます。AWSの既製のマシンイメージとコンテナには、あらかじめインストールされています。

開発者にとっての意味

TorchNeuronは、Trainiumを特別なターゲットから、PyTorchのデバイスの1つに変えます。これによって、一般的なチームでもGPUとの価格比較が現実的になります。

• まず自分のモデルを試してください。 89%のカバー率でも、人気のある設計のおよそ10個に1個は作業が必要です。移行を計画する前に、短い学習または推論のジョブを実行してください。
• すべての環境でバージョンを固定してください。 requirementsファイルでtorch-neuronxのバージョンを固定しないと、古いXLA製品が入る可能性があります。CIでcollect_envの出力を確認してください。
• チップの世代を確認してください。 サポートされているのは、Trainium 2と3のインスタンスだけです。計画がFP8での学習に依存しているなら、Trainium 3が必要です。
• ベータ版として扱ってください。 3つの情報源のどれも、本番環境で使える状態だとは述べていません。期限のある作業には、GPUへの切り替え手段を残しておいてください。
• 時間単位ではなく、ジョブ単位でコストを比較してください。 2.4倍という数字は、比較対象が示されていない1つのカーネルについてのものです。予算を投じる前に、1ドルあたりのトークン数を自分で測ってください。