Googleは、PyTorchモデルや操作をGoogleのTensor Processing Units（TPU）上で直接実行できるように設計された専用バックエンド「TorchTPU」をオープンソース化しました。これまでPyTorchはTPUハードウェアとのインターフェースに様々な回避策に依存してきましたが、TorchTPUは、TPUアーキテクチャ向けに特別に実装されたPyTorch操作の基本構成要素であるカスタムATenカーネルを含めることで、より統合された体験を提供します。このネイティブサポートは、非ネイティブな統合による摩擦なしにGoogleの特殊なAIアクセラレータを活用したい研究者やエンジニアの開発ワークフローを簡素化することを目的としています。
この画像は、PyTorchモデルをGoogleのTensor Processing Unitsで実行するためのネイティブバックエンドを提供するTorchTPUプロジェクトを特定するものです。
このバックエンドは、パフォーマンスとスケーラビリティを最適化するためにいくつかの技術的機能を導入しています。これには、TPUが機械学習ワークロードを処理する中心的な仕組みであるXLA（Accelerated Linear Algebra）計算の実行を高速化するためのコンパイルキャッシュが含まれます。さらに、TorchTPUは、LlamaやQwenのような大規模モデルのトレーニングに不可欠なデータ並列化（Data Parallelism）やテンソル並列化（Tensor Parallelism）を含む分散トレーニング技術をサポートしています。また、このプロジェクトは`torch.compile()`とも統合されており、ユーザーはトレーニングと推論の両方においてスループットを大幅に向上させ、レイテンシを低減できるグラフモード実行の最適化の恩恵を受けることができます。
スター履歴チャートは、この新しいTPU向けネイティブPyTorchバックエンドに対する開発者の関心の高まりを例示しています。
ネイティブサポートへの移行にもかかわらず、このプロジェクトは最適な使用のための特定のハードウェアおよびソフトウェア要件を強調しています。例えば、ドキュメントでは、様々なモデルをサポートしている一方で、インタラクティブなノートブックやユーザーガイドは、v6e GCEインスタンスなどのTPU VMで実行するのが最適であると記されています。さらに、このプロジェクトはXLA/MLIRバックエンドに依存しているため、よりシームレスなPyTorch体験を提供しつつも、依然として広範なXLAエコシステム内で動作します。この開発は、PyTorchエコシステムにおけるTPUの第一パーティサポートへの移行を意味しており、JAXのハードウェア中心の最適化と、PyTorchの広範なソフトウェア採用との間のギャップを縮小させる可能性があります。