AIAIExplainer NO. ---

TorchTPUとTorchNeuronを解説：Nvidia CUDAの先を目指すPyTorch

GoogleとAmazonが同じ週に、自社AIチップ向けのネイティブなPyTorchバックエンドを公開しました。バックエンドの役割、開発者が今NvidiaのGPUから移行できるもの、そしてチップの仕様よりソフトウェアが重要になり得る理由を解説します。

By Rowan Ashby Published 21.10.26 12:31 UTC 9 min read

チップとヒートシンクを備えたGoogle Tensor Processing Unit 3.0ボードのクローズアップ — Zinskauf / Wikimedia Commons / CC BY-SA 4.0 · Source

要点

GoogleはTorchTPUをGitHub上でオープンソースプロジェクトとして公開しました。自社のTensor Processing Unit（TPU）向けのネイティブなPyTorchスタックで、MetaのPyTorchチームと共同で開発されました。
その前日、AmazonはTrainium2とTrainium3向けの、Apache 2.0ライセンスのネイティブPyTorchバックエンド「TorchNeuron」を発表しました。10月20日にパブリックベータが始まりました。
どちらのツールも、新しいハードウェア向けのPyTorch公式拡張ポイントであるPrivateUse1に接続します。実際には、デバイス設定を変えるだけでモデルを移行できる場合もあります。
どちらもまだ完成していません。Amazonは未対応の演算子や、動的形状を含む非対応機能を挙げています。Google自身のエンジニアも、入力サイズが変わると再コンパイルが発生すると述べています。
競争の焦点はシリコンだけでなくソフトウェアにもあります。Nvidiaの優位は、PyTorchがおおむねその上に構築されてきたプログラミング基盤CUDAに大きく支えられています。

過去10年近く、PyTorchでモデルを書くAI開発者は、実質的にNvidiaのハードウェア向けにコードを書いていました。その状況が変わり始めています。今週のわずか2日間に、二大クラウド事業者が、PyTorchのコードをほとんど、あるいはまったく書き換えずに自社製AIチップで動かせるソフトウェアを公開しました。先行したのはAmazonで、Trainiumチップ向けのTorchNeuronを発表しました。続いてGoogleが、Tensor Processing Unit向けのTorchTPUをオープンソースで公開しました。TechCrunchはこの2つを、AIアプリ開発者の間でのNvidiaの「牙城」に対する共同攻勢だと表現しました。

この解説記事では、PyTorchのバックエンドとは何か、各社が何を公開したのか、開発者が今移行できるもの・できないもの、そしてAIハードウェア競争の行方をチップの仕様以上にソフトウェア層が左右し得る理由を取り上げます。

まず、PyTorchとは何か、なぜ重要なのか

PyTorchは、ニューラルネットワークを構築・学習するためのオープンソースのフレームワークです。TechCrunchによると、Meta（当時はFacebook）が2016年に開発し、2022年にLinux Foundationへ移管しました。現在の研究で使われる大規模言語モデル、画像生成モデル、音声システムの大半はPyTorchのコードで定義されています。Amazonは、AI研究者・開発者の85%がPyTorchでコードを書いているとしています。この数字はAmazon自身の発表によるもので、独立した集計ではありません。

問題は、PyTorchがNvidiaのGPUの上で育ってきたことです。その機能、テスト、性能向上の工夫の多くは、NvidiaのプログラミングプラットフォームであるCUDAを前提としていました。NvidiaはCUDAを「GPUコンピューティングの基盤」と説明しており、コンパイラ、ライブラリ、デバッグツールをまとめたツールキットです。PyTorchを別のチップで動かしたいチームは、通常、専用のラッパーを使うか、コードを修正するか、性能低下を受け入れる必要がありました。

「バックエンド」とは、わかりやすく言うと何か

わかりやすいたとえがあります。PyTorchが車だとすれば、バックエンドはその下にあるエンジンです。開発者は、レイヤー、テンソル、オプティマイザー、学習ループといった慣れた操作系で運転します。バックエンドはその指示を、特定のチップが実行できる処理に変換します。Amazonも同様のたとえを使い、TorchNeuronを「車を分解する」ことなくエンジンを載せ替える方法だと説明しています。

技術的な仕掛けはPrivateUse1と呼ばれる仕組みです。PyTorchのドキュメントによると、これはハードウェアメーカーがPyTorchのコアファイルを編集せずに新しいデバイスを登録できる予約枠です。これを使うバックエンドは、PyTorchの基本演算の独自実装を登録します。まだ実装していない演算はCPUにフォールバックできます。デバイス名も変更できるため、開発者は汎用的な名前ではなくdevice="neuron"やTPUのデバイス名を書くことになります。

新しい2つのツールはいずれもこの方式を採用しています。これが重要なのは、Googleの従来のPyTorch/XLAプロジェクトなど、以前の取り組みがPyTorchの内部ではなく横に並ぶ形だったためです。Googleのエンジニアは2026年4月の開発者向けブログ記事で、TorchTPUの目標を「PyTorchらしく感じられること」という簡潔な言葉で表しました。

Amazonが公開したもの：TorchNeuron

AmazonのエンジニアはPyTorchの公式ブログでTorchNeuronを説明しました。最も単純な場合、変更は1行だけで、モデルを'neuron'デバイスに移すことです。PyTorchが推奨するデバイス非依存のインターフェースtorch.acceleratorで書かれたコードは、変更なしで動作します。コンパイルするモデルでは、Neuronコンパイラのバックエンドを一度指定するだけで済みます。

この記事とAWS Neuronのドキュメントによると、TorchNeuronは次に対応しています。

演算を即時に実行するイーガーモード（研究・デバッグ向け）と、演算を融合してオーバーヘッドを減らす試験的な「adaptive eager」モード
モデルをグラフとして取り込み、本番環境向けの性能を得るためにAmazonのNeuronコンパイラに渡すtorch.compile
DDP、FSDP2、テンソル並列・パイプライン並列、DTensorなどの標準的な分散学習ツール
モデルの最も負荷の高い部分を手作業で調整するためのカーネル言語NKIと、既製カーネルをそろえたオープンソースのライブラリ

Hugging Faceによると、Trainiumは同社のTransformersとDiffusersライブラリのネイティブバックエンドになりました。Amazonはまた、「ハードコードされたCUDA前提」を取り除くため、PyTorch本体にプルリクエストを送ったと報告しています。TechCrunchはこれを、PyTorchを標準でよりハードウェア中立にする取り組みだと説明しました。

Googleが公開したもの：TorchTPU

GoogleはPyTorch Conference 2026で、Google Cloudのブログ記事を通じてオープンソース化を発表しました。執筆者は、AI・コンピューティングインフラ担当バイスプレジデント兼ゼネラルマネージャーのMark Lohmeyer氏と、エンジニアリング担当バイスプレジデントのBill Jia氏です。コードはGitHubのgoogle-pytorch組織で公開されています。Googleは、既存のPyTorchワークロードを「使い慣れたAPIを使い、最小限のコード変更で」TPU上で実行できるとしています。

TorchNeuronと同様に、TorchTPUもデバッグ用のイーガーモードと本番用のtorch.compile経路を備えています。Googleの4月の技術記事によると、コンパイル経路ではPyTorch標準のInductorコンパイラではなく、Googleが長年使ってきたTPU用コンパイラXLAを使います。記事では3種類のイーガーモードも紹介されており、その一つである「fused eager」モードは、Googleによると、1演算ずつ実行する厳格モードより50%から100%超高速に動作しました。

Googleの発表は、より広いエコシステムに重点を置いています。推論エンジンのvLLMとSGLang、学習ライブラリのTorchTitan、オーケストレーションフレームワークのRayでのネイティブ対応を挙げています。Hugging Face Transformers v5では、device_map="tpu"と設定するだけでTPUを指定できます。Googleはまた、導入、診断、チューニングの自動化を支援する「TPU agent suite」をプレビュー版で公開しました。さらに、オンデバイス向けランタイムのLiteRTを更新し、PyTorchモデルをスマートフォンやノートPC向けに変換できるようにしました。

比較

項目	Amazon TorchNeuron	Google TorchTPU
対象チップ	Trainium2、Trainium3	GoogleのTPU（大規模テストではIronwood/v7xに言及）
状況	2026年10月20日からパブリックベータ	GitHubでオープンソース公開、agent suiteはプレビュー
ライセンス	Apache 2.0	Googleの発表には記載なし
統合方式	PyTorch PrivateUse1	PyTorch PrivateUse1
コンパイラ経路	TorchDynamo → StableHLO → Neuronコンパイラ	TorchDynamo → StableHLO → XLA
ベンダーの主な数値	2026年8月の実行でアップストリームのテスト結果約15万4,000件、うち約6万9,000件がスキップ（そのうち約2万2,000件は演算子が未実装のため）	Qwen 3 8Bでモデル FLOPs 使用率最大52%、Ironwood TPU 8,192基までほぼ線形にスケール

出典：PyTorchブログ（TorchNeuron、2026年10月20日更新）、AWS Neuronドキュメント、Google Cloudブログ（TorchTPU、2026年10月）、Google Developers Blog（2026年4月）。性能とテストの数値はすべて企業が公表したもので、独立したベンチマークによる検証は行われていません。

開発者が今移行できるもの、できないもの

両社のうたい文句は「書き換え不要」です。しかし、各社自身のドキュメントに基づく実態は「多くの一般的なケースでは書き換え不要」というものです。

少ない手間で移行できそうなもの

Hugging Faceのライブラリで読み込む標準的なTransformerモデル（両社ともアップストリームで対応済み）
TorchTitanや一般的なFSDP方式のシャーディングで構築された学習ジョブ
両クラウドが対応しているvLLMによるLLMの推論
cudaをハードコードせず、PyTorchのデバイス非依存APIをすでに使っているコード

まだ難しい、または未対応のもの

独自のCUDAカーネル。手書きのNvidia向けカーネルは他のチップでは動きません。ベンダーのカーネルライブラリに頼るか、NKI（Amazon）やPallas（Google）で書き直す必要があります。
動的形状。Amazonは「動的形状とシンボリック形状」を、TorchNeuronが最初のリリースで対応しないことにした機能に挙げています。Googleのエンジニアは、系列長やバッチサイズが変わるとTPUで再コンパイルが発生すると書き、動的形状への対応強化をロードマップの項目に挙げました。
あまり使われない演算。Amazonの非対応機能一覧には、スパーステンソル、NestedTensor、フォワードモードの自動微分、一部のアップサンプリングモードが含まれます。イーガーモードでは、Neuron版のない演算は警告付きでCPU上で実行され、モデルが大幅に遅くなることがあります。
CUDA固有のツール。CUDA graphsやNCCL固有の集団通信などの機能は、その性質上Nvidia専用です。

また、「動く」ことと「速く動く」ことは同じではありません。両社とも、本番レベルの速度にはコンパイルモードと調整済みカーネルが必要だと強調しています。顧客の主張はワークロードによって異なります。Deepgramの最高技術責任者（CTO）であるAdam Sypniewski氏は、Amazonの記事の中で、Trainiumは同社の既存のGPU基準と比べて1ドルあたり約60%多くの入力音声秒数を処理したと述べ、それは「当社が評価したベンチマークにおいて」だとしています。これは一顧客の測定結果であり、一般的な結果として受け取るべきではありません。

チップの仕様よりソフトウェア層が重要な理由

チップの単純な性能値は公表も比較も容易です。一方、乗り換えコストは見えにくいものです。企業の本当の投資は、モデルのコード、学習レシピ、デバッグの習慣、監視ツール、そしてそれらに精通したエンジニアにあります。AmazonのPyTorchブログ記事はこの点を率直に述べています。AIの運用コストには、ワークロードを動かし、ライブラリが変わっても動かし続けるために必要なエンジニアリング時間が含まれる、というものです。

Nvidiaの地位は、そのエンジニアリング時間を自社ハードウェア上で最も小さくすることに支えられてきました。TechCrunchが指摘するように、PyTorchの多くの機能は、NvidiaのGPUがすでにAI研究の標準だったため、それに合わせて作られてきました。仕様で勝るライバルのチップでも、移行に数カ月の移植作業がかかるなら、負けることがあり得ます。

ネイティブバックエンドは、そのコストに直接切り込みます。あらゆるベンチマークでNvidiaを上回る必要はありません。価格、入手性、供給能力が決め手になるほど、乗り換えを安くすればよいのです。Amazonにとっての利害は大きく、同社は自社製シリコン事業の年間収益ランレートが250億ドルを超えたとしています。一方、GoogleはGoogle Cloudを通じて外部顧客にTPUを提供しています。Metaは、TorchTPUによって自社のPyTorchワークロードをこれらのチップに展開できるとしています。

PyTorch自体の内部でも、目立たない変化が起きています。隠れたCUDA前提を取り除くアップストリームの変更はどれも、AmazonやGoogleのチップだけでなく、Nvidia以外のすべてのチップに役立ちます。MetaのPyTorchエンジニアリング担当ディレクターであるJana van Greunen氏は、両社の発表で発言が引用されています。Amazonの記事では、開発者は「ハードウェアプラットフォームをまたいで一貫した体験を得られるべきだ」と述べています。

影響と意味

AI開発者にとっての実際的な効果は、選択肢が増えることです。チームは同じPyTorchモデルをNvidiaのGPU、Trainium、TPUでより簡単に試し、自社のワークロードでコストと速度を比較できます。最初に試すのに適しているのは、大幅にカスタマイズされた研究用コードではなく、対応ライブラリ上の標準的なモデルです。

Nvidiaにとって、今週の動きがソフトウェア面での優位を終わらせるわけではありません。TechCrunchは、開発者が自動的にNvidiaのハードウェアから離れるわけではないと指摘しました。Nvidiaのツールは成熟しており、CUDA向けに書かれた独自カーネルは依然として手作業で移植する必要があります。AmazonとGoogleのどちらのツールも、まだ汎用的な代替手段ではありません。Amazonは今回のリリースをベータ版と呼び、Googleはスタックの一部をプレビュー版としています。

変わったのは方向性です。2つのハイパースケーラーが、自社独自のフレームワークではなくPyTorchを自社チップへの入り口として位置づけるようになりました。両社とも、PyTorch本体に修正を還元しているとしています。こうした取り組みが続けば、AIチームにとっての重要な問いは「このチップで自分のコードは動くか」から「どのチップが自分のコードを最も安く動かせるか」へと移るかもしれません。その変化は、AIハードウェア市場全体の価格と利益率に圧力をかけることになります。残る課題がどれだけ早く解消されるか、独立したベンチマークでどう比較されるか、Nvidiaがどう対応するかなど、まだわからない点がいくつもあります。

よくある質問

PyTorchのバックエンドとは何ですか？

バックエンドとは、PyTorchの演算を特定のチップが実行できる処理に変換する層です。TorchTPUとTorchNeuronは、新しいハードウェア向けのPyTorch公式拡張ポイントであるPrivateUse1を通じてGoogleのTPUとAWS Trainiumを登録するため、開発者はデバイス設定を変えるだけでこれらを利用できます。