# Integrating TorchTPU into TorchTitan and Optimizing Models for TPU > Source: > Published: 2026-10-20 15:17:51.385342+00:00 Over the past years, Meta’s **TorchTitan** has emerged as the PyTorch ecosystem’s premier clean-room reference platform for large-scale generative AI model training. Built entirely on PyTorch-native scaling primitives, TorchTitan provides an unencumbered, highly modular showcase for pre-training flagship dense and Mixture-of-Experts (MoE) architectures like Llama 3, Qwen 3, and FLUX. Historically, developers looking to train PyTorch models on **Google Cloud TPUs** faced some challenges: either undertake a complete codebase rewrite into JAX, or grapple with legacy torch_xla. That paradigm has officially shifted. With the launch of **TorchTPU**, Cloud TPUs now feature a native backend (`torch.device("tpu")`) integrated directly into PyTorch's core interface. Enabling a training script to use TPUs now only requires minimal changes to your device initialization code. Providing device compatibility, however, is only half the battle, developers also need high performance production-grade, distributed training recipes. This is where **TorchTitan on TPU** comes in. **TorchTitan TPU** is a fork of Meta’s TorchTitan implementation that is adapted to enable model training on Google Cloud TPUs. The fork is available at [https://github.com/google-pytorch/torchtitan/](https://github.com/google-pytorch/torchtitan/) ork is available at [https://github.com/google-pytorch/torchtitan/](https://github.com/google-pytorch/torchtitan/) Rather than requiring engineering teams to invent custom TPU sharding strategies or hand-craft low-level XLA communication passes from scratch, Google and Meta co-developed TPU-native execution passes directly inside TorchTitan. TorchTitan serves as the definitive vehicle for scaling PyTorch on Google Cloud TPUs. TorchTitan provides official, battle-tested recipes for pre-training large dense models (such as Llama 3 70B/405B and Qwen 3 8B/32B), Mixture-of-Experts (such as Qwen 3 235B MoE), and diffusion transformers (such as FLUX.1 12B). Beyond initial pre-training, TorchTitan serves as the scalable trainer inside post-training pipelines like VeRL for Group Relative Policy Optimization (GRPO) and Reinforcement Learning (RL) alignment. In these workflows, TorchTitan handles the parameter update step while synchronizing weights with high-throughput serving engines like vLLM-TorchTPU. Model definitions in TorchTitan are written in pure, single-device PyTorch (nn.Module) with no CUDA-specific code. Collectives are inserted either by a composable framework layer (e.g. FSDP), by declarative sharding plans (e.g. tensor-parallelism), or as backend-specific plugins (e.g. MoE layers). TorchTPU supports C10D collective APIs and the PyTorch DeviceMesh abstraction. Together, this allows the TPU path to choose which parallelisms to apply for a model, and which implementations to use. TorchTitan builds a model from a configuration tree, and that tree can be rewritten before the model is constructed. This offers a clean way to replace individual blocks of a model with versions that run better on TPU, without touching the model definition or the shared trainer loop. We use this actively. The TPU recipes replace standard attention with Splash Attention, a Pallas kernel that handles causal, sliding-window and packed-document masks without materializing the attention matrix; replace the final projection and cross-entropy loss with a fused Pallas kernel that never writes the full logits tensor to HBM; and replace MoE expert matmuls with grouped Pallas kernels. Each kernel is a JAX function turned into a PyTorch custom op by TorchTPU's `torch.tpu.pallas.jax_op` decorator, which is all the glue a new kernel needs. Each replacement is a config flag, so a recipe can turn it on, and a test can compare it against the plain PyTorch path it replaces. Kernels can come from [tokamax](https://github.com/openxla/tokamax), the OpenXLA kernel library for TPU and GPU. FSDP2 is built around the GPU execution model. Its all-gathers and reduce-scatters are issued from module hooks as explicit collective calls, and overlap with compute comes from running them on dedicated CUDA streams synchronized with events. None of that is visible to a compiler: the collectives sit outside the traced region, so on TPU a torch.compile region breaks at every FSDP hook and a whole-step graph cannot be captured at all. TorchTPU provides no-op stream and event objects only so that FSDP2 runs; XLA has no use for them, since it schedules work from data dependencies. [SimpleFSDP](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/graph_trainer/simple_fsdp.py) takes the compiler's view instead. Parameter sharding is a DTensor parametrization: reading a module's weight performs a redistribute from sharded to replicated, which is the all-gather, and its autograd backward is the reduce-scatter. There are no hooks and no explicit collective calls, so both collectives are ordinary nodes in the traced graph, whether that graph is a per-block torch.compile region or the Graph Trainer's single make_fx step. The XLA compiler schedules the all-gathers ahead of the matmuls that need them and overlaps the reduce-scatters with the backward compute, inside one program. Mixed precision is part of the same mechanism, a policy that sets the gather dtype (bf16 parameters) and the reduce dtype (fp32 or bf16 gradients) on that redistribute. Hybrid sharding for multislice is the same code on a 2-D mesh: shard within a slice over ICI, replicate across slices over DCN. On FLUX this is what brought TorchTitan to 99% performance parity with JAX MaxDiffusion: the per-layer HLO for a SimpleFSDP block is a single module with the weight all-gathers, matmuls, attention and norms fused together, which is the shape of program XLA was built to optimize. Many PyTorch trainers either run in eager mode with CUDA graphs to cut launch overhead, or apply torch.compile to isolated regions, typically one transformer block at a time. TorchTPU does not yet support the PyTorch `torch.Stream` abstraction, which limits the performance of mixing eager-based FSDP2 with layer-wise torch.compile. SimpleFSDP helps by expressing parameter sharding in a way that is compatible with compilation, whether using layer-wise torch.compile or full-graph capture. TorchTitan's Graph Trainer provides full training-step capture (forward, loss and backward), producing one FX graph ahead of time with PyTorch's make_fx tracer, skipping AOTAutograd partitioning, so the whole computation is visible to XLA). GraphTrainer also runs optimization passes: to implement graph-based activation memory policy that decides what to save, recompute or offload. The train step then executes as a single XLA program. This is the model XLA was built for, and it makes PyTorch training on TPU conceptually the same as JAX, where the whole train step is one jitted function. Owning the whole step as one graph is also what makes the large-model compilation techniques possible: scan over layers, which runs the decoder layers as a single loop body so that compile time and memory stop growing with depth, and the graph splitter, which cuts the step into separately compiled forward and backward programs with activations offloaded to host memory in between. On Qwen3 235B this path reaches 29%+ MFU. The optimizer can join the graph as well: While the fused_adam optimizer is very performant, a functional AdamW update can be traced into the train-step (functional_adamw_step) and ensures an iteration has no host-side work between the first forward op and the weight write-back. On Qwen3-8B this path reaches 51-52% MFU. Every lever below is described in the [TorchTitan TPU optimization guide](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/torch-tpu-optimization-guide.md), and nearly all of them are already wired into [the recipes](https://github.com/google-pytorch/torchtitan/tree/main/torchtitan/experiments/tpu/recipes). For most models, getting a good MFU on TPU is a sequence of configuration changes on an existing recipe rather than new code. The order matters, because each step changes what the next one measures. A lot of optimizations in TorchTitan are available as ready-to-use building blocks that can be enabled by a flag flip (see illustration below). **1. Establish a Baseline** **2. Tune Compute and Kernals** **3. Optimize Memory and Batch Size** **4. Tune Compiler, Communication and Profile** To check that TorchTitan runs efficiently across different classes of models, we tuned and scaled several of them on TPU v7x (Ironwood). At scale, Llama 3 405B (large dense model) holds about 57% MFU from 128 to 1,024 chips. Qwen3 235B (large MoE model) reaches 32.42% MFU on a single 1,024-chip slice and 30.51% on 8,192 chips across eight slices. Muse Glimmer 30B (modern dense model) reaches 43.4% MFU on 128 chips (4x4x8). FLUX.1 (diffusion model) reaches 19.4% MFU on 1,024 chips (8x8x16). For all models we run The following case studies show how we got there for three of these models, each representing a different class: [FLUX.1](https://github.com/google-pytorch/torchtitan/tree/main/torchtitan/experiments/tpu/flux) is a 12-billion-parameter diffusion transformer, which makes it a useful check that the playbook above is not specific to LLMs. We trained it at 1024x1024 resolution on a TPU v6e-8 with FSDP in bf16 and a batch of three samples. Out of the box, in eager mode with FSDP2 and full activation checkpointing, a step took 11.49 seconds. Applying the playbook in order brought that down one stage at a time. [Deferred-and-fused eager execution](https://github.com/google-pytorch/torch_tpu/blob/main/docs/markdown/2_in_depth_4_torchtpu_execution_modes.md) cut the number of HLO module invocations by 28% and the step to 9.95 s (15.5% improvement). Replacing FSDP2 with [SimpleFSDP](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/graph_trainer/simple_fsdp.py) reduced the step to 6.52 s (52.6% improvement): the trace was dominated by FSDP2 copying parameter shards into a flat staging buffer before each all-gather and back out into per-parameter tensors afterwards; SimpleFSDP all-gathers each parameter directly, so the staging buffer and both copies are gone (see the SimpleFSDP section above). The [splash attention kernel](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/kernels/splash_attention.py), with sequence packing and segment ids so that no MXU cycles go to padding, took the step to 5.55 s (17.5% improvement), and [per-block torch.compile](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/torch-tpu-optimization-guide.md), which hands XLA large sub-graphs instead of fine-grained modules and moves control flow out of the Python interpreter, took it to 3.02 s (83.8% improvement). The final step came from [three small, local rewrites in the model code](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/models/flux/model/layers.py) that the profile pointed at: splitting a fused projection into two separate linear layers so its output no longer had to be sliced apart, simplifying the rotary embedding to plain element-wise math, and unpacking Q, K and V with a view instead of a transpose. None of them changes what the model computes; each removes copies and lets XLA fuse more of the graph. Together they brought the step to 2.78 s (8.6% improvement), 4.1x faster than the starting point and about 99% of the throughput of the JAX-based [MaxDiffusion](https://github.com/AI-Hypercomputer/maxdiffusion) on the same hardware. We also scaled FLUX from 8 to 1,024 v7x chips, keeping three images per TPU. The best configuration went from 2.53 s per step (20.76% MFU) on 8 chips to 2.71 s (19.4% MFU) on 1,024, which is 93% weak-scaling efficiency. Past one 64-chip cube, pure FSDP slows down due to the model’s compact size. Hybrid sharding via HSDP fixes this: it shards within each 4x4x4 cube and replicates across cubes. We used the topology-aware device mesh, which maps the parallelism optimally to the physical 8x8x16 torus. [Muse Glimmer 30B](https://github.com/google-pytorch/torchtitan/tree/main/torchtitan/experiments/tpu/muse_glimmer) is a modern dense decoder: 52 layers, grouped-query attention with 32 query and 2 key-value heads, a 202k-token vocabulary with a soft-capped output head, and a layer pattern that mixes sliding-window and global attention, with RoPE skipped on every fourth layer. It is the case that shows how much works out of the box. The model definition, the parallelization and the training loop ran on TPU unchanged; the port came down to two model-specific pieces. Upstream implements the sliding-window pattern with FlexAttention, which has no TPU backend, so the first step was a [windowed SDPA inner attention](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/muse_glimmer/attention.py) that reproduces the window exactly and gives a correctness baseline, checked against upstream on CPU. The second piece was the [fused cross-entropy kernel](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/kernels/linear_softmax_ce_softcap_kernel.py), which we extended to apply the soft-cap to each vocabulary tile as it computes the loss, so the 202k-wide logits never have to be materialized; without that the LM head alone dominated memory. Everything else, including the [Graph Trainer wrapper](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/graph_trainer/muse_glimmer.py), was the standard onboarding ladder, with each feature behind a flag and its own test. The resulting [recipe](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/recipes/ironwood/muse-glimmer-30b/tpu7x-4x4x4/README.md) trains the 30B on 64 v7x chips (a 4x4x4 slice) at sequence length 16,384 with one sequence per device, Graph Trainer plus SimpleFSDP and fp32 gradient reductions, and reaches 5,848 tokens per second per chip at ~44% MFU. [Qwen3-235B-A22B](https://github.com/google-pytorch/torchtitan/tree/main/torchtitan/experiments/tpu/qwen3) is the largest MOE model in this set: 94 layers, 128 experts with 8 active per token, 235 billion parameters of which 22 billion are active per token. It is the case that exercises the large-model half of the playbook. At this depth compiling the whole unrolled step as one program is not an option: XLA plans every buffer of a program at compile time, and a 94-layer graph laid out flat asks for more HBM than the chip has before a step can run, while its compile time grows with every layer. Getting it to work at all meant enabling one of the two Graph Trainer paths for large models. Scan over layers is the XLA-native answer, the same construct JAX and MaxText rely on: the 94 identical layers become one compiled loop body, so compile time and program size stop growing with depth. The graph splitter is the more flexible alternative: it keeps the layers unrolled but cuts the captured step into per-layer executables, which works no matter what kind of layers you have (and does not require layers to be the same). We run both for this model and they land within a few percent of each other. Either way the activations do not fit in HBM at sequence length 4096 with 16 sequences per device, so both paths offload layer activations to host memory and prefetch them back for the backward pass. From there the playbook applied, one change at a time, each measured on the same 64-chip slice. We raised the scoped VMEM budget so the compiler could keep more kernel intermediates on chip, which also removed an exposed update at the tail of the optimizer step. We enabled a set of [libtpu scheduling flags](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/tpu_flags.py) that let XLA's latency-hiding scheduler overlap the FSDP all-gathers and reduce-scatters with the surrounding compute and offload the reduce-scatters to the SparseCores, so most of the collective time is hidden behind compute, leaving only a small exposed reduce-scatter tail in the backward pass. We switched attention to the [tokamax splash kernel](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/kernels/tokamax_splash_attention.py), and replaced the expert matmuls with specialized [paired grouped-matmul kernels](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/kernels/paired_segment_matmul.py) backed by a K-tiled weight-gradient kernel, with hierarchical reduce-scatter and the SparseCore latency-hiding scheduler to hide the expert-gradient collectives. A [fused route-score backward](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/kernels/route_score.py) removed a large fp32 copy from the MoE backward pass. The last round of fixes came from profiling at 1,024 chips: keeping the blocked activation layout under scan offload, raising the backward layer-scheduler threshold so the expert reduce-scatters overlap with compute again, and pipelining small reduce-scatters across the scan loop's back-edge. The recipe runs at ~32% MFU on 256 v7x chips with fp32 gradient reductions kept throughout. We also used this model to test how far TorchTPU scales, from 64 chips up to 8,192-chip v7x pod scale, in both single-slice and multi-slice configurations. A TPU slice is a set of chips connected directly over the [inter-chip interconnect (ICI)](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm); [multislice](https://cloud.google.com/tpu/docs/multislice-introduction) lets one training job span several slices connected over the [data-center network (DCN)](https://cloud.google.com/tpu/docs/multislice-introduction), which is how a job grows past the largest single slice. Within a slice we run FSDP over ICI, sharding parameters, gradients and optimizer state across the slice's chips; across slices the model is replicated and the gradients are all-reduced over DCN once per step, so the only traffic on the slower network is a single gradient reduction, which overlaps with the backward pass. That is [hybrid sharded data parallelism (HSDP)](https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html), and it is what the 2,048- to 8,192-chip runs use, with 1,024-chip slices as the sharding group. The framework behaved as expected throughout: per-chip throughput dropped by 8% going from 64 to 8,192 chips (4,568 to 4,208 tokens per second per chip), near-linear scaling across 128x more chips. One caveat: for a model of this size, data parallelism beyond about 1K chips is a scaling experiment rather than a training recipe. This is a weak-scaling run: the batch per device stays fixed as chips are added, so the global batch grows with the chip count and at 8k chips is larger than one would use to train this model. What the run shows is that the framework, the compiler and the collectives hold up at that scale, not that this is the configuration to train with. The three case studies above are deep dives: one model at a time, with a parity test and a profile behind every change. A team arriving on TPU usually has a different problem. They have a model they already train, often a Hugging Face one, and quite likely it won’t be the model that TorchTitan ships for TPU. They may well want every last percent out of it eventually, but first they need an easy way to get it running at all. That is the problem the [Transformers modeling backend](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/transformers_modeling_backend/README.md) solves: it builds the model from the Transformers library by its Hub id and runs it through TorchTitan's trainer and parallelism instead of its own. On TPU the backend has a [TorchTPU lane](https://github.com/google-pytorch/torchtitan/tree/main/torchtitan/experiments/tpu/transformers_modeling_backend) with a [single configuration](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/transformers_modeling_backend/config_registry.py) that applies the first rungs of the playbook optimizations to any supported checkpoint, with no change to the model code: SimpleFSDP in place of FSDP2, per-block torch.compile, the splash attention kernel [registered as an attention implementation](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/transformers_modeling_backend/attention.py) with the transformers attention interface so the model's own attention calls land on it, the fused Pallas cross entropy [fed the final hidden states and the head weight](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/transformers_modeling_backend/model.py) instead of materialized logits, and, for mixture-of-experts checkpoints, the TPU grouped-matmul experts substituted for the portable implementation. Most of the dense architectures the backend supports train this way unmodified; mixture-of-experts checkpoints run through the same path but without expert parallelism yet. None of this reaches the hand-tuned numbers above, but it is a working, reasonably efficient first run for a model that has no TPU recipe yet, and a way to find out whether it deserves one. Sooner or later a model deserves more than a generic path: it is the one you will train for months, it has an architecture the backend cannot express, or the out-of-the-box number is not the one you need. Porting a model to a new accelerator used to be a specialist project measured in weeks, with the hard-won lessons living in one engineer's head. On TorchTitan TPU it is a repeatable procedure. The repository ships a [model-onboarding skill](https://github.com/google-pytorch/torchtitan/tree/main/torchtitan/experiments/tpu/.claude/skills/tpu_model_onboarding) for coding agents that captures how the case studies above were done and walks an engineer, or the agent itself, through the same ladder, with the tests, the acceptance bars and the traps built in. The roadmap runs along five lines: production readiness, precision, hardware, parallelism and models. On production readiness, the recipes gain the operational tooling that long runs on Google Cloud expect: integration with the [ML Goodput Measurement library](https://docs.cloud.google.com/tpu/docs/goodput), so a run reports how much of its chip time turns into preserved training progress and where the rest goes, and with [Multi-Tier Checkpointing](https://docs.cloud.google.com/kubernetes-engine/docs/how-to/machine-learning/training/multi-tier-checkpointing) on GKE, which keeps checkpoints in local memory and in-cluster replicas with Cloud Storage as the durable tier, so large jobs save often and resume quickly. On precision, the [float8 recipe](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/float8_converter.py) that today ships for Llama 3 70B, built on torchao's Float8Linear and v7x's native fp8 matmuls, extends to the rest of the model set. On hardware, Google has announced its [eighth-generation TPUs](https://blog.google/innovation-and-ai/infrastructure-and-cloud/google-cloud/eighth-generation-tpu-agentic-era/), TPU 8t for pre-training and TPU 8i for inference and reinforcement learning, and TorchTitan will support them as they become available on Google Cloud, with the same recipes carried forward. On parallelism, expert parallelism for the MoE kernels is the open item named earlier, and long-context training is next after it, now that million-token context windows are the norm for frontier models. On models, the targets are the next size up from Qwen3-235B: [Qwen3.8-Max](https://www.alibabagroup.com/document-2021044032125272064), a 2.4-trillion-parameter mixture of experts, and [Kimi K3](https://huggingface.co/moonshotai/Kimi-K3), 2.8 trillion parameters with Kimi Delta Attention in place of standard attention on most layers, which will exercise the kernel plug-in path as much as the scaling one. TorchTitan on TPU is no longer an experiment. The same codebase and plain-PyTorch model definitions train Llama 3 405B at about 57% MFU on up to 1,024 chips, and run Qwen3 235B, a mixture-of-experts model, on up to 1,024 chips in a single slice. In a weak-scaling test with a fixed per-chip batch, Qwen3 235B kept 92% of its per-chip throughput from 64 to 8,192 chips. That shows the framework, compiler and collectives hold up at pod scale, even though a model of this size would not be trained at that scale. The pieces that make this possible are not buried in one recipe. SimpleFSDP, the Graph Trainer with scan and the graph splitter, host offload, the [Pallas kernel library](https://github.com/google-pytorch/torchtitan/tree/main/torchtitan/experiments/tpu/kernels) and the [optimization playbook](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/tpu/torch-tpu-optimization-guide.md) are building blocks, each behind a flag and each documented with the number it is worth. For models that TorchTitan does not ship, the [Transformers modeling backend](https://github.com/google-pytorch/torchtitan/blob/main/torchtitan/experiments/transformers_modeling_backend/README.md) runs 500+ Hugging Face models through the same trainer with the first optimizations already applied, and for teams that want more than that, the [onboarding skill](https://github.com/google-pytorch/torchtitan/tree/main/torchtitan/experiments/tpu/.claude/skills/tpu_model_onboarding) turns porting a model into a repeatable procedure that an engineer or an agent can follow. That leaves two ways to use it. The first is to adopt TorchTitan as the training framework: pick a [recipe](https://github.com/google-pytorch/torchtitan/tree/main/torchtitan/experiments/tpu/recipes), launch it on a slice, and build on a configuration whose performance is already known. The second is to keep the training code you have and take what you need: the kernels are standard PyTorch custom ops, SimpleFSDP and the Graph Trainer are modules you can import, the compiler flags and the playbook apply to any TorchTPU workload, and every recipe spells out what it enables and why. Either way, the fastest way to find out what this does for your model is to run it.