We’re excited to share that TorchTPU—our native PyTorch stack for Google Cloud TPUs—is now open source on GitHub, alongside the early preview of the TPU Agent Suite!
Co-developed with the PyTorch team at Meta, TorchTPU brings Cloud TPUs (v5e, v5p, v6e, and v7 / Ironwood) alongside GPUs and CPUs as a native backend option for PyTorch. You can now train, fine-tune, and serve PyTorch models on TPUs with minimal code changes using the standard APIs and open-source tools you already rely on.
⚡ Quickstart: Installing TorchTPU
You can install the latest torch_tpu release directly via pip (Python 3.11+):
python -m venv .venv
source .venv/bin/activate
pip install torch_tpu --extra-index-url https://download.pytorch.org/whl/cpu
🛠️ What’s Included in TorchTPU?
⚬ Native PyTorch Experience: Debug models effortlessly in eager mode (.to("tpu")) and transition to production-grade performance with torch.compile(), backed by custom ATen kernels and an XLA compilation cache.
⚬ Out-of-the-Box Ecosystem Integrations:
⚬ Hugging Face: Deploy popular open-weight models natively in Transformers v5 simply by passing device_map="tpu", plus native Diffusers support.
⚬ High-Throughput Serving: Native PyTorch backends for vLLM and SGLang, paired with llm-d for prefix-cache-aware Kubernetes fleet routing across GKE.
⚬ Distributed Training & RL at Scale: Production-grade integrations with TorchTitan (SimpleFSDP, custom Pallas kernels), VeRL for post-training/RL alignment, and Ray for heterogeneous compute orchestration.
🤖 Introducing the TPU Agent Suite (Early Preview)
To make onboarding and performance tuning even faster, we’ve also open-sourced the TPU Agent Suite in preview. It provides declarative bundles of multi-agent skills, routing protocols, and tools that compile and install natively across coding agent harnesses (including Claude Code and Antigravity, with Codex and Cursor support coming soon).
With the tpu-agent skill enabled, you can use natural language prompts directly in your coding environment to automate:
⚬ Model Onboarding & Benchmarking: Spin up vLLM inference servers on torch_tpu, run test completions, and benchmark TTFT, ITL, and throughput.
⚬ Profiling & Diagnostics: Analyze XProf profiles to pinpoint memory bottlenecks and TPU compute utilization.
⚬ Kernel Optimization: Optimize custom Pallas kernels and run automated self-refinement benchmark loops.
🔗 Resources & Links
⚬ TorchTPU GitHub Repository: https://github.com/google-pytorch/torch_tpu
⚬ TPU Agent Suite GitHub Repository: https://github.com/google-pytorch/tpu-agent-suite
⚬ Official Google Cloud Documentation: https://docs.cloud.google.com/tpu/docs/pytorch-on-tpu
⚬ Launch Blog Post: https://cloud.google.com/blog/topics/ai-infrastructure/torchtpu-is-now-open-source-run-native-pytorch-on-tpus
We’d love for you to try out both projects! If you have questions about migrating a workload, feedback on the agent skills, or run into any issues, drop a reply below or open an issue directly on the GitHub trackers. What PyTorch workloads are you most excited to test on TPUs?
Top comments (0)