PyTorch/XLA

// Overview

PyTorch/XLA is the package that lets PyTorch code run on Google Cloud TPUs. It uses the XLA deep-learning compiler as the bridge: tensors are moved to an `xla` device, the training step is traced into an XLA graph, and the compiler turns that graph into TPU execution. In day-to-day use the change to an existing script is small — move the model and inputs to `'xla'`, and wrap the step in `torch_xla.step()`.

The distributed story is what most users come for. Single-process and multi-process training are both supported, along with FSDP for sharded parameters, `DistributedDataParallel`, and SPMD, where you describe how tensors are sharded and the compiler works out the collectives. Dynamic shapes and eager mode are available for the cases where full tracing is awkward, and the project ships profiling and performance-analysis tooling because compiled execution makes "why is this slow?" a different question than it is on GPU.

Google has since announced TorchTPU, a more native PyTorch-on-TPU path built on the PrivateUse1 device extension mechanism rather than on a compiler bridge, and has said it will replace PyTorch/XLA once its public repository ships. Until that transition happens, PyTorch/XLA remains the supported way to run PyTorch on TPU hardware, and it is the codebase the TPU ecosystem is currently built on.

// What it does

• Runs PyTorch models on Cloud TPUs by moving tensors to an `xla` device
• Single-process and multi-process training modes
• SPMD sharding, FSDP and DistributedDataParallel for large-model training
• Dynamic shape and eager mode support alongside traced execution
• Profiling and performance-analysis tooling for compiled graphs
• Versioned in step with PyTorch releases (`torch` and `torch_xla` share a version)

// Getting started

Install `torch_xla` matched to your `torch` version, then adapt the training loop. The TPU extra pulls in the runtime.

pip install torch==2.8.0 'torch_xla[tpu]==2.8.0'

import torch_xla

model.to('xla')

with torch_xla.step():
    inputs, labels = inputs.to('xla'), labels.to('xla')
    # forward, loss, backward, optimizer step

Commands and code are distilled from the project's own documentation — always check the official repo for the latest.

// When to use it

• Train an existing PyTorch model on Cloud TPU pods without rewriting it for JAX
• Shard a large model across TPU cores with SPMD or FSDP
• Fine-tune on TPU capacity when GPUs are unavailable or more expensive
• Profile and debug TPU execution of a PyTorch training loop

// How PyTorch/XLA compares

PyTorch/XLA alongside other open-source efficient training tools AI/TLDR tracks, ranked by GitHub stars.