将 TorchTPU 集成到 TorchTitan 并针对 TPU 优化模型

2026年10月20日

简介

过去几年里，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/ 获取

为什么选择 TorchTitan on TPU

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

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

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

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

模型定义与模型分区解耦

TorchTitan 中的模型定义使用纯单设备 PyTorch（nn.Module）编写，不含任何 CUDA 专用代码。集合通信（collectives）的插入方式包括：通过可组合的框架层（如 FSDP）、通过声明式分片计划（如张量并行），或作为后端特定的插件（如 MoE 层）。TorchTPU 支持 C10D 集合通信 API 和 PyTorch DeviceMesh 抽象。这些能力共同使得 TPU 路径可以选择为模型应用哪些并行方式，以及使用哪些实现。

可插拔架构

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

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

编译友好的 FSDP 实现

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

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

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

Graph Trainer：捕获整个训练步

许多 PyTorch 训练器要么在急切模式下运行并使用 CUDA 图来减少启动开销，要么对孤立区域应用 torch.compile，通常一次处理一个 Transformer 块。TorchTPU 尚不支持 PyTorch 的 torch.Stream 抽象，这限制了将基于急切模式的 FSDP2 与逐层 torch.compile 混合使用的性能。SimpleFSDP 通过以与编译兼容的方式表达参数分片来提供帮助，无论使用逐层 torch.compile 还是全图捕获。

TorchTitan 的 Graph Trainer 提供完整的训练步捕获（前向、损失和反向），使用 PyTorch 的 make_fx 追踪器提前生成一个 FX 图，跳过 AOTAutograd 分区，因此整个计算对 XLA 可见。GraphTrainer 还运行优化 pass：用于实现基于图的激活内存策略，决定保存、重计算或卸载什么。训练步随后作为单个 XLA 程序执行。

这正是 XLA 为之构建的模型，它使 PyTorch 在 TPU 上的训练在概念上与 JAX 相同，即整个训练步是一个 jit 过的函数。将整个步作为一个图来拥有，也是使大模型编译技术成为可能的原因：对层进行 scan，将解码器层作为单个循环体运行，从而使编译时间和内存不再随深度增长；以及图切分器，将步切分为单独编译的前向和反向程序，其间激活卸载到主机内存。在 Qwen3 235B 上，这条路径达到 29%+ MFU。

优化器也可以加入图中：虽然 fused_adam 优化器性能很高，但函数式 AdamW 更新可以被追踪进训练步（functional_adamw_step），并确保一次迭代在第一个前向算子和权重写回之间没有主机端工作。在 Qwen3-8B 上，这条路径达到 51-52% MFU。

TPU 模型优化手册

下面的每个杠杆都在 TorchTitan TPU 优化指南 中有描述，而且几乎全部已经接入 配方 中。对于大多数模型，在 TPU 上获得良好 MFU 是对现有配方的一系列配置更改，而非新代码。顺序很重要，因为每一步都会改变下一步所测量的内容。TorchTitan 中的许多优化都可以作为现成的构建块使用，只需翻转一个标志即可启用（见下图）。

1. 建立基线

• 选择步的执行方式：eager、compile 或 Graph Trainer。 Eager 模式几乎无需改动即可在 TPU 上运行 PyTorch 模型，这使它成为让模型跑起来并进行调试的正确场所。它不是性能配置：每个算子都是独立的程序，TPU 在启动之间处于空闲状态。TorchTPU 的 延迟并融合的 eager 模式（--tpu_config.eager_mode=DEFER_AND_FUSE）通过让 XLA 跨算子边界融合，在不改动代码的情况下恢复了部分性能，配方中保持开启以覆盖编译区域之外的任何内容。启用 逐块 torch.compile（--compile.enable）才是将调试配置转变为高性能配置的关键：每个 Transformer 块变成单个 XLA 程序，因此其算子被融合、一次性启动而非逐个启动，并与该块的集合通信一起调度。该块必须能够追踪，这排除了块内依赖数据的 Python 代码，而这就是模型所放弃的全部。Graph Trainer 更进一步（--tpu_config.use_graph_trainer），让整个训练步——前向、损失、反向和优化器更新——被捕获为一个程序。整个步必须能够追踪，你得到的是一个完全没有 eager 胶水代码的步，外加第 5 步中的大模型技术。
• 使用 SimpleFSDP 和混合精度策略的分片。 SimpleFSDP 将分片表示为追踪图中的集合通信（--parallelism.use_simple_fsdp），因此编译器可以将 all-gather 和 reduce-scatter 与周围的矩阵乘法一起调度。FSDP2 本身是为 eager 性能优化的，但这也使得它难以被追踪到图中。将其与 bf16 参数和 fp32 梯度归约配对：计算以 bf16 运行，而主权重和归约保持在 fp32，这是每个 TPU 配方所使用的设置。
• 对于大型模型（~100B+ 参数），启用 scan 或拆分图，并将激活值卸载。 单个展开的步进图随着深度增加会停止编译，并且它会将每层保存的输入一直保持活跃直到反向传播。对层进行 scan（--tpu_config.use_scan_layers）将解码器作为单个循环体运行，因此编译时间和内存不再随深度增长，并将堆叠的激活值卸载到主机内存；Qwen3 235B 和 Llama 3 405B 采用的就是这条路径。图拆分器（--tpu_config.use_graph_split）保持层展开，但将前向和反向（以及可选的每层组）编译为独立的程序，并在它们之间卸载激活值。两者都建立在 Graph Trainer 之上。

2. 调优计算和内核

• 开启 Pallas 内核。 Splash attention（--splash_attention_kernel.use_splash_attention_kernel）从不物化注意力矩阵，并处理因果、滑动窗口和打包文档掩码；它在长序列长度下最为重要。融合的线性交叉熵损失（配方中的 PallasCrossEntropyLoss）一次计算一个词表分块的 LM-head 矩阵乘法、softmax 和损失，因此完整的 logits 张量从不写入 HBM。它的主要价值在于内存，而下一步将其转化为吞吐量。对于 MoE 模型，分组专家矩阵乘法和配对的 gate/up 内核（use_pallas_paired_grouped_mm）用每层一个内核取代了对专家的循环。这些中的每一个都是配方中的一个开关，并且每一个都针对它所取代的普通 PyTorch 路径进行了测试。
• 调优内核分块大小。 Splash attention、融合损失和 MoE 内核各自暴露了它们的块大小（sa_block_* 等），默认值是起点而非最优值。在作用域 VMEM 预算（xla_tpu_scoped_vmem_limit_kib）内，针对模型形状和 TPU 代次进行调优，并预期在 Trillium 上调优的值不会迁移到 Ironwood。配方中包含了每个模型和拓扑的调优值。

3. 优化内存和批量大小

• 提高批量。 在小规模下，每芯片批量是最大的单一杠杆，因为它提高了每个启动程序所完成的工作量。持续推高直到模型受 HBM 限制；融合损失和注意力内核正是释放内存以实现这一点的关键。在大型切片上，请记住全局批量等于每芯片批量乘以芯片数量，因此应根据模型能利用的量来设定，而不是根据芯片能容纳的量

4. 调优编译器、通信和性能分析

• 调整 XLA 和 libtpu 标志。每个配方都带有自己的标志组合，例如 SparseCore 延迟隐藏调度器和分层 reduce-scatter，并附有注释解释每个标志，而 --libtpu-init-flags 允许你从命令行设置标志。标志从不改变模型计算的内容，但可能改变计算的方式：纯调度标志会使损失保持位相同，而改变融合或归约顺序的标志（分层 reduce-scatter 就是其中之一）会改变舍入。因此，每次扫描后都要检查第 1 步的损失和梯度范数。位相同意味着该标志是安全的。微小差异意味着舍入顺序发生了变化，只有当与基线配方的损失曲线对比表明安全时，这种变化才是安全的。
• 进行性能分析，找到瓶颈，然后回到解决该瓶颈的杠杆。 捕获 xprof 跟踪，覆盖几个稳态步骤（--profiler.enable_profiling）；TorchTPU 性能分析指南介绍了在 TorchTPU 上的捕获方法，xprof 文档介绍了相关工具。跟踪回答一个问题：步骤时间花在了哪里？有两个参考资料可以继续深入：Cloud TPU 性能指南解释了 XLA 级别的杠杆（填充、张量维度、批大小、融合）以及每个杠杆何时适用，而扩展书籍的章节 如何分析 TPU 程序展示了如何逐个算子地阅读 TPU 跟踪。

案例研究

为了验证 TorchTitan 在不同类别的模型上都能高效运行，我们在 TPU v7x（Ironwood）上对其中几个模型进行了调优和扩展。在大规模下，Llama 3 405B（大型稠密模型）从 128 到 1,024 个芯片保持约 57% 的 MFU。Qwen3 235B（大型 MoE 模型）在单个 1,024 芯片切片上达到 32.42% MFU，在跨八个切片的 8,192 个芯片上达到 30.51%。Muse Glimmer 30B（现代稠密模型）在 128 个芯片（4x4x8）上达到 43.4% MFU。FLUX.1（扩散模型）在 1,024 个芯片（8x8x16）上达到 19.4% MFU。

对于所有我们运行的模型

以下案例研究展示了我们如何为其中三个模型实现这些成果，每个模型代表一个不同的类别：

FLUX.1

FLUX.1 是一个 120 亿参数的扩散 Transformer，这使它成为检验上述操作手册并非仅适用于 LLM 的有用案例。我们在 TPU v6e-8 上以 1024x1024 分辨率对其进行了训练，使用 FSDP、bf16 精度和 3 个样本的批大小。开箱即用，在 eager 模式下使用 FSDP2 和完整激活检查点，一个步骤耗时 11.49 秒。按顺序应用操作手册，逐步将其降低。延迟并融合的 eager 执行将 HLO 模块调用次数减少了 28%，步骤时间降至 9.95 秒（提升 15.5%）。将 FSDP2 替换为 SimpleFSDP 后，步骤时间降至 6.52 秒（提升 52.6%）：跟踪显示 FSDP2 在每次 all-gather 之前将参数分片复制到一个扁平的暂存缓冲区中，之后再复制回各个参数张量，这占据了主导地位；SimpleFSDP 直接对每个参数进行 all-gather，因此暂存缓冲区和两次复制都消失了（参见上文的 SimpleFSDP 部分）。splash attention 内核配合序列打包和 segment ids，使 MXU 周期不浪费在填充上，将步骤时间降至 5.55 秒（提升 17.5%），而 逐块 torch.compile 向 XLA 传递大型子图而非细粒度模块，并将控制流移出 Python 解释器，将其降至 3.02 秒（提升 83.8%）。

最后一步来自性能分析指向的模型代码中的三处小规模局部重写：将一个融合的投影拆分为两个独立的线性层，使其输出不再需要被切分；将旋转位置编码简化为普通的逐元素数学运算；以及用 view 而非转置来解包 Q、K 和 V。这些改动都不改变模型的计算内容；每一处都消除了拷贝，让 XLA 能够融合更多的计算图。它们共同将单步耗时降至 2.78 秒（提升 8.6%），比起点快 4.1 倍，达到同一硬件上基于 JAX 的 MaxDiffusion 约 99% 的吞吐量。

我们还将 FLUX 从 8 扩展到 1,024 个 v7x 芯片，每个 TPU 保持三张图像。最佳配置从 8 芯片上的每步 2.53 秒（20.76% MFU）变为 1,024 芯片上的 2.71 秒（19.4% MFU），弱扩展效率为 93%。超过一个 64 芯片立方体后，纯 FSDP 会因模型体积紧凑而变慢。通过 HSDP 的混合分片解决了这个问题：它在每个 4x4x4 立方体内分片，并在立方体之间复制。我们使用了拓扑感知的设备网格，将并行性最优地映射到物理的 8x8x16 环面上。

MuseGlimmer

Muse Glimmer 30B 是一个现代的稠密解码器：52 层，分组查询注意力（32 个查询头和 2 个键值头），202k 词表并带有软上限输出头，层模式混合了滑动窗口注意力和全局注意力，每四层跳过一次 RoPE。这个案例展示了有多少功能可以开箱即用。模型定义、并行化和训练循环在 TPU 上原封不动地运行；移植工作归结为两个特定于模型的部分。上游使用 FlexAttention 实现滑动窗口模式，而 FlexAttention 没有 TPU 后端，因此第一步是实现一个带窗口的 SDPA 内部注意力，它精确复现窗口并提供正确性基线，在 CPU 上与上游进行了校验。第二部分是融合的交叉熵内核，我们对其进行了扩展，使其在计算损失时对每个词表分块应用软上限，这样 202k 宽的 logits 就永远不需要被完整物化；没有这一点，仅 LM 头就会主导内存占用。其他一切，包括 Graph Trainer 包装器，都是标准的接入流程，每个功能都有一个标志位和自己的测试。最终的配方在 64 个 v7x 芯片（一个 4x4x4 切片）上以序列长度 16,384、每设备一个序列训练 30B 模型，使用 Graph Trainer 加 SimpleFSDP 和 fp32 梯度归约，达到每芯片每秒 5,848 个 token，MFU 约 44%。

Qwen3 235B

Qwen3-235B-A22B 是这组模型中最大的 MOE 模型：94 层，128 个专家，每个 token 激活 8 个，2350 亿参数，其中每个 token 激活 220 亿。这是检验大模型那一半方案的案例。在这个深度，将整个展开的步骤编译为一个程序是行不通的：XLA 在编译时规划程序的每个缓冲区，而一个平铺展开的 94 层图在一步运行之前就需要超过芯片所拥有的 HBM，同时其编译时间随每一层递增。要让它跑起来，就必须启用 Graph Trainer 针对大模型的两条路径之一。按层扫描（Scan over layers）是 XLA 原生的方案，也是 JAX 和 MaxText 所依赖的相同构造：94 个相同的层变成一个编译后的循环体，因此编译时间和程序大小不再随深度增长。图分割器（graph splitter）是更灵活的替代方案：它保持层展开，但将捕获的步骤切割成按层的可执行文件，无论你有什么类型的层都能工作（并且不要求各层相同）。我们对这个模型运行了这两种方式，结果相差几个百分点。无论哪种方式，在序列长度为 4096、每设备 16 条序列时，激活值都放不进 HBM，因此两条路径都将层激活值卸载到主机内存，并为反向传播预取回来。

接下来按方案逐步实施，每次只改一处，每次都在同一 64 芯片切片上测量。我们提高了作用域内的 VMEM 预算，使编译器能在芯片上保留更多内核中间结果，这也消除了优化器步骤尾部一次显式的更新。我们启用了一组 libtpu 调度标志，让 XLA 的延迟隐藏调度器将 FSDP 的 all-gather 和 reduce-scatter 与周围的计算重叠，并将 reduce-scatter 卸载到 SparseCore 上，从而将大部分集合通信时间隐藏在计算之后，只在反向传播中留下一小段显式的 reduce-scatter 尾部。我们将注意力切换到 tokamax splash 内核，并用专门的 成对分组矩阵乘法内核替换专家矩阵乘法，这些内核由 K 分块的权重梯度内核支撑，配合分层 reduce-scatter 和 SparseCore 延迟隐藏调度器来隐藏专家梯度的集合通信。一个 融合的路由分数反向传播消除了 MoE 反向传播中一次大的 fp32 拷贝。最后一轮修复来自在 1,024 芯片上的性能分析：在扫描卸载下保持分块激活布局，提高反向层调度阈值使专家 reduce-scatter 再次与计算重叠，以及在扫描循环的回边之间流水化小的 reduce-scatter。该方案在 256 个 v7x 芯片上以约 32% 的 MFU 运行，全程保持 fp32 梯度归约。

我们还用这个模型测试了 TorchTPU 的扩展能力，从 64 颗芯片扩展到 8,192 颗芯片的 v7x pod 规模，涵盖单切片和多切片两种配置。TPU 切片是一组通过 芯片间互连（ICI）直接连接的芯片；多切片功能让一个训练任务跨越多个通过数据中心网络（DCN）连接的切片，这就是任务规模超越最大单一切片的方式。在单个切片内，我们通过 ICI 运行 FSDP，将参数、梯度和优化器状态分片到该切片的芯片上；跨切片时，模型被复制，梯度每步通过 DCN 进行一次 all-reduce，因此较慢网络上唯一的流量就是这一次梯度归约，且它与反向传播重叠。这就是混合分片数据并行（HSDP），2,048 至 8,192 颗芯片的运行采用的就是这种方式，以 1,024 颗芯片的切片作为分片组。整个过程中框架表现符合预期：从 64 颗芯片扩展到 8,192 颗芯片，单芯片吞吐量下降了 8%（每芯片每秒从 4,568 个 token 降至 4,208 个 token），在芯片数量增加 128 倍的情况下实现了近线性扩展。有一点需要注意：对于这种规模的模型，超过约 1K 颗芯片的数据并行更像是一次扩展性实验，而非训练配方。这是一次弱扩展运行：随着芯片增加，每设备的批量保持固定，因此全局批量随芯片数量增长，在 8k 芯片时已超过训练该模型通常会使用的批量。这次运行表明的是框架、编译器和集合通信在该规模下能够撑住，而不是说这就是训练该模型的配置。

HuggingFace Transformers 建模后端用于 TorchTitan TPU

上面的三个案例研究是深度剖析：一次只针对一个模型，每次改动背后都有一次一致性测试和一份性能剖析。一个刚上手 TPU 的团队通常面临的是不同的问题。他们有一个已经在训练的模型，往往是 Hugging Face 的模型，而且很可能不是 TorchTitan 为 TPU 发布的那个模型。他们最终可能希望榨干它的每一分性能，但首先需要一种简单的方法让它跑起来。这正是 Transformers 建模后端所解决的问题：它通过 Hub id 从 Transformers 库构建模型，并让它通过 TorchTitan 的训练器和并行机制运行，而不是用模型自带的。在 TPU 上，该后端有一条 TorchTPU 通道，通过单一配置即可将优化手册中的前几步优化应用于任何受支持的检查点，且无需修改模型代码：用 SimpleFSDP 替代 FSDP2、对每个 block 进行 torch.compile、将 splash attention 内核注册为一种注意力实现并接入 transformers 的注意力接口，使模型自身的注意力调用落到该内核上、将融合的 Pallas 交叉熵直接喂入最终隐藏状态和分类头权重而非物化的 logits，以及对于混合专家（mixture-of-experts）检查点，用 TPU 的 grouped-matmul 专家实现替代可移植实现。该后端支持的大多数稠密架构无需修改即可这样训练；混合专家检查点也走同一条路径，但尚未支持专家并行。这些都达不到上面手工调优的数字，但对于一个还没有 TPU 训练配方的模型来说，这是一个可用且效率尚可的首次运行，也是一种判断它是否值得拥有专属配方的方法。

基于智能体辅助的模型接入 TorchTitan TPU

迟早，某个模型会需要超越通用路径的方案：它可能是你要训练数月的模型，其架构是后端无法表达的，或者开箱即用的数字并非你所需。将模型移植到新加速器过去是一个以周计的专家项目，来之不易的经验只存在于某位工程师的脑海中。在 TorchTitan TPU 上，这已成为一个可重复的流程。该仓库为编码代理提供了一个模型接入技能，它记录了上述案例研究的完成方式，并引导工程师（或代理本身）走完同样的阶梯，其中内置了测试、验收标准和已知的坑。

路线图

路线图沿五条线展开：生产就绪、精度、硬件、并行和模型。在生产就绪方面，配方获得了 Google Cloud 上长时间运行所需的运维工具：与ML Goodput Measurement 库集成，使一次运行能报告其芯片时间有多少转化为保留的训练进度、其余流向何处；以及与 GKE 上的多层检查点集成，它将检查点保存在本地内存和集群内副本中，以 Cloud Storage 作为持久层，使大型作业能频繁保存并快速恢复。在精度方面，目前为 Llama 3 70B 提供的float8 配方（基于 torchao 的 Float8Linear 和 v7x 的原生 fp8 矩阵乘）将扩展到其余模型集。在硬件方面，Google 已宣布其第八代 TPU——用于预训练的 TPU 8t 和用于推理与强化学习的 TPU 8i——TorchTitan 将在其登陆 Google Cloud 时提供支持，并沿用相同的配方。在并行方面，MoE 内核的专家并行是前文提到的待办项，其后是长上下文训练，因为百万 token 的上下文窗口如今已是前沿模型的常态。在模型方面，目标是比 Qwen3-235B 更大一号的模型：Qwen3.8-Max，一个 2.4 万亿参数的混合专家模型；以及Kimi K3，2.8 万亿参数，在大多数层上用 Kimi Delta Attention 取代标准注意力，这将同时考验内核插件路径和扩展路径。

总结

TPU 上的 TorchTitan 已不再是一项实验。同一代码库和纯 PyTorch 模型定义可在最多 1,024 片芯片上以约 57% 的 MFU 训练 Llama 3 405B，并可在单个切片的最多 1,024 片芯片上运行混合专家模型 Qwen3 235B。在每芯片批量固定的弱扩展测试中，Qwen3 235B 从 64 片到 8,192 片芯片保持了 92% 的每芯片吞吐量。这表明框架、编译器和集合通信在 pod 规模下依然稳健，尽管这种规模的模型并不会真的在该规模下训练。使这一切成为可能的组件并未埋没在某个配方中。SimpleFSDP、带 scan 的 Graph Trainer 与图分割器、主机卸载、Pallas 内核库和优化手册都是构建模块，每个都由一个开关控制，并各自注明了其价值所在。对于 TorchTitan 未提供的模型，Transformers 建模后端可让 500 多个 Hugging Face 模型通过同一训练器运行，且首批优化已应用；对于想要更多的团队，接入技能将模型移植变成工程师或代理可遵循的可重复流程。

这样就剩下两种使用方式。第一种是采用 TorchTitan 作为训练框架：选择一个 配方，在切片上启动它，并基于一个性能已知的配置进行构建。第二种是保留你现有的训练代码，按需取用：这些内核是标准的 PyTorch 自定义算子，SimpleFSDP 和 Graph Trainer 是可以导入的模块，编译器标志和操作手册适用于任何 TorchTPU 工作负载，而且每个配方都明确说明了它启用了什么以及为什么。无论哪种方式，了解它对你的模型有什么效果的最快方法就是运行它。

致谢

^{Alexander Brauckmann Jack Oh Jerry Hsu Jialei Chen Tanvi Bajpai Wenjun Gu Yaoshiang Ho Tushar Agarwal Karlo Basioli Parker Holloway Huangze Tang Lee Howes Claudio Basile Freya Azad Kat Ko Andi Gavrilescu Chintan Shah}