DEV Community

Mikhail Savchenko
Mikhail Savchenko

Posted on Originally published at inite.ai

Hugging Face libraries now run natively on TPUs via TorchTPU

Hugging Face announced on October 20, 2026 that its main open-source libraries now run natively on TPUs, Google's AI accelerator chips available on Google Cloud. The change rests on TorchTPU, a native PyTorch backend that makes TPU "just another device".

What changed

PyTorch users who wanted TPUs previously had to go through PyTorch/XLA (torch_xla) or learn JAX. The post describes the old route as workable but awkward: lazy tensors meant code did not always run when expected, small changes could trigger graph recompilation, and debugging was hard to reason about.

With TorchTPU, you write regular PyTorch. It ships with Pallas kernels tuned for TPUs, distributed training support and a torch.compile integration. Installation on a TPU VM is a single pip command, using the CPU build of torch to avoid pulling the larger CUDA build.

What runs on it

  • Transformers: pass device_map="tpu" when loading a model. It runs in eager mode by default, with torch.compile or distributed setups as options.
  • Diffusers: TPU is now a first-class device. Eager mode needs only pipe.to("tpu"), and compiled mode adds one compile call. Large models can be sharded across chips with a TensorParallelConfig.
  • Kernels: the Kernels Hub has ready-made TPU kernels written in Pallas, and kernels from the Tokamax library can be used too.

Optimum TPU, Hugging Face's earlier TPU package, is no longer required to run models with Transformers and Diffusers.

The numbers, and their limits

Hugging Face reports that at batch size 64, a Hub ragged-dot kernel made decoding 21.5 times faster on Qwen3-30B-A3B and 12.8 times faster on Gemma 4 26B-A4B than the default implementation, measured on a TPU v6e-8. It also says the gain grows with batch size. These are the publisher's own benchmarks.

Support is still maturing. Of the tests that run on TPU, 96.5% pass. Hugging Face has reported 34 issues to the TorchTPU repository, 19 of which are fixed. A nightly CI run on TPU is something the team says it is working to add, not something that exists yet.

What the post does not say

The source gives no pricing, no cost comparison with other accelerators and no guidance on which workloads suit TPUs better than GPUs. Those questions need your own testing.

Top comments (0)