PyTorch 更新 TorchTitan：支持 Qwen3.5、Kimi K3、DeepSeek-V4 并推出 TitanRL

PyTorch：Blog· The PyTorch Team at Meta·2026-10-22 03:45· 1 小时前AI 评分52

TorchTitan 更新：全新 SOTA 模型、大规模性能表现、TitanRL

AI 导读

PyTorch 团队发布 TorchTitan 年度更新，新增 Muse Spark Open、Qwen3.5、Kimi K3、DeepSeek-V4 等 SOTA 模型的完整训练栈支持，并推出强化学习方案 TitanRL，从预训练到 vLLM 推理共用同一模型代码。

正文 · 原文

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 集群上评估了 Deepseek V3 671b+MTP1 模块在 MLPerf 训练基准上的性能和数值，仅扩展 PP、DP 和 EP 轴。使用的量化方案是 MXFP8。

结果

Muse Spark 开放

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

结果

生态系统集成

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

• TorchTitan 已作为训练后端集成到 veRL 中（路线图，代码）
• TorchTitan 也已添加到 Miles 中（readme，文章）
• 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 Update: New SOTA Models, Performance at Scale, TitanRL 首次出现在 PyTorch。

来源：PyTorch：Blog · pytorch.org