TorchTPU is Google’s open-source PyTorch backend for its TPUs, announced on October 20, 2026, with Meta’s PyTorch team. A basic run can start with installing the package and selecting the tpu device. Distributed jobs need extra TPU-specific launch configuration, and a device change alone does not guarantee peak performance.
Start with the device and execution mode
Google’s Cloud TPU instructions use pip install torch_tpu and torch.device("tpu"). The PyTorch project’s example changes a model placement from .to("cuda") to .to("tpu"). TorchTPU runs PyTorch’s eager mode by default; torch.compile is also available to compile a model for TPU execution.
That is a practical starting point, not a promise that every GPU workload will run unchanged. The model’s operations and shape can affect compatibility and performance. The PyTorch project reports support for about 450 ATen operators, excluding structural or composite delegates from that count. In one GPU-origin Muse Glimmer 30B example, the project replaced FlexAttention, which was unsupported in that test, with F.scaled_dot_product_attention.
Distributed jobs need TPU-specific setup
For distributed execution, TorchTPU uses the tpu_dist process-group backend and a one-process-per-device arrangement. A distributed launch also needs TPU topology and slice-builder configuration. The difference is concrete: selecting a TPU device addresses where the model runs, while the process group and launch settings coordinate work across devices.
| Workload | Device or backend setting | Additional distributed setup |
| Basic run | Select torch.device("tpu") or move the model to .to("tpu") | Not required for a single-device run |
| Distributed run | Select the TPU device and use tpu_dist | One process per device, plus TPU topology and slice-builder configuration |
Google documents eager execution as the default and torch.compile as an option. Compilation and workload-specific tuning can matter when the goal is peak performance; switching devices by itself does not ensure that result.
What the published tests measure
The PyTorch project reports three compatibility evaluations with different pass conditions. In its Tier 1 set of 4,917 popularity-selected models, 4,821 passed a runtime test defined as no crash, reported as 98%. In a separate Tier 2 subset of 657 representative architectures, 638 achieved numerical parity against a GPU baseline, reported as 97%. For bounded-dynamism correctness, the project reports passing results for 834 of 839 evaluated causal LLMs, or 99%.
Those numbers describe the tested sets and their specific criteria. A no-crash result is not the same as numerical parity, and neither test establishes how every workload performs. The parity evaluation compares numerical results with a GPU baseline; it is not a standardized speed comparison between TPUs and GPUs.
The project also reports a Qwen3 MoE 235B training run across as many as 8,192 TPU v7 chips. In that configuration, with a batch of 262,144 and sequence length 4,096, it reports 2,104 tokens per second per TensorCore, 27.1% model FLOPs utilization (MFU), and a 31.2-second step time. These figures describe that reported configuration, not a general performance estimate for other models or hardware.
Loading saved model weights
Compatible PyTorch model weights can be loaded for TPU execution when they are restored into a compatible model definition. PyTorch’s state_dict and model-loading guidance describes saving and restoring model parameters; torch.load supports remapping saved tensors to a target device. TorchTPU’s PyTorch-native semantics and TPU device make that a route to using compatible weights on a TPU.
Resuming training can involve more than model weights. Optimizer, scheduler, and other training state are separate checkpoint components, so a saved model’s parameters alone do not restore them.
