Research· Important
Google and Meta bring native TPU support to PyTorch's TorchTitan training framework
Google introduced TorchTPU, a native PyTorch backend for Cloud TPUs (torch.device("tpu")), and co-developed a fork of Meta's TorchTitan trainer adapted for TPU training. Previously, running PyTorch training on TPUs meant rewriting code in JAX or relying on the older torch_xla path; now only minimal device-initialization changes are needed. The fork provides tested recipes for pre-training large dense models (Llama 3 70B/405B, Qwen 3 8B/32B), Mixture-of-Experts models (Qwen 3 235B), and diffusion transformers (FLUX.1 12B), plus a 'Graph Trainer' that captures an entire training step as one XLA program for better compiler optimization.