TorchTPU is an open-source, native PyTorch backend for Google Cloud TPUs, built by Google in collaboration with the PyTorch team at Meta. With TorchTPU, you can run standard PyTorch code on TPUs by setting device = torch.device("tpu") (or device = torch.accelerator.current_accelerator()), then debug in eager mode, compile with torch.compile, and scale from a single chip to a multi-host slice with torch.distributed.
In this video, we walk through how TorchTPU integrates with the PyTorch dispatcher, StableHLO, and the XLA compiler, and run a live demo on a Cloud TPU v6e (Trillium) chip in both eager and compiled modes. We also cover multi-chip scaling through the standard PyTorch process group API (tpu_dist) and integrations across the open-source ecosystem, including Hugging Face Transformers, TorchTitan, vLLM, SGLang, Ray, and llm-d.
Chapters:
0:00 – Introduction
0:17 – Cloud TPUs for machine learning
0:32 – What is TorchTPU?
0:53 – Setting the TPU device in PyTorch
1:29 – How it works: PyTorch dispatcher, StableHLO, and XLA
1:59 – Demo: Eager mode and torch.compile on Trillium (v6e)
3:08 – Scaling with torch.distributed
3:31 – PyTorch ecosystem integrations
3:57 – Custom Pallas kernels and profiling tools
4:36 – Getting started
🔗 Resources:
* Explore TorchTPU on GitHub → https://g.dev/cloud/torch-tpu
* PyTorch on Cloud TPU → https://g.dev/cloud/pytorch-on-tpu
* TorchTitan on TPU → https://g.dev/cloud/torchtitan
* Learn more about Google Cloud TPU → https://g.dev/cloud/tensor-processing-units
Speaker: Rob Mulla
Products Mentioned: Tensor Processing Units
