TorchTPU: Running PyTorch Natively on TPUs at Google Scale (developers.googleblog.com)

Google has announced TorchTPU, a significant advance in running PyTorch natively and efficiently on its Tensor Processing Units (TPUs). This new integration aims to meet the growing demands of AI infrastructure, which often requires scaling models across thousands of accelerators. TorchTPU allows developers to run existing PyTorch workloads with minimal code changes, providing a familiar experience while maximizing the performance capabilities of TPUs. The innovative architecture incorporates an “Eager First” philosophy, supports diverse execution modes, and integrates with PyTorch's distributed APIs for enhanced usability and reliability. Key technical features of TorchTPU include three eager execution modes—Debug Eager, Strict Eager, and a breakthrough Fused Eager mode—which improve performance by fusing operations into larger, efficient computational chunks on the fly. TorchTPU also utilizes XLA as its backend compiler to optimize computation and communication effectively. Importantly, it supports divergent executions, addressing limitations faced by its predecessor, PyTorch/XLA. By strategically optimizing both regular and custom kernels, TorchTPU aims to streamline the development process, ensuring a seamless transition for developers while leveraging the unique capabilities of TPUs. This initiative is poised to enable higher performance for next-generation AI models and workloads in the PyTorch ecosystem.