PyTorchコードをそのままAWS専用AIチップで動かせる「TorchNeuron」公開 — コード書き直しなしでNVIDIA以外の選択肢が現実的に

10月19日、About Amazonが「Amazon opens its AWS Trainium AI chips to every PyTorch developer, no rewrite needed」と題した記事を公開した。この記事では、AmazonがPyTorchコードをそのままAWS Trainiumチップ上で動かせるオープンソースライブラリ「TorchNeuron」を発表したことについて詳しく紹介されている。

コード変更なしでTrainiumが使える、その意味

AI開発者の85%がPyTorchでコードを書いている（元記事が引用するAmazon調査による数字）。だが、そのコードを動かせるハードウェアは長らく限られており、新しいチップに移行しようとすると数カ月かけて整備したコードの書き直しと再検証が必要だった。

Amazonが発表した**TorchNeuron**は、この問題を「設定の1行変更」で解決する。既存のPyTorchコードがそのままAWS Trainiumチップ上でネイティブに動作し、コードの書き直しは不要だ。

仕組みはシンプルで、PyTorchにはもともと新しいハードウェアが自身を登録するための標準的なプラグイン機構が備わっている。具体的には、PyTorchが提供するPrivateUse1バックエンド（カスタムデバイスをファーストクラスのtorch.deviceとして登録できる拡張API）を活用しており、TorchNeuronはこの標準機構に乗っかる形で実装されている。そのため、テスト・デバッグ・スケーリングに使っている既存ツールがそのまま動く。

Amazonのエンジニアが同一のAI学習スクリプトを従来チップとTrainiumで並べて実行したデモでは、どちらのチップを使うかを指定する設定のみが異なり、モデルの学習結果（精度・速度）は同等だったと報告されている。

なお、GitHubリポジトリ名はtorch-neuronxとなっているが、これはNeuron SDK向けPyTorchバックエンドの実装リポジトリであり、製品としての呼称「TorchNeuron」とは区別して理解するとよい。

TorchNeuronは10月20日のPyTorch Conference North Americaにて、Apache 2.0ライセンスのオープンソースとして公開される。

エンジニアにとって何が変わるか

コスト面での影響は大きい。スタートアップにとって最も高い費用は往々にしてエンジニアリング工数であり、新しいハードウェアへの適応に費やす1週間は、プロダクト改善に使えない1週間だ。

大企業にとっても恩恵がある。推論（inference）と学習（training）に同一ハードウェアを使い回せるため、オフピーク時に推論用の空きキャパシティをそのまま学習ジョブに転用できる。これまで専用チップが遊んでいた時間帯を有効活用できる。

さらに本質的な変化として、記事は「チップ選択の自由」に言及している。開発者がチップ間をペナルティなく移動できる環境が整えば、チップメーカーは開発者に選ばれるために競争しなければならなくなる。NVIDIAが事実上独占してきたAIチップ市場において、NVIDIA以外の選択肢を現実的なものにする取り組みだ。

すでに本番で使っている企業

TorchNeuronの早期採用企業として記事が挙げるのは以下だ：

Cartesia（リアルタイム音声AI）：標準の推論パイプラインをPyTorchコードの最小限の変更のみでTrainium上で動作させた
Lieber Institute for Brain Development（神経精神疾患の創薬研究機関）：統合失調症などの疾患向けに生物・化学・言語AIモデルを活用。計算コストが極めて高い分子動力学シミュレーションをTrainium上で実行している

記事によると、ほかにもロボティクス・映像・自動運転・LLM・金融など幅広い分野のスタートアップ・企業がTorchNeuronを使って開発中だとしている。

エコシステムへの統合状況

単体ライブラリとしての完成度だけでなく、周辺エコシステムへの統合も進んでいる：

**Hugging Face**：人気のAI開発ツール群でTrainiumのネイティブサポートが有効化され、Amazon固有のコードなしで利用可能
**TorchTitan**（PyTorch公式の分散学習フレームワーク）：FSDP（Fully Sharded Data Parallelism）を使った複数Trainiumインスタンスへのスケールアウトに対応
**vLLM**（推論高速化ツール）：Amazonがオープンソースへの主要コントリビューションを実施

また、PyTorch本体への上流コントリビューションとして、フレームワークを特定チップに縛るコードの除去にも取り組んでいる。Trainium2・Trainium3に対応し、数十万規模のAIモデルを毎ビルドごとにテストして品質を担保しているとのことだ。

背景：Amazonの10年間のカスタムシリコン投資

TorchNeuronはAmazonが単発で出してきたツールではない。2015年にチップ設計会社Annapurna Labsを買収して以来、Amazonはクラウドワークロードに特化した4つのチップファミリーを内製してきた。TorchNeuronはその成果を、PyTorchを書くすべての開発者に開放するための出口に位置づけられる。PyTorchエコシステムとの深い統合を通じて、AWS独自チップの採用障壁を下げるという戦略的な意図が、このライブラリには込められている。

詳細はAmazon opens its AWS Trainium AI chips to every PyTorch developer, no rewrite neededを参照していただきたい。