Meta 的 PyTorch 团队发布 TorchTitan 更新,新增对最新开放权重模型和强化学习的支持,并降低大规模训练的通信与存储负担。生成式人工智能开发团队可使用同一模型实现,从预训练衔接到强化学习,并通过多 GPU 扩大实验规模。

1. 减轻单独维护各模型代码的负担

TorchTitan 是面向生成式人工智能预训练和后训练的平台,后训练指预训练后针对特定目的进行调整。本次新增对 Qwen3.5 和 DeepSeek-V4 等最新模型的支持,并提供配置体系,让开发者更换模型组件、并行方式及执行硬件特定 GPU 运算的内核时,无须单独维护各模型代码。

并行化是将训练任务分配到多个 GPU 的方式。开发团队可基于同一模型,更换任务划分方法或计算组件。这有助于减轻每次接入新模型时都要重建整个训练基础设施的负担。

2. 通过 TitanRL 统一训练与推理的模型实现

TitanRL 将用于预训练和后训练的模型直接用于 vLLM 内部的推理,减少修改模型结构时必须同时修改训练与推理实现的重复工作。用于训练的回复生成以无须等待完成的异步方式进行,训练和生成均可通过多种并行方式扩展。

团队表示,已在按输入仅激活部分计算的专家模块模型,以及对所有输入使用全部参数的稠密模型上,测试了数百个训练步的稳定训练,任务涉及上下文与代码执行等。上下文指模型一次可参考的输入范围,训练步则是训练数据处理的一次迭代单位。公开案例之一是使用 40 台 GB300、64K 上下文,对 Qwen3.5-35B-A3B 基础模型进行训练和验证的软件问题解决实验。

团队表示,在使用批次不变模式(batch-invariant mode)的受支持 on-policy 配置中,训练器与生成器的对数概率可达到逐比特一致。on-policy 指使用当前策略生成的数据进行训练,对数概率则是对概率取对数的值。团队将其用于研究数学和终端智能体强化学习中的数值不一致问题。

3. 大规模吞吐量需结合运行条件解读

据公司公布的数据,在内部 GB300 集群上,DeepSeek V3 671B 搭配辅助下一词元预测的 MTP1 模块,并采用 8 位浮点数格式 MXFP8 进行低精度量化。在衡量机器学习性能的公开 MLPerf 训练基准中,256 台 GPU 配置达到每 GPU 每秒 5,324.21 词元。将计算拆分为多个阶段的流水线阶段数为 2,划分数据的并行度为 128,划分专家模块计算的并行度为 64。单设备一次处理的数据批次,即本地批次,以及将批次细分后的微批次数,分别为 120。

在公司公布的 Muse Spark 1.2 OSS 结果中,96K 上下文、256 台 GPU 配置下,学习式路由的每 GPU 吞吐量为每秒 1,500.61 词元,强制均衡专家负载的路由为每秒 1,839.52 词元。192K、512 台配置下,两者分别为每秒 1,335.86 和 1,588.22 词元。上下文、GPU 数量及并行配置不同的结果,不能作为相同条件下的性能排名进行比较。

4. 降低通信量的关键在于模型分片边界

FlexShard 能更精细地确定 GPU 之间的模型状态分片边界。据公司公布的数据,在 256 台 H100 上运行 DeepSeek V3 时,为跨 GPU 更新训练状态的分布式 Muon 优化器采用 Owned·BlockShard 布局,可消除更新前的 All Gather 通信。相较仅使用传统全分片数据并行训练方式 FSDP 的等效实现,用于更新模型参数的优化器阶段耗时最多降低 79%。

收益与成本取决于数据划分方式。传统方式在分片边界与量化块不对齐时,需要用更大的格式传输数据,或额外交换压缩传输所需的校正值。新方式将边界与量化块对齐,从而消除额外通信。据公司公布的数据,在 Llama3 405B、256 台 H100 配置下,完整训练步耗时最多降低 20%。

5. 进一步缩短检查点保存等待时间

新的检查点库在首次保存前共享描述数据的信息,即元数据,并在后续保存中省去多个进程共同交换数据的集合通信。这能减轻每次保存训练状态时多个进程同步的负担。据公司公布的数据,在 512 台 H100 上,与传统方式相比,保存等待及保存、加载耗时均有所缩短,具体如下。

表:传统方式与新检查点方式的耗时

条件:512 台 H100 GPU,DSv3 671B,检查点间隔 50(据公司公布数据)

单位:训练器等待时间为 ms,保存、加载时间为 s

进一步缩短检查点保存等待时间
项目DCPv1 平均值DCPv1 P99torch_checkpointing 平均值torch_checkpointing P99
训练器等待时间841.2941.4634.136.6
检查点保存时间384.4403.113.113.2
无须重新分片的加载时间194.7198.27982.3

P99 是将测量值从小到大排序后位于 99% 位置的值。因保存而暂停训练的时间缩短后,就有余地更频繁地保存检查点。应用时,可先克隆 TorchTitan 仓库,再按照指南文档操作。

本次更新的价值在于同时减少模型实现的重复工作和分布式训练的等待时间,完善从新模型实验衔接到大规模强化学习的开发路径。