TorchTitan on Cloud TPU: Configure and Run Training Recipes

What TorchTitan on TPU Actually Is#

TorchTitan TPU is a fork of Meta's TorchTitan, maintained by Google at github.com/google-pytorch/torchtitan. It adapts TorchTitan's training platform to run on Google Cloud TPUs through TorchTPU, which provides a native PyTorch backend via torch.device("tpu"). No JAX rewrite, no legacy torch_xla. The fork plugs into PyTorch's core device interface.

Meta and Google co-developed the TPU-native execution passes inside TorchTitan, so the sharding strategies and XLA communication passes are already wired in rather than left for you to build. The shipped recipes cover Llama 3 (70B, 405B), Qwen3 (8B, 32B, 235B MoE), and FLUX.1 12B. As of October 20, 2026, the case studies run on TPU v7x (Ironwood).

On that hardware, Llama 3 405B holds about 57% MFU from 128 to 1,024 chips. Qwen3 235B reaches 32.42% MFU on a single 1,024-chip slice and 30.51% across 8,192 chips in eight slices. Muse Glimmer 30B hits 43.4% on 128 chips (4x4x8). FLUX.1 reaches 19.4% on 1,024 chips (8x8x16). 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.

Enabling TPU in Your Training Script#

The minimal change is device initialization. Instead of torch.device("cuda"), use torch.device("tpu"). TorchTPU integrates this directly into PyTorch's core interface, so most model code stays untouched.

The rest of the configuration happens through the fork's flags. Here is the progression from a bare eager run to a fully optimized one:

# Eager mode with deferred fusion: bring the model up, debug it
python train.py --tpu_config.eager_mode=DEFER_AND_FUSE

# Per-block torch.compile: each transformer block becomes one XLA program
python train.py --compile.enable

# Graph Trainer: forward, loss, backward, and optimizer as one program
python train.py --tpu_config.use_graph_trainer

Eager mode with DEFER_AND_FUSE lets XLA fuse across op boundaries without code changes. It is the right starting point for debugging, but every op is still its own program and the TPU idles between launches. Per-block torch.compile turns each transformer block into a single XLA program, so ops fuse and launch once. Graph Trainer captures the whole training step into one program with no eager glue at all.

Sharding With SimpleFSDP#

The fork uses SimpleFSDP instead of FSDP2, and the reason is specific. FSDP2 is optimized for eager performance, which makes it difficult to trace into a graph. SimpleFSDP expresses sharding as collectives inside the traced graph, so the compiler can schedule all-gathers and reduce-scatters against the surrounding matrix multiplies.

python train.py --parallelism.use_simple_fsdp

Every TPU recipe pairs SimpleFSDP with bf16 parameters and fp32 gradient reduction. Compute runs in bf16; master weights and reductions stay in fp32. This is the default across all shipped recipes.

Large Model Graph Techniques#

For models around 100B parameters and above, like Llama 3 405B or Qwen3 235B MoE, an unrolled step graph stops compiling as depth grows, and every layer's saved input stays live until backward. Two techniques address this, both building on Graph Trainer:

Scan over layers runs the decoder as a single loop body. Compile time and memory stop growing with depth, and the one stacked activation gets offloaded to host memory. This is the path Qwen3 235B and Llama 3 405B take.

python train.py --tpu_config.use_scan_layers

Graph split keeps layers unrolled but compiles forward and backward as separate programs, optionally splitting each layer group. Activations get offloaded between groups.

python train.py --tpu_config.use_graph_split

Both require --tpu_config.use_graph_trainer to be enabled.

Pallas Kernels for Compute and Memory#

Three kernels matter, each toggled by a flag:

Splash attention never materializes the attention matrix. It handles causal, sliding-window, and packed-document masks, and matters most at long sequence lengths.

python train.py --splash_attention_kernel.use_splash_attention_kernel

Fused linear cross-entropy loss (PallasCrossEntropyLoss in the recipe) computes the LM-head matmul, softmax, and loss one vocabulary tile at a time. The full logits tensor never hits HBM. Its main value is memory, which you then convert to throughput by raising the batch.

For MoE models, the paired group matmul replaces a loop over experts with one kernel per layer:

python train.py --use_pallas_paired_grouped_mm

Each kernel exposes block-size knobs (sa_block_* and similar). The defaults are a starting point, not an optimum. Values tuned on Trillium don't carry over to Ironwood. Tune them for your model shape and TPU generation within the VMEM budget (xla_tpu_scoped_vmem_limit_kib).

Batch Size, Compiler Flags, and Profiling#

Once the kernels free memory, raise the per-chip batch until the model is HBM-bound. The fused loss and splash attention are what create the headroom. On a large slice, the global batch equals per-chip batch times chip count, so size it to what the model can use, not what the chips can hold.

Each recipe carries a bundle of XLA and libtpu flags with comments explaining each one. You can set flags from the command line:

python train.py --libtpu-init-flags="your_flags_here"

Flags don't change what the model computes, but they can change how. A pure scheduling flag leaves the loss bit-identical. A flag that alters fusion or reduction order, hierarchical reduce-scatter being one, shifts rounding. Check the step-1 loss and gradient norm after every flag sweep. Bit-identical means the flag is free. A small difference means a rounding-order change, which is only safe once a loss-curve comparison against the baseline recipe confirms it.

To profile, capture a trace over a few steady-state steps:

python train.py --profiler.enable_profiling

The trace answers one question: where does the step time go? Find the bottleneck, then go back to the lever that addresses it.

HuggingFace Backend, Agent Onboarding, and What Is Next#

If your model isn't one TorchTitan ships for, the Transformers modeling backend builds it from the Hugging Face library by its Hub ID and runs it through TorchTitan's trainer and parallelism. On TPU, the backend applies the first rungs of the playbook with no model code changes: SimpleFSDP in place of FSDP2, per-block torch.compile, splash attention registered through the transformers attention interface, fused Pallas cross-entropy fed the final hidden states and head weight, and grouped-matmul experts for MoE checkpoints. The backend supports 500+ Hugging Face models. Most dense architectures train unmodified; MoE checkpoints run through the same path but without expert parallelism yet.

For models that need more than the generic path, the repository ships a model-onboarding skill for coding agents. It captures how the case studies were done and walks an engineer or the agent itself through the same ladder, with tests, acceptance bars, and traps built in.

On the roadmap, Google has announced 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. The float8 recipe that ships for Llama 3 70B extends to the rest of the model set. Next model targets are Qwen3.8-Max (2.4 trillion parameters) and Kimi K3 (2.8 trillion parameters with Kimi Delta Attention on most layers). On the production readiness side, recipes gain ML Goodput Measurement library integration and Multi-Tier Checkpointing on GKE, so large jobs save often and resume quickly.

Gotchas#

• FSDP2 and Graph Trainer don't mix. FSDP2 is optimized for eager performance, which makes it hard to trace into a graph. Use SimpleFSDP (--parallelism.use_simple_fsdp) when you want compilation or Graph Trainer.
• Kernel tile sizes don't transfer between TPU generations. Splash attention, fused loss, and MoE kernel block sizes need tuning per generation. What is optimal on Trillium isn't on Ironwood.
• Some flags change rounding. Hierarchical reduce-scatter and fusion-altering flags shift the rounding of reductions. Check the step-1 loss and gradient norm after changing compiler flags. A small difference is a rounding-order change, not a bug.
• Per-block compile requires the block to trace. Data-dependent Python inside a transformer block breaks compilation. The block has to trace end to end; that is all the model gives up.
• Expert parallelism isn't available for MoE yet. MoE checkpoints run through the Transformers backend with grouped-matmul experts but without expert parallelism. The roadmap names this as the open item.
• Scan layers and graph split both require Graph Trainer. They won't work in eager or per-block compile mode alone.
• Graph Trainer requires the whole step to trace. Forward, loss, backward, and optimizer update all have to be capturable. Anything with data-dependent Python control flow outside the transformer blocks will block it.

Quick Reference#

# Minimal eager run (debug)
--tpu_config.eager_mode=DEFER_AND_FUSE

# Per-block compile
--compile.enable

# Graph Trainer (full step as one program)
--tpu_config.use_graph_trainer

# Sharding
--parallelism.use_simple_fsdp

# Large model graph techniques (requires Graph Trainer)
--tpu_config.use_scan_layers
--tpu_config.use_graph_split

# Kernels
--splash_attention_kernel.use_splash_attention_kernel
--use_pallas_paired_grouped_mm

# Profiling
--profiler.enable_profiling

# Compiler flags from the command line
--libtpu-init-flags="..."

The order matters. Start eager, add compile, add Graph Trainer, add SimpleFSDP, add kernels, raise the batch, tune flags, profile, repeat.

Quick answers