AWSが、自社のAI向けチップTrainiumをPyTorchの標準デバイスとして扱えるTorchNeuronを公開しました。10月20日のPyTorch Conference North AmericaでApache 2.0のオープンソースとして公開ベータに入っています。既存コードの.to('neuron')の1行を変えるだけで学習ループが動く点が売りです。Nvidia GPUとCUDAが事実上の標準となっている状況に、AWSがソフトウェア面から挑む動きになります。
背景と文脈
AIの開発現場では、ハードウェアよりソフトウェアの慣れが乗り換えを阻んできました。NvidiaはCUDA（同社GPUを汎用計算に使うためのプログラミング基盤）とPyTorchの連携を長年磨いてきました。その結果、PyTorchで書いたコードはNvidia GPUで動かすのが「何もしなくても動く」選択肢になっています。
AWSもTrainiumを以前から提供していました。ただ、従来はNeuron SDK固有の書き方や設定が必要でした。Hugging Faceも、これまではモデルをNeuron向けに作り替える手間があったと認めています。チップが安くても、移植の工数が割に合わないと判断する開発チームは少なくありませんでした。
Amazonによると、AIアプリの約85%はPyTorchで書かれています。この85%をそのまま取り込めれば、Trainiumの価格面の強みが初めて生きてきます。Amazonの発表でも、PyTorchを特定のチップに縛られない形にすること自体を目標に掲げています。
技術/ビジネス面
TorchNeuronの核は、PyTorch公式の拡張口であるPrivateUse1を使ってTrainiumを登録している点です。PrivateUse1は、PyTorch本体を改造せずに新しいハードウェアを「デバイス」として追加するための仕組みです。これによりtorch.nnやtorch.autocast、torch.distributed、torch.profilerがそのままneuronデバイス上で動きます。
実行モードは2つあります。1つはeager（命令を1つずつ即座に実行する通常モード）で、printやブレークポイントがそのまま使えるため研究やデバッグ向きです。もう1つはtorch.compile（モデル全体をグラフにまとめて最適化してから実行する仕組み）で、本番の速度を狙う場面向けです。eagerの遅さを補う実験的な「adaptive eager」も用意されました。
周辺ライブラリの対応も広めです。分散学習のFSDP2（モデルの重みを複数チップに分割して持たせる手法）やテンソル並列、パイプライン並列が標準APIのまま使えます。Hugging Face TransformersではTrainiumが正式なデバイスとなり、device_map="neuron"でモデルを載せられます。推論ではvLLM向けのプラグインも提供されます。
対応チップはTrainium2とTrainium3です。Hugging Face上の10万モデルで動作を試験し続けているとしています。一方、疎テンソルや動的な形状など一部機能は初版で未対応です。音声AIのDeepgramは、GPUの従来構成と比べて1ドルあたり約60%多くの音声を処理できたと述べています。
これからどうなるか
注目点は、AWSがPyTorch本体にもCUDA前提のコードを外す修正を送っていることです。これが進めば、AMDやAppleなど他社チップにも恩恵が及びます。逆に言えば、「どのチップでも同じコードが動く」状態に近づくほど、勝負は価格と調達のしやすさに移っていきます。
ただし、設定の切り替えだけで本当に同等の性能が出るかは、モデルごとに検証が必要です。演算子の未対応部分はCPUに処理が戻る場合があり、速度が落ちる原因になります。公開ベータの段階なので、本番投入は自分のモデルで試してから判断するのが安全でしょう。
開発者にとっての実利は、GPU不足やクラウド費用の見直しの際に「コードを書き直さずに試せる選択肢」が増えることです。torch.acceleratorのようなデバイス非依存のAPIで書いておけば、将来のチップ移行コストを小さく保てます。
まとめ
TorchNeuronは、TrainiumをPyTorchの普通のデバイスとして扱える公式バックエンドです。eagerとtorch.compileの両方に対応し、Hugging FaceやvLLMとも連携します。CUDA一強の構図がすぐ崩れるわけではありません。それでも、乗り換えの壁をソフトウェア側から下げる一歩になりそうです。
参考リンク
- Introducing TorchNeuron: Native PyTorch backend for AWS Trainium（PyTorch Blog）
- TorchNeuron announcement（About Amazon）
- aws-neuron/torch-neuronx（GitHub）
- Amazon is taking its biggest shot at Nvidia yet with new TorchNeuron dev tool（TechCrunch）
アイキャッチ画像: Photo by Akshar Dave🌻 on Unsplash
