Google and Meta introduced TorchTPU, a native PyTorch backend for Google TPUs. It replaces PyTorch/XLA and supports eager execution out of the box. Models can migrate from GPU to TPU by changing `.to("cuda")` to `.to("tpu")`. It integrates with torch.compile and PyTorch’s distributed primitives. TorchTPU supports ~450 Aten ops and runs 5,000 Hugging Face models without modification. It offers three precision modes, defaulting to bfloat16 with FP32 accumulation. The project is open source and used in production by Google Cloud customers. It enables custom kernels via JAX/Pallas and Helion DSL. TorchTPU validates DTensor-based sharding and multi-host inference. It is designed for seamless PyTorch ecosystem integration.
What changed
PyTorch on TPUs required PyTorch/XLA with Lazy Tensors, intrusive code changes, and lacked eager mode support.
What it unlocks
Running PyTorch models on Google TPUs with minimal code changes and native eager execution.
- ~450 Aten ops supported
- 5,000 Hugging Face models validated
- 288 GB HBM in TPU 8i
- 19.2 TBps inter-chip interconnect in TPU 8t/8i
What you need to act on it
- Google TPU hardware
- TorchTPU wheel installed
- PyTorch codebase
Sources