Google 发布 TorchTitan TPU 分支并给出 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 现在拥有原生后端（torch.device("tpu")），直接集成到 PyTorch 的核心接口中。现在，只需对设备初始化代码做极少改动，即可让训练脚本使用 TPU。然而，提供设备兼容性只是成功的一半，开发者还需要高性能、生产级的分布式训练配方。这正是 TPU 上的 TorchTitan 的用武之地。TorchTitan TPU 是 Meta TorchTitan 实现的一个分支，经过调整后可在 Google Cloud TPU 上进行模型训练。该分支可在 https://github.com/google-pytorch/torchtitan/ 获取，工作内容可在 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 装饰器将 JAX 函数转换为 PyTorch 自定义算子，这就是新内核所需的全部胶水代码。每个替换都是一个配置标志，因此配方可以将其开启，测试可以将其与其所替换的普通 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，而其 autograd 反向则是 reduce-scatter。没有钩子，也没有显式的集合调用，因此这两个集合操作都是被追踪图中的普通节点，无论该图是逐块的 torch.compile 区域，还是 Graph Trainer 的单个 make_fx 步骤。XLA 编译器在需要它们的矩阵乘法之前调度 all-gather，并将 reduce-scatter 与反向计算重叠，全部在一个程序内完成。混合精度是同一机制的一部分，一种策略在该 redistribute 上设置 gather dtype（bf16 参数）和 reduce dtype（fp32 或 bf16 梯度）。多切片混合分片是在二维网格上的相同代码：在一个切片内通过 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 还运行优化通道：实现基于图的激活内存策略，决定保存、重计算或卸载什么。训练步骤随后作为单个 XLA 程序执行。

这正是 XLA 为之构建的模型，它使 TPU 上的 PyTorch 训练在概念上与 JAX 相同，整个训练步骤就是一个 jit 函数。将整个步骤作为一个图拥有也是使大模型编译技术成为可能的原因：层扫描（scan over layers），将解码器层作为单个循环体运行，使编译时间和内存不再随深度增长；以及图分割器（graph splitter），将步骤切分为单独编译的前向和反向程序，其间将激活卸载到主机内存。在 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 trace，覆盖几个稳态步骤（--profiler.enable_profiling）；TorchTPU 分析指南介绍了在 TorchTPU 上捕获 trace，xprof 文档介绍了相关工具。trace 回答一个问题：步骤时间花在哪里？接下来有两份参考资料：Cloud TPU 性能指南解释了 XLA 层面的杠杆（填充、张量维度、批次大小、融合）以及各自适用的时机，而 scaling book 的 How to Profile TPU Programs 一章展示了如何逐操作读取 TPU trace。

案例研究

为了验证 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 分辨率训练它，使用 bf16 的 FSDP 和三个样本的批次。开箱即用，在带有 FSDP2 和完整激活检查点的 eager 模式下，一个步骤耗时 11.49 秒。按顺序应用该方案，一步一步将其降低。延迟且融合的 eager 执行将 HLO 模块调用次数减少了 28%，步骤时间降至 9.95 秒（提升 15.5%）。将 FSDP2 替换为 SimpleFSDP 将步骤时间降至 6.52 秒（提升 52.6%）：trace 中主要耗时是 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 而非 transpose 来解包 Q、K 和 V。这些改动都没有改变模型所计算的内容；每一处都消除了拷贝，并让 XLA 能融合更多计算图。它们共同将每步耗时降至 2.78 秒（提升 8.6%），比起点快 4.1 倍，并且在相同硬件上达到了基于 JAX 的 MaxDiffusion 约 99% 的吞吐量。

我们还将 FLUX 从 8 块 v7x 芯片扩展到 1,024 块，每个 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 token 词表并带有软上限输出头，以及一种混合滑动窗口和全局注意力的层模式，其中每第四层跳过 RoPE。这是一个展示了大量开箱即用效果的案例。模型定义、并行化和训练循环在 TPU 上未加改动即可运行；移植工作归结为两个模型特定的部分。上游使用 FlexAttention 实现滑动窗口模式，而它没有 TPU 后端，因此第一步是窗口化 SDPA 内部注意力，它精确复现窗口并给出正确性基线，已在 CPU 上与上游进行核对。第二个部分是融合交叉熵内核，我们对其进行了扩展，在计算损失时对每个词表分块应用软上限，因此 202k 宽的 logits 永远不必被实体化；否则仅 LM 头就会占用主要内存。其他一切，包括 Graph Trainer 封装器，都是标准的接入流程，每个功能都在开关之后并配有各自的测试。由此产生的配方在 64 块 v7x 芯片（一个 4x4x4 切片）上训练 30B 模型，序列长度为 16,384，每设备一个序列，使用 Graph Trainer 加 SimpleFSDP 和 fp32 梯度归约，达到每芯片每秒 5,848 tokens，约 44% MFU。

Qwen3 235B

Qwen3-235B-A22B 是这一组中最大的 MOE 模型：94 层，128 个专家，每个 token 激活 8 个，共 2350 亿参数，其中每个 token 激活 220 亿。它是检验该方案中大型模型那一半的案例。在这样的深度下，将整个展开步编译为一个程序并不可行：XLA 在编译时会规划程序的每一个缓冲区，而一个平铺展开的 94 层计算图在单步运行之前所需的 HBM 就已超过芯片容量，同时其编译时间还随每一层增长。要让它跑起来，意味着必须启用两条面向大型模型的 Graph Trainer 路径之一。分层 Scan 是 XLA 原生的答案，也是 JAX 和 MaxText 所依赖的同一构造：94 个相同的层变为一个编译后的循环体，因此编译时间和程序大小不再随深度增长。图切分器则是更灵活的替代方案：它保持层展开，但将捕获的步切分为逐层的可执行文件，这适用于任何类型的层（且不要求各层相同）。我们对该模型同时运行了两种方案，结果相差仅几个百分点。无论哪种方案，在每设备 16 条序列、序列长度 4096 时激活值都无法放入 HBM，因此两条路径都将层激活卸载到主机内存，并在反向传播时预取回来。

从那里开始，逐一应用该方案中的各项改动，每项都在同一个 64 芯片切片上测量。我们提高了限定范围的 VMEM 预算，使编译器能将更多内核中间结果保留在片上，这也消除了优化器步末尾一处暴露的更新。我们启用了一组 libtpu 调度标志，让 XLA 的延迟隐藏调度器将 FSDP 的 all-gather 和 reduce-scatter 与周围的计算重叠，并将 reduce-scatter 卸载到 SparseCore，因此大部分集合通信时间被隐藏在计算之后，反向传播中只留下少量暴露的 reduce-scatter 尾部。我们将注意力切换到 tokamax splash 内核，并用专门的 成对分组矩阵乘内核替换了专家矩阵乘，后者由 K 分块的权重梯度内核支撑，并配合分层 reduce-scatter 和 SparseCore 延迟隐藏调度器来隐藏专家梯度的集合通信。一个融合的 route-score 反向消除了 MoE 反向传播中的一次大规模 fp32 拷贝。最后一轮修复来自 1024 芯片上的性能分析：在 scan 卸载下保持分块激活布局，提高反向层调度器阈值以使专家 reduce-scatter 再次与计算重叠，以及在 scan 循环的回边处对小型 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 个），在芯片数增加 128 倍的情况下实现了接近线性的扩展。有一点需要注意：对于这种规模的模型，超过约 1K 芯片的数据并行是扩展实验而非训练配方。这是一次弱扩展运行：随着芯片增加，每个设备的批次大小保持不变，因此全局批次随芯片数量增长，在 8k 芯片时比训练该模型通常使用的批次更大。这次运行表明的是框架、编译器和集合通信在该规模下能够支撑，而不是说这就是应该采用的训练配置。

面向 TorchTitan TPU 的 HuggingFace Transformers 建模后端

以上三个案例研究都是深度剖析：一次一个模型，每次改动背后都有对等性测试和性能分析。一个刚接触 TPU 的团队通常面临不同的问题。他们有一个已经在训练的模型，往往是 Hugging Face 的模型，而且很可能不是 TorchTitan 为 TPU 提供的模型。他们最终可能希望榨干它的每一分性能，但首先他们需要一种简单的方法让它先跑起来。这就是 Transformers 建模后端所解决的问题：它根据 Hub id 从 Transformers 库构建模型，并通过 TorchTitan 的训练器和并行机制运行它，而不是用模型自己的。在 TPU 上，该后端有一条 TorchTPU 通道，带有一个单一配置，可将手册优化中最初几级的优化应用于任何受支持的检查点，且无需更改模型代码：用 SimpleFSDP 替代 FSDP2、逐块 torch.compile、将 splash attention 内核注册为注意力实现并接入 transformers 注意力接口，使模型自身的注意力调用落到该内核上、将融合 Pallas 交叉熵输入最终隐藏状态和 head 权重而非物化 logits，以及对于混合专家检查点，用 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 工作负载，而且每个配方都明确说明了它启用了什么以及为什么。无论哪种方式，了解这对你的模型有什么作用的最快方法就是运行它。

上一页

下一页