TorchTPU brings a native device path to PyTorch
Transformers and Diffusers can target tpu through familiar APIs, while full-pipeline compatibility remains the engineering check.
device_map="tpu" is the loading path Hugging Face demonstrates for Transformers with TorchTPU. On a TPU VM, its documented installation is:
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
The additional index supplies the CPU distribution of PyTorch, avoiding an unnecessary CUDA distribution for this environment. The article was published October 20, 2026.
TorchTPU provides a native PyTorch backend with eager execution, torch.compile integration and distributed training support. The Diffusers example moves a pipeline using pipe.to("tpu"), with compilation as a subsequent option.
Keeping model definitions in the project's existing framework is the architectural benefit I would prioritize. Compatibility still needs attention: the authors explicitly report remaining failures in the tests they run on TPU.
My starting point would be eager execution with output validation, followed by a comparison against the compiled path. That separates getting the workload to run from deciding how to optimize it. The engineering boundary is whether the entire pipeline works, rather than whether its tensors reach the device.
Source: Alvaro Moran and coauthors at Hugging Face, Hugging Face ❤️ TorchTPU: TPUs, Natively.