TorchTitan 更新：新增 SOTA 模型支持、规模化性能提升与 TitanRL

TorchTitan 是一个 PyTorch 原生平台，用于生成式 AI 模型的快速实验以及大规模预训练和后训练。它被设计为 PyTorch 原生技术的最小化、洁净室实现，并带有灵活扩展点以满足用户需求。这使得 TorchTitan 在学术界和工业界的研究与生产中得到广泛采用。

过去一年，TorchTitan 团队取得了大量改进，同时没有失去其最小化、可扩展的理念。增强包括对 SOTA 模型的支持、大规模开箱即用性能的显著提升、RL 支持以及新的通信特性。

TorchTitan 适合所有人——易于使用，提供 SOTA 性能，并允许用户从小规模开始并无缝扩展。我们将在以下章节中介绍所有这些新改进，并期待您的反馈！

新模型支持

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

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

TitanRL

TorchTitan 现在支持强化学习！许多 RL 栈为模型保留单独的训练和推理实现。架构变更必须在两者之间保持对齐，而当运行出现偏差时，很难判断罪魁祸首是你的奖励、你的算法，还是两个副本在悄悄不一致。TitanRL 从预训练、后训练到 vLLM 内部都保持同一模型，因此你可以在一个代码库中完成工作。

大规模智能体 RL

TitanRL 支持训练智能体工作负载所需的功能。Rollout 异步运行——训练和生成各自以多种并行形式扩展。

我们在多个工作负载上进行了测试，如下所示，证明它可以在 MoE 和稠密模型上稳定训练数百步，任务需要长上下文、代码执行、压缩、harness、自博弈等。

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

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

正确性与调试

TitanRL 提供强大的正确性保证！假设你刚刚启用了一个新算法或内核，启动了一次运行，而奖励与你的基线不同。这是你的改动造成的，还是噪声？批不变模式消除了一个噪声来源：在支持的 on-policy 配置中，训练器和生成器产生逐位相同的对数概率，因此每一步它们之间的差距恰好为零。与 TorchTitan 团队合作的一项近期研究利用它来检查数学和终端智能体 RL 中的数值不匹配问题。这帮助用户专注于他们的研究，而不是基础设施。

开箱即用的性能提升

过去一年，TorchTitan 为 MoE 模型带来了巨大的性能提升——大致围绕三个重点领域：用于提高计算利用率的更快内核、最大化通信与计算重叠，以及细粒度的内存控制。

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

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

在分布式计算和扩展方面，Dist-MoE 为 NVIDIA SM100+ GPU 提供了融合的通信-计算内核，并通过激活打包处理动态张量，使其保持完全可 CUDA 图化。同时，一个全面革新的流水线并行包提供了可 CUDA 图化、确定性的执行，并带有用户控制旋钮，以便在组合 FSDP 和流水线并行时最大化重叠。

最后，TorchTitan overrides 能够干净地集成特定模型的专用内核——例如融合 MLA、融合 SwiGLU 和融合 Complex RoPE——而不会牺牲代码的简洁性和可维护性。

DSV3 671B 结果

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

结果

MLPerf 基准上的数值

如下图所示，TorchTitan 在 DSv3 MLPerf 基准上确实与 Megatron 的损失收敛相匹配。

Muse Spark Open

对于我们的 MuseSpark 1.2 OSS 发布，我们展示了我们的高精度长上下文训练效率，展示了我们的 TP、EP 和 CP 扩展能力。

结果

生态系统集成

TorchTitan 的优势之一是其简洁性和可复用性。因此，生态系统中的多个项目选择集成 TorchTitan 并将其作为训练器提供。具体来说：

• TorchTitan 已作为训练后端集成到 veRL 中（路线图、代码）
• TorchTitan 也已加入 Miles（readme，post）
• TorchTitan 是 Meshy 的默认训练后端，并为 Areal 的训练引擎 Archon 的核心实现提供了灵感

FlexShard

TorchTitan 利用 FSDP2 对模型状态（参数、梯度和优化器）进行按行或按列的分片。然而，分布式 Muon 优化器和参数分块量化等较新的特性需要或受益于更特殊的切分策略，而这些策略在 FSDP2 中无法实现。为此，TorchTitan 现已集成 FlexShard——一个用于切分模型和优化器状态的新数据并行库。FlexShard 提供以桶为中心的 API，每个参数桶都可以按照自定义的 Placement 策略进行切分，并具有优化的运行时，可将分桶通信与计算重叠。FlexShard 提供与 FSDP1 和 FSDP2 类似的放置方式，但还提供新的放置方式，例如 Owned（其中一个 rank 拥有整个参数）以及 BlockShard（允许二维 tile 对齐的切分）。

FlexShard 使 TorchTitan 能够以高效的方式支持分布式 Muon 优化器，其中使用 Muon 优化的参数通过 Owned 和 BlockShard 放置方式进行切分——从而无需在执行优化器更新之前进行 All Gather 集合通信。在 256 块 H100 上运行 DeepSeekV3 时，与仅使用 FSDP 的等效实现相比，这将优化器耗时最多缩短了 79%。

FlexShard 还使 TorchTitan 能够支持分块量化的参数 AllGather：由集合通信收集 FP8 量化的参数，将通信量减少约 50%。使用 FSDP2 时，由于切分边界不一定与分块量化边界对齐——我们被迫要么以 BF16 收集参数（增加通信量），要么执行 FP8 all-gather，但需要为量化缩放因子增加额外的集合通信。FlexShard 的 BlockShard 放置方式使 TorchTitan 能够兼得两者之长——收集量化参数而无需为缩放因子进行任何额外通信，因为切分策略在定义上就与量化方案对齐：在 256 块 H100 上运行 LLama3 405B 时，这将端到端训练步耗时最多缩短了 20%。

GraphTrainer

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

TorchTitan 同时支持基于图和基于 eager 的训练路径，且数值结果逐位一致。基于图的训练器支持在 eager 执行中难以表达的高级图级优化，包括跨模块融合、显式图级调试，以及对异构 AI 硬件加速器的定制支持。

分布式基础设施

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

新的 NCCL 后端，已在 TorchTitan 上验证：在 PyTorch 2.14 中，来自 torchcomms 的 NCCL 实现被移植到 c10d 中，作为一个新的、可选启用的 NCCL 后端。这使我们能够更快地启用新功能。它被设计为现有后端的直接替代品，并以标准的 "nccl" 后端名称注册。与现有后端相比，它新增了：

• 单向 RMA 窗口
• 用于容错的就地进程组重配置
• 急切通信器初始化（在 PyTorch 2.15 中）
• 对 NCCL shrink/grow 和 suspend/resume API 的原始访问

对称内存：对称内存仍然是 TorchTitan 的重要性能杠杆。TorchTitan 的 MoE 层现在可以运行在 Dist-MoE 上，这是一个新开源的分布式 Mixture-of-Experts 库，为 NVIDIA Blackwell GPU 提供了开箱即用的 DeepSeek V3 配方。Dist-MoE 通过对称内存在专家之间移动 token，并将该通信与专家计算融合。TorchTitan 中的 FSDP 也可以从对称内存池中分配其通信缓冲区。这使得 NCCL 能够使用其对称内存内核实现，包括在启用 NCCL 的 zero-CTA 策略时使用 copy-engine 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 的规模下，我们测量到以下改进：

这些改进减少了阻塞时间和检查点写入时间，并为在 Titan 配方中采用 torch_checkpointing 的用户解锁了更频繁的检查点，而不会冒着 TFLOP 下降的风险。

开始使用

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

文章 TorchTitan 更新：新的 SOTA 模型、大规模性能、TitanRL 首次出现在 PyTorch 上。