For those who ever tried to run PyTorch on a TPU, you in all probability ported your code to torch_xla, otherwise you ended up having to study JAX. With TorchTPU, a local PyTorch backend, TPU is now simply one other system. It opens up native help of TPU within the Hugging Face ecosystem. Transformers, Diffusers and Kernels run natively on TPU. Decide a mannequin from the Hub, ship it to "tpu", and also you’re executed.
Wait, what’s a TPU? 🤔
TPU stands for Tensor Processing Unit. TPUs are Google’s personal AI accelerators. They’re chips constructed for one job: the matrix math that neural networks spend most of their time on. Google has been designing them for about ten years, and makes use of them to coach and serve its personal fashions, together with Gemini.
TPUs can be found on Google Cloud, from a single chip as much as pods of hundreds of chips. They shine while you want quite a lot of compute, for coaching or for serving at scale.
The catch has all the time been software program. The TPU stack grew up round TensorFlow and JAX, and PyTorch customers needed to work tougher to get in.
The outdated manner
Till now, working PyTorch on TPUs meant going by way of PyTorch/XLA (torch_xla). It labored, however it had its quirks. Tensors have been lazy (the tensors, not you), so your code did not all the time run while you thought it did. Small modifications might set off graph recompilations, and debugging was arduous to purpose about.
To make life simpler, we constructed Optimum TPU. It gave Hugging Face fashions inference and coaching help on TPUs. Now that TorchTPU brings native help, Optimum TPU will not be required anymore to run Hugging Face fashions with Transformers and Diffusers.
Enter TorchTPU 🚀
TorchTPU provides TPU as a brand new backend to PyTorch. You write common PyTorch code. It comes with Pallas kernels tuned for TPUs, help for distributed coaching, and a torch.compile integration for while you need extra velocity. Putting in it on a TPU VM takes one command:
pip set up torch-tpu --extra-index-url https://obtain.pytorch.org/whl/cpu
The additional index provides you the CPU construct of torch. With out it, pip pulls the a lot greater CUDA construct, which you do not want on a TPU.
Your favorite libraries, now on TPU
Transformers
Transformers v5 was constructed round one concept: be the supply of reality for mannequin definitions, to allow them to run anyplace. “Wherever” now consists of TPUs. To load a mannequin on a TPU, cross device_map="tpu". That is it.
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3.5-9B"
mannequin = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="tpu",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The way forward for AI is", return_tensors="pt").to("tpu")
out = mannequin.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(out[0], skip_special_tokens=True))
The mannequin runs in keen mode by default. You too can use torch.compile for higher efficiency, or run it in a distributed setup.
Pictures? Movies? Diffusers paints them on TPUs 🎨
TPU is now a first-class system in Diffusers. For prototyping or debugging, keen mode solely wants pipe.to("tpu"). For deployment, compiled mode provides one compile name. Your pipelines and code keep the identical, and there is not any TPU fork to take care of.
import torch
from diffusers import StableDiffusion3Pipeline
pipe = StableDiffusion3Pipeline.from_pretrained(
"stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16
)
pipe.to("tpu")
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)
picture = pipe("a corgi astronaut planting a flag on a pizza moon").photographs[0]
The primary name compiles, and each name after that with the identical picture dimension reuses the graph.
Fashions too massive for one chip? Cross a TensorParallelConfig to from_pretrained and the transformer is sharded throughout chips because it hundreds, so the complete mannequin by no means has to slot in one place. Many widespread pipelines help it out of the field, and those that do not are only a tp_plan away. The Diffusers TPU information has the complete recipe.
Listed here are some examples run on a Cloud TPU v6e-4:
TP=4 (every mannequin break up throughout 4 chips) · bf16 · batch 1
| MiniMax-H3 (video + audio) | Qwen-Picture-2.1 | |
|---|---|---|
 |
 |
|
| “The fluffy plush Hugging Face mascot from the image, a spherical yellow smiling face with rosy cheeks and two little open palms, stands on the prow of an historical Greek trireme […] Sound of crashing waves, creaking timber, wind within the rigging, distant seagulls, and a low epic orchestral swell.” · 1344×768, 5.2 s, 49 steps · 188 s compiled, 262 s keen | “A capybara sporting a wizard hat, studying a ebook by candlelight, oil portray” · 1024×1024, 40 steps · 3.5 s compiled, 8.0 s keen | “A neon store signal that reads ‘QWEN IMAGE 2.1’, wet night time, reflections on moist pavement” · 2048×2048, 40 steps |
Kernels
The Kernels Hub has ready-to-use TPU kernels. Load one with the Kernels library and provides your mannequin a velocity increase.
TPU kernels are written in Pallas, which provides you low-level management over the {hardware}. You too can usher in kernels from the Tokamax library. Here’s a ragged-dot kernel from Tokamax plugged right into a Combination-of-Consultants mannequin:
from kernels import get_kernel
from transformers import AutoModelForCausalLM
from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
repo_id = "tengomucho/ragged-dot-tpu"
kernel = get_kernel(repo_id, model=1, trust_remote_code=[repo_id])
ALL_EXPERTS_FUNCTIONS.register("tokamax_ragged_dot", kernel.experts_forward)
mannequin = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-30B-A3B-Instruct-2507",
experts_implementation="tokamax_ragged_dot",
)
The larger the batch, the larger the win. At batch dimension 64, decoding is 21.5× sooner on Qwen3-30B-A3B and 12.8× sooner on Gemma 4 26B-A4B.
Decode throughput on a TPU v6e-8, Hub kernel vs. the default implementation.
We check it so you do not have to 🧪
We’re constantly enhancing TPU help in transformers by working the library check suite. We’re working so as to add a nightly CI on TPU and make the outcomes obtainable within the (transformers dashboard)[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard].
The Transformers CI dashboard, with TPU outcomes subsequent to different accelerators. Work in progress!
It is not all inexperienced but, and we’re effective with that. Of the checks that run on TPU, 96.5% cross. Each failure we discover turns into a small repro script and a problem on the TorchTPU repo. Up to now we’ve reported 34 points, and 19 are already mounted. Working on TPU additionally caught a number of bugs in Transformers itself, like in StaticCache and the Qwen3 tensor-parallel plans. We mounted these upstream.
If one thing breaks on TPU, we need to discover out earlier than you do.
Get began
TPUs are not a aspect quest for PyTorch customers. Set up TorchTPU, seize a mannequin from the Hub, and ship it to "tpu".
We will not wait to see what you construct 🤗. For those who hit an issue, open a problem and inform us about it.
