Discover TorchTPU, Google’s new engineering stack designed to run PyTorch natively on TPU infrastructure with peak efficiency. Learn how its "Eager First" philosophy and XLA integration simplify model migration while unlocking massive scalability for the next generation of AI.

GoogleDevs' platform is a central hub for developers interested in Google technologies, APIs, and developer tools. Through articles, tutorials, and documentation, GoogleDevs offers insights into building applications using Google Cloud Platform, Android, Chrome, and other Google services. Developers can learn about cloud computing, machine learning, and mobile app development with Google's developer tools and platforms.

Google Developers

Google's TorchTPU is a new engineering stack enabling PyTorch to run natively on TPU hardware with minimal code changes. Built on an 'Eager First' philosophy using PyTorch's PrivateUse1 interface, it offers three execution modes: Debug Eager (synchronous, for debugging), Strict Eager (asynchronous, mirrors standard PyTorch), and Fused Eager (auto-fuses operations for 50–100%+ performance gains). For peak performance, it integrates with torch.compile via XLA as the backend compiler, translating PyTorch operators to StableHLO IR. TorchTPU supports DDP, FSDPv2, and DTensor for distributed workloads, and adds MPMD support to handle divergent execution across ranks — a key limitation of its predecessor PyTorch/XLA. Custom kernels via Pallas and JAX are supported. The 2026 roadmap includes bounded dynamism for dynamic shapes, precompiled kernel libraries, and broader ecosystem integration.

TorchTPU: Running PyTorch Natively on TPUs at Google Scale

Architecting for Usability, Portability, and Performance

Engineering the TorchTPU Stack: The Technical Reality