Googleは、Google独自のTensor Processing Units（TPU）上でPyTorchのモデルや演算を実行できるPyTorchバックエンド「TorchTPU」をオープンソース化しました。このプロジェクトには、多数のPyTorch演算向けのカスタムATenカーネル、XLA計算の実行を最適化するコンパイルキャッシュ、分散学習およびモデル推論の例が含まれています。
主な機能として、Data ParallelismとTensor Parallelismによる分散学習のサポート、グラフモード実行のためのtorch.compile()との統合、デバッグおよびベンチマーク用のユーティリティがあります。サンプルモデルには、Llama、Qwen、ResNet、DLRM、minGPT、UNetが含まれます。
ユーザーガイドと、Jupyterに似たインタラクティブなMarimoノートブックが提供されています。ノートブックはTPU VMから実行する必要があり、最も簡単な選択肢としてv6e GCEインスタンスが推奨されています。現在、Python 3.11以降がサポートされており、これはPyTorchのリリース互換性マトリックスに合わせたものです。
最新の公開版である0.2.0は、追加インデックスからPyTorchのCPUビルドを使用し、pipでPyPIからインストールできます。開発者はBazelを使ってソースからwheelをビルドすることもでき、最適化ビルドはデフォルトで有効になっています。認証情報を持たない場合に備えて、Remote Build Executionを無効にするオプションも用意されています。
出典: Hacker News · 要約：HeadlinesBriefing