Integrating TorchTPU into TorchTitan and Optimizing Models for TPU

将 TorchTPU 集成到 TorchTitan 并为 TPU 优化模型 - Google 开发者博客

Google Tag Manager (noscript)

End Google Tag Manager (noscript)

HTML

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

2026 年 10 月 20 日

Aleksey Vlasenko

软件工程师

Rob Torres

高级产品经理

Will Constable

Meta

分享

• Facebook

• Twitter

• LinkedIn

• 邮件

介绍

近年来，Meta 的 TorchTitan 已成为 PyTorch 生态系统中用于大规模生成式 AI 模型训练的首选独立开发参考平台。TorchTitan 完全基于 PyTorch 原生扩展原语构建，为预训练主流密集模型和专家混合模型（MoE）架构（如 Llama 3、Qwen 3 和 FLUX）提供了一个无限制、高度模块化的展示平台。

历史上，开发者在 Google Cloud TPU 上训练 PyTorch 模型时面临一些挑战：要么将整个代码库重写为 JAX，要么应对过时的 torch_xla。这一范式现已正式转变。随着 TorchTPU 的发布，Cloud TPU 现在具备了直接集成到 PyTorch 核心接口的原生后端（torch.device("tpu")）。现在只需对设备初始化代码进行最小修改，即可启用使用 TPU 的训练脚本。然而，提供设备兼容性只是解决问题的一半，开发者还需要高性能的生产级分布式训练方案。这就是 TorchTitan TPU 的用武之地。TorchTitan TPU 是 Meta TorchTitan 实现的一个分支，专门针对 Google Cloud TPU 进行了适配以支持模型训练。该分支可在 https://github.com/google-pytorch/torchtitan/ 获取。

为何选择在 TPU 上运行 TorchTitan

Google 与 Meta 联合开发了 TPU 原生执行流程，直接内置于 TorchTitan 中，避免了工程团队需要从零开始设计自定义 TPU 分片策略或手动编写低级 XLA 通信流程的麻烦。TorchTitan 成为了在 Google Cloud TPU 上扩展 PyTorch 的标准方案。

专为大规模预训练和后训练设计

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

除了初始预训练，TorchTitan 还作为后训练流水线（如 VeRL 的组相对策略优化（GRPO）和强化学习（RL）对齐）中的可扩展训练器。在这些工作流中，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；还将 MoE 专家矩阵乘法替换为分组的 Pallas 内核。每个内核都是通过 TorchTPU 的 torch.tpu.pallas.jax_op 装饰器将 JAX 函数转换为 PyTorch 自定义操作实现的，这正是新内核所需的所有接口。每个替换都通过配置标志控制，因此配方可以启用该功能，而测试可以将其与被替换的纯 PyTorch 实现进行对比。内核可以来自 tokamax，这是为 TPU 和 GPU 设计的 OpenXLA 内核库。

适合编译的 FSDP 实现

FSDP2 围绕 GPU 执行模型构建。其 all-gathers 和 reduce-scatters 作为显式集体调用通过模块钩子发出，通过在专用 CUDA 流上运行并使用事件同步实现计算重叠。这些对编译器不可见：集体操作位于追踪区域之外，因此在 TPU 上，torch.compile 区域会在每个 FSDP 钩子处中断，整个步骤的图完全无法捕获。TorchTPU 仅提供无操作的流和事件对象以使 FSDP2 运行；XLA 不需要这些对象，因为它根据数据依赖关系调度工作。

SimpleFSDP 采用编译器视角。参数分片是 DTensor 参数化：读取模块权重时会从分片状态重新分布到复制状态（这即为 all-gather），其 autograd 反向传播即为 reduce-scatter。没有钩子和显式集体调用，因此无论图是 per-block torch.compile 区域还是 Graph Trainer 的单次 make_fx 步骤，集体操作都是追踪图中的普通节点。XLA 编译器会提前安排所有需要的 all-gathers，并在单个程序中将 reduce-scatters 与反向计算重叠。混合精度是同一机制的一部分：该策略设置重新分布时的 gather 类型（bf16 参数）和 reduce 类型（fp32 或 bf16 梯度）。多切片的混合分片在 2-D mesh 上使用相同代码：在切片内按 ICI 分片，在切片间按 DCN 复制。

在 FLUX 上，这使 TorchTitan 实现了与 JAX MaxDiffusion 99% 的性能对齐：SimpleFSDP 块的每层 HLO 是一个包含权重 all-gather、矩阵乘法、注意力和归一化操作的单一模块，这正是 XLA 最擅长优化的程序结构。

图训练器：捕获完整的训练步骤

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

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

这就是 XLA 设计的模型，它使 TPU 上的 PyTorch 训练在概念上与 JAX 相同，整个训练步骤是一个 jit 编译函数。将整个步骤视为一个图也是实现大模型编译技术的基础：通过层扫描技术，将解码器层作为单个循环体运行，使编译时间和内存不再随深度增长；以及图分割器，将步骤分割为单独编译的前向和反向程序，在两者之间将激活数据卸载到主机内存。在 Qwen3 235B 上，这条路径实现了 29%+ 的 MFU。

优化器也可以加入图中：虽然融合的 fused_adam 优化器性能优异，但可以将函数式 AdamW 更新（functional_adamw_step）追踪到训练步骤中，确保从第一次前向操作到权重写回之间没有主机端工作。在 Qwen3-8B 上，这条路径实现了 51-52% 的 MFU。

TPU 模型优化指南

以下每个杠杆都在 TorchTitan TPU 优化指南中有详细描述，几乎所有内容都已集成到配方中。对于大多数模型，在 TPU 上获得良好的 MFU 是对现有配方进行一系列配置更改，而不是编写新代码。顺序非常重要，因为每一步都会改变下一步的测量基准。TorchTitan 中的许多优化都作为即用型构建模块提供，只需通过标志切换即可启用（见下图说明）。

1. 建立基准线

• 选择步骤的执行方式：即时执行、编译或 Graph Trainer。即时模式几乎无需修改即可在 TPU 上运行 PyTorch 模型，这使其成为模型引入和调试的合适选择。这不是性能配置：每个操作都是独立的程序，TPU 在操作启动之间处于空闲状态。TorchTPU 的延迟融合即时模式（--tpu_config.eager_mode=DEFER_AND_FUSE）通过让 XLA 跨操作边界融合，无需代码更改即可恢复部分性能，配方中默认启用该模式以覆盖编译区域外的所有内容。启用按块编译（--compile.enable）将调试配置转换为高性能配置：每个 Transformer 块成为一个 XLA 程序，其操作被融合、一次性启动，并与块的集合操作一起调度。该块需要进行追踪，这排除了块内的数据依赖 Python 代码，这也是模型需要放弃的全部内容。Graph Trainer 进一步扩展（--tpu_config.use_graph_trainer），使整个训练步骤（前向传播、损失计算、反向传播和优化器更新）被捕捉为一个程序。整个步骤需要追踪，最终得到的步骤完全没有即时执行的粘合代码，同时包含第 5 步中的大模型技术。

• 使用 SimpleFSDP 和混合精度策略进行分片。SimpleFSDP 通过在追踪图中使用集体操作来表达分片 (--parallelism.use_simple_fsdp)，这使编译器能够将 all-gather 和 reduce-scatter 操作与矩阵乘法进行调度。FSDP2 本身针对 eager 模式进行了优化，但这也使得难以将其追踪到图中。结合 bf16 参数和 fp32 梯度减少：计算在 bf16 中进行，而主权重和减少操作保持在 fp32 中，这是所有 TPU 配方使用的设置。

• 对于大型模型（约 1000 亿参数以上），启用 scan 或拆分图，并卸载激活。展开的单步图随着深度增加会停止编译，且会一直保留每层的保存输入直到反向传播。通过层扫描 (--tpu_config.use_scan_layers) 将解码器作为单个循环体运行，这样编译时间和内存不会随深度增长，同时将堆叠的激活卸载到主机内存；这是 Qwen3 235B 和 Llama 3 405B 采用的路径。图拆分器 (--tpu_config.use_graph_split) 保持层展开，但将前向和反向传播（以及可选的每层组）编译为独立的程序，并在中间卸载激活。两者都基于 Graph Trainer。

1. 调整计算和内核

• 启用 Pallas 内核。Splash 注意力 (--splash_attention_kernel.use_splash_attention_kernel) 从不显式生成注意力矩阵，处理因果、滑动窗口和打包文档掩码；在长序列长度时效果最显著。融合线性交叉熵损失（配方中的 PallasCrossEntropyLoss）按词汇块逐步计算 LM 头矩阵乘法、softmax 和损失，因此完整的 logits 张量永远不会写入 HBM。其主要价值是内存，下一步会将其转化为吞吐量。对于 MoE 模型，分组专家矩阵乘法和配对门控/上行内核 (use_pallas_paired_grouped_mm) 用每个层的一个内核替代专家循环。这些都只是配方中的开关选项，每个都与它替代的纯 PyTorch 路径进行了测试。

• 调整内核块大小。Splash 注意力、融合损失和 MoE 内核各自暴露了它们的块大小（sa_block_* 和相关参数），默认值只是一个起点而非最优值。根据模型形状和 TPU 代次，在限定的 VMEM 预算（xla_tpu_scoped_vmem_limit_kib）内调整它们，并预期在 Trillium 上调优的值可能无法直接迁移到 Ironwood。配方中为每个模型和拓扑携带了调优后的值。

1. 优化内存和批量大小

• 提高批量大小。在小规模时，每芯片的批量大小是最重要的杠杆，因为它能提升每个启动程序的工作量。推动它直到模型受 HBM 限制；融合损失和注意力内核正是释放内存进行此操作的关键。在大规模切片中，请记住全局批量大小是每芯片批量大小乘以芯片数量，因此应根据模型能使用的规模来设置，而非芯片能容纳的规模。

1. 调整编译器、通信和性能分析

• 调整 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 到 1024 个芯片上实现了约 57% 的 MFU。Qwen3 235B（大型 MoE 模型）在单个 1024 芯片切片上达到 32.42% 的 MFU，在 8 个切片的 8192 芯片上达到 30.51%。Muse Glimmer 30B（现代密集模型）在 128 芯片（4x4x8）上达到 43.4% 的 MFU。FLUX.1（扩散模型）在 1024 芯片（8x8x16）上达到 19.4% 的 MFU。

对于所有运行的模型

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

FLUX.1

FLUX.1 是一个 120 亿参数的扩散变换器，这使其成为验证上述方法论不仅适用于大语言模型的有用案例。我们在 TPU v6e-8 上以 1024x1024 分辨率训练它，使用 FSDP 在 bf16 精度下进行批量大小为 3 的样本训练。开箱即用时，以 FSDP2 和完整激活检查点的急切模式运行，每步耗时 11.49 秒。按顺序应用上述方法论逐步优化，每阶段减少一次耗时。延迟融合的急切执行减少了 28% 的 HLO 模块调用次数，将步骤时间降至 9.95 秒（提升 15.5%）。将 FSDP2 替换为 SimpleFSDP 将步骤时间降至 6.52 秒（提升 52.6%）：跟踪记录显示 FSDP2 在每次 all-gather 操作前后将参数分片复制到平面 staging 缓冲区再复制回每个参数张量，占用了大量时间；SimpleFSDP 直接对每个参数执行 all-gather，消除了 staging 缓冲区和两次复制（见上文 SimpleFSDP 部分）。带序列打包和段 ID 的 splash 注意力内核使步骤时间降至 5.55 秒（提升 17.5%），而每个块的 torch.compile（将大型子图而非细粒度模块传递给 XLA，并将控制流移出 Python 解释器）进一步将时间降至 3.02 秒（提升 83.8%）。

最后一步来自模型代码中三个小的本地重写：将融合的投影拆分为两个独立的线性层，使其输出不再需要被分割；将旋转嵌入简化为普通的逐元素数学运算；使用视图而非转置来解包Q、K和V。这些改动并未改变模型的计算内容，但每个改动都消除了复制操作，使XLA能够融合更多计算图部分。这些优化共同将单步耗时降至2.78秒（提升8.6%），比初始基准快4.1倍，达到同硬件下基于JAX的MaxDiffusion吞吐量的99%。

我们还将FLUX从8块v7x芯片扩展到1,024块，每块TPU保持3张图像。最佳配置从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上无需修改即可运行；移植工作仅涉及两个模型特定组件。上游实现滑动窗口模式时使用了无TPU后端的FlexAttention，因此第一步是实现精确复现窗口的窗口化SDPA内注意力，通过CPU与上游进行正确性验证。第二部分是融合交叉熵内核，我们扩展其实现，在计算损失时对每个词汇块应用软截断，使202k宽度的logits无需显式生成；缺少此优化时LM头本身就会主导内存占用。其余所有内容（包括Graph Trainer封装器）均为标准上架流程，每个功能均通过标志控制并配有独立测试。最终方案在64块v7x芯片（4x4x4切片）上以16,384序列长度、每设备单序列、Graph Trainer+SimpleFSDP和fp32梯度聚合方式训练30B模型，达到约44% MFU时每芯片5,848 tokens/秒。

Qwen3 235B

Qwen3-235B-A22B 是该系列中最大的 MOE 模型：94 层结构，128 个专家（每个 token 激活 8 个），总参数量 2350 亿，其中每个 token 激活 220 亿参数。这正是大型模型训练策略中所采用的方案。在如此深度的模型中，将整个展开步骤编译为一个程序已不可行：XLA 在编译时为每个程序规划所有缓冲区，而将 94 层结构完全展开后，其所需的 HBM 内存会超过芯片容量，且编译时间会随层数增加而增长。要使该模型正常运行，必须启用图训练器路径中的其中一个方案。逐层扫描是 XLA 原生的解决方案，JAX 和 MaxText 也依赖这一结构：94 个相同层被编译为一个循环体，使编译时间和程序规模不再随深度增长。图分割器是更灵活的替代方案：它保持层的展开状态，但将捕获的步骤拆分为每层可执行模块，无论层的类型如何都能正常工作（且不要求各层相同）。我们对这两种方案都进行了测试，其性能差异在百分之几以内。无论如何，在序列长度 4096、每设备 16 个序列的情况下，激活值无法装入 HBM，因此两种方案都会将层激活值卸载到主机内存，并在反向传播时预取回来。

从那里开始，训练策略逐步应用，每次只改变一个因素，并在相同的 64 芯片切片上进行测量。我们提高了 VMEM 的预算范围，使编译器能够保留更多内核中间结果在芯片上，这也消除了优化器步骤末尾的一个显式更新。我们启用了 libtpu 调度标志集，使 XLA 的延迟隐藏调度器能够重叠 FSDP 的 all-gathers 和 reduce-scatters 操作，并将 reduce-scatters 卸载到 SparseCores，从而使大部分集合操作时间被计算隐藏，仅在反向传播中保留少量的 reduce-scatter 尾部。我们转向了 tokamax splash 内核，将专家 matmuls 替换为由 K 分块权重梯度内核支持的专用配对分组 matmul 内核，并结合分层 reduce-scatter 和 SparseCore 延迟隐藏调度器来隐藏专家梯度集合操作。一个融合的路由评分反向传播消除了 MoE 反向传播中的一个大 fp32 复制操作。最后一轮优化来自 1024 芯片的性能分析：保持扫描卸载下的阻塞激活布局，提高反向传播层调度器阈值使专家 reduce-scatters 再次与计算重叠，并在扫描循环的回边中对小型 reduce-scatters 进行流水线处理。该方案在 256 个 v7x 芯片上运行时，保持 fp32 梯度缩减全程启用的情况下，MFU 约为 32%。

我们还使用该模型测试了TorchTPU的扩展能力，从64个芯片扩展到8,192芯片的v7x机柜规模，分别在单切片和多切片配置下进行验证。TPU切片是指通过芯片间互连（ICI）直接连接的一组芯片；多切片配置允许一个训练任务跨多个通过数据中心网络（DCN）连接的切片运行，这是任务规模突破单个最大切片的实现方式。在单个切片内，我们通过ICI在芯片间对参数、梯度和优化器状态进行分片处理；跨切片时模型会进行复制，每个步骤通过DCN进行一次梯度all-reduce操作，这样在较慢的网络上仅产生单次梯度归约流量，且该流量可与反向传播过程重叠。这种混合分片数据并行（HSDP）方法被用于2,048到8,192芯片规模的运行，其中以1,024芯片切片作为分片组。框架表现符合预期：从64芯片扩展到8,192芯片时，每芯片吞吐量下降8%（从4,568到4,208 tokens/秒），在芯片数量增加128倍的情况下实现了近线性扩展。需要说明的是：对于如此大规模的模型，超过约1,000芯片的数据并行属于扩展性实验而非标准训练方案。本次运行属于弱扩展测试：每个设备的批量大小保持固定，随着芯片数量增加全局批量同步增长，当达到8,000芯片时全局批量已超过实际训练该模型时的常规规模。该运行验证了框架、编译器和通信操作在如此规模下的稳定性，而非推荐的训练配置。

TorchTitan TPU的HuggingFace Transformers模型后端

上述三个案例研究属于深度分析：每次仅针对一个模型，每个变更都包含等效性测试和性能分析。通常使用TPU的团队面临的问题不同。他们已有正在训练的模型（通常是Hugging Face模型），且该模型很可能不是TorchTitan为TPU预置的标准模型。他们最终可能希望榨取模型的每个百分比性能，但首要需求是快速实现模型在TPU上的运行。这就是Transformers模型后端要解决的问题：该后端通过模型在Hub上的ID从Transformers库构建模型，并通过TorchTitan的训练器和并行化框架运行，而非使用模型自带的实现。在TPU上，后端包含TorchTPU专用通道，采用统一配置对所有支持的检查点应用优化方案，无需修改模型代码：使用SimpleFSDP替代FSDP2，按块启用torch.compile，将splash注意力内核注册为transformers注意力接口的实现方案，使模型的注意力调用直接命中该内核，将融合的Pallas交叉熵直接作用于最终隐藏状态和头部权重而非显式计算logits，对于专家混合检查点则使用TPU分组矩阵乘法专家替代可移植实现。后端支持的大多数密集型架构均可直接使用该方式训练；专家混合检查点遵循相同路径但尚未启用专家并行。虽然这些优化尚未达到手动调优的性能指标，但为尚未适配TPU的模型提供了可用且效率合理的首次运行方案，并可判断该模型是否需要专门的TPU训练方案。

基于代理的模型迁移至TorchTitan TPU

最终，模型需要的已不仅仅是通用路径：这可能是你训练数月的模型，其架构无法被后端表达，或者现成的参数数量不符合需求。在TorchTitan TPU上，将模型移植到新加速器已成为可重复的流程。该仓库为编码代理提供了模型接入技能，完整记录了上述案例研究的操作方式，并引导工程师或代理自身逐步完成相同流程，内置测试、验收标准和陷阱规避机制。

路线图

路线图围绕五个方向展开：生产就绪性、精度、硬件、并行性和模型。在生产就绪性方面，配方集将集成Google Cloud长期运行所需的运维工具：与ML Goodput Measurement库整合，使运行能报告芯片时间转化为训练进度保留量及剩余时间去向；并支持GKE多层级检查点，将检查点存储在本地内存和集群副本中，以云存储作为持久化层级，使大规模任务能频繁保存并快速恢复。在精度方面，当前为Llama 3 70B提供的float8配方（基于torchao的Float8Linear和v7x原生fp8矩阵乘法）将扩展到整个模型集。在硬件方面，Google已宣布其第八代TPU：用于预训练的TPU 8t和用于推理及强化学习的TPU 8i，TorchTitan将在这些TPU于Google Cloud上线时提供支持，延续现有配方。在并行性方面，MoE内核的专家并行性是前述待解决的关键点，长上下文训练将在其之后推进，因为百万级token上下文窗口已成为前沿模型的标配。在模型方面，目标是超越Qwen3-235B的下一个级别：Qwen3.8-Max，一个2.4万亿参数的专家混合模型，以及Kimi K3，一个采用Kimi Delta Attention替代标准注意力机制（在大部分层）的2.8万亿参数模型，这将同时考验内核插件路径和扩展性路径。

总结

TorchTitan在TPU上的应用已不再是实验。相同代码库和纯PyTorch模型定义可在最多1024个芯片上以约57%的MFU训练Llama 3 405B，并在单个切片中以最多1024个芯片运行Qwen3 235B（专家混合模型）。在固定每芯片批量的弱扩展性测试中，Qwen3 235B从64个芯片扩展到8192个芯片时保持了92%的每芯片吞吐量。这表明框架、编译器和集合通信在Pod规模下仍能稳定运行，尽管这种规模的模型通常不会在如此规模下进行训练。实现这一目标的要素并未隐藏在单一配方中：SimpleFSDP、带扫描和图分割器的图训练器、主机卸载、Pallas内核库和优化策略手册都是构建模块，每个模块都通过标志启用并附有其价值的文档说明。对于TorchTitan未提供的模型，Transformers建模后端可通过相同训练器运行500多个Hugging Face模型，已预先应用首批优化；对于需要更多功能的团队，接入技能可将模型移植转化为工程师或代理可重复遵循的标准流程。

这为使用它提供了两种方式。第一种是将 TorchTitan 作为训练框架：选择一个配方，在一个切片上启动它，并基于已知性能的配置进行构建。第二种是保留你已有的训练代码，按需获取所需内容：内核是标准的 PyTorch 自定义操作，SimpleFSDP 和图训练器是你可以导入的模块，编译器标志和操作指南适用于任何 TorchTPU 工作负载，每个配方都会明确说明其启用的功能及原因。无论哪种方式，最快了解它对你的模型有何影响的方法就是直接运行。

posted in:

• AI

• Cloud

• Announcements

• Best Practices

• Documentation

• Community

• Learn

• Influence

Previous

Next

Related Posts

List

AI

Cloud

Announcements

Best Practices

宣布在 TPU 上推出新的 vLLM 和 SGLang 推理引擎

Mobile

Industry Trends

统一边缘 AI：Coral NPU 的开放标准如何将可移植、能耗高效的智能从模型带到硅芯片

Case Studies

Learn

通过 Model Explorer 为自定义边缘编译器带来可视化图调试功能

Navigation dots