PyTorch now runs natively on Google Cloud TPUs. Google released TorchTPU to fix a long-standing problem. Teams previously had to rewrite their code in JAX to use TPUs. Or they had to use the old torch_xla library which was hard to use. This new backend lives directly inside PyTorch. It uses the standard torch.device("tpu") command. Developers can now train on TPUs with minimal code changes.
Meta's TorchTitan platform is now optimized for this new backend. TorchTitan was built to show how to scale PyTorch models. It handles large dense models and Mixture-of-Experts architectures. Google adapted TorchTitan to work specifically with Cloud TPUs. This fork allows teams to use familiar PyTorch tools. They no longer need to learn a completely new programming language.
How TorchTitan TPU trains models on Cloud TPUs
TorchTitan provides battle-tested recipes for training massive models. These recipes cover pre-training large language models like Llama 3 and Qwen 3. They also support diffusion models like FLUX.1. The platform handles both dense models and sparse Mixture-of-Experts designs. It works for models with hundreds of billions of parameters.
Beyond initial training, the platform supports post-training alignment. It runs Group Relative Policy Optimization inside the training loop. It also handles Reinforcement Learning tasks for safety alignment. The training scripts synchronize weights with high-throughput serving engines. This ensures the model is ready for immediate deployment after training finishes.
Model definition decoupled from model partitioning
Model definitions in TorchTitan are written in pure PyTorch. The code uses standard nn.Module classes without any CUDA-specific instructions. Collective operations are added by a composable framework layer. These layers include FSDP for distributed training or tensor-parallelism for sharding. Backend-specific plugins handle special cases like MoE layers.
This separation allows the TPU path to choose its own parallelism strategy. The system decides which implementation to use for a specific model. Teams can replace blocks of code without touching the main model definition. They can swap in versions that run better on TPU hardware. This keeps the shared trainer loop unchanged while optimizing performance.