Google 将 TorchTPU 集成进 TorchTitan 并发布 TPU 模型优化方案

OCT. 20, 2026

Introduction

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/ ork is available at https://github.com/google-pytorch/torchtitan/

Why TorchTitan on TPU

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.

Designed to Demonstrate Pre- and Post-Training at Scale

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 Definition Decoupled from Model Partitioning

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.

Pluggable Architecture

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, the OpenXLA kernel library for TPU and GPU.

Compile-Friendly FSDP Implementation

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 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.

Graph Trainer: Capturing the Entire Training Step

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.

TPU Model Optimization Playbook

Every lever below is described in the TorchTitan TPU optimization guide, and nearly all of them are already wired into the 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

• Choose how the step executes: eager, compile, or Graph Trainer. Eager mode runs PyTorch models on TPU with virtually no changes, which makes it the right place to bring a model up and debug it. It is not a performance configuration: every op is its own program and the TPU idles between launches. TorchTPU's deferred-and-fused eager mode (--tpu_config.eager_mode=DEFER_AND_FUSE) recovers some of that with no code change by letting XLA fuse across op boundaries, and the recipes keep it on to cover whatever sits outside compiled regions. Enabling per-block torch.compile (--compile.enable) is what turns a debug configuration into a performant one: each transformer block becomes a single XLA program, so its ops are fused, launched once instead of one at a time, and scheduled together with the block's collectives. The block has to trace, which rules out data-dependent Python inside it, and that is all the model gives up. Graph Trainer goes further (--tpu_config.use_graph_trainer) and lets the whole training step, forward, loss, backward and optimizer update, be captured as one program. The whole step has to trace, and what you get back is a step with no eager glue at all, plus the large-model techniques in step 5.
• Shard with SimpleFSDP and a mixed-precision policy. SimpleFSDP expresses sharding as collectives inside the traced graph (--parallelism.use_simple_fsdp), so the compiler can schedule all-gathers and reduce-scatters against the matrix multiplies around them. FSDP2 itself is optimized for eager performance, but that also makes it difficult to trace into a graph. Pair it with bf16 parameters and fp32 gradient reduction: compute runs in bf16 while master weights and reductions stay in fp32, which is the setting every TPU recipe uses.
• For large models (~100B+ parameters), enable scan or split the graph, and offload activations. One unrolled step graph stops compiling as depth grows, and it keeps every layer's saved input live until the backward pass. Scan over layers (--tpu_config.use_scan_layers) runs the decoder as a single loop body, so compile time and memory stop growing with depth, and offloads the one stacked activation to host memory; this is the path Qwen3 235B and Llama 3 405B take. The graph splitter (--tpu_config.use_graph_split) keeps the layers unrolled but compiles forward and backward, and optionally each layer group, as separate programs with activations offloaded in between. Both build on Graph Trainer.

2. Tune Compute and Kernals

• Turn on the Pallas kernels. Splash attention (--splash_attention_kernel.use_splash_attention_kernel) never materializes the attention matrix and handles causal, sliding-window and packed-document masks; it matters most at long sequence lengths. The fused linear cross-entropy loss (PallasCrossEntropyLoss in the recipe) computes the LM-head matmul, softmax and loss one vocabulary tile at a time, so the full logits tensor is never written to HBM. Its main value is memory, which the next step turns into throughput. For MoE models, the grouped expert matmul and the paired gate/up kernel (use_pallas_paired_grouped_mm) replace a loop over experts with one kernel per layer. Each of these is a switch in the recipe, and each is tested against the plain PyTorch path it replaces.
• Tune the kernel tile sizes. Splash attention, the fused loss and the MoE kernels each expose their block sizes (sa_block_* and friends), and the defaults are a starting point rather than an optimum. Tune them for the model shape and the TPU generation, within the scoped VMEM budget (xla_tpu_scoped_vmem_limit_kib), and expect values tuned on Trillium not to carry over to Ironwood. The recipes carry the tuned values for each model and topology.

3. Optimize Memory and Batch Size

• Raise the batch. At small scale the per-chip batch is the single largest lever, because it raises the work done per launched program. Push it until the model is HBM-bound; the fused loss and the attention kernel are what free the memory to do so. On a large slice, remember that the global batch is the per-chip batch times the chip count, so size it to what the model can use, not to what the chips can hold

4. Tune Compiler, Communication and Profile

• Tune the XLA and libtpu flags. Each recipe carries its flag bundle, such as the SparseCore latency-hiding scheduler and hierarchical reduce-scatter, with a comment explaining every flag, and --libtpu-init-flags lets you set flags from the command line. Flags never change what the model computes, but they can change how: a pure scheduling flag leaves the loss bit-identical, while a flag that alters fusion or the order of a reduction, hierarchical reduce-scatter being one, shifts rounding. So check the step-1 loss and gradient norm after every 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 says so.
• Profile, find the bottleneck, and go back to the lever that addresses it. Capture an xprof trace over a few steady-state steps (--profiler.enable_profiling); the TorchTPU profiling guide covers capture on TorchTPU and the xprof documentation covers the tools. The trace answers one question: where does the step time go? Two references take it from there: the Cloud TPU performance guide explains the XLA-level levers (padding, tensor dimensions, batch size, fusion) and when each applies, and the scaling book's chapter How to Profile TPU Programs shows how to read a TPU trace op by op.

Case Studies

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

FLUX.1 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 cut the number of HLO module invocations by 28% and the step to 9.95 s (15.5% improvement). Replacing FSDP2 with SimpleFSDP 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, 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, 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 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 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.

MuseGlimmer

Muse Glimmer 30B 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 that reproduces the window exactly and gives a correctness baseline, checked against upstream on CPU. The second piece was the fused cross-entropy kernel, 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, was the standard onboarding ladder, with each feature behind a flag and its own test. The resulting recipe 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

Qwen3-235B-A22B 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 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, and replaced the expert matmuls with specialized paired grouped-matmul kernels 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 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); multislice lets one training job span several slices connected over the data-center network (DCN), 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), 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.

HuggingFace Transformers modeling backend for TorchTitan TPU

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 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 with a single configuration 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 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 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.

Agent-assisted Model Onboarding to TorchTitan TPU

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 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.

Roadmap

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, 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 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 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, 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, a 2.4-trillion-parameter mixture of experts, and 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.

Summary

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 and the optimization playbook 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 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 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, 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.

Previous

Next