Google open-sources TorchTPU to run PyTorch AI workloads on its chips

Developers using PyTorch to build and run AI models now have an open-source route onto Google’s Tensor Processing Units, or TPUs. Google Cloud released TorchTPU on GitHub on October 20, 2026. Co-developed with Meta’s PyTorch team, the stack promises to move existing workloads onto Google’s chips with minimal code changes, while retaining familiar tools.

Keep the PyTorch workflow, change the hardware

TorchTPU is a software stack for running PyTorch workloads natively on TPUs, rather than a new model or chip. Google says it supports training, serving and deployment through standard PyTorch programming interfaces. The intended change is where the work runs, without requiring developers to replace the framework they use to write it.

The development path has two stages. Google says developers can debug models using eager mode, then use torch.compile to optimize execution for production. The company says that compilation step improves TPU use and reduces memory overhead. TorchTPU also supplies diagnostics and built-in support for distributing work across devices.

The migration promise needs precise wording. TechCrunch describes moving existing applications with a configuration change and no rewrite. Google’s own announcement uses the more qualified phrase “minimal code changes.” That leaves room for adaptation; it is not a blanket assurance that every existing workload will move unchanged.

At Meta, advancing the frontier of AI requires immense compute scale and deep hardware diversity. We cannot be constrained by a single hardware architecture.

Jana van Greunen, Director of PyTorch Engineering, Meta, in Google Cloud’s announcement

The connection extends beyond model code

Google worked with Hugging Face to bring native TorchTPU execution into its libraries. In Transformers v5, Google says developers can select TPUs by passing device_map="tpu". Its Diffusers integration supports eager mode, compilation and splitting large generative models across multiple accelerators. Google also built dedicated pipelines to test the Transformers suite on TPU infrastructure every night.

The release also connects to tools that serve models and coordinate computing resources. These integrations address different parts of an AI system, rather than asking one library to handle the whole workflow:

• vLLM: Built with Inferact, the TPU backend supports continuous batching, quantization and parallel execution across chips. Google says it preserves the existing workers, schedulers and API servers used to serve models.
• SGLang: Developed with RadixArk, its backend targets structured and agent-based workloads. It reuses cached model state across requests, letting repeated prompt beginnings skip recomputation while retaining upstream serving components.
• Ray: Native TorchTPU support lets developers coordinate data processing, training, reinforcement learning and serving across CPUs, GPUs and TPUs within one workflow, assigning stages through resource configuration.

Training scale is a separate test

For training, TorchTPU integrates with TorchTitan, PyTorch’s reference implementation for distributed training. Google says the combination uses compilation, specialized Pallas kernels and graph-level optimizations. TorchTitan supplies training recipes for large dense models, mixture-of-experts models and diffusion transformers, giving the release a training path alongside its model-serving integrations.

Google reports that this training path achieved up to 52% Model Flops Utilization, a measure of training compute use, on Qwen 3 8B. It also reports near-linear scaling and stability across as many as 8,192 Ironwood TPUs. Those are company-reported results for a named model and hardware setup, not a demonstrated performance advantage for every PyTorch workload.