TorchTPU

TorchTPU is an open-source, native PyTorch backend for Google Cloud TPUs, developed in collaboration between Google and the PyTorch team at Meta. With TorchTPU, developers can run existing PyTorch workloads on TPU hardware with minimal code changes and scale training and inference workloads from a single chip to multi-slice TPU pods.

Why use TorchTPU?

Google TPUs are application-specific integrated circuits (ASICs) with an architecture purpose-built for the matrix and vector operations that power deep learning, offering high throughput, energy efficiency, and performance per dollar at scale. While TPUs have historically been paired with the JAX and TensorFlow ecosystems, TorchTPU brings native TPU support directly to the PyTorch ecosystem, designed intentionally around the following three guiding principles:

• Native PyTorch citizenship: TorchTPU integrates directly with PyTorch with no JAX runtime dependencies. It preserves standard PyTorch tensor view and in-place mutation semantics, adopts PyTorch’s one-process-per-device execution model, and supports torch.compile through Dynamo.
• Out-of-the-box usability: Porting a workload from a GPU to a TPU starts with changing device="cuda" to device="tpu" and backend="nccl" to backend="tpu_dist". Your unmodified PyTorch model runs immediately in eager mode, with clear error messages when diagnosing runtime or compiler issues.
• Production-scale performance: Beyond op-by-op eager execution, TorchTPU introduces a deferred-fusion eager mode, a Dynamo-to-XLA compiler pipeline, bounded dynamism to avoid shape-driven recompilations, and custom kernel APIs to help you reach peak hardware utilization.

How TorchTPU works

TPUs are programmed with the XLA compiler stack through two different entry points: StableHLO and Mosaic. The actual TPU ISA is proprietary, subject to change, and not publicly available. Effectively, one can think of StableHLO and (even more so) of Mosaic as equivalent to PTX. See picture below for a high-level sketch:

A C++ program can access a local TPU by instantiating a PJRT client, whose API provides methods to compile MLIR (StableHLO or Mosaic) into a TPU executable, to launch that executable, and to copy data to and from the TPU. Note that TPU execution is asynchronous and the CPU can wait for a TPU result to be ready by waiting on the associated PjRtBuffer. StableHLO follows the SSA Paradigm and, hence, doesn’t allow for in-place mutations, i.e., a write operation always produces a new value (at least conceptually). This is particularly significant for TorchTPU, which provides familiar PyTorch “view” semantics, including mutations, through view decomposition and inversion. The XLA compiler assumes that tensor shapes are known at compile time. This constraint is relaxed in TorchTPU by bounded dynamism. Recognizing this as a limitation for PyTorch workloads, we have been adding support for bounded dynamism, i.e., when the maximum bounds of a shape’s dimensions are known at compile time, but the actual dimensions are known only at runtime.

TorchTPU: An out of tree backend with autoload support

Following current best practice where additional backends are constructed out of tree and registered at runtime, we extend the PyTorch dispatcher with a new C++ backend, registered as PrivateUse1. Additionally, we have used the standard at::Tensor mechanism to define a device-specific implementation of tensor storage (by using the c10::UniqueVoidPtr inside c10::StorageImpl).

We designed TorchTPU with usability in mind. When the user imports PyTorch on a TPU machine where the TorchTPU wheel is installed, Python will autoload the TorchTPU wheel. Autoloading the wheel sets up the backend and replaces the PrivateUse1 name with tpu. This makes a TorchTPU program no different in simple example from the in-tree backends, and a PyTorch program can run on the TPU with little-to-no change:

import torch
import torch.distributed as dist

def main():
    device = torch.device("tpu")
    dist.init_process_group(backend="tpu_dist")
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

Because the TPU becomes the current accelerator and tpu_dist is the default backend for TorchTPU, they don’t need to be specified explicitly:

import torch
import torch.distributed as dist

def main():
    device = torch.accelerator.current_accelerator()
    dist.init_process_group()
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

Capturing the graph

TorchTPU traces eager ops at the ATen level, relying on the PyTorch op dispatch hooks that are part of the backend infrastructure. TorchTPU ops are constructed as StableHLO programs, and passed to XLA singly, or in groups, allowing the compiler to perform significant fusion of operations to more effectively utilize the hardware. When TorchTPU receives a callback invocation from the PyTorch dispatcher the runtime:

• Creates a Deferred Op instance that contains:
  • A StableHLO builder object, which when invoked returns a StableHLO function implementing the semantics of aten::mul, as shown below (note that any implicit broadcasting and dtype promotion are handled elsewhere)

func.func @main(%arg0: tensor<4x16xbfloat16>, %arg1: tensor<4x16xbfloat16>) -> tensor<4x16xbfloat16> {
  %0 = stablehlo.multiply %arg0, %arg1 : tensor<4x16xbfloat16>
  return %0 : tensor<4x16xbfloat16>
}

• An optional map of static op parameters (e.g., training and momentum for aten::native_batch_norm) to be used for the computation of a cache key;
• Return a new at::Tensor with c10::StorageImpl pointing (through indirection) to that Deferred Op instance.

Upon reaching an appropriate synchronization point, dependent on the execution mode, TorchTPU will:

• Analyze the region of ops;
• Determine which ops to materialize into executable code;
• Create a compilation cache key for the graph;
• If it already exists in the cache, fetch and dispatch the executable on the TPU;
• If it does not exist in the cache, asynchronously compile the object, add it to the cache and dispatch for execution;
• Potentially wait for the value to be ready and copy it back, if that was appropriate for the synchronization point reached.

We discuss more details on the various modes of execution in TorchTPU execution model.

The TorchTPU compilation cache

To avoid unnecessary compilations, TorchTPU employs a multi-tiered compilation cache:

• Tier 1: Process-local in-memory cache for sharing compilation within the same process. This is the fastest of the three tiers. It is always enabled.
• Tier 2: Machine-local shared memory that enables cross-process sharing on the same host. This is useful for multi-device training/inference. Tier-2 is enabled by default for distributed programs, but for single-device programs it offers no benefit so is disabled.
• Tier 3: Persistent cache backed by a remote file system (e.g., Google Cloud Storage), allowing reuse across hosts, programs, and job restarts. This is useful when running TPU programs in execution environments like Colab and data centers. This is disabled by default but can be enabled by specifying a storage location.

Each entry is a compiled TPU executable representing a StableHLO computation graph generated from an ATen op or a sequence of ATen ops. The StableHLO graphs are represented by fingerprints combining the structure of the graph, the tensor shapes of the inputs/outputs, and the compiler options including the versions of the TPU hardware, libtpu, and the PJRT runtime.

torch.compile support

To maintain a user experience consistent with other hardware supporting torch.compile() we register a new TorchDynamo backend, which can be used as follows: compiled_model = torch.compile(model, backend="tpu").

Dynamo invokes the TorchTPU backend by passing torch.fx.GraphModule plus a list of example input tensors. At that point the runtime traces the provided graph as if we were in eager mode, but with all ops deferred and no execution on the TPU. The runtime generates and compiles a single StableHLO program for the entire FX graph and returns a Python callable that wraps the invocation. Note that TorchTPU does not use TorchInductor, instead reusing the mature XLA compiler stack that has developed alongside the TPU architecture.

More detail

In the following documents we’ll go into more detail about the:

• Memory model supported by the TPU and what its implications are on PyTorch;
• The Execution Modes TorchTPU supports;
• Support for custom kernels;
• Distributed execution on clusters of TPUs.