powered by TechFeed
表示モード
Amazon

AWSがPyTorchコードを「ほぼ書き換えなし」で自社AIチップに移植できる「TorchNeuron」を発表 — `.to("cuda")`を`.to("neuron")`に変えるだけ、GoogleのTPU対応バックエンドとの競争も激化

11月9日、InfoQが「AWS Launches TorchNeuron to Run PyTorch Models on AWS Trainium Without Modifications」と題した記事を公開した。AWSがPyTorchコードをほぼ無修正でTrainiumチップ上で実行できるネイティブバックエンド「TorchNeuron」を発表したことを報じている。

11月9日、InfoQが「AWS Launches TorchNeuron to Run PyTorch Models on AWS Trainium Without Modifications」と題した記事を公開した。AWSがPyTorchコードをほぼ無修正でTrainiumチップ上で実行できるネイティブバックエンド「TorchNeuron」を発表したことを報じている。


.to('cuda') を .to('neuron') に書き換えるだけ

AWSはPyTorch Conference North America 2026にて、TorchNeuronを発表した。既存のPyTorchコードをAWS Trainiumチップ上でほぼ無修正で動かすことができるネイティブPyTorchバックエンドだ。

AWSによれば、ハードウェアの切り替えに必要な変更は設定1箇所のみ。PyTorchの自動デバイス検出(torch.accelerator)を利用している場合、必要な変更は.to('cuda')を.to('neuron')に書き換えるだけで済む。

現在、AIの研究者や開発者の85%がPyTorchでコードを書いている。これまで、そのコードを動かせるハードウェアはコード自体よりもはるかに限られており、新しいチップへの対応には、何ヶ月もかけて完成させた成果物の書き直しと再検証が必要だった。

新しいチップへの移行コストが開発者の大きな負担となっていた背景を踏まえると、この「1行の変更」というアプローチは実用上の意味が大きい。

分散学習・カスタムカーネルにも対応

TorchNeuronは単なるデバイス置き換えにとどまらない。DDP、FSDP2、DTensorなど、PyTorchの分散学習・並列化APIをネイティブサポートしている。また、torch_neuronx.Streamおよびtorch_neuronx.Eventを通じて、Torchのストリームとイベントもネイティブに扱える。

さらに、Neuron Kernel Interface(NKI)を使ったカスタムオペレータの開発・最適化・実行も可能だ。

Adaptive EagerやTorch.compileを超えたパフォーマンス制御として、標準PyTorchオペレーション(ATen ops)の代わりにNKIカーネルをモデルコード内で直接呼び出すことができる。

TorchNeuronのドキュメントによれば、バックエンドはすでに実用段階にあり、Amazon Annapurna Labsの副社長Colin Braceが率いるチームが最も人気の高い10万モデルと800以上のMLアーキテクチャを検証済みだという。

訓練と推論のリソースを柔軟に共有できる

TorchNeuronが提供するもう一つの価値は、ハードウェアリソースの共有だ。同一のTrainiumインスタンス上で、訓練ワークロードと推論ワークロードを混在させることができる。ピークオフ時には推論用リソースを訓練に回すといった柔軟な運用が可能になる。

すでにいくつかの組織が採用しており、音声AI推論にCartesia、ドラッグターゲット特定と創薬探索の研究にLieber Institute for Brain Developmentが活用している。

「本当のテストはプロダクションでのコスト」——業界の声

LinkedInでのAWSの発表に対し、Liner Notes ConsultingのCEO Yasuo Maedaはこう指摘した。

本当のテストは、開発者がプロダクションで競争力のあるパフォーマンスとTCO(総所有コスト)を達成できるかどうかだ。AIインフラの競争は、最高のチップを作ることだけではない。いかにスケールで使いやすくするかでもある。

別のユーザー(元記事では「Pratham pratham」と表記されており、同一人物の名前が重複している可能性がある)も同様に、マイグレーションの摩擦を減らすことはアクセラレータのロックイン回避に重要としつつも、ワークロード固有のパフォーマンス、トークンあたりのコスト、再現性が真の評価基準になると強調している。

Googleの「TorchTPU」との競争

TorchNeuronの発表は、GoogleがTPU向けネイティブPyTorchバックエンドとしてTorchTPUを発表した直後のタイミングだ。PyTorchエコシステムを自社ハードウェアに取り込む動きが、クラウド大手間で加速している。


詳細はAWS Launches TorchNeuron to Run PyTorch Models on AWS Trainium Without Modificationsを参照していただきたい。