Google、PyTorchのTPU版「TorchTPU」公開 Nvidia依存の緩和を狙う

TorchTPUとは何か

Googleは火曜日、かねて待たれていたPyTorchのGoogle版「TorchTPU」をオープンソースプロジェクトとして公開したと発表しました。PyTorchはAIモデルの構築・学習・実行に使われる代表的なオープンソースのツールキットで、一部の見方ではAIアプリの85%で使われているとされます。2016年にMeta（当時のFacebook）が開発を始め、2022年にLinux Foundationへ移管されました。

PyTorchが生まれた当時、AI研究の標準ハードウェアはすでにNvidiaのGPUでした。そのためPyTorchの機能の多くはNvidiaのGPUを前提に作られており、これがNvidiaの強さの一因になっています。TorchTPUはこの前提を崩し、GoogleのAIチップであるTPU上でAIアプリを動かせるようにするものです。

設定変更だけでTPUへ移行

TorchTPUの特徴は、既存のPyTorchアプリをTPU、ひいてはGoogle Cloudへ移しやすくする点にあります。Googleによれば、開発者はアプリを書き直すことなく、簡単な設定変更だけで移行できるとしています。クラウドの選択肢を広げる効果が期待される仕組みです。

Googleは移行を確実にするため、MetaやHugging Face、vLLMを手がけるInferact、SGLangを手がけるRadixArkなどと直接協力したと述べています。これらの主要なAI開発ツール上で、本家PyTorchと同様にTorchTPUの機能やコードライブラリが動作するようにしたということです。あわせて、移行作業やエラー修正を支援する「TPU agent suite」を用意し、スマートフォンやノートPC、組み込み機器でAIを動かすLiteRTもPyTorch製アプリとの相性を高めるよう更新しました。

Amazonも先行して対抗

Googleより前の月曜日、Amazonも独自のPyTorch版「TorchNeuron」を発表しました。TorchNeuronはAIアプリをAmazon自社開発のAIチップTrainiumとAWSへ、書き換えなしで移植できるようにするものです。Amazonはさらに踏み込み、PyTorch本体のプロジェクトに対して、各種機能をNvidiaのチップから切り離し、既定で任意のハードウェアに対応しやすくするコードを投稿していると説明しています。

両社の狙いは共通しており、PyTorchが実質的にNvidia向けに書かれている状態を解消することにあります。ただしGoogleは協業先や周辺ツールの整備を前面に出しており、単なる移植手段ではなく開発体験ごと取り込む構えに見えます。

Nvidia優位はすぐ崩れない

もっとも、移行が容易になったからといって、AIアプリの開発者が一斉にNvidiaのハードウェアを捨てるわけではありません。原文もその点を明示しており、今回の発表は選択肢を増やす段階と位置づけられます。Nvidiaの優位性がどれだけ揺らぐかは、移行後の性能やコストが見えてこないと判断できません。

また、原文にはTPUやTrainium上での実際の性能差や、移行に伴う制約、提供時期などの具体的な情報は含まれていません。TorchTPUがどの範囲のPyTorch機能に対応するのかも、この発表だけでは分かりません。