TorchTPU brings native PyTorch support for Google TPUs
PyTorch developers can now target TPU hardware through the same APIs they already use, reducing the migration cost that historically kept most teams on NVIDIA GPUs. Teams running training or inference pipelines with TorchTitan, vLLM, or SGLang can evaluate TPUs without rewriting code. The use of agents for kernel generation and sharding also hints at how optimization work itself may be automated.
Google's Bill Jia and Meta's Zongwei Zhou introduced TorchTPU at PyTorchCon 2026, an open source PyTorch backend for Google TPUs built with PyTorch core maintainers. The backend integrates with Hugging Face Transformers for modeling, TorchTitan for training, and vLLM and SGLang for serving. The keynote also demonstrated how long-horizon agentic workflows can migrate complex model workloads from GPUs to TPUs. Agents were additionally shown handling performance optimization tasks such as quantization, custom kernel generation, and sharding strategies.
Bill Jia (@GoogleCloud) and Zongwei Zhou (@Meta) introduce TorchTPU, a native PyTorch backend for Google TPUs, and agentic developer workflows at #PyTorchCon 2026.
TorchTPU is an open source collaboration between Google and Meta, developed in close partnership with PyTorch core maintainers. It integrates with Hugging Face Transformers for modeling, TorchTitan for training, and vLLM and SGLang for serving.
They demonstrate how long-horizon agentic workflows can migrate complex model workloads from GPUs to TPUs. They explore using agents for performance optimization, including quantization, custom kernel generation, and advanced sharding strategies.
Watch the keynote livestream: https://t.co/6TtaEkvsHz
- AWS Machine Learning Blog10-21 20:00原文
