Google and Meta bring TorchTitan to TPUs with PyTorch-native training recipes

A native TPU backend for PyTorch

Until now, teams that wanted to train PyTorch models on Cloud TPUs had two main options, according to the Google Developers Blog: rewrite their codebase in JAX, or work with the legacy torch_xla. Google says that has changed with the launch of TorchTPU, which gives Cloud TPUs a native backend, torch.device("tpu"), integrated into PyTorch's core interface. Enabling a training script to use TPUs now only requires minimal changes to device initialization code.

Device support alone is not enough, the post argues, because developers also need high-performance, production-grade distributed training recipes. That is the role of TorchTitan TPU, a fork of Meta's TorchTitan adapted for Google Cloud TPUs. It is available at github.com/google-pytorch/torchtitan. Google and Meta co-developed the TPU-native execution passes inside TorchTitan, so engineering teams do not need to invent their own TPU sharding strategies or hand-write low-level XLA communication passes.

Recipes, decoupled models and swappable kernels

TorchTitan provides recipes for pre-training large dense models such as Llama 3 70B/405B and Qwen 3 8B/32B, Mixture-of-Experts models such as Qwen 3 235B MoE, and diffusion transformers such as FLUX.1 12B. It also acts as the trainer in post-training pipelines like VeRL for Group Relative Policy Optimization (GRPO) and reinforcement learning alignment. In those workflows it handles the parameter update and synchronizes weights with serving engines such as vLLM-TorchTPU.

Model definitions are written in pure, single-device PyTorch with no CUDA-specific code. Collectives are added by a framework layer such as FSDP, by declarative sharding plans such as tensor parallelism, or by backend plugins such as MoE layers. TorchTPU supports C10D collective APIs and the PyTorch DeviceMesh abstraction.

The configuration tree that builds a model can be rewritten before construction, which lets the TPU recipes swap in TPU-friendly blocks without touching the model definition or the shared trainer loop. Google lists three replacements:

• Splash Attention, a Pallas kernel that handles causal, sliding-window and packed-document masks without materializing the attention matrix.
• A fused Pallas kernel for the final projection and cross-entropy loss that never writes the full logits tensor to HBM.
• Grouped Pallas kernels for MoE expert matmuls.

Each kernel is a JAX function turned into a PyTorch custom op by the torch.tpu.pallas.jax_op decorator, and each replacement is a config flag. Kernels can come from tokamax, the OpenXLA kernel library for TPU and GPU.

SimpleFSDP and Graph Trainer

The post explains that FSDP2 is built around the GPU execution model, with collectives issued from module hooks and overlapped using CUDA streams. A compiler cannot see that, so on TPU a torch.compile region breaks at every FSDP hook and a whole-step graph cannot be captured. TorchTPU provides no-op stream and event objects only so FSDP2 runs.

SimpleFSDP takes the compiler's view instead. Parameter sharding becomes a DTensor parametrization, so the all-gather and reduce-scatter are ordinary nodes in the traced graph. XLA can then schedule all-gathers ahead of the matmuls that need them and overlap reduce-scatters with backward compute. Mixed precision is handled by a policy on the same mechanism, and hybrid sharding for multislice uses a 2-D mesh. On FLUX, Google says this brought TorchTitan to 99% performance parity with JAX MaxDiffusion.

Graph Trainer goes further by capturing the full training step (forward, loss and backward) as one FX graph with PyTorch's make_fx tracer, so the step runs as a single XLA program. The post says this makes PyTorch training on TPU conceptually the same as JAX. It also enables scan over layers and a graph splitter that offloads activations to host memory between separately compiled forward and backward programs. According to Google, this path reaches 29%+ MFU on Qwen3 235B. With a functional AdamW update traced into the step, Qwen3-8B reaches 51-52% MFU.

A playbook of configuration changes

Google also published a TPU model optimization playbook, described in the TorchTitan TPU optimization guide, with nearly all levers already wired into the recipes. The post says that for most models, getting good MFU on TPU is "a sequence of configuration changes on an existing recipe rather than new code." The order matters, since each step changes what the next one measures.

The first step is choosing a baseline: eager, compile, or Graph Trainer. Eager mode runs PyTorch models on TPU with virtually no changes, which suits bringing up and debugging a model, but it is not a performance configuration because every op is its own program. TorchTPU's deferred-and-fused eager mode recovers some of that by letting XLA fuse across op boundaries. For PyTorch teams, the practical message is that TPU training no longer has to start with a JAX rewrite. The available text is cut off before the remaining playbook steps.