Skip to content

Hugging Face adds native PyTorch support for TPUs

TorchTPU, built by Google with Meta's PyTorch team, drops the torch_xla dependency so PyTorch code runs on TPUs unmodified. Hugging Face reports decoding up to 21.5 times faster on some models.

Source: Hugging Face

Hugging Face confirmed on 20 October that TorchTPU, a native PyTorch backend for Google’s TPUs built with Meta’s PyTorch team, now runs Transformers, Diffusers and Kernels directly, with no porting step required.

Until now, PyTorch on TPUs meant going through PyTorch/XLA’s torch_xla, or Hugging Face’s own Optimum TPU wrapper. With TorchTPU, a TPU is just another device: Transformers takes device_map="tpu", Diffusers takes pipe.to("tpu"). The torch-tpu 0.2.0 package landed on PyPI on 19 October. Hugging Face says that at batch size 64, decoding is 21.5 times faster on Qwen3-30B-A3B and 12.8 times faster on Gemma 4 26B-A4B than the previous default, measured on a TPU v6e-8.

model = AutoModelForCausalLM.from_pretrained(model_id, device_map="tpu")

Why it matters: the main friction in trying TPUs for inference or training has always been rewriting code for torch_xla. If device_map="tpu" is really a drop-in swap for device_map="cuda", that removes the main reason teams default to GPUs out of familiarity rather than cost or availability.

The caveat: Hugging Face’s own numbers put it at 96.5% of its TPU test suite passing, with 34 known issues logged and 19 already fixed, so this is early and some rough edges remain.

Share

Related reading

ML for Web Devs

Running a model in the browser, and when that is a good idea

Transformers.js and WebGPU put real models on the client. The demo is easy; knowing which features belong there is the part worth thinking about.

ML for Web Devs

Machine learning for web developers: the only mental model you need

You do not need calculus to work with models. You need one idea, that a model is a function you fit instead of write, and the rest follows from things you already know.

Engineering

The N+1 you cannot see

Everyone knows to eager load. The queries that actually take sites down are the ones hiding in an accessor, a Blade partial, or a policy. There is a one-line switch that finds all of them.

More from The Wire

All briefs