Google and Meta Introduce TorchTPU, a Native PyTorch Backend for Google TPUs

TorchTPU is an open source collaboration between Google and Meta that introduces a native PyTorch backend for Google TPUs. It enables out-of-the-box model execution and TPU migration while targeting industry-competitive inference and training performance through eager execution optimizations, torch.compile integration, and native hardware acceleration.

Google TPUs have been used to train and deploy large ML models at massive scale. While TPUs have historically been tied to TensorFlow and JAX, this effort started at Google in mid 2025 with the PyTorch Team at Meta. The earlier PyTorch/XLA integration, introduced in 2020 and based on Lazy Tensors, required intrusive source-code modifications, suffered from severe compile-time latency, lacked eager execution support, and lagged in distributed primitives.

Design principles and usability

TorchTPU was built as a principled TPU backend for PyTorch rather than an adaptation of JAX. Its principles include PyTorch citizenship (no JAX dependencies in the runtime, solid eager mode, torch.compile based on Torch Dynamo, and support for DDP, FSDP and DTensor), seamless ecosystem integration, and community engagement. The focus on users is porting simplicity: changing model.to("cuda") to model.to("tpu") should let a model run out of the box, though peak performance may need further tuning. The project also pursues bounded dynamism to avoid XLA recompilation from shape changes and explores dynamic TPU kernels that can be precompiled ahead of time. Importing PyTorch on a TPU machine with the TorchTPU wheel installed requires only trivial changes through PyTorch's autoload mechanism.

Eager execution and compilation

TorchTPU extends the PyTorch dispatcher with a C++ backend registered as PrivateUse1 and defines TPU tensor storage via the standard at::Tensor mechanism. It offers three eager modalities: Strict (default) dispatches one Aten op at a time asynchronously; Debug (TORCH_TPU_LAUNCH_BLOCKING=1) is synchronous with added bound checks; and DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1) defers and fuses multiple ops, reported to be 50% or more faster than Strict eager while reducing peak HBM memory. At a materialization point such as a .to("cpu") call, TorchTPU applies liveness analysis, builds a cache key for the op graph, and dispatches a cached executable or compiles asynchronously. A Repeated Ops Heuristic detects repeated sequences, such as transformer block stacks, and triggers materialization early. The three-tier compilation cache spans process-local memory, machine-local shared memory, and a persistent remote file system cache for reuse across hosts and restarts.

Supported workloads and precision

TorchTPU supports about 450 Aten ops, including forward and backward passes for Core ATen ops required by major architectures, plus custom kernels via JAX/Pallas or the Helion DSL. It has been validated out of the box on the 5,000 most downloaded Hugging Face models, covering LLMs, audio and speech, vision, and text encoders. DTensor-based sharding, multi-host inference with vLLM and SGLang, fine-tuning and training with TorchTitan, and integration with Ray and TorchRec have also been validated. TPU matrix multiply units use bfloat16 inputs with FP32 accumulation, while non-contracting ops run in FP32 on the vector unit. TorchTPU maps PyTorch precision APIs to this hardware: medium, the default, uses BF16 multiplication with FP32 accumulation; high uses a 3-pass algorithm at roughly 3x compute cost; highest uses 6 passes at roughly 6x overhead. Selecting high or highest triggers a performance warning.

Sources: PyTorch · Google Cloud Blog · Hugging Face Blog