PyTorch/XLA: Running PyTorch on TPUs, and the TorchTPU Transition

The gap PyTorch/XLA fills between PyTorch and Cloud TPU

PyTorch's standard device story covers CPU and CUDA. Cloud TPUs are neither. PyTorch/XLA exists to close that gap: the README describes it as a Python package that uses the XLA deep learning compiler to connect PyTorch and Cloud TPUs. The audience is therefore narrow and specific. It is for engineers who already have a PyTorch training loop and want it to execute on TPU hardware without rewriting the model in a different framework. The README also points at Kaggle as a zero-cost entry point, offering a single Cloud TPU VM and linking two notebooks, one on Stable Diffusion with PyTorch/XLA 2.0 and one on distributed PyTorch/XLA basics with PJRT. That is a deliberate on-ramp: you can check whether your loop survives the XLA path before renting anything.

The repository itself is C++, with Python as the user-facing surface. That split matters when you debug. Errors you see in Python often originate in the tracing and compilation layers underneath, and the documentation map in the README reflects that, splitting material across learn, accelerators, perf, features and contribute directories, plus a separate torchax tree with its own docs and examples. The scope is broader than a single training script: the perf docs cover AMP, DDP, Dynamo, Fori loop, FSDP, quantization, recompilation and SPMD, which tells you the project is aimed at people running real distributed workloads, not toy single-device experiments.

Lazy tensors, torch_xla.step and the sync boundary

The mechanism visible in the README is a lazy execution model. The single-process example shows the minimum diff to an existing training loop: import torch_xla, wrap the batch loop in a with torch_xla.step() block, move inputs and labels with .to('xla'), move the model parameters with model.to('xla'), and call torch_xla.sync() after the loop. Nothing in that diff changes the model definition or the optimizer. The work is relocated, not rewritten.

That shape implies a deferred graph. Operations issued inside the step block are not dispatched one at a time to the device; they accumulate and are handed to the XLA compiler, which is why an explicit sync call appears at the end. The step context is the boundary that tells the runtime where one iteration of the graph begins and ends. This is the design decision that shapes everything else in the project. It is also the source of the most common class of problem: code that reads a tensor value mid-step, prints it, or branches on it forces a materialization and breaks the batching the compiler was trying to build. The README does not enumerate those cases in the section quoted here, so treat the troubleshooting material under docs/source/learn as required reading rather than optional.

The README also names three execution modes and is explicit that they are not interchangeable. Single process means one Python interpreter controlling one TPU. Multi process means N interpreters for N TPUs. SPMD means one interpreter controlling all N TPUs. The README states plainly that multiprocessing is more complex and is not compatible with SPMD, and that the tutorial does not cover SPMD, pointing instead at a separate SPMD guide. Pick your mode before you write the launcher, because the choice is not reversible by a flag.

Installing wheels: version pinning, Python range and the C++11 ABI split

Installation is a pip command with a hard version coupling. The stable TPU path in the README is pip install torch==2.8.0 'torch_xla[tpu]==2.8.0'. The two versions are written to match, and the install note says builds are available for Python 3.8 to 3.11, so a 3.12 interpreter will not work on that path. A later note in the same README says that as of 07/16/2025 and starting from the 2.8 release, nightly and release wheels are provided for Python 3.11 to 3.13. Those two statements sit next to each other and do not obviously agree; verify against the wheel index for the exact version you intend to run rather than trusting either line alone.

The nightly path is a direct wheel URL, not a package name: pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.9.0.dev-cp312-cp312-linux_x86_64.whl' with -f pointing at the libtpu wheel index. The README's own comment tells you to edit cp310-cp310 to match your Python version, which is a reminder that these URLs are built per interpreter version and per ABI. If you use custom kernels, there is a separate optional install for pallas dependencies, pulled from a JAX wheel index with --pre.

The C++11 ABI story is the most consequential packaging detail. The README states that as of 03/18/2025 and starting from the 2.7 release, C++11 ABI builds are the default and pre-C++11 ABI wheels are no longer provided. It also gives the reason the choice existed: C++11 ABI wheels and docker images have better lazy tensor tracing performance. The README offers a diagnostic rather than a blanket recommendation. If your model is tracing bound, meaning the host CPU is busy tracing while TPUs sit idle, switching ABI can help. It cites a Mixtral 8x7B result on v5p-256 at global batch size 1024, going from 33% MFU pre-C++11 to 39% with the C++11 ABI. That is a single configuration from the project's own material, not a general claim, and the pre-C++11 option is no longer shipped for current releases.

Where PyTorch/XLA is the wrong tool

The first limitation is stated by the project itself, at the top of the README. A note dated 4/22/2026 says that once TorchTPU is public it will replace PyTorch/XLA, and an earlier note from 10/2025 points to an RFC at issue #9684 proposing a more native direction for PyTorch on TPU, based on community feedback. For anyone choosing a stack to maintain for years, that is the single most important fact on the page. The repository is not archived and the last push is recent, but the stated direction is replacement, not indefinite development of this surface.

The second limitation is hardware scope. The package connects PyTorch to Cloud TPUs. The README's plugin documentation lists a CPU PJRT plugin, which is useful for local iteration, but the install commands, the wheel index URLs, the docker image names and the performance discussion all assume TPU VMs. If your accelerator is something else, this is not the project for you, and the CPU plugin is a development convenience rather than a deployment target.

The third limitation is the lazy model itself. Any training loop that depends on per-step host-side control flow, dynamic shapes that change every iteration, or frequent synchronization with tensor values will fight the compiler. The README acknowledges this space by pointing at a dedicated dynamic shape document and a recompilation document under the perf tree. Recompilation is listed as a first-class topic, which tells you it is a normal occurrence rather than an error condition. Budget engineering time for it.

Finally, the packaging surface is brittle in a way that pure-Python libraries are not. You are matching a torch version, a torch_xla version, a Python minor version, a C++ ABI flavor and a wheel index URL. A mismatch produces failures that look like environment problems rather than dependency problems.

Alternatives and the difference in approach

The obvious alternative is JAX, which also targets TPUs and also compiles through XLA. The difference is not the backend, it is who owns the graph. JAX is built around explicit functional transformations and a tracing model from the ground up; the user writes code that is already shaped for compilation. PyTorch/XLA takes the opposite route, keeping the imperative PyTorch loop and inserting a compilation boundary with torch_xla.step() and torch_xla.sync(). You keep your nn.Module, your optimizer and your data loader, and you accept that some patterns will not survive the boundary. If your team already knows PyTorch and the model is not written in a way that fights tracing, that trade favors PyTorch/XLA. If you are starting fresh on TPU and are willing to learn a different programming model, JAX removes the impedance mismatch rather than managing it.

Within the same project there is a second option worth noting: torchax. The README's documentation map gives it its own docs and examples directories, and the 10/2025 note about a more native direction for PyTorch on TPU is consistent with that line of work. For anyone evaluating the project today, torchax is the part of the repository most likely to resemble where things are heading, and it is documented separately from the main PyTorch/XLA path.

The third comparison is against not moving at all. If your workload runs acceptably on GPUs or on CPU and you have no TPU allocation, the cost of adopting a lazy execution model, a version-coupled wheel stack and an announced migration is hard to justify. PyTorch/XLA makes sense when TPU capacity is the thing you already have.

Maintenance cost, licensing and what the repository does not tell you

The maintenance cost is dominated by version coupling. The README pins torch and torch_xla together in every install example, and the nightly example uses a literal wheel URL containing both the version (2.9.0.dev) and the interpreter tag (cp312-cp312). Upgrading PyTorch means upgrading torch_xla in lockstep, and possibly moving Python versions, since the supported range has shifted across releases. The ABI default changed at 2.7 and pre-C++11 wheels stopped being published, so any environment pinned to an older ABI needs a migration rather than a patch. Docker users have an equivalent constraint: the README gives an image tag of the form r2.6.0_3.10_tpuvm_cxx11, encoding release, Python version and ABI in one string.

On licensing, the repository metadata reports NOASSERTION, which means GitHub could not map the licence file to a known identifier. The README in the supplied material does not state a licence. Treat the licence as something to confirm from the LICENSE file and the PyTorch project's own terms before you ship a product that links against these wheels, and route that question to whoever handles legal review rather than assuming a permissive default. Nothing here is legal advice.

The README also leaves gaps a reader should notice. The multi-process section is truncated in the material available, so the launcher details are not verifiable from what is quoted. The Kaggle notebooks are linked but their contents are not summarized. The performance figures are limited to the one Mixtral 8x7B configuration cited for the ABI comparison. And the TorchTPU note gives no timeline beyond once it is public. Anyone planning a migration should read issue #9684 directly rather than inferring scope from the note.

Who should adopt PyTorch/XLA now

The case for adopting is strongest when the TPU allocation already exists and the training loop is conventional. If your code is a standard loop over a data loader with backward and optimizer step, the README's diff is small: import torch_xla, wrap the loop in torch_xla.step(), move tensors and parameters to 'xla', call torch_xla.sync() at the end. That is a day of work, not a rewrite, and the Kaggle notebooks let you validate it before provisioning anything.

The case against is equally clear. If you need a stack with a stated multi-year roadmap, the README's own note says TorchTPU will replace this once public, so the honest position is that you are adopting a bridge with a named successor. If your accelerator is not a Cloud TPU, stop here. If your model is tracing bound, the README points at the C++11 ABI wheels as the lever, and that lever is now the default rather than an option, so the remaining tuning work sits in the tracing and recompilation docs under docs/source/perf rather than in packaging.

What to verify first, concretely: your Python version against the range your target release supports, the exact torch and torch_xla pair from the release you intend to run, and whether the multi-process or SPMD mode matches how you will launch. The README states those two modes are incompatible, and that is not a detail you want to discover after writing the launcher.

Editorial conclusion

Adopt PyTorch/XLA if you are already training on Cloud TPU VMs and your loop maps onto torch_xla.step(), torch_xla.sync() and .to('xla'), pinning torch and torch_xla to the same version (for example torch==2.8.0 with torch_xla[tpu]==2.8.0). Do not start a new long-lived TPU stack on it without reading the TorchTPU note at the top of the README and issue #9684, since the README states TorchTPU will replace PyTorch/XLA once public. Before committing, verify your Python version against the supported range, check whether your model is tracing bound (the README ties that symptom to the C++11 ABI wheel choice), and confirm your accelerator is a Cloud TPU rather than an XLA backend the project does not ship wheels for.

Official sources

1. Issues
2. Project website
3. pytorch/xla on GitHub
4. README
5. Releases