研究团队在欧洲公共基础设施上训练1T参数MoE模型

如前所述，原始 GPU FLOPS 并不是限制性资源。在预训练一个 1T MoE 时，挑战存在于系统的另外两个部分。

节点之间的 InfiniBand 流量比同一节点上 GPU 之间的流量慢 6 倍，比 GPU 与 CPU 之间的流量慢 36 倍。在我们的设置中，训练步骤分布在每个数据并行副本的 384 个 GPU 上，这一链路在一个训练步骤中会被多次跨越。因此，我们需要 InfiniBand 上的流量要么很小，要么与 GPU 上发生的计算重叠。

在 Jupiter 的登录节点上，GH200 每个 Grace 最多拥有 480 GB 的 LPDDR5X，带宽为 384 GB/s。然而，计算节点拥有 120 GB 内存，带宽为 512 GB/s，这更偏向带宽而非容量。虽然这种配置适合受内存带宽限制的仿真工作负载，但对于训练而言，这限制了每个节点在 HBM 和 Grace 内存上能够处理和存储的状态量。

可用内存量实际上少于 120 GB。在 GH200 上，Linux 将每个 Hopper 的 HBM 暴露为一个没有 CPU 的 NUMA 节点。每个节点报告 858 GiB 内存，其中 456 GiB 是 Grace 内存。在总量中，内核保留约 24 GiB，每个 GPU 剩余约 108 GiB 的可用主机内存。更多损失在主机内存损失一节中描述。

对于 1T MoE，为硬件架构选择正确的并行布局对于实现生产吞吐量至关重要。我们最初考虑了流水线并行，但鉴于以下原则，我们最终聚焦于 FSDP 和专家并行，并投入精力优化它们，最终取得了回报。由于每种布局决定了所有后续的下游优化，我们无法断言基于流水线并行的布局是否会更快。这一决定基于三条原则：

根据上述考虑，我们最终采用了 384 路全分片数据并行（FSDP），每个节点内采用 4 路专家并行。本文中的所有运行都使用单个数据并行副本（DP = 1），因此一个 FSDP 组跨越整个作业。在实践中，我们可以在 1536 个 GPU 上运行最多 4 个数据并行副本，而对吞吐量没有明显影响。

FSDP 仅重叠了部分通信。一层上的计算常常在下一层的权重到达之前就已完成，这导致 GPU 保持空闲。在基线配置（每个序列 8,192 个 token）中，rank 0 上一个被分析的步骤总共持续 83.6 秒，其中计算耗时 14.3 秒，而有 68.3 秒用于等待通信发生。

由于 FSDP 中的通信仅限于权重，且每层收集的权重体积不依赖于 token 数量，因此在更长的序列上计算会增加每字节通信所对应的计算量。因此，我们在内存允许的范围内增加了序列长度。然而，即使采用基线的激活检查点技术，12,288 个 token 的序列也已经超出了 HBM 容量。

更长的序列需要更多的激活内存。首先，我们仅在层边界处存储激活，并在反向传播期间重新计算层内的一切。完全重计算增加了计算量，但 GPU 在步骤的大部分时间处于空闲状态，因此增加的计算并未延长步骤时间。通过这一更改，序列长度达到了 32,768 个 token，吞吐量达到每 GPU 370 tokens/s，是基线的 4 倍。

在 65,536 个 token 时，仅边界激活就占用 74 层 × 65,536 tokens × 6,144 × 2 字节 ≈ 每 GPU 55 GiB，这无法与权重和剩余激活一起放入 HBM。因此，我们决定通过 900 GB/s 的 NVLink-C2C 链路将它们卸载。每层的边界激活在前向传播后复制到 Grace，并在反向传播期间重新计算之前复制回 GPU。我们将 offload_fraction 参数设置为 0.86，该参数设置执行卸载的层比例，并以 Grace 内存换取 HBM。通过卸载，序列长度达到了 49,152 个 token，吞吐量达到每 GPU 508 tokens/s。

我们将激活检查点卸载的钩子贡献给了 TorchTitan，并附带了一个最小化的 CPU 卸载器作为文档示例。然后，我们在贡献的钩子之上构建了此处描述的卸载策略。

优化器状态需要相当大的内存份额。我们对除路由器之外的每个权重矩阵使用 Muon（Jordan et al., 2024），约占参数的 99.5%，并对嵌入、LM head、路由器和归一化使用 AdamW。Muon 为每个参数存储一个 fp32 状态（动量），而 Adam 存储两个。移除第二个状态使得 65,536 个 token 的序列长度能够放入内存。据报道，Muon 训练大型模型比 AdamW 更高效（Liu et al., 2025），并且已被用于万亿参数 MoE 模型（Kimi Team, 2025; Kimi Team et al., 2026）。

使用 Muon 时，每次更新通过几次 Newton-Schulz 迭代进行正交化，这些迭代作用于整个矩阵。然而，在我们的 FSDP 设置中，矩阵被分割成 384 个分片。一个直接的实现会在一个 rank 上收集每个矩阵，在那里运行 Newton-Schulz，然后将结果分散回去，这会显著增加通信量。

FlexShard 现在是 TorchTitan 的一部分，它使我们能够选择哪个 rank 拥有哪些参数。借此，我们按整个专家对专家进行分片，因此每个 rank 拥有完整的专家矩阵，并在无需通信的情况下运行其 Newton-Schulz 迭代。密集矩阵（如注意力投影）仍然分片并每个 rank 收集。这种设置使我们能够进一步利用计算-通信重叠，因为每个 rank 可以在密集矩阵传输期间处理自己的专家，从而实现完全的通信重叠。

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）