PyTorch on TPUs: Deep Dive into TorchTPU, XLA and Cloud TPU

Home » Python » PyTorch on TPUs: Deep Dive into TorchTPU, XLA and Cloud TPU

Running PyTorch on Google TPUs is no longer a niche, experimental path reserved for a handful of experts. Between Google’s new TorchTPU stack, the battle-tested PyTorch/XLA project, and a growing ecosystem of tooling and frameworks, training and serving models on TPUs is rapidly becoming as natural as working on NVIDIA GPUs. The big shift is that you can now aim for high performance, huge scale, and a much smoother developer experience at the same time.

This article dives deep into how PyTorch leverages TPUs today and where the stack is going: we will unpack the TorchTPU architecture, the differences versus traditional PyTorch/XLA, how distributed training, compilation, and hardware specifics work, and what this means in practice if you are migrating GPU-centric PyTorch workflows. If you live in the world of LLMs, diffusion, or large-scale recommendation systems, the details below are exactly the kind of low-level reality that will decide whether your TPU runs fly or crawl.

Why PyTorch on TPUs matters right now

Modern AI workloads have outgrown the simple “one machine, few GPUs” era. State-of-the-art models now stretch across clusters holding tens of thousands of accelerators, pushing software to handle extreme scale, reliable distributed execution, and portable performance across different chips and vendors in AI infrastructure.

Google’s Tensor Processing Units (TPUs) sit at the heart of this frontier. They power internal systems like Gemini and Veo, as well as a large fraction of Google Cloud customers’ training and inference workloads. Historically, TPUs were closely paired with JAX and TensorFlow, but the broader ecosystem has standardized heavily on PyTorch, which created a painful split: GPUs meant “PyTorch + CUDA”, TPUs meant “JAX + XLA”.

Google’s answer is a full-throttle effort to make TPUs feel like a first-class PyTorch target. TorchTPU aims to give you native, eager PyTorch semantics with top-tier performance, while PyTorch/XLA remains a powerful, lazily-compiled path that’s already widely adopted in production. Around these stacks, Cloud TPU, GKE, Vertex AI and community frameworks like easy-torch-tpu are turning TPU clusters into straightforward, scriptable infrastructure for anything from 1B to 70B+ parameter models.

Inside TPU hardware: more than just a faster chip

A TPU system is fundamentally a tightly integrated fabric of chips, hosts and interconnects, not just a single accelerator card. Understanding this layout is essential to making sense of TorchTPU’s design and why its compiler choices differ from pure GPU stacks.

Each TPU host connects to multiple TPU chips via an Inter-Chip Interconnect (ICI). The ICI forms a high-bandwidth 2D or 3D torus topology, which allows large-scale pods to behave like a single logical accelerator. Instead of shuttling gradients through traditional networking stacks, collectives ride directly on this torus, making scale-out far more efficient once your software knows how to express those collectives correctly.

Within a TPU chip, compute is split between TensorCores and SparseCores. TensorCores are specialized, single-threaded engines that excel at dense matrix math—exactly what powers transformers, CNNs, and most standard deep learning layers. SparseCores are designed for workloads with irregular memory access patterns, such as embeddings, gathers/scatters, and offloaded collective operations.

This architecture is fantastic for deep learning, but it is picky about how you feed it. For example, many transformer implementations hardcode attention head dimensions of 64. Current TPU generations tend to hit their sweetest spot at 128-256, which means simply doubling the head dimension can dramatically improve matrix multiplication efficiency and TensorCore utilization. Portability does not erase these hardware realities; it just makes it easier to reach them.

From PyTorch/XLA to TorchTPU: two complementary ways to run PyTorch on TPUs

PyTorch can already run on TPUs today via PyTorch/XLA (torch_xla), which presents TPUs as standard PyTorch devices and compiles lazy XLA graphs under the hood. However, many researchers have found that while the changes to their code are minor on paper, the behavior difference versus GPU eager execution can feel jarring.

TorchTPU is Google’s new, native PyTorch backend designed to feel like “real” PyTorch, not a wrapper. Instead of forcing PyTorch into a JAX-like model with Lazy Tensors everywhere, TorchTPU leans into PyTorch’s eager execution and modern compilation APIs like torch.compile. It uses the PrivateUse1 device mechanism in PyTorch, so from your perspective you are just working with regular torch.Tensor objects that happen to live on a TPU.

The key difference between the two approaches is execution style. PyTorch/XLA defaults to lazy execution: operations build up a graph, which then triggers an XLA compilation when you hit a sync barrier such as a step in your training loop. TorchTPU, by contrast, is architected as “Eager First”, with additional modes that progressively fuse operations and hand off optimized subgraphs to XLA without asking you to abandon the standard PyTorch mental model.

Cloud TPU, GKE, and Vertex AI: the infrastructure backbone

Underneath any PyTorch-on-TPU stack you choose is the Cloud TPU platform, which exposes custom ASICs as scalable cloud resources tuned for both training and inference. These accelerators are used for a wide variety of workloads: conversational agents, code generation, image and media models, speech, recommendation systems, and personalization engines.

Cloud TPUs are tightly integrated with Google Kubernetes Engine (GKE), so you can schedule large-scale PyTorch jobs using standard Kubernetes primitives. The Dynamic Workload Scheduler lets you request the entire fleet of accelerators you need in one go, ensuring that thousands of TPU chips come online together to train or serve a model without manual orchestration.

For teams that want the simplest on-ramp, Vertex AI abstracts away most of the cluster management. You can target TPUs from managed training and serving workflows, including when you are using PyTorch-based models. Google Cloud positions this flexibility—TPUs or GPUs, managed or DIY Kubernetes—as a direct answer to the exploding demand for AI infrastructure from enterprises and research labs alike.

TorchTPU’s core philosophy: “PyTorch Citizenship”

The central design goal of TorchTPU is blunt: it should feel like PyTorch, not like a foreign framework. If you already know how to train a model on CUDA GPUs, you should be able to port that same training script to TPUs with minimal code edits and without rewriting your mental model.

In practical terms, the ideal migration looks almost comically simple. Where you would normally write device = torch.device(‘cuda’), you instead obtain a TPU device from the TorchTPU module—conceptually something like device = tpu.get_device()—and call model.to(device) just as you would on GPU. Your forward pass, optimizer logic, and the way you call into Hugging Face models can remain unchanged.

Previous TPU integrations often pushed PyTorch to imitate JAX: they relied heavily on Lazy Tensors and forced you into static-graph thinking. That broke one of PyTorch’s biggest strengths: you couldn’t just insert a print in the middle of your forward pass to inspect shapes or values. TorchTPU rejects that trade-off. It keeps eager behavior as the baseline and builds performance around it, rather than asking you to abandon it.

This “PyTorch Citizenship” principle extends to error handling too. Instead of cryptic, 500-line C++ stack traces buried deep in the XLA stack, the goal is to surface clean Python tracebacks that point directly to the offending line in your training loop or model definition. When you are juggling multi-billion-parameter models and thousands of TPUs, that quality-of-life improvement is the difference between an afternoon fix and days of aimless debugging.

Eager modes in TorchTPU: Debug, Strict, and Fused

Delivering a native eager experience on hardware built for large fused graphs is non-trivial. TorchTPU solves this by offering several eager modes backed by a shared compilation and execution pipeline, so you can move smoothly from “make it work” to “make it fast”.

Debug Eager is the slowest but most transparent mode. It dispatches one operation at a time to the TPU and synchronizes with the CPU after each op. Performance is intentionally sacrificed so you can easily track down NaNs, shape mismatches, or out-of-memory errors with immediate feedback and clear stack traces.

Strict Eager keeps this single-op dispatch semantics but executes asynchronously. The TPU and CPU can run in parallel until the user code hits a synchronization point, providing an experience much closer to standard GPU-backed eager PyTorch, but still without heavy graph compilation requirements.

Fused Eager is where things get really interesting from a performance standpoint. TorchTPU observes the stream of operations you execute and automatically fuses them into larger, denser computation chunks before sending them to the TPU via XLA. This dynamic fusion step significantly boosts TensorCore utilization and cuts down memory bandwidth overhead, routinely yielding 50-100%+ speedups over Strict Eager without any model code changes.

All three eager modes share a common Compilation Cache that can live on a single host or be made persistent across multiple hosts in a distributed setup. Over time, as your training loop stabilizes and the system sees the same patterns, compilation cost drops and you spend more wall-clock time crunching tensors instead of building executables.

Static compilation: torch.compile, XLA, and StableHLO

When you need absolute peak performance on TPUs, TorchTPU hooks directly into the modern PyTorch compilation pipeline. You can wrap models or functions with torch.compile(), which captures an FX graph using Torch Dynamo, then bypasses the usual TorchInductor backend and hands control to XLA instead.

Choosing XLA as the primary backend is a deliberate decision rooted in TPU reality. XLA has been hardened over years of deployment across TPU pods, and it deeply understands the intersection of dense math and collective communication over the ICI torus. TorchTPU maps PyTorch operators directly into StableHLO, the tensor IR understood by OpenXLA, then lets XLA’s lowering passes generate optimized TPU binaries, reusing the same runtime paths as the eager modes wherever possible.

Extensibility for custom operators is not an afterthought. TorchTPU supports custom kernels defined in Pallas and JAX: by decorating a JAX function with something like @torch_tpu.pallas.custom_jax_kernel, you can inject low-level hardware-tuned code into the compilation path without losing the benefits of the global optimizer. Work is also underway to support additional DSLs such as Helion for even more flexible kernel authoring.

Distributed PyTorch on TPUs: DDP, FSDP, DTensor and MPMD

Massive models do not train on a single accelerator, and TorchTPU is built with that reality front and center. It integrates directly with PyTorch’s standard distributed APIs, including DistributedDataParallel (DDP), FSDPv2, and DTensor, and has been validated with third-party libraries that build on those abstractions.

One of the big historical pain points with PyTorch/XLA was its strict SPMD (Single Program, Multiple Data) bias. Many real-world PyTorch training scripts have small divergences between ranks—rank 0 might handle logging, checkpointing, or metrics, while other ranks do pure compute. For XLA’s global-graph view, this kind of behavior was awkward and often forced developers to rewrite code to avoid divergence.

TorchTPU explicitly embraces MPMD (Multiple Program, Multiple Data) scenarios. It carefully isolates and scopes communication primitives so that divergent behavior does not break correctness or kill performance. Wherever possible, it still lets XLA see a global picture of the distributed computation to overlap communication with compute, but it no longer forces you into an unrealistically pure SPMD style.

The way this meshes with existing PyTorch Distributed paradigms is especially important. Frameworks like FSDP, DTensor, and ecosystem tools such as TorchTitan rely on the ProcessGroup API for collectives like all-reduce, all-gather, and broadcast. On GPUs, those calls typically resolve to NCCL. TorchTPU intercepts these collectives at the ProcessGroup layer and lowers them into StableHLO collective ops, which the TPU hardware and ICI torus execute natively. From the perspective of FSDP or DTensor, nothing has changed—they simply see a different backend.

PyTorch/XLA: lazy execution, sync points, and practical tips

While TorchTPU is the long-term, fully native path, PyTorch/XLA remains a key tool for running PyTorch on TPUs today. If you are used to CUDA’s eager execution, the biggest conceptual shift with PyTorch/XLA is that tensors are lazy. Operations record a graph; actual execution and compilation happen when you explicitly or implicitly synchronize.

Synchronization points are where PyTorch/XLA hands the built-up graph to XLA for compilation and execution. Typical barriers include calls like torch_xla.sync() or higher-level utilities such as xm.optimizer_step(optimizer), which both step your optimizer and synchronize gradients across devices when you are in a distributed setup.

This lazy model has major performance implications. The first time a given graph (or a graph with new input shapes) executes, you pay a compilation cost, but subsequent iterations run much faster as long as the structure remains stable. That’s why shape stability—fixed sequence lengths, consistent batch sizes—matters so much for PyTorch/XLA workloads, and why padding inputs to fixed sizes is such a common pattern.

Multi-process training on PyTorch/XLA uses its own convenience tools. You typically wrap your core training function (for example, _mp_mnist_fn) and launch it across devices with torch_xla.launch. Data loading is managed via torch_xla.distributed.parallel_loader.MpDeviceLoader, which takes a standard PyTorch DataLoader and ensures each process sees a unique shard of data while prefetching batches to the appropriate TPU device.

Data loading, distributed execution, and AMP on TPUs

Efficient input pipelines are just as critical on TPUs as they are on GPUs. On PyTorch/XLA, MpDeviceLoader overlaps host-side data loading and device-side execution, feeding batches directly to the TPU and helping you avoid prolonged idle periods while the accelerator waits for new data.

For distributed training, xm.optimizer_step(optimizer) is doing more than a vanilla optimizer step. It performs gradient all-reduces across devices, averages them, applies the weight updates, and handles the necessary synchronization, so you typically do not need a separate explicit sync call in each iteration. Logging helpers like xm.is_master_ordinal(local=False) ensure only one process handles metrics and checkpointing to avoid duplication.

Automatic Mixed Precision (AMP) looks a bit different on TPUs than on GPUs. TPUs natively support bfloat16 (BF16), which offers a much larger exponent range than float16 and usually does not require explicit loss scaling for stability. PyTorch/XLA extends PyTorch AMP to map automatically between BF16 and FP32 where needed, making mixed-precision training on TPUs both straightforward and robust.

Saving models also has a TPU-specific best practice. While you can call torch.save from device tensors, it is generally recommended to move state dicts to CPU before serialization when using PyTorch/XLA, which makes them easier to reload on non-TPU hardware such as standard GPU machines.

Easy-torch-tpu and real-world TPU training frameworks

On top of the official stacks, the community is building higher-level frameworks to make TPUs easier to adopt. One example is aklein4/easy-torch-tpu, a lightweight training framework created specifically to simplify PyTorch/XLA workflows on Google Cloud TPU clusters.

Easy-torch-tpu positions itself as a simpler, more flexible alternative to large, rigid codebases like Hypercomputer/torchprime. Its design priorities are clear: easy setup, straightforward customizability, and smooth integration with gcloud ssh-driven cluster workflows. It deliberately targets “academic scale” experiments—models in the 1-10B parameter range over roughly 32-64 TPU chips.

Extensibility is handled via subclassing and configuration files. By adding new subclasses, you can plug in your own architectures, training loops, optimizers, data loaders, and even custom sharding and rematerialization strategies. This lets you experiment freely while reusing the framework’s distributed and logging scaffolding.

The framework integrates tightly with key ecosystem tools. Weights & Biases support makes experiment tracking trivial, while Hugging Face integration simplifies loading datasets, pulling pretrained checkpoints, and saving models that can later be run on standard GPU-based PyTorch. The repository includes installation docs, starter examples, and is actively evolving with community feedback.

Limitations, debugging, and performance pitfalls

Even with all these improvements, running PyTorch on TPUs is not completely frictionless yet. Understanding where things can go wrong will save you a lot of time when you are pushing large models or dynamic workloads.

Graph recompilations remain one of the biggest hidden performance killers. Any time your computation graph or input shapes change between sync points, XLA may need to recompile, which introduces noticeable pauses. This is especially common with variable-length sequences or adaptive batch sizes, which are common in language modeling and generation workloads.

Unsupported or partially supported operators can silently drag performance down. While PyTorch/XLA and TorchTPU aim for broad operator coverage, some ATen ops may not have native XLA lowerings yet. In those cases, execution may fall back to CPU, which is technically correct but can be orders of magnitude slower. Built-in debugging utilities and metrics (such as torch_xla.debug.metrics) help you spot where CPU fallbacks or unexpected recompilations are happening.

Classic GPU profiling tools like Nsight and nvprof do not see inside TPU kernels. Instead, you rely on XLA-specific profiling hooks, TPU runtime metrics, and higher-level logging to understand bottlenecks. Many teams find that once they adopt best practices (static-ish shapes, careful data loading, monitoring recompilations), they quickly converge on predictable performance.

Google’s compiler roadmap is explicitly targeting these pain points. Work on advanced bounded dynamism within XLA is meant to let models handle varying sequence lengths and batch sizes without triggering fresh compiles. A growing library of precompiled TPU kernels aims to slash cold-start latency on the first iteration of new graphs.

Roadmap and ecosystem: toward frictionless PyTorch on TPUs

Looking ahead, Google’s TorchTPU roadmap is ambitious and tightly aligned with the wider PyTorch ecosystem. A public GitHub repository is planned, complete with extensive documentation, architecture tutorials, and reproducible examples spanning both training and serving scenarios.

Integration with PyTorch’s Helion DSL is on the horizon, which should expand developer options for writing custom TPU kernels without diving into the deepest layers of XLA or hardware-specific code. Native, first-class support for dynamic shapes via torch.compile is also a priority, reflecting the realities of modern sequence-based models.

Multi-queue support is another key focus area. Many production PyTorch codebases rely heavily on asynchronous execution patterns and decoupled memory/compute streams. Making those idioms map cleanly to TPUs without major refactors will significantly lower migration friction for large, mature projects.

Deep ecosystem integrations are already in motion. Efforts are underway to validate strong scaling to full TPU Pod size, and to hook into major PyTorch-based systems like vLLM and TorchTitan. At the same time, Google is collaborating closely with Meta and the PyTorch community, and is exploring open-sourcing key parts of TorchTPU to accelerate adoption and transparency.

All of this sits against a larger business backdrop where TPU capacity is scaling dramatically. Google Cloud is signing more multi-billion-dollar AI infrastructure deals, Anthropic is planning for access to up to a million TPUs (on the order of a gigawatt of capacity), and Google is even selling TPUs directly for on-prem data centers. The days when TPUs were a niche, internal Google-only resource are long gone.

Putting it all together, the PyTorch-on-TPU story is moving from “quirky side path” to “standard option” remarkably fast. Between TorchTPU’s native eager experience, PyTorch/XLA’s battle-tested lazy execution, frameworks like easy-torch-tpu, and the rich Cloud TPU infrastructure around them, you can now take mainstream PyTorch models—often with little more than a device string change—and run them efficiently on some of the largest AI supercomputers available. The more the stack converges on familiar PyTorch idioms instead of forcing new mental models, the more realistic it becomes to treat hardware choice as an implementation detail rather than a fundamental design constraint.