Googleは2026年10月20日、AI開発フレームワーク「PyTorch」を自社のAIチップ「TPU」で動作させるための実行基盤「TorchTPU」をオープンソースで公開した。PyTorch Conference 2026で発表されたもので、既存のPyTorchコードを大幅に書き換えることなく、TPUで学習や推論を実行できるようにする。時期を同じくして、Amazon Web Services（AWS）も、独自のAIチップ「Trainium」向けに開発した「TorchNeuron」の公開ベータ版を提供開始した。両社の取り組みに共通するのは、開発者が使い慣れたAIモデルやソフトウェアを維持しながら、NVIDIA製GPU以外のハードウェアを選びやすくすることだ。ただし、既存のコードを動かせることと、実際のサービスで十分な性能を引き出せることは別であり、コンパイルや対応演算に関する制約を理解する必要がある。

AD

GoogleがTorchTPUを公開、既存のPyTorchコードをTPUで実行可能に

TorchTPUは、GoogleがMetaのPyTorch開発チームと共同開発した、TPU向けのネイティブバックエンドだ。PyTorchでは、開発者がモデルの構造や計算処理を記述し、バックエンドが実際のハードウェアで演算を実行する。Googleは10月20日の公式ブログで、従来のPyTorch APIを使いながら、わずかなコード変更で既存の処理をTPUに移行できると説明した。これにより、開発者がTPU専用のプログラミング手法を一から習得する負担を減らす狙いがある。

TorchTPUは、開発やデバッグに適した逐次実行モード（Eager Mode）と、性能を重視したコンパイル実行の両方に対応する。逐次実行では個々の演算を順番に実行し、モデルの動作を確認できる。一方、本番運用ではtorch.compileを使って計算処理をまとめて最適化し、TPUの演算能力を効率よく利用できるようにする。開発時の柔軟性を維持しながら、運用段階ではより高い処理性能を目指せる構成だ。

主要なAI開発ライブラリとの連携も進んでいる。Hugging FaceのTransformers v5では、モデルの配置先としてdevice_map="tpu"を指定することで、TPUを利用できるようになった。画像生成などに使われるDiffusersもTPU上での逐次実行やtorch.compile、テンソル並列に対応する。Googleはさらに、TransformersのテストをTPU上で毎晩自動実行する仕組みを構築した。これは単に新しいハードウェアへの対応を追加するだけでなく、Transformersの更新後も互換性を継続的に確認するための取り組みだ。

AIモデルの推論を担うソフトウェアにもTorchTPUの対応が広がっている。GoogleはInferactと協力してvLLM向けのTPUバックエンドを開発し、RadixArkとはSGLang向けのバックエンドを整備した。どちらも既存のワーカー構成やスケジューラー、APIサーバーなどを維持しながら、ハードウェアに近い演算処理をTPU向けに最適化する構成を採る。これにより、AIモデルを多数の利用者へ配信するための仕組みを全面的に作り直さずに、TPUを導入しやすくなる。

学習分野では、大規模分散学習向けのTorchTitanと連携し、複数のTPUを使ったモデルの学習に対応する。分散処理フレームワークのRayにもTorchTPUが統合され、CPU、GPU、TPUを組み合わせたワークフローを一元的に管理できるようになった。さらにGoogleは、TPUへの移行支援や動作診断、推論性能の調整を補助する「TPU agent suite」もGitHubで早期プレビューとして公開した。ただし、こちらは正式な一般提供とは異なる開発段階にあり、利用できる機能や成熟度を分けて評価する必要がある。

AWSもTorchNeuronを公開、TrainiumをPyTorchの標準デバイスとして利用

AWSも同様に、独自のAIチップをPyTorchから直接利用できる環境の整備を進めている。2026年10月19日付のPyTorch公式ブログでの発表によると、AWSは翌20日からTrainium2とTrainium3向けの「TorchNeuron」を公開ベータ版として提供開始した。TorchNeuronはApache License 2.0で公開されるオープンソースのバックエンドで、TrainiumをPyTorchの標準的なデバイスとして扱えるようにする。既存のPyTorchコードを大きく変更することなく、TrainiumでAIモデルを実行することを目指している。

TorchNeuronの基盤となるのは、PyTorchが独自ハードウェア向けに用意した拡張機構「PrivateUse1」だ。TorchNeuronはこの機構を使ってneuronというデバイスを登録し、PyTorchの標準APIからTrainiumを利用できるようにする。対応するモデルであれば、モデルや入力データを.to("neuron")で配置することで、逐次実行による学習や推論を行える。コンパイルによって性能を高める場合は、torch.compile(model, backend="neuron")のようにNeuron用バックエンドを指定するか、あらかじめ既定のバックエンドとして登録する。

これまでもTrainiumでPyTorchモデルを動かす方法は存在したが、利用には特有の制約があった。Hugging Faceの解説によると、従来のOptimum NeuronはPyTorch XLAを基盤としており、遅延実行や固定された入力形状を前提とするなど、開発者がコンパイラの動作を意識する必要があった。Optimum Neuronはこうした複雑さを吸収するため、モデル専用の実装や学習コードを独自に保守していた。TorchNeuronではTrainiumをPyTorchの標準デバイスとして扱えるため、TransformersやDiffusers、Accelerate、PEFT、TRLなどを直接利用できるようになり、Hugging Faceも今後はOptimum Neuronを介さずに各ライブラリを使うことを推奨している。

この変更の意義は、デバイスの指定方法が簡単になることだけではない。AIモデルごとにTrainium専用の実装を維持する必要が減れば、普段利用しているライブラリの更新を取り込みながら、別のハードウェアで動作を試しやすくなる。AWSはさらに、PyTorch本体のテストなどに残るCUDAを前提とした処理についても、特定のハードウェアに依存しない形へ改める修正を提案している。ただし、TorchNeuronによってCUDA専用のコードがすべて自動的にTrainiumへ移植されるわけではない。PyTorchの標準APIを利用する部分を増やし、ハードウェアごとの専用実装を減らしていく取り組みといえる。

AD

GoogleとAWSの違いは、本番運用での制約にある

TorchTPUとTorchNeuronは、どちらもPyTorchの標準APIを通じて独自のAIチップを利用しやすくする技術だ。ただし、対応するハードウェアや実行方法、現時点で公表されている制約には違いがある。両社の発表とAWS Neuronのリリースノートを基に、主な特徴を比較すると次のようになる。

比較項目 Google「TorchTPU」 AWS「TorchNeuron」
公開状況と対象 TPU向けにオープンソース公開 Trainium2・Trainium3向けの公開ベータ
基本的な利用方法 既存のPyTorch APIを利用。Transformers v5ではdevice_map="tpu"を指定可能 モデルと入力をneuronデバイスに配置し、PyTorchの標準APIを利用
開発・デバッグ 逐次実行モードに対応 逐次実行モードに対応
性能最適化 torch.compileに対応。Pallasによる専用カーネルも利用可能 torch.compileに対応。NKIによる専用カーネルも作成可能
現時点で明示された制約・条件 TPU agent suiteは早期プレビュー 初版のコンパイルは固定された入力形状が前提。逐次実行では未対応演算をCPUで代替実行可能

この比較は、両社が2026年10月に公表した機能と提供条件を整理したものであり、同じモデルを使った性能や価格の比較ではない。また、AWSのTorchNeuronで明示されている制約が、そのままGoogleのTorchTPUにも当てはまるわけではない。両社に共通するのは、アプリケーション側で利用するAPIをできるだけ維持し、ハードウェア固有の処理をバックエンドや専用カーネルに任せるという設計方針だ。これにより別のAIチップを試す際の障壁は下がるが、本番環境で十分な性能を引き出すための調整まで不要になるわけではない。

対応モデルは89％に拡大、ただし性能評価には注意が必要

TorchNeuronによる改善を示す具体的な数字として、Hugging Faceは対応モデルの拡大を挙げている。同社によると、利用頻度の高いテキスト生成モデルの上位100アーキテクチャのうち、追加の専用実装なしで利用できる割合は、Optimum Neuron v0.4.5の46％からTorchNeuronでは89％へ上昇した。これは従来より幅広いモデルをTrainiumで試せるようになったことを意味する。ただし、対象は特定の100アーキテクチャに限られており、あらゆるAIモデルの89％が動作するという意味ではない。また、対応することと、高い処理速度で動作することも区別する必要がある。

特に注意が必要なのが、コンパイル時の入力形状に関する制約だ。AWSのTorchNeuronリリースノートによると、初版のtorch.compileではdynamic=Falseを指定する必要があり、動的な入力形状にはまだネイティブ対応していない。そのため、入力データの形状が変わると新たなコンパイルが必要になる場合がある。AWSは対策として、入力をパディングで一定の長さにそろえたり、近いサイズの入力をグループ化したりして、コンパイル対象となる形状の種類を減らすことを推奨している。

例えば、長さの異なる文章を大量に処理するAIサービスでは、文章ごとに入力形状が変化しやすい。そのたびにコンパイルが発生すれば、推論処理そのものが高速でも、利用者が応答を受け取るまでの時間は長くなり得る。Hugging Faceが公開した推論のサンプルでも、頻繁な再コンパイルを避けるために固定サイズのキャッシュを使用する構成が推奨されている。学習用のサンプルでは、PyTorch Dynamoによる入力形状の自動的な動的処理を無効にし、現時点で対応していない非同期の重み読み込みも停止している。標準のモデル定義や学習コードを使えるようになっても、Trainiumで安定して動かすには実行時の設定を確認する必要がある。

対応していない演算の扱いも、性能評価で見落とせない点だ。TorchNeuronの逐次実行モードでは、Trainium上で実行できない演算をCPUに切り替えて処理する仕組みが用意されている。これによってモデルを最後まで実行できる場合がある一方、CPUへの切り替えが発生すれば、データ転送や処理待ちが性能に影響する可能性がある。AWSは、代替実行された演算をtorch_neuronx.get_fallback_ops()で確認できるとしている。モデルが正常に動作したとしても、すべての演算がTrainiumで実行されたとは限らないため、実行先を確認せずに得られた速度をAIチップ自体の性能として評価するのは適切ではない。

さらに、最初の実行で発生するコンパイル時間と、コンパイル結果をキャッシュから再利用できる2回目以降の実行時間も分けて測定する必要がある。GoogleはPallas、AWSはNeuron Kernel Interface（NKI）を通じて、それぞれのAIチップに最適化された専用カーネルを利用できるようにしている。つまり、PyTorchの共通APIを使っていても、高い性能を追求する段階ではハードウェア固有の最適化が依然として重要になる。導入を検討する開発者は、対応する演算の範囲に加え、コンパイルの頻度やキャッシュの再利用状況、実際の処理時間を確認する必要がある。

GoogleとAWSによる今回の取り組みは、AI開発におけるハードウェア選択の自由度を広げる動きとして注目される。これまではNVIDIAのCUDAを中心に構築されてきたソフトウェア環境が大きな強みとなっていたが、TorchTPUとTorchNeuronは、広く利用されているPyTorchのコードやライブラリを維持したまま、別のAIチップを試せる環境を整えつつある。ただし、導入の容易さだけで、実際の性能や運用コストにおける優劣が決まるわけではない。今後は、同じモデルと負荷条件での処理速度、コンパイル時間、実行可能な演算の範囲、運用コストを比較することで、TPUやTrainiumが既存のGPU環境に対してどこまで有力な選択肢となるかが明らかになっていくだろう。