Google 将 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。然而，提供设备兼容性只是成功的一半，开发者还需要高性能、生产级的分布式训练配方。这正是 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 recipes 用 Splash Attention 替代标准注意力——这是一种 Pallas 内核，能够处理因果、滑动窗口和打包文档掩码，而无需物化注意力矩阵；用融合的 Pallas 内核替代最终投影和交叉熵损失，该内核从不将完整的 logits 张量写入 HBM；并用分组的 Pallas 内核替代 MoE 专家矩阵乘法。每个内核都是通过 TorchTPU 的 torch.tpu.pallas.jax_op 装饰器转换为 PyTorch 自定义算子的 JAX 函数，这就是新内核所需的全部粘合代码。每个替换都是一个配置标志，因此 recipe 可以将其开启，测试也可以将其与被替换的普通 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 dtype（bf16 参数）和 reduce dtype（fp32 或 bf16 梯度）的策略。多切片场景下的混合分片是在二维网格上的相同代码：在切片内通过 ICI 进行分片，在切片间通过 DCN 进行复制。

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

Graph Trainer：捕获整个训练步骤

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

TorchTitan 的 Graph Trainer 提供完整的训练步骤捕获（前向、损失和反向），借助 PyTorch 的 make_fx 追踪器提前生成一个 FX 图，跳过 AOTAutograd 分区，从而使整个计算对 XLA 可见）。GraphTrainer 还会运行优化 pass：实现基于图的激活内存策略，决定保存、重算或卸载哪些内容。随后训练步骤作为单个 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 层面的杠杆（填充、张量维度、批大小、融合）以及各自何时适用，而扩展手册中的如何对 TPU 程序做性能分析一章展示了如何逐算子阅读 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 和批大小为三个样本对它进行训练。开箱即用，在 eager 模式下配合 FSDP2 和完整激活检查点，一步耗时 11.49 秒。按顺序应用手册，一步步将其降低。延迟且融合的 eager 执行将 HLO 模块调用次数减少了 28%，步骤时间降至 9.95 秒（提升 15.5%）。用 SimpleFSDP 替换 FSDP2 将步骤时间降至 6.52 秒（提升 52.6%）：trace 显示主要耗时在 FSDP2 每次 all-gather 前将参数分片复制到扁平暂存缓冲区，之后再复制回各参数张量；SimpleFSDP 直接对每个参数做 all-gather，因此暂存缓冲区和两次复制都消失了（见上文 SimpleFSDP 部分）。splash attention 内核配合序列打包和 segment id，使没有 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 切片）上以序列长度 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 亿。它正是用来检验 playbook 中“大模型”那一半的案例。在这样深的模型上，把整个展开的 step 当作一个程序来编译并不可行：XLA 在编译期会为程序的每个 buffer 做规划，而一个 94 层的图平铺展开后，在一步能运行之前所需的 HBM 就超过了芯片容量，同时编译时间还会随层数增长。要让它跑起来，就必须启用大模型两条 Graph Trainer 路径中的一条。按层 Scan 是 XLA 原生的方案，也是 JAX 和 MaxText 依赖的同一构造：94 个相同的层变成一个编译后的循环体，于是编译时间和程序大小不再随深度增长。图切分器则是更灵活的替代方案：它保持层展开，但把捕获的 step 切成逐层的可执行文件，这对任何类型的层都适用（也不要求各层相同）。我们对该模型两条路径都跑了，结果相差在几个百分点以内。无论哪条路径，在序列长度 4096、每设备 16 条序列时激活都放不进 HBM，因此两条路径都把层激活卸载到主机内存，并在反向传播时预取回来。

此后 playbook 便逐项应用，每一项都在同一个 64 芯片切片上测量。我们提高了 scoped VMEM 预算，让编译器能把更多 kernel 中间结果留在片上，这也消除了优化器 step 尾部一处暴露的 update。我们启用了一组 libtpu 调度标志，让 XLA 的延迟隐藏调度器能把 FSDP 的 all-gather 和 reduce-scatter 与周围的计算重叠起来，并把 reduce-scatter 卸载到 SparseCore，于是大部分集合通信时间被隐藏在计算之后，反向传播中只剩下一小段暴露的 reduce-scatter 尾部。我们把注意力切换到了 tokamax splash kernel，并用专门的 paired grouped-matmul kernel 替换了专家 matmul，其背后是一个 K-tiled 权重梯度 kernel，配合分层 reduce-scatter 和 SparseCore 延迟隐藏调度器来隐藏专家梯度的集合通信。一个 融合的 route-score 反向 消除了 MoE 反向传播中一次大规模的 fp32 拷贝。最后一轮修复来自 1024 芯片上的 profiling：在 scan 卸载下保持分块的激活布局，提高反向层调度器阈值，使专家 reduce-scatter 再次与计算重叠，以及在 scan 循环的回边（back-edge）上对小规模 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 交叉熵直接接收最终隐藏状态和输出头权重而非物化 logits，并且对于混合专家检查点，用 TPU 分组矩阵乘专家替代可移植实现。该后端支持的大多数稠密架构都能以这种方式未经修改地训练；混合专家检查点走同一路径，但尚不支持专家并行。这些都无法达到上文手工调优的数字，但对于尚无 TPU 配方的模型来说，这是一个可运行且相当高效的首次运行方案，也是判断它是否值得专门配方的途径。

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

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

路线图

路线图沿五条线展开：生产就绪、精度、硬件、并行和模型。在生产就绪方面，配方获得了在 Google Cloud 上长时间运行所需的运维工具：与 ML Goodput Measurement 库集成，使一次运行能报告其芯片时间有多少转化为保留的训练进度，以及其余部分去了哪里；以及与 GKE 上的 Multi-Tier Checkpointing 集成，它将检查点保存在本地内存和集群内副本中，以 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 splitter 的 Graph Trainer、主机卸载、Pallas 内核库和优化手册是构建模块，每个都在一个标志之后，每个都记录了其价值数字。对于 TorchTitan 不提供的模型，Transformers 建模后端通过同一个训练器运行 500+ Hugging Face 模型，并已应用首批优化；对于想要更多的团队，接入技能将移植模型转变为工程师或代理可以遵循的可重复流程。

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

上一页

下一页