10月19日、PyTorchが「Introducing TorchNeuron: Native PyTorch backend for AWS Trainium」と題した記事を公開した。AWS Trainiumを既存のPyTorchコードからほぼそのまま使えるようにするネイティブバックエンド「TorchNeuron」のオープンソース公開を発表した内容だ。
NVIDIA GPU依存からの脱却は、AI開発コストの観点から長年の課題だ。AmazonはNVIDIA対抗として機械学習専用チップ「AWS Trainium」を開発してきたが、普及の壁となっていたのはエコシステムの問題だった。NVIDIA向けに書かれたPyTorchコードをTrainiumに移植するには相応の書き換えが必要で、移行コストが高かった。TorchNeuronはその障壁をほぼ取り除く試みであり、PyTorch Conference North America（2025年10月20日）にあわせてGitHubでApache 2.0ライセンスのオープンソースとして公開された。
.to('neuron') の一行だけでTrainiumが使える
TorchNeuronの核心は、デバイス指定を .to('neuron') に変えるだけで学習ループがTrainium上で動くという単純さにある。
# Run a training loop on Trainium
model = MyModel().to('neuron') # the only change
optimizer = torch.optim.AdamW(model.parameters())
for batch in dataloader:
optimizer.zero_grad()
loss = criterion(model(batch), targets)
loss.backward()
optimizer.step()
torch.accelerator（PyTorchが推奨するデバイス非依存API）を使って書かれたコードであれば、修正ゼロでそのまま動作する。torch.compile を使う場合は torch.compile(model, backend="neuron") と指定するか、torch.compiler.set_default_backend("neuron") を一度呼ぶだけだ。
この統合は、PyTorchコミュニティが整備した PrivateUse1デバイスバックエンド（サードパーティがデバイスをPyTorchに登録するための公式拡張ポイント）と torch.accelerator API の上に直接構築されている。torch.nn、torch.tensor、torch.autocast、torch.distributed、torch.profiler といった標準APIがそのまま neuron デバイス上で動作する。
2つの実行モード：研究向けEagerと本番向けtorch.compile
Eagerモードはオペレーションを即座に実行するため、Pythonの制御フロー、print文、ブレークポイントがコンパイルやトレースのステップなしで機能する。新たに導入された Adaptive Eagerモード（実験的機能）では、オペレーションのフュージョン機会を増やし、コード変更なしでコンパイル済みグラフに近いパフォーマンスを狙える。
torch.compile はモデルをグラフとしてキャプチャしてNeuronコンパイラが最適化する本番向けモードだ。内部的には TorchDynamo → AOTAutograd → Torch-MLIR → StableHLO → NEFF（Neuron Executable File Format：Trainium向けの実行バイナリ形式）というパイプラインで処理される。FlexAttention もネイティブサポートされている。
推論スタックとして vLLM Neuron（LLM向け）と vLLM Omni Neuron（画像・動画生成の拡散モデル向け）も同じ torch.compile サポート上に構築されている。
分散学習もそのまま動く
大規模モデルの学習に不可欠な分散APIも、標準PyTorchのものがそのまま使える。対応しているのは以下だ。
- DDP、FSDP2、SimpleFSDP
- テンソル並列、パイプライン並列
- FSDP2 + テンソル並列の組み合わせ、ZeRO
- DTensor、DeviceMesh（シャーディングと多次元プロセスグループ）
- コレクティブAPI（
all_gather、reduce_scatter、all_reduce）とP2P API（send、recv）
コレクティブはノード内ではTrainiumの専用コレクティブコア上で、ノード間では**Elastic Fabric Adapter（EFA）**（AWSの高スループット低レイテンシネットワークインターフェース）上で実行される。PyTorch APIへの変更は不要だ。
カスタムカーネルをPyTorchスタイルで書ける：NKI
Neuron Kernel Interface（NKI）を使うと、EagerモードとTorch.compileの両方で動くカスタムオペレーターを開発できる。カスタムオペレーターの登録はPyTorchの register_autograd スタイルで行い、バックワードパスも含めた定義がPyTorchコードとして完結する。
NKI-lib としてTrainium最適化済みの高性能オペレーターもオープンソースで提供されており、FlashAttentionやMoEブロックワイズカーネルなどLLM・マルチモーダルシナリオをカバーしている。
プロファイリング：フレームワークからハードウェアまで一気通貫
標準の torch.profiler.profile APIに NeuronConfig オブジェクトを組み合わせることで、CPUレベルのPyTorchディスパッチからTrainiumのハードウェアアクティビティまで、エンドツーエンドのトレースが取得できる。取得したプロファイルは Neuron Explorer UIで可視化でき、任意のイベントからソースコードの該当行へジャンプしたり、依存チェーンをたどってストールの原因を追ったりできる。さらにParquet形式での出力にも対応しており、AIエージェントが直接クエリする自動ボトルネック分析の用途にも活用できる。
Meta PyTorch EngineeringディレクターのJana van Greunen氏は次のようにコメントしている。
「PyTorchのビジョンはシンプルだ。開発者はハードウェアプラットフォームをまたいで一貫した体験を持てるべきである。AWS Trainiumのネイティブサポートはその選択肢を広げ、研究者とエンジニアが既に知っているEager実行と分散APIを使って実験・反復できるようにする」
TorchNeuronはTrainium2およびTrainium3をサポートしており、github.com/aws-neuron/torch-neuronxで公開されている。NVIDIA GPU向けのPyTorchコードを大規模な書き直しなしにAWS Trainiumへ持ち込める選択肢として、AI開発コスト最適化を検討しているチームにとって試す価値のある選択肢となりそうだ。
詳細はIntroducing TorchNeuron: Native PyTorch backend for AWS Trainiumを参照していただきたい。
