Hugging Face ❤️ TorchTPU: TPUs, Natively

TorchTPU: Bringing Native PyTorch to Google TPUs

Google’s Tensor Processing Units (TPUs) have long been the domain of TensorFlow and JAX users, leaving PyTorch developers to wrestle with the `torch_xla` bridge or rewrite code in JAX. Hugging Face’s TorchTPU changes that narrative by adding a first‑class TPU backend to PyTorch itself. The result? Hugging Face models—Transformers, Diffusers, and custom kernels—can be run on TPUs with the same code you’d use on CPUs or GPUs, simply by targeting the device name "tpu".

---

Why TPUs Matter for Large‑Scale AI

TPUs are Google‑designed AI accelerators optimized for the matrix‑heavy workloads that dominate modern deep learning. Available on Google Cloud from single‑chip VMs to massive pods, they excel at both training and inference at scale. Historically, the TPU software stack revolved around TensorFlow and JAX, making it cumbersome for PyTorch users to tap into this raw compute power. The old workflow required `torch_xla`, which introduced lazy tensors, frequent graph recompilations, and a steep debugging curve.

---

TorchTPU: A Seamless Backend

TorchTPU integrates directly with the PyTorch runtime, exposing TPUs as a native device type. Installation is a single pip command:

```bash pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu ```

The extra index ensures you get the lightweight CPU build rather than the larger CUDA package. Once installed, you write ordinary PyTorch code and move tensors or models to the TPU with `to("tpu")` or by using `device_map="tpu"` in Hugging Face utilities.

Transformers on TPUs

Loading a model from the Hub onto a TPU is as simple as:

```python from transformers import AutoModelForCausalLM, AutoTokenizer modelid = "Qwen/Qwen3.5-9B" model = AutoModelForCausalLM.frompretrained(modelid, devicemap="tpu") tokenizer = AutoTokenizer.frompretrained(modelid) inputs = tokenizer("The future of AI is", returntensors="pt").to("tpu") output = model.generate(**inputs, maxnewtokens=64) print(tokenizer.decode(output[0], skipspecial_tokens=True)) ```

The model runs eagerly by default, but you can boost performance with `torch.compile` or enable distributed training across multiple TPU chips using the built‑in TensorParallel configuration.

Diffusers and Media Generation

Diffusers now treat TPUs as a first‑class device. An eager‑mode pipeline requires only a single `pipe.to("tpu")` call. For production workloads, you can compile the transformer component:

```python import torch from diffusers import StableDiffusion3Pipeline pipe = StableDiffusion3Pipeline.from_pretrained( "stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16 ) pipe.to("tpu") # eager mode

optional compiled mode

pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False) image = pipe("a corgi astronaut planting a flag on a pizza moon").images[0] ```

Large models that exceed a single chip’s memory can be sharded across chips at load time via `TensorParallelConfig`, eliminating the need for a separate TPU‑specific fork.

---

TPU‑Optimized Kernels

The Hugging Face Kernels Hub supplies pre‑written Pallas kernels tuned for TPU hardware. You can pull a kernel with the `kernels` library and plug it into a model, for example integrating a ragged‑dot kernel from the Tokamax collection into a Mixture‑of‑Experts transformer. These kernels often deliver order‑of‑magnitude speedups, especially at larger batch sizes (e.g., 21× faster decoding for Qwen‑3‑30B at batch‑64 on a v6e‑8 TPU).

---

Quality Assurance and Community Support

Hugging Face runs the full Transformers test suite on TPUs, reporting a 96.5 % pass rate on the current CI. Failures are turned into reproducible issue reports on the TorchTPU repository, many of which have already been fixed. This proactive testing not only stabilizes the TPU backend but also surfaces bugs in the core Transformers library that are then patched upstream.

If you encounter a problem, the recommended workflow is to open an issue on the TorchTPU GitHub repo and share details of your environment and code. The community and Hugging Face engineers are quick to respond, ensuring that TPU support continues to improve.

---

Getting Started

1. Spin up a Cloud TPU VM (e.g., a v6e‑4 or v6e‑8 instance).
2. Install TorchTPU with the pip command above.
3. Pick a model from the Hugging Face Hub and load it using `device_map="tpu"` or `pipe.to("tpu")`.
4. Optional: Enable `torch.compile` or distributed TensorParallel for maximum throughput.

With TorchTPU, TPUs are no longer a niche platform for PyTorch users—they’re just another device in your toolbox. Whether you’re prototyping a new LLM, generating high‑resolution images, or scaling inference across a pod of chips, the workflow stays familiar, and the performance gains are substantial.

Happy building on TPUs!