Google 将 TorchTPU 集成进 TorchTitan，为 TPU 优化模型训练

Google 跟踪代码管理器（noscript）结束 Google 跟踪代码管理器（noscript） HTML 将 TorchTPU 集成到 TorchTitan 并为 TPU 优化模型 2026 年 10 月 20 日 Aleksey Vlasenko 软件工程师 Rob Torres 高级产品经理 Will Constable Meta

引言

过去几年，Meta 的 TorchTitan 已成为 PyTorch 生态系统中大规模生成式 AI 模型训练的首要净室参考平台。TorchTitan 完全基于 PyTorch 原生扩展原语构建，为 Llama 3、Qwen 3 和 FLUX 等旗舰级稠密和混合专家（MoE）架构的预训练提供了一个不受束缚、高度模块化的展示平台。

过去，希望在 Google Cloud TPU 上训练 PyTorch 模型的开发者面临一些挑战：要么将整个代码库重写为 JAX，要么费力应对遗留的 torch_xla。这种局面已正式改变。随着 TorchTPU的发布，Cloud TPU 现在拥有直接集成到 PyTorch 核心接口中的原生后端（torch.device("tpu")）。现在，要让训练脚本使用 TPU，只需对设备初始化代码做最小改动。然而，提供设备兼容性只是成功的一半，开发者还需要高性能、生产级的分布式训练配方。这正是 TorchTitan on TPU 的用武之地。 TorchTitan TPU 是 Meta TorchTitan 实现的一个分支，经过调整后可在 Google Cloud TPU 上训练模型。该分支可在 https://github.com/google-pytorch/torchtitan/ ork 可在 https://github.com/google-pytorch/torchtitan/

为什么要在 TPU 上使用 TorchTitan

Google 和 Meta 没有要求工程团队自行发明自定义 TPU 分片策略或从头手工编写底层 XLA 通信传递，而是直接在 TorchTitan 内共同开发了 TPU 原生执行传递。TorchTitan 是在 Google Cloud TPU 上扩展 PyTorch 的权威载体。

旨在展示大规模预训练和后训练

TorchTitan 为大型稠密模型（如 Llama 3 70B/405B 和 Qwen 3 8B/32B）、混合专家模型（如 Qwen 3 235B MoE）以及扩散 Transformer（如 FLUX.1 12B）的预训练提供了官方且经过实战检验的配方。

除了初始预训练之外，TorchTitan 还作为后训练流水线中的可扩展训练器，例如用于组相对策略优化（GRPO）和强化学习（RL）对齐的 VeRL。在这些工作流中，TorchTitan 负责参数更新步骤，同时与 vLLM-TorchTPU 等高吞吐服务引擎同步权重。

模型定义与模型分区解耦

TorchTitan 中的模型定义使用纯单设备 PyTorch（nn.Module）编写，不含 CUDA 专用代码。集合通信要么由可组合框架层（例如 FSDP）插入，要么由声明式分片计划（例如张量并行）插入，要么作为后端专用插件（例如 MoE 层）插入。TorchTPU 支持 C10D 集合通信 API 和 PyTorch DeviceMesh 抽象。二者结合，使 TPU 路径能够选择对模型应用哪些并行方式，以及使用哪些实现。

可插拔架构

TorchTitan 从配置树构建模型，并且该配置树可以在构建模型之前被重写。这提供了一种简洁的方式，将模型的各个模块替换为在 TPU 上运行更好的版本，而无需改动模型定义或共享的训练器循环。

我们正在积极使用这一机制。TPU 配方用 Splash Attention 替换标准注意力，这是一个 Pallas 内核，能够处理因果、滑动窗口和打包文档掩码，而无需物化注意力矩阵；用融合 Pallas 内核替换最终投影和交叉熵损失，该内核从不将完整 logits 张量写入 HBM；并用分组 Pallas 内核替换 MoE 专家矩阵乘法。每个内核都是通过 TorchTPU 的 torch.tpu.pallas.jax_op 装饰器转换为 PyTorch 自定义算子的 JAX 函数，这就是新内核所需的全部粘合代码。每次替换都是一个配置标志，因此配方可以将其开启，测试也可以将其与被替换的普通 PyTorch 路径进行比较。内核可以来自 tokamax，这是面向 TPU 和 GPU 的 OpenXLA 内核库。

对编译友好的 FSDP 实现

FSDP2 围绕 GPU 执行模型构建。它的 all-gather 和 reduce-scatter 从模块钩子中作为显式集合通信调用发出，与计算的重叠来自在专用 CUDA 流上运行它们并与事件同步。这些对编译器都不可见：集合通信位于被追踪区域之外，因此在 TPU 上，torch.compile 区域会在每个 FSDP 钩子处中断，整个步骤的图根本无法捕获。TorchTPU 仅提供空操作的流和事件对象，以便 FSDP2 能够运行；XLA 用不到它们，因为它根据数据依赖来调度工作。

SimpleFSDP 则采用编译器的视角。参数分片是一种 DTensor 参数化：读取模块权重会执行从分片到复制的 redistribute，这就是 all-gather，而其自动微分反向就是 reduce-scatter。没有钩子，也没有显式集合通信调用，因此这两个集合通信都是被追踪图中的普通节点，无论该图是逐块的 torch.compile 区域，还是 Graph Trainer 的单个 make_fx 步骤。XLA 编译器会在需要它们的矩阵乘法之前调度 all-gather，并在一个程序内将 reduce-scatter 与反向计算重叠。混合精度是同一机制的一部分，即在该 redistribute 上设置 gather 数据类型（bf16 参数）和 reduce 数据类型（fp32 或 bf16 梯度）的策略。多切片混合分片是在二维网格上的相同代码：在切片内通过 ICI 分片，在切片间通过 DCN 复制。

在 FLUX 上，正是这一点使 TorchTitan 达到与 JAX MaxDiffusion 99% 的性能持平：SimpleFSDP 块的逐层 HLO 是一个单一模块，其中权重 all-gather、矩阵乘法、注意力和归一化融合在一起，这正是 XLA 旨在优化的程序形态。

Graph Trainer：捕获整个训练步骤

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.

posted in: