hn.today

Google open-sources TorchTPU, a native PyTorch back end for TPUs

github.com375 points88 comments
Screenshot of Google open-sources TorchTPU, a native PyTorch back end for TPUs

TorchTPU is a native PyTorch back end for Google TPUs that implements PyTorch ATen kernels and a compilation cache to run and optimize XLA computations on TPU devices. It supports distributed training (data and tensor parallelism), integrates with torch.compile() for graph-mode execution, and provides utilities for debugging and benchmarking. The project includes end-to-end examples for large models and common benchmarks (Llama, Qwen, ResNet, DLRM, minGPT, UNet), user-facing APIs, and interactive Marimo notebooks (static copies available in docs) that are intended to be run from a TPU VM (a v6e GCE instance is recommended).

Installation supports Python 3.11+ and a public pip release (torch_tpu==0.2.0) with the CPU torch index URL; building from source uses Bazel/Bazelisk with an explicit wheel build workflow (bazel build config=wheel_common //ci/wheel:torch_tpu_wheel) and defaults to optimized (-c opt) builds. The repo documents nuances for wheel builds (shared libxla_base.so handling), options to disable remote build execution (config=no_rbe), and VS Code/clangd setup for C++ development. Contributions are welcome under an Apache 2.0 license; docs, examples, and build scripts are included to help developers reproduce, extend, and deploy PyTorch models on TPUs.

Read on github.com88 comments on Hacker News

Summary generated by AI from the linked article. hn.today is not affiliated with Hacker News or Y Combinator.

More in AI

Claude Duet

Claude Duet

Claude Duet is a new AI app for professionals and small business owners that acts as a proactive assistant, managing tasks and priorities. It is available for iOS, Android, and web, and offers a research preview for Pro and Max subscribers. (claude.com)

The daily digest

Today's best Hacker News stories, summarized and screenshotted, one email a day.