# Introducing TorchTPU: Native PyTorch Backend for Google TPUs A collaborative effort between Google and Meta to enable seamless PyTorch execution on TPUs. - **Source:** [PyTorch Blog](https://pytorch.org/blog) - **Published:** 2026-10-20 - **Topic:** [Software & Engineering](https://brifing.si/c/engineering) - **Original article:** - **On Brifing:** https://brifing.si/s/385cd96d-3196-4375-a33b-0e622b80d0ef ## New Native TPU Integration TorchTPU is an open-source collaboration between Google and Meta, launched in mid-2025 to provide a principled PyTorch backend for TPUs. It moves away from previous approaches that shoehorned PyTorch into JAX, focusing instead on native hardware acceleration and seamless integration within the existing PyTorch ecosystem for training and inference. ## Addressing Previous Integration Limitations Earlier efforts like PyTorch/XLA relied on Lazy Tensors and required intrusive source-code modifications to define compilable regions. These legacy systems suffered from significant compile-time latency and lacked support for PyTorch's native imperative eager execution mode. TorchTPU solves these issues by prioritizing eager execution and providing a more robust, developer-friendly architecture. ## Core Design Principles The project adheres to strict principles, including PyTorch citizenship by avoiding JAX dependencies in the runtime and supporting native distributed idioms like DDP and FSDP. Developers can now port workloads by simply changing device targets, such as switching from cuda to tpu, while maintaining access to familiar tools like torch.profiler. ## Advanced Eager Execution Modalities TorchTPU introduces flexible dispatch modalities, allowing Aten operations to execute immediately or be deferred to a materialization point. This DeferAndFuse approach aggregates operations for compilation, bridging the performance gap between standard op-by-op execution and fully compiled workloads. This flexibility is managed through the PyTorch dispatcher and a dedicated C++ backend. ## Hardware Architecture Overview Google TPUs utilize specialized ASICs featuring TensorCores with systolic arrays of matrix-multiply units and dedicated SparseCores for gather and embedding operations. High Bandwidth Memory and bidirectional inter-chip interconnects facilitate massive scale operations. The architecture is designed to handle large-scale tensor operations efficiently, provided that model shapes are sufficiently wide for utilization. ## Compiler Stack and PjRT Programming TPUs involves the XLA compiler stack using StableHLO and Mosaic intermediate representations. A C++ program interacts with the hardware by instantiating a PjRT Client, which manages the compilation of MLIR into executable code. TPU execution is asynchronous, with the CPU waiting on PjRTBuffer results to ensure proper synchronization during model training. ## Optimizing Numerical Precision TorchTPU defaults to medium precision, utilizing native bfloat16 for matrix multiplications while maintaining FP32 accumulation for stability. Users can select high or highest precision modes, which employ multi-pass algorithms to approximate full IEEE FP32 accuracy. These modes incur significant compute overhead, making them suitable only for specific, numerically sensitive deep learning layers. ## Extensive Model Validation The development team has validated out-of-box execution for the 5,000 most downloaded models on Hugging Face. This includes large language models like Llama and Gemma, vision models such as DinoV2 and ViT, and various audio processing architectures. The framework supports custom kernel authoring via JAX/Pallas and the Helion DSL for specialized needs. ## Distributed Computing Support TorchTPU implements a one-process-per-device model to align with PyTorch ecosystem standards, differing from JAX’s one-process-per-host approach. It supports DTensor-based sharding and multi-host inference using vLLM and SGLang. Ray-based orchestration is fully integrated, enabling developers to scale training and fine-tuning workloads across large TPU Pods with minimal configuration changes. ## Bounded Dynamism Support Recognizing that XLA traditionally requires static shapes, TorchTPU introduces bounded dynamism. This allows the compiler to handle scenarios where the maximum dimensions of a shape are known at compile time, even if the actual dimensions vary during runtime. This feature significantly reduces the need for frequent XLA recompilation during model execution. ## Usability and Autoloading The project emphasizes developer experience through an autoload mechanism that simplifies the transition from GPU to TPU. By importing the TorchTPU wheel, existing PyTorch codebases can be adapted with trivial changes to initialization and device assignment. Actionable error messages are provided to help developers debug runtime or compiler issues during the porting process. ## Future Outlook and Production Currently in active development, TorchTPU is already utilized in production for both inference and training by major Google Cloud customers. By providing a stable, high-performance bridge between PyTorch and TPU hardware, the project aims to standardize TPU usage for the broader AI research and development community, ensuring long-term compatibility and performance. ## Sources covering this story - https://pytorch.org/blog/introducing-torchtpu/ — --- _Summary written by Brifing from the original reporting. Copyright in the original article belongs to its publisher; read the full piece at the link above._