Hugging Face ❤️ TorchTPU: TPUs, Natively

Hugging Face says TPUs are no longer a side quest for PyTorch users. With TorchTPU, the company is adding TPU support as a native PyTorch backend, so users can treat a TPU like any other device instead of routing everything through torch_xla or switching to JAX. That opens native TPU support across its ecosystem, including Transformers, Diffusers and Kernels.

The pitch is simple: pick a model from the Hub, send it to "tpu", and get on with your day. TorchTPU also ships Pallas kernels tuned for TPUs, distributed training support, and torch.compile integration for people who want more speed. Installing it on a TPU VM is a one-liner, though Hugging Face notes you need the CPU build of torch in the install command so pip doesn’t drag in a much bigger CUDA build you won’t use.

For Transformers, the change is almost absurdly small. Pass device_map="tpu" when loading a model and it comes up directly on the accelerator. The model runs in eager mode by default, but you can also compile it or run distributed. Diffusers gets the same treatment: pipe.to("tpu") is enough for eager mode, and compiled mode adds one more compile call. And if the model is too large for one chip, TensorParallelConfig can shard it across chips as it loads.

Hugging Face says it has already been testing the stack on Cloud TPU hardware, including a v6e-4 setup for example runs. The company says the TPU test coverage is not perfect yet — 96.5% of the tests that run on TPU pass — but the failures are being turned into small repros and issues. So far, it says, 34 issues have been reported and 19 fixed, including some bugs found in Transformers itself.

There’s a broader point here. TPU support has always been less about raw silicon and more about whether mainstream PyTorch tooling would bother to meet it halfway. Hugging Face just did the boring, necessary work that makes hardware matter: make it easy, make it native, and make the weird edge cases somebody else’s problem.