PyTorchをTPUとTrainiumで：Nvidia依存はどこまで減るか

TorchTPUとTorchNeuronの登場で、PyTorchをTPUやTrainiumで動かすのはデバイス指定の変更だけになりました。数値、制約、実際のコスト、誰が先に試すべきかを分析します。

本文

これまでPyTorchをGoogleのTPUやAWSのTrainiumで動かすには、中間ライブラリとかなりの修正が必要でした。2026年10月19〜20日、PyTorch Conferenceの直前に、両社はネイティブなPyTorchバックエンドをオープンソースで公開しました。TorchTPU（Google、MetaのPyTorchチームと共同開発）とTorchNeuron（AWS、パブリックベータ）です。共通の狙いは、Nvidia GPU向けに書いたPyTorchコードを、ごくわずかな変更で自社チップ上で動かすことです。

• 書き直しではなくデバイスの切り替え。 TorchNeuronについてAWSは、必要な変更はmodel.to('neuron')だけだとしています。TorchTPUについてGoogleは「最小限のコード変更」とし、Hugging Face Transformers v5ではdevice_map="tpu"を指定するだけです。
• 2つの実行モードに対応。 どちらもeagerモード（1演算ずつ実行し、デバッグしやすい）とtorch.compile（本番向けにグラフ全体をコンパイル）をサポートします。
• 分散学習を標準で提供。 DDP、FSDP2、DTensorは両方で動作し、TorchNeuronはテンソル並列とパイプライン並列にも対応します。
• 推論エコシステムと連携。 TorchTPUはvLLMとSGLang向けのバックエンドを備え、TorchNeuronはvLLM NeuronのほかTransformers v5、Diffusers、trlに対応します。
• 両プロジェクトともApache 2.0でGitHubに公開されています。

分析

分析：コード1行でNvidia GPUから離れられるのか

結論から言えば、試すには十分ですが、移行するにはまだ足りません。本質的な変化は統合の仕組みにあり、実際の節約額は動かすモデルによって大きく変わります。

従来と何が違うのか

最大の違いはPrivateUse1です。これは新しいハードウェア向けにPyTorchが用意した公式の拡張ポイントです。TPUやTrainium上のテンソルはラッパーのない通常のPyTorchテンソルになります。そのため、Hugging Face、vLLM、TorchTitanなどのライブラリはチップごとに別の分岐を持つ必要がなくなります。

なぜ書き直さずに動くのか

torch.compileを呼ぶと、PyTorchはTorchDynamoで計算グラフを取得します。TorchTPUはそれをGPU向けコンパイラではなくStableHLOに変換し、XLAがTPU向けに最適化します。TorchNeuronもTorch-MLIR、StableHLO、Neuronコンパイラという似た経路をたどります。開発者は使い慣れたPyTorchを書き続け、チップへの変換は下の層で行われます。

TorchTPUの「Fused Eager」モードは、多数の小さな演算を自動的に大きなまとまりにしてからチップに送ります。Googleは1演算ずつの実行と比べて50%から100%超の向上を報告しています。ただしこれはメーカーの数値で、独立した測定はまだありません。

誰にどんな影響があるのか

• MLエンジニア： 同じ学習ループをGPU、TPU、Trainiumで動かして比較できます。独自カーネルには引き続きPallas（TPU）やNKI（Trainium）の習得が必要です。
• 技術マネージャー： GPU価格交渉での選択肢が増えます。モデルが複数のチップで動けば、クラウド契約がCUDAに縛られなくなります。
• 中小企業： 多くはAPI経由でモデルを使うため、影響は間接的です。Google CloudやAWSでオープンモデル（Qwen、Llama）を自社運用するチームは直接の恩恵を受けます。
• マーケターや製品チーム： すぐには何も変わりません。恩恵があるとすれば、これらのチップを使う提供者のAPI価格低下という形になります。

実際のコストはいくらか

どちらの企業も新しい料金は公表していません。唯一のコスト数値はDeepgramのもので、既存のGPUと比べて1ドルあたり約60%多く音声を処理できたというものです。下表はこの数値を2つの規模に当てはめた仮定の例で、移行コストを3万ドル（エンジニア2人で1〜2か月程度）と仮定しています。

表からわかるのは、費用が小さければチップの切り替えはほとんど割に合わないということです。費用が大きく、モデルが安定していれば、採算が合い始めます。60%という数値は1社・1種類の処理（音声）の結果であり、皆さんのモデルではそれより良い場合も悪い場合も、まったく効果がない場合もあります。

リスクと限界

• 演算子の対応範囲が不十分です。 2026年8月のPyTorchテストでは、TorchNeuronに対応する演算子がないため約2万2,000件のテストがスキップされました。初版は動的形状、スパーステンソル、NestedTensorに未対応で、FP8はプライベートベータです。
• 動的形状は依然としてTPUの弱点です。 Googleは、系列長やバッチサイズが変わると再コンパイルが発生し、初回の反復は遅延が大きいと認めています。プロンプトの長さが絶えず変わるチャットボットは、まさに影響を受ける処理です。
• ロックインは場所を変えるだけです。 TPUはGoogle Cloudにしかなく、TrainiumはAWSにしかありません。CUDAから離れることが、特定のクラウドへのより強い依存を意味する場合があります。
• CUDA専用のコードは動きません。 CUDA Graphs、NCCL固有の呼び出し、自作のCUDAカーネルは書き直しが必要です。

検証が必要な点

• 52%のMFUと8,192チップはGoogleによるTorchTitanの社内テストで、同価格帯のGPUとの独立した比較はありません。
• Zenteiq（約4倍）とDeepgram（1ドルあたり+60%）の数値は、メーカーが選んで掲載した顧客の声です。
• TechCrunchが引用した「AIアプリの85%がPyTorchを使用」という数値は出典が明示されておらず、TBN Growは確認できていません。
• 自分で確かめる方法：本番で動いているモデルを1つ選び、現在のGPUでトークン/秒と100万トークンあたりのコストを測り、同じ入力長でTPU v6eまたはTrainium2上でも実行してください。

TBN Growの見解

これはインフラソフトウェアの変化であり、Google CloudやAWSでオープンモデルを学習・自社運用するチームにとって実際の価値があります。月に数万ドル規模のGPU費用がかかっているチームは、安定したモデルで今すぐ試す価値があります。小規模チーム、API利用者、入力が大きく変動する処理では、演算子の対応範囲と動的形状のサポートが成熟するまで、1〜2回のリリースを待つのが妥当です。

今後の注目点

PyTorchをTPUとTrainiumで動かす際に注目すべき点

今すぐできること

• [ ] コード内でcudaを直接呼んでいる箇所を探し、torch.accelerator APIに置き換えてチップの切り替えを容易にします。
• [ ] TorchTPUを試す：TPU v6eのVMでpip install torch_tpu==0.2.0を実行します（READMEではこれが最も簡単な始め方とされています。CUDA版を入れないよう--extra-index-url https://download.pytorch.org/whl/cpuを付けます）。
• [ ] TorchNeuronを試す：trn2インスタンスでpip install torch-neuronxを実行し、デバイスをneuronに切り替えます。
• [ ] コストを見積もる前に、自分のプロジェクトのテストを走らせて未対応の演算子を洗い出します。
• [ ] 速度だけでなく、100万トークンあたりや音声1時間あたりのコストで比較します。

スケジュールと未発表の事項

• AWSは次のリリースでTorchAOとFP8を一般提供にし、演算子の不足を順次解消するとしていますが、日付は未定です。
• Googleはtorch.compileでの動的形状対応とコンパイル済みカーネルのライブラリを目標に掲げていますが、時期は示していません。
• GoogleはTorchTPUをベータとも一般提供とも位置づけておらず、対応するTPU世代の完全な一覧も公表していません。
• Nvidiaは公式な反応を示していません。

関連記事：中国のAIチップ市場が転換、Nvidiaは34%、Huaweiは28%に、AIインフラ投資の回収という課題。

検証に使った情報源：Google Cloud Blog、Google Developers Blog、GitHubのTorchTPU、PyTorch Blog – TorchNeuron、AWS What's New、TechCrunch。