PyTorchのコードを書き換えずに、TPUで動かせるのでしょうか。
Hugging Faceが公開したHugging Face ❤️ TorchTPU: TPUs, Nativelyは、その答えがYesに近づいたことを伝えています。PyTorchにTPUを新しいバックエンドとして加えるTorchTPUにより、Transformers、Diffusers、Kernelsが TPU上でネイティブに動くようになりました。
デバイス名を変えるだけ、という手触り
使い方は単純です。Hubからモデルを選び、デバイスに”tpu”を指定します。GPUで”cuda”を指定するのと同じ感覚で、TPUが「もう一つのデバイス」になります。
移植が前提だった時代との違い
これまでPyTorchでTPUを使うには、torch_xlaへ移植するか、JAXを学ぶ必要がありました。torch_xlaではテンソルが遅延評価されるため、コードが思ったタイミングで実行されません。小さな変更で再コンパイルが走り、デバッグも難しいという癖がありました。
Hugging Faceはこの橋渡しとしてOptimum TPUを作っていましたが、TorchTPUの登場で、TransformersとDiffusersの実行にOptimum TPUは不要になったとしています。
- Before: torch_xlaへの移植、またはJAXの習得が必要
- After: 通常のPyTorchコードのまま、デバイスに”tpu”を指定
計算資源を選べるようになるという機会
TPUはGoogleが約10年かけて設計してきたAIアクセラレータです。Google Cloudから、単一チップ、数千チップ規模のポッドまで利用できます。ソフトウェアの壁が下がれば、PyTorch資産を持つチームも、GPU一択ではなくTPUを選択肢に入れて比較できます。
ただし、元記事の抜粋からは性能やコストの実測値までは確認できません。自社のワークロードで小さく試し、数字で判断する段階です。
冒頭の問いへの答えは、Yesです。少なくともHugging Faceのエコシステム上では、TPUは「移植先」ではなく「選べるデバイス」になりました。
関連記事
- ALTK-Evolve: How learned memory tops the τ²-bench leaderboard on customer service tasks without a frontier model
- FairStress: A deeper evaluation of fairness in LLM decision-making
- NVIDIA Cosmos3-H-Dreams: Bringing Clinical Surgery to Life Through Generative Simulation
参考文献
コメント