TorchTitan Update: New SOTA Models, Performance at Scale, TitanRL

TorchTitan 更新：支持最先进模型、扩展性能、TitanRL - PyTorch

特色项目

TorchTitan 是一个原生 PyTorch 平台，支持生成式 AI 模型的快速实验和大规模预训练与后训练。它以最小化、干净的实现方式构建 PyTorch 原生技术，提供灵活的扩展接口以满足用户需求。这使得 TorchTitan 在学术界和工业界的研究与生产环境中得到了广泛应用。

过去一年，TorchTitan 团队在保持其最小化、可扩展理念的同时实现了重大改进。改进包括支持最先进模型、显著提升开箱即用的扩展性能、强化学习支持以及新增通信功能。

TorchTitan 适用于所有人 – 易于使用，提供最先进性能，并允许用户从小规模开始无缝扩展。以下部分将详细介绍所有这些新改进，期待您的反馈！

新模型支持

最先进模型支持对于训练框架的成功至关重要，TorchTitan 始终致力于在开放权重架构方面提供最新进展。TorchTitan 新增了对 Muse Spark Open、Muse Glimmer、Qwen3-VL、Qwen3.5、Kimi K2.7、Kimi K3 和 DeepSeek-V4 的原生 PyTorch 支持。这些模型实现了与 TorchTitan 五维并行架构栈的完整训练堆栈集成。该工作还引入了基于类型化 Python 配方、模型配置转换和目标组件覆盖的可组合配置系统。用户可以在不维护独立模型分支的情况下，修改并行策略、模型组件或硬件特定内核。相同的接口集成了来自 Attention Gym 的优化内核（包括 Gated DeltaNet、Kimi Delta Attention 和稀疏注意力内核），以及可选的 Helion 内核（如融合 RoPE），同时保持与 torch.compile 的兼容性。

每个前沿模型都扩展了 TorchTitan 的可复用工具包，从多模态数据管道、样本打包、共享视觉组件、Hugging Face 检查点互操作性，到可组合的五维并行、声明式分片、优化的 Attention Gym 和 Helion 内核，以及严格的数值验证。开发者可以基于此基础，更快地将新发布的架构以原生 PyTorch 实现带入验证过的可扩展训练流程。

TitanRL

TorchTitan 现在支持强化学习！许多强化学习框架会为模型保持独立的训练和推理实现。架构变更需要在两者之间保持一致，当训练出现异常时，很难判断问题是出在奖励函数、算法本身，还是两个副本在悄然分歧。TitanRL 从预训练到 vLLM 内部使用相同的模型，使您可以在一个代码库中完成所有工作。

扩展规模的智能体强化学习

TitanRL 支持您训练智能体工作负载所需的一切。Rollouts 异步运行 – 训练和生成各自通过多种并行形式扩展。

我们在多个工作负载上进行了测试，如下图所示，证明其能够在需要长上下文、代码执行、压缩、Harness、自我对弈等任务上，对 MoE 和密集模型进行数百步的稳定训练。

图1：SWE-Rebench easy 上的训练/验证，模型 Qwen3.5-35B-A3B base，64K 上下文，40 GB300

图2：Tmax 15k 上训练，终端基准 2.1 上评估，模型 Qwen3.5-35B-A3B base，128K 上下文，20 GB300

正确性与调试

TitanRL 提供了强大的正确性保障！假设你刚刚启用了一种新算法或内核，启动了一次运行，但奖励值与基线存在差异。这种差异是你的修改导致的，还是噪声？批处理不变模式消除了一个噪声来源：在支持的策略梯度配置中，训练器和生成器产生的对数概率在位级别完全一致，因此它们之间的差距在每一步都精确为零。TorchTitan 团队最近的一项研究利用该特性检查了数学和终端代理强化学习中的数值不匹配问题。这有助于用户专注于研究本身，而非基础设施问题。

开箱即用的性能改进

过去一年，TorchTitan 为 MoE 模型带来了显著的性能提升，主要集中在三个领域：更高计算利用率的加速内核、最大化通信与计算重叠、以及细粒度的内存控制。

local_compile 在无需承诺全图或全块追踪的情况下，带来了原生 torch.compile 的优势，允许用户选择性优化关键区域，同时保持数值精度而无需维护自定义内核。

在内存管理方面，torch_remat 提供了现代激活检查点 API，支持对需要重新计算的区域进行细粒度控制，避免了 CPU 开销和自定义操作包装器的需求。

在分布式计算和扩展方面，Dist-MoE 为 NVIDIA SM100+ GPU 提供了融合通信-计算内核，结合动态张量的激活打包技术，使张量保持完全 CUDA 图表兼容性。同时，全新升级的流水线并行工具包提供了 CUDA 图表兼容、确定性执行功能，并通过用户控制参数最大化 FSDP 与流水线并行的组合重叠效果。

最后，TorchTitan 的覆盖功能实现了专用模型内核（如融合 MLA、融合 SwiGLU 和融合 Complex RoPE）的干净集成，同时不牺牲代码的简洁性和可维护性。

DSV3 671B 结果

通过整合上述所有性能特性，我们在内部 GB300 集群上评估了 Deepseek V3 671B+MTP1 模块在 MLPerf 训练基准测试中的性能和数值表现，仅扩展了 PP、DP 和 EP 轴。使用的量化方案为 MXFP8。

#### 结果

运行

GPU 数量

PP / VPP / DP / EP

LBS / 微批次

TPS/GPU

全局 TPS

TFLOPs/GPU

峰值 HBM

PP2-256

256

2 / 8 / 128 / 64

120 / 120

5,324.21

1,362,998

1,553.89

243.86 GiB

PP2-512

512

2 / 8 / 256 / 64

60 / 60

5,262.68

2,694,493

1,535.93

224.65 GiB

PP4-256

4 / 4 / 64 / 64

240 / 240

5,178.22

1,325,624

1,511.28

238.08 GiB

PP4-512

4 / 4 / 128 / 64

5,087.99

2,605,049

1,484.95

221.14 GiB

PP8-512

8 / 2 / 64 / 64

4,780.79

2,447,763

1,395.29

233.96 GiB

Muse Spark Open

在 MuseSpark 1.2 开源版本中，我们展示了高精度长上下文训练效率，突出了 TP、EP 和 CP 缩放能力。

路由

上下文长度

并行性

MFU

学习

96K

PP1/DP16/TP4/CP4/EP64

1,500.61

384,156

646.03

25.84%

253.66 GiB

强制平衡

1,839.52

470,918

791.95

31.68%

192K

PP1/DP16/TP4/CP8/EP64

1,335.86

683,961

781.66

31.27%

235.64 GiB

1,588.22

813,168

929.33

37.17%

235.65 GiB

生态系统集成

TorchTitan 的优势之一在于其简洁性和可复用性。因此，生态系统中的多个项目选择将其作为训练器进行集成。特别是：

• TorchTitan 已集成到 veRL 中作为训练后端（roadmap，code）

• TorchTitan 已添加到 Miles ( readme , post )

• TorchTitan 是 Meshy 的默认训练后端，并启发了 Areal 训练引擎 Archon 的核心实现

FlexShard

TorchTitan 采用 FSDP2 实现模型状态（参数、梯度和优化器）的按行或按列分片。然而，分布式 Muon 优化器等新特性以及参数的块级量化需要或受益于 FSDP2 无法实现的更复杂的分片策略。为解决这一问题，TorchTitan 现在集成了 FlexShard – 一个用于分片模型和优化器状态的新数据并行库。FlexShard 提供以桶为中心的 API，每个参数桶可根据自定义的放置策略进行分片，并具有优化的运行时，可将分片通信与计算重叠。FlexShard 提供与 FSDP1 和 FSDP2 类似的放置方式，但还新增了 Owned（一个 rank 拥有完整参数）和 BlockShard（允许 2D 瓷砖对齐分片）等放置策略。

FlexShard 使 TorchTitan 能够以高效方式支持分布式 Muon 优化器，其中使用 Muon 优化的参数通过 Owned 和 BlockShard 放置进行分片 – 消除了在执行优化器更新前进行 All Gather 集合操作的需求。在 256 个 H100 上运行 DeepSeekV3 时，与等效的 FSDP 独立实现相比，优化器时间缩短了高达 79%。

FlexShard 还使 TorchTitan 能够支持块级量化参数的 AllGather 操作：通过集合操作将 FP8 量化参数进行聚合，通信量减少约 50%。使用 FSDP2 时，由于分片边界可能与块级量化边界不一致 – 我们被迫选择在 BF16 中聚合参数（增加通信量），或执行 FP8 AllGather 但需要额外的量化缩放因子集合操作。FlexShard 的 BlockShard 放置使 TorchTitan 能够兼得二者优势 – 在无需额外通信处理缩放因子的情况下聚合量化参数，因为分片策略定义上与量化方案对齐：在 256 个 H100 上运行 LLama3 405B 时，端到端训练步骤时间缩短高达 20%。

GraphTrainer

GraphTrainer 是 TorchTitan 的图模式训练方案，提供围绕显式训练图构建的替代执行模型。

TorchTitan 支持基于图和基于即时执行的训练路径，且数值结果位对位匹配。基于图的训练器能够实现即时执行难以表达的高级图级优化，包括跨模块融合、显式图级调试以及针对异构 AI 硬件加速器的定制支持。

分布式基础设施

TorchTitan 是 PyTorch 新通信特性发布时的重要验证阶段。大规模训练会对通信堆栈（初始化、集合操作、故障处理和内存注册）的每个部分施加压力，这使 TorchTitan 成为一个严苛的测试环境。

新的 NCCL 后端（已在 TorchTitan 验证）：在 PyTorch 2.14 中，我们从 torchcomms 将 NCCL 实现移植到 c10d，作为新的可选 NCCL 后端。这使我们能够更快地启用新功能。该后端设计为现有后端的即插即用替代方案，并在标准 "nccl" 后端名称下注册。与现有后端相比，新后端新增了以下特性：

• 单边 RMA 窗口
• 用于容错的原地进程组重新配置
• 立即通信器初始化（PyTorch 2.15 中新增）
• 对 NCCL 缩放/扩展和挂起/恢复 API 的原始访问能力

对称内存：对称内存仍然是 TorchTitan 的重要性能优化手段。TorchTitan 的 MoE 层现在可以运行在新推出的分布式 Mixture-of-Experts 库 Dist-MoE 上，并提供针对 NVIDIA Blackwell GPU 的 DeepSeek V3 配方。Dist-MoE 通过对称内存在专家之间传输 token，并将通信与专家计算融合。TorchTitan 的 FSDP 也可以从对称内存池分配通信缓冲区。这使 NCCL 能够使用其对称内存内核实现，包括在 NCCL 启用零 CTA 策略时的复制引擎 all-gather 操作。

检查点

TorchTitan 已集成支持新的分布式检查点库 torch_checkpointing。该新库消除了默认 torch.distributed.checkpointing（DCPv1）保存中的每次保存同步操作。我们在第一次保存前仅执行一次集体操作（在默认进程组上）以共享元数据；第一次保存后不再执行任何集体操作。

随着预训练规模超过 16k 个 rank，使用 GLOO 集体操作实现同步的旧方法在节点数量增加时扩展性较差。torch_checkpointing 每次保存仅运行轻量级屏障的方法在大规模场景下表现更优。

我们还改进了 DCPv1 的其他方面，例如在 torch.save 中一次性保存整个检查点项（如模型、优化器等），而不是对每个张量单独调用 torch.save，并使自定义组件可插拔（例如编写针对您基础设施优化的自定义存储后端）。

在 512 个 H100 GPU 运行 DSv3 671B 模型且检查点间隔为 50 的测试中，我们测得以下改进：

训练器阻塞时间

平均值

P99

DCPv1

841.2 毫秒

941.46 毫秒

torch_checkpointing

34.1 毫秒

36.6 毫秒

检查点写入时间

384.4 秒

403.1 秒

13.1 秒

13.2 秒

非重新分片加载时间

194.7 秒

198.2 秒

79 秒

82.3 秒

这些改进减少了阻塞时间和检查点写入时间，使用户在采用 torch_checkpointing 的 Titan 配方时能够更频繁地进行检查点保存而不会导致 TFLOP 下降。

入门指南

克隆 github.com/pytorch/torchtitan，并按照此处列出的说明操作。如果您有任何反馈，我们非常乐意听取！