TorchTPUでPyTorchの対応先を拡大
Googleは、PyTorchを自社のAI向けチップ「TPU」で動かすソフトウェア基盤「TorchTPU」をオープンソースとして公開しました。PyTorchは、AIモデルの開発、学習、実行に使われるオープンソースのツールキットです。AIアプリの約85％で使われているとの推計もあり、対応ハードウェアが増えると、開発者が実行環境を選びやすくなります。
TorchTPUは、既存のPyTorchアプリを大幅に書き換えずにTPUへ移すことを目指しています。GoogleはMetaのPyTorchチームと共同開発し、開発時の実行モードから本番向けの最適化まで、PyTorchの標準APIを使えるようにしました。
Googleは、Hugging Faceのライブラリや、モデル提供に使われるvLLM、SGLangなどとの連携も進めています。Transformers v5では設定で「tpu」を指定してモデルをTPUに割り当てられます。画像生成向けのDiffusersもTorchTPUに対応します。
AmazonもTrainium向け基盤を公開
Googleの発表に先立ち、AmazonはPyTorchモデルを自社チップ「Trainium」で動かすTorchNeuronを発表しました。両基盤は、既存のPyTorchアプリを大幅に書き換えずに、自社チップとクラウドへ移すことを目指しています。
AmazonはTorchNeuronの公開に加え、PyTorch本体にNVIDIA製チップへの依存を減らすコードを提出するとしています。PyTorchの機能を特定のハードウェアから切り離す取り組みが進めば、複数のチップを選べるソフトウェア基盤につながります。
TorchTPUの開発にはMetaが参加し、GoogleはHugging Face、vLLMを開発するInferact、SGLangを開発するRadixArkなどとも連携しました。
CUDAの蓄積が乗り換えの壁に
NVIDIAのGPU向け開発では、並列計算基盤のCUDAが長く使われてきました。CUDAは2006年に発表され、C言語の既存知識を生かせる形で開発者に広がりました。
NVIDIAのIan Buck副社長は、初期世代のGPU向けに書かれたCUDAコードが新しい世代でも動く互換性を、普及を支えた要素として挙げています。CUDAはC以外の言語にも対応し、データ処理や画像処理などに使う1,000を超えるCUDA-Xライブラリを備えています。
NVIDIAのcuDNNは、ニューラルネットワークで使われる演算のGPU向け実装や、複数の演算をまとめて処理する仕組みを提供します。こうした開発基盤やライブラリが蓄積するなか、GoogleとAmazonは既存のPyTorchアプリを大幅に書き換えずに自社チップへ移す方法を用意しています。
開発支援から配備までを整備
GoogleはTorchTPUに加えて「TPU agent suite」も公開し、TPUへの導入支援、プロファイリング、ハードウェアの確認などを行うエージェント群を提供しています。リクエストに応じて専門のエージェントへ振り分ける構成で、複数の開発環境で動かす仕組みも含みます。
本番環境向けには、Googleは推論サーバーのvLLMやSGLang、複数の計算機に処理を分散するllm-dとの連携を進めています。llm-dでは、Google Kubernetes Engine上でGPUまたはTPUを選び、共通の接続口から推論処理を運用できるとしています.
