Hugging Face ships TorchTPU, making TPUs a native PyTorch device

Hugging Face published a blog post on October 20 introducing TorchTPU, a native PyTorch backend for TPUs. According to the post, Transformers, Diffusers and Kernels now run natively on TPU: pick a model from the Hub, send it to "tpu", and that is the whole change.

Until now, running PyTorch on TPUs meant porting code to torch_xla or learning JAX instead. The XLA path worked but had quirks: tensors were lazy, small edits could trigger graph recompilations, and debugging was hard to reason about. Hugging Face built Optimum TPU to fill the gap with inference and training support.

TorchTPU changes that premise. It adds TPU as a new PyTorch backend, so users write ordinary PyTorch code, with Pallas kernels tuned for TPUs, distributed training support and an optional torch.compile integration. Installation on a TPU VM is a single pip command pointed at the CPU build of torch.

The usage surface is deliberately small. In Transformers v5, device_map="tpu" loads a model straight onto the accelerator; in Diffusers, pipe.to("tpu") runs inference, with an added compile call when more speed is needed. With native support in place, the post says Optimum TPU is no longer required for Transformers and Diffusers.

For models too large for one chip, TorchTPU offers TensorParallelConfig, which shards the transformer across chips as it loads so the full model never has to fit in one place. Many popular pipelines support it out of the box, and the rest are a tp_plan away.

The published measurements used a Cloud TPU v6e-4 with four-way tensor parallelism, bf16 and batch size 1. MiniMax-H3, a video-and-audio model, produced a 1344x768, 5.2-second, 49-step clip in 188 seconds compiled versus 262 seconds eager; Qwen-Image-2.1 rendered 1024x1024 at 40 steps in 3.5 seconds versus 8.0.

Those numbers show a steady gain from compilation rather than a blanket TPU victory over GPUs. The structural change is at the abstraction layer: TPU stops being a porting target and becomes just another device, so model definitions, pipelines and kernels can land on a different accelerator without a fork.

For infrastructure teams, that lowers the cost of multi-hardware sourcing and makes TPUs easier to evaluate for inference and smaller fine-tuning jobs. What to watch: behaviour at larger distributed training scales, how quickly the Kernels Hub fills out, and the operational complexity that graph compilation brings in production.