![\"The](\"https://huggingface.co/buckets/huggingface/hf-torchtpu-blog-assets/resolve/images/banner.jpg\")
If you ever tried to run PyTorch on a TPU, you probably ported your code to torch_xla, or you ended up having to learn JAX. With TorchTPU, a native PyTorch backend, TPU is now just another device. It opens up native support of TPU in the Hugging Face ecosystem. Transformers, Diffusers and Kernels run natively on TPU. Pick a model from the Hub, send it to "tpu", and you're done.
TPU stands for Tensor Processing Unit. TPUs are Google's own AI accelerators. They are chips built 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 uses them to train and serve its own models, including Gemini.
\nTPUs are available on Google Cloud, from a single chip up to pods of thousands of chips. They shine when you need a lot of compute, for training or for serving at scale.
\nThe catch has always been software. The TPU stack grew up around TensorFlow and JAX, and PyTorch users had to work harder to get in.
\nUntil now, running PyTorch on TPUs meant going through PyTorch/XLA (torch_xla). It worked, but it had its quirks. Tensors were lazy (the tensors, not you), so your code didn't always run when you thought it did. Small changes could trigger graph recompilations, and debugging was hard to reason about.
To make life easier, we built Optimum TPU. It gave Hugging Face models inference and training support on TPUs. Now that TorchTPU brings native support, Optimum TPU is not required anymore to run Hugging Face models with Transformers and Diffusers.
\nTorchTPU adds TPU as a new backend to PyTorch. You write regular PyTorch code. It comes with Pallas kernels tuned for TPUs, support for distributed training, and a torch.compile integration for when you want more speed. Installing it on a TPU VM takes one command:
\nThe extra index gives you the CPU build of torch. Without it, pip pulls the much bigger CUDA build, which you don't need on a TPU.
Transformers v5 was built around one idea: be the source of truth for model definitions, so they can run anywhere. "Anywhere" now includes TPUs. To load a model on a TPU, pass device_map="tpu". That's it.
\n\n = \n = \n = \n = \n = \n\nThe model runs in eager mode by default. You can also use torch.compile for better performance, or run it in a distributed setup.
TPU is now a first-class device in Diffusers. For prototyping or debugging, eager mode only needs pipe.to("tpu"). For deployment, compiled mode adds one compile call. Your pipelines and code stay the same, and there's no TPU fork to maintain.
\n\n\n = \n # eager mode: that's the whole change\n# optional: compiled mode\n\n\n = .\nThe first call compiles, and every call after that with the same image size reuses the graph.
\nModels too big for one chip? Pass a TensorParallelConfig to from_pretrained and the transformer is sharded 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 ones that don't are just a tp_plan away. The Diffusers TPU guide has the full recipe.
Here are some examples run on a Cloud TPU v6e-4:
\nTP=4 (each model split across four chips) · bf16 · batch 1
\n| MiniMax-H3 (video + audio) | \nQwen-Image-2.1 | \n\n |
|---|---|---|
| \n | ![\"A](\"https://huggingface.co/buckets/huggingface/hf-torchtpu-blog-assets/resolve/images/qwen-image-21-capybara-1024.jpg\") | \n![\"Neon](\"https://huggingface.co/buckets/huggingface/hf-torchtpu-blog-assets/resolve/images/qwen-image-21-neon-2048.jpg\") | \n
| "The fluffy plush Hugging Face mascot from the picture, a round yellow smiling face with rosy cheeks and two little open hands, stands at the prow of an ancient Greek trireme […] Sound of crashing waves, creaking timber, wind in the rigging, distant seagulls, and a low epic orchestral swell." · 1344×768, 5.2 s, 49 steps · 188 s compiled, 262 s eager | \n"A capybara wearing a wizard hat, reading a book by candlelight, oil painting" · 1024×1024, 40 steps · 3.5 s compiled, 8.0 s eager | \n"A neon shop sign that reads 'QWEN IMAGE 2.1', rainy night, reflections on wet pavement" · 2048×2048, 40 steps | \n
The Kernels Hub has ready-to-use TPU kernels. Load one with the Kernels library and give your model a speed boost.
\nTPU kernels are written in Pallas, which gives you low-level control over the hardware. You can also bring in kernels from the Tokamax library. Here is a ragged-dot kernel from Tokamax plugged into a Mixture-of-Experts model:
\n\n\n\n\n = \n = \n\n = \nThe bigger the batch, the bigger the win. At batch size 64, decoding is 21.5× faster on Qwen3-30B-A3B and 12.8× faster on Gemma 4 26B-A4B.
\n![\"Decode](\"https://huggingface.co/buckets/huggingface/hf-torchtpu-blog-assets/resolve/images/kernel-ragged-dot-speedup.png\")
We're continuously improving TPU support in transformers by running the library test suite. We are working to add a nightly CI on TPU and make the results available in the (transformers dashboard)[[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard]](https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard%5D).
\n![\"Transformers](\"https://huggingface.co/buckets/huggingface/hf-torchtpu-blog-assets/resolve/images/tpu-ci-dashboard.png\")
It's not all green yet, and we're fine with that. Of the tests that run on TPU, 96.5% pass. Every failure we find becomes a small repro script and an issue on the TorchTPU repo. So far we have reported 34 issues, and 19 are already fixed. Running on TPU also caught a few bugs in Transformers itself, like in StaticCache and the Qwen3 tensor-parallel plans. We fixed those upstream.
If something breaks on TPU, we want to find out before you do.
\nTPUs are no longer a side quest for PyTorch users. Install TorchTPU, grab a model from the Hub, and send it to "tpu".
We can't wait to see what you build 🤗. If you hit a problem, open an issue and tell us about it.
\n", "content_markdown": "\n\nIf you ever tried to run PyTorch on a TPU, you probably ported your code to `torch_xla`, or you ended up having to learn JAX. With [TorchTPU](https://github.com/google-pytorch/torch_tpu), a native PyTorch backend, TPU is now just another device. It opens up native support of TPU in the Hugging Face ecosystem. Transformers, Diffusers and Kernels run natively on TPU. Pick a model from the Hub, send it to `\"tpu\"`, and you're done.\n\n## Wait, what's a TPU? 🤔\n\nTPU stands for [Tensor Processing Unit](https://cloud.google.com/tpu). TPUs are Google's own AI accelerators. They are chips built 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 uses them to train and serve its own models, including Gemini.\n\nTPUs are available on [Google Cloud](https://docs.cloud.google.com/tpu/docs/intro-to-tpu), from a single chip up to pods of thousands of chips. They shine when you need a lot of compute, for training or for serving at scale.\n\nThe catch has always been software. The TPU stack grew up around TensorFlow and JAX, and PyTorch users had to work harder to get in.\n\n## The old way\n\nUntil now, running PyTorch on TPUs meant going through PyTorch/XLA (`torch_xla`). It worked, but it had its quirks. Tensors were lazy (the tensors, not you), so your code didn't always run when you thought it did. Small changes could trigger graph recompilations, and debugging was hard to reason about.\n\nTo make life easier, we built Optimum TPU. It gave Hugging Face models inference and training support on TPUs. Now that TorchTPU brings native support, Optimum TPU is not required anymore to run Hugging Face models with Transformers and Diffusers.\n\n## Enter TorchTPU 🚀\n\n[TorchTPU](https://github.com/google-pytorch/torch_tpu) adds TPU as a new backend to PyTorch. You write regular PyTorch code. It comes with Pallas kernels tuned for TPUs, support for distributed training, and a `torch.compile` integration for when you want more speed. Installing it on a TPU VM takes one command:\n\n```shell\npip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu\n```\n\nThe extra index gives you the CPU build of `torch`. Without it, pip pulls the much bigger CUDA build, which you don't need on a TPU.\n\n## Your favourite libraries, now on TPU\n\n### Transformers\n\nTransformers v5 was built around one idea: be the source of truth for model definitions, so they can run anywhere. \"Anywhere\" now includes TPUs. To load a model on a TPU, pass `device_map=\"tpu\"`. That's it.\n\n```py\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nmodel_id = \"Qwen/Qwen3.5-9B\"\nmodel = AutoModelForCausalLM.from_pretrained(\n model_id,\n device_map=\"tpu\", # will load model directly on the TPU\n)\ntokenizer = AutoTokenizer.from_pretrained(model_id)\ninputs = tokenizer(\"The future of AI is\", return_tensors=\"pt\").to(\"tpu\")\nout = model.generate(**inputs, max_new_tokens=64)\nprint(tokenizer.decode(out[0], skip_special_tokens=True))\n```\n\nThe model runs in eager mode by default. You can also use `torch.compile` for better performance, or run it in a distributed setup.\n\n### Images? Videos? Diffusers paints them on TPUs 🎨\n\nTPU is now a first-class device in Diffusers. For prototyping or debugging, eager mode only needs `pipe.to(\"tpu\")`. For deployment, compiled mode adds one `compile` call. Your pipelines and code stay the same, and there's no TPU fork to maintain.\n\n```py\nimport torch\nfrom diffusers import StableDiffusion3Pipeline\n\npipe = StableDiffusion3Pipeline.from_pretrained(\n \"stabilityai/stable-diffusion-3.5-medium\", dtype=torch.bfloat16\n)\npipe.to(\"tpu\") # eager mode: that's the whole change\n# optional: compiled mode\npipe.transformer.compile(backend=\"tpu\", fullgraph=True, dynamic=False)\n\nimage = pipe(\"a corgi astronaut planting a flag on a pizza moon\").images[0]\n```\n\nThe first call compiles, and every call after that with the same image size reuses the graph.\n\nModels too big for one chip? Pass a `TensorParallelConfig` to `from_pretrained` and the transformer is sharded 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 ones that don't are just a `tp_plan` away. The [Diffusers TPU guide](https://huggingface.co/docs/diffusers/main/optimization/tpu) has the full recipe.\n\nHere are some examples run on a Cloud TPU v6e-4:\n\n*TP=4 (each model split across four chips) · bf16 · batch 1*\n\n| [MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) (video + audio) | [Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) | |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| |  |  |\n| *\"The fluffy plush Hugging Face mascot from the picture, a round yellow smiling face with rosy cheeks and two little open hands, stands at the prow of an ancient Greek trireme \\[…\\] Sound of crashing waves, creaking timber, wind in the rigging, distant seagulls, and a low epic orchestral swell.\"* · 1344×768, 5.2 s, 49 steps · 188 s compiled, 262 s eager | *\"A capybara wearing a wizard hat, reading a book by candlelight, oil painting\"* · 1024×1024, 40 steps · 3.5 s compiled, 8.0 s eager | *\"A neon shop sign that reads 'QWEN IMAGE 2.1', rainy night, reflections on wet pavement\"* · 2048×2048, 40 steps |\n\n### Kernels\n\nThe [Kernels Hub](https://huggingface.co/kernels?sort=trending&search=tpu) has ready-to-use TPU kernels. Load one with the Kernels library and give your model a speed boost.\n\nTPU kernels are written in [Pallas](https://docs.jax.dev/en/latest/pallas/tpu/), which gives you low-level control over the hardware. You can also bring in kernels from the [Tokamax](https://github.com/openxla/tokamax) library. Here is a ragged-dot kernel from Tokamax plugged into a Mixture-of-Experts model:\n\n```py\nfrom kernels import get_kernel\nfrom transformers import AutoModelForCausalLM\nfrom transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS\n\nrepo_id = \"tengomucho/ragged-dot-tpu\"\nkernel = get_kernel(repo_id, version=1, trust_remote_code=[repo_id])\nALL_EXPERTS_FUNCTIONS.register(\"tokamax_ragged_dot\", kernel.experts_forward)\nmodel = AutoModelForCausalLM.from_pretrained(\n \"Qwen/Qwen3-30B-A3B-Instruct-2507\",\n experts_implementation=\"tokamax_ragged_dot\",\n)\n```\n\nThe bigger the batch, the bigger the win. At batch size 64, decoding is 21.5× faster on Qwen3-30B-A3B and 12.8× faster on Gemma 4 26B-A4B.\n\n[](https://huggingface.co/buckets/huggingface/hf-torchtpu-blog-assets/resolve/images/kernel-ragged-dot-speedup.png) *Decode throughput on a TPU v6e-8, Hub kernel vs. the default implementation.*\n\n## We test it so you don't have to 🧪\n\nWe're continuously improving TPU support in transformers by running the library test suite. We are working to add a nightly CI on TPU and make the results available in the (transformers dashboard)\\[[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard]](https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard%5D).\n\n[](https://huggingface.co/buckets/huggingface/hf-torchtpu-blog-assets/resolve/images/tpu-ci-dashboard.png) *The Transformers CI dashboard, with TPU results next to other accelerators. Work in progress!*\n\nIt's not all green yet, and we're fine with that. Of the tests that run on TPU, 96.5% pass. Every failure we find becomes a small repro script and an issue on the TorchTPU repo. So far we have reported 34 issues, and 19 are already fixed. Running on TPU also caught a few bugs in Transformers itself, like in `StaticCache` and the Qwen3 tensor-parallel plans. We fixed those upstream.\n\nIf something breaks on TPU, we want to find out before you do.\n\n## Get started\n\nTPUs are no longer a side quest for PyTorch users. Install TorchTPU, grab a model from the Hub, and send it to `\"tpu\"`.\n\n- [TorchTPU on GitHub](https://github.com/google-pytorch/torch_tpu)\n- [TPU kernels on the Hub](https://huggingface.co/kernels?sort=trending&search=tpu)\n- [Cloud TPU documentation](https://docs.cloud.google.com/tpu/docs)\n\nWe can't wait to see what you build 🤗. If you hit a problem, open an issue and tell us about it.\n", "dateCreated": "2026-10-20T17:55:58.869157456+00:00", "date_created": "2026-10-20T00:00:00+00:00", "date_modified": "2026-10-20T00:00:00+00:00", "date_published": "2026-10-20T00:00:00+00:00", "description": "If you ever tried to run PyTorch on a TPU, you probably ported your code to torch_xla, or you ended up having to learn JAX. With TorchTPU, a native PyTorch backend, TPU is now just another device. It opens up native support of TPU in the…", "external_url": "https://huggingface.co/blog/hf-torchtpu", "git": { "commit": "ebb0941c5c06a2ab2377342d607535262a3e969e", "path": "items/huggingface-co/2026/10/2026-10-20-hugging-face-heart-torchtpu-tpus-natively.md", "permalink": "https://github.com/aymericbeaumet/aggr-instance/blob/ebb0941c5c06a2ab2377342d607535262a3e969e/items/huggingface-co/2026/10/2026-10-20-hugging-face-heart-torchtpu-tpus-natively.md" }, "id": "urn:aggr:item:ff289941cd8f1e94d5fcedee41061ada5750027e", "isBasedOn": { "@id": "https://huggingface.co/blog/hf-torchtpu" }, "mainEntity": { "@id": "https://huggingface.co/blog/hf-torchtpu", "@type": "CreativeWork", "archivedAt": { "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-10-20-hugging-face-heart-torchtpu-tpus-natively/#webpage" }, "dateModified": "2026-10-20T00:00:00+00:00", "datePublished": "2026-10-20T00:00:00+00:00", "headline": "Hugging Face TorchTPU: TPUs, Natively", "image": { "@type": "ImageObject", "caption": null, "height": 90, "url": "https://huggingface.co/blog/assets/hf-torchtpu/thumbnail.png", "width": 160 }, "timeRequired": "PT6M", "url": "https://huggingface.co/blog/hf-torchtpu", "wordCount": 1129 }, "name": "Archived snapshot: Hugging Face TorchTPU: TPUs, Natively", "preview": { "alt": null, "color": null, "height": 90, "placeholder": { "data_url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=", "hash": "" }, "url": "https://huggingface.co/blog/assets/hf-torchtpu/thumbnail.png", "width": 160 }, "reading_minutes": 6, "replicated_at": null, "source": "huggingface-co", "source_name": "Hugging Face - Blog", "summary": null, "tags": [], "temporalCoverage": "2026-10-20T17:55:58.869157456+00:00", "title": "Hugging Face TorchTPU: TPUs, Natively", "url": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-10-20-hugging-face-heart-torchtpu-tpus-natively/", "word_count": 1129 }