EUROPA 用 PyTorch 与 TorchTitan 在 JUPITER 公有机群训练 1T 参数 MoE

PyTorch：Blog· EUROPA 联盟（Domyn、Fraunhofer IAIS 和 IIS）·2026-11-12 20:00· 1 小时前AI 评分40

在欧盟公共基础设施上使用 PyTorch 和 TorchTitan 训练 1T 参数 MoE

AI 导读

EUROPA 联盟借助 PyTorch 与 TorchTitan，在 EuroHPC 的 JUPITER 超算上把约 1T 参数 MoE 的预训练吞吐从每 GPU 90 tokens/s 提升到约 1,300 tokens/s，使用 384 块 NVIDIA GH200（96 节点）。

正文 · 原文

TL;DR

此前，在共享的公共基础设施上对约 1 万亿参数的混合专家（MoE）模型进行预训练扩展被认为不可行，原因在于其独特的架构和互连瓶颈。然而，通过利用 PyTorch 和 TorchTitan 的细粒度灵活性，EUROPA 联盟成功在 JUPITER 上实现了前沿规模训练，作为 EuroHPC Frontier AI Grand Challenge 的一部分，在每 DP 副本 384 块 NVIDIA GH200 GPU（96 节点）上，将吞吐从每 GPU 90 tokens/s 提升到约 1,300 tokens/s。

EuroHPC Frontier AI Grand Challenge 旨在证明主权开源 AI 能够利用共享公共资源在前沿水平上竞争。在此过程中，我们——由 Domyn、Fraunhofer IAIS 和 IIS 组成的 EUROPA 联盟——被选中开发一个开放的欧洲混合专家（MoE）模型，参数至少 400B，支持全部 24 种欧盟官方语言，并严格遵循 EU AI Act。

在本文中，我们介绍如何将预训练扩展到 JUPITER 上的 1T MoE，JUPITER 是托管于 Forschungszentrum Jülich 的 EuroHPC 百亿亿次超级计算机，从而证明公共基础设施确实能够支撑前沿规模。

该表列出了文中实验所用的配置。

这种规模的模型通常在其开发者运营或独家租用的集群上训练。而 JUPITER 是由公共资金资助、由数百个项目共享的研究基础设施。它还被设计用于支持广泛的科学研究，这导致了一些在 AI 领域不太典型的架构选择，我们将在后文详细描述。因此，要在保持可行吞吐的同时达到 1T 参数规模，需要超越标准分布式技术的精细优化。

为了获得对预训练设置进行细粒度优化所需的灵活性，我们选择使用 PyTorch 和 TorchTitan。在这些条件下，我们在 384 块 GH200 GPU（96 节点）上将吞吐从每 GPU 90 tokens/s 提升到约 1,300 tokens/s。我们希望分享实现这一目标的主要里程碑，每一个都建立在 TorchTitan 的设计之上，并促成了对 TorchTitan 和 PyTorch 的贡献：

1. 为 Muon 优化器实现 FlexShard 放置，消除了该优化器的大部分通信。
2. 将激活检查点卸载到 Grace CPU。
3. 在计算流上运行 FSDP 的 all-gather 设置（unshard_async_op），从而每 GPU 释放约 14 GiB HBM。

完整报告详细介绍了其余更改，并附有交互式图表和 profiler 跟踪。

机器及其约束

每个 JUPITER Booster 节点配备四个 NVIDIA GH200 Grace Hopper 超级芯片。每个超级芯片将一个 Hopper GPU（96 GB HBM3）与一个 72 核 Grace CPU（120 GB LPDDR5X）配对，通过 900 GB/s 的 NVLink-C2C 链路连接。节点内的 GPU 通过 NVLink 4 连接，每个方向 150 GB/s，每个节点有四个 InfiniBand NDR200 适配器，每 GPU 约 25 GB/s。

图 2. JUPITER 节点内的带宽，按比例绘制。所有跨节点流量都通过细窄的 InfiniBand 链路离开。

除了加速器之外，系统的两个关键部分也与设计相关：

• 跨节点带宽。 InfiniBand 比 GPU 与其自身 CPU 之间的链路慢 36 倍，而将 1T 参数模型分布到数百个加速器上，需要在训练过程中持续跨越这一边界。因此，InfiniBand 上的流量必须与计算积极重叠。
• 主机内存容量。 GH200 每个 Grace 最高可配 480 GB LPDDR5X，但 JUPITER 的计算节点实际使用 120 GB 配置（登录节点则为 480GB），这更倾向于带宽而非容量。扣除内核占用后，每 GPU 约 108 GiB 可用，而我们发现实践中有些节点启动作业时可用内存要少得多。这限制了可以从 HBM 中移出的状态量。

这些约束决定了并行布局：384 路 FSDP，每个节点内 4 路专家并行，且不使用流水线并行。MoE token 的 all-to-all 难以重叠，因此我们将其保留在 NVLink 上。没有流水线并行时，步骤往往更短，这有利于在大规模下提供快速反馈，而固定开销通过 384 条 65,536 token 序列的大批次（每步约 25M token）来摊销。FSDP 已经为我们提供了第 n+1 层 gather 与第 n 层计算的重叠，因此剩余的大部分工作都用于让重叠变得有效。我们通过三个维度实现这一点：转向更长的序列以计算更长时间、减少通过 InfiniBand 传输的字节数，以及利用 Grace 内存提供足够的总容量来容纳模型、梯度和优化器状态。

FlexShard：无需通信的 Muon

除路由器和嵌入（约占参数的 99.5%）外，我们对每个权重矩阵使用 Muon，其余使用 AdamW。Muon 每个参数需要一个 fp32 状态，而 Adam 需要两个。去掉第二个状态对于让 65,536 token 的序列放入内存至关重要。

Muon 通过几次 Newton-Schulz 迭代对每个更新进行正交化，这些迭代作用于整个矩阵。由于 FSDP 将每个矩阵拆分为 384 个分片，简单实现会在一个 rank 上收集每个矩阵，在那里运行 Newton-Schulz，然后将结果分散回去。在我们的互连上，这种通信过于昂贵。

FlexShard 现已成为 TorchTitan 的一部分，让我们可以控制哪个 rank 拥有哪些参数。我们按整个专家对专家进行分片，因此每个 rank 拥有完整的专家矩阵，并在完全无通信的情况下运行其 Newton-Schulz 迭代。稠密矩阵（如注意力投影）仍收集在拥有它的 rank 上，每个 rank 在稠密矩阵传输期间处理自己的专家，从而重叠通信。

在 FSDP 的 CPU offload 下，优化器运行在 Grace 上，参数也存储在那里，而 Newton-Schulz 在 72 个 Arm 核心上每步约需 35 秒。我们为 TorchTitan 贡献了支持，使其能够在参数保持 offload 的同时在 GPU 上运行 Newton-Schulz。优化器步骤从 41.7 秒降至 2.2 秒，每 GPU 吞吐量从 441 提升至 597 tokens/s（+35%）。

图 3. Rank 0 从 80 秒开始的性能分析步骤。上：Newton-Schulz 在 Grace 上运行，反向传播后 GPU 空闲约 34 秒，同时 CPU 进行计算。下：Newton-Schulz 在 Hopper 上运行，优化器只占该步骤的一小部分。

激活 offload 到 Grace

在我们的设置中，每层收集的权重体积不依赖于 token 数量，因此更长的序列会增加计算时间，但不影响通信。我们最初只在层边界存储激活，并在反向传播期间重新计算层内所有内容。在 TorchTitan 中，这被称为 FullAC。由于 GPU 在基线步骤的大部分时间里处于空闲状态，额外的计算并未延长该步骤。结果，在 32,768 tokens 时，每 GPU 吞吐量从 90 提升至 370 tokens/s。

在 65,536 tokens 时，仅边界激活就占用 74 层 × 65,536 tokens × 6,144 × 2 字节 ≈ 每 GPU 55 GiB，这无法与权重一起放入 HBM。GH200 具有 900 GB/s 的链路连接到 120 GB 的 CPU 内存，因此我们决定将它们 offload。每个边界激活在前向传播产生后立即复制到 Grace，并在反向传播重新计算它的前一层复制回 Hopper。一个 offload 参数设置 offload 的层比例，以此用 Grace 内存换取 HBM。

TorchTitan 的激活检查点包装器使这一更改变得方便。我们向 TorchTitan 贡献了激活检查点 offload 的钩子，并在其之上构建了我们的 offload 策略。通过 offload，每 GPU 吞吐量从 370 提升至 508 tokens/s（+37%），序列长度增长到 49,152 tokens，然后借助 Muon 增长到 65,536。

图 4. Offload 调度，用 16 层代替 74 层绘制。在前向传播期间，每个边界激活移动到 Grace。在反向传播期间，它在重新计算的前一层返回 HBM。

消除 FSDP 流的碎片

在升级我们的软件栈后，CUDA 缓存分配器开始在约 70% 内存利用率时在每个 rank 上重试分配。在重试时，有 25-30 GiB 的 HBM 空闲，但分配器无法使用它。

在调查根本原因后，我们在 FSDP 重叠通信与计算的方式中找到了它。FSDP 使用一个计算流和四个侧流，分配器为每个流维护单独的池。一个流释放的内存只能由该流重用。FSDP 在写入它的流上分配每个缓冲区，而我们其他一些通信优化在 FSDP 的调用内部分配了临时缓冲区，因此这些缓冲区也落在了侧流上。在峰值内存时重放分配器的事件日志显示，侧流持有 27.3 GiB，其中只有 6.4 GiB 在使用中。剩余的 20.9 GiB 是计算流无法触及的，因此计算流无法为 MoE 缓冲区找到 4 GiB 的块。

图 5. 在峰值活跃内存时每个分配器流缓存持有的 HBM，从 rank 0 上分配器的事件跟踪重建。实心条是活跃张量，虚线条是只有所属流可以重用的缓存内存。

幸运的是，FSDP2 有一个私有标志 unshard_async_op，它会在计算流上运行 all-gather 设置，并让 NCCL 在其内部流上运行。随后，gather 缓冲区从计算流的池中分配，并归还到该池。我们调整了自定义 gather，使其异步运行，因此计算流只在需要权重时才等待网络。这将发生分配器重试的 rank 数量从 384 降至零。峰值预留内存从 88.1 GiB 降至 74.6 GiB，吞吐量从每 GPU 996 tokens/s 提升至 1,210 tokens/s。

我们提出了一个 TorchTitan 选项来启用该标志。经过评审，unshard_async_op 正被移入公开的 FSDP API（PyTorch #202343），以便其他 FSDP 用户也能依赖它。

向上游贡献与结语

我们刚刚描述的配置使得在公共欧洲基础设施上以可用的吞吐量将预训练运行扩展到 1T 参数成为可能。为了让其他人能够复现这些结果（更多细节见配套文章），我们将所有修复都上游到了 PyTorch 和 TorchTitan。

除了我们已经提到的三项贡献外，我们还向 TorchTitan 贡献了共享 RoPE 缓存（在我们的序列长度下每 GPU 2.56 GiB）、在反向传播中跳过 embedding 重新 gather，以及可配置的 seed-offset mesh 轴。在 PyTorch 中，我们修复了CPU 初始化重置 RNG 状态和trunc_normal_ all-gathering DTensors 的问题。我们还向于利希超级计算中心报告了作业启动时未归属的主机内存问题，因为它会影响该机器的每一位用户。

得益于 PyTorch 和 TorchTitan 提供的控制水平，训练现在在每 DDP 副本 384 块 GH200 上以每 GPU 约 1,300 tokens/s 的持续吞吐量运行，4 个 DDP 副本共 1536 块 GPU，并通过充分利用 HBM 和 Grace 架构而适配内存。配套文章涵盖了图 1 中其余路径：数据和专家负载均衡、减少通过 InfiniBand 发送的字节数、节点本地重分片、SonicMoE 中的动态 tile 调度，以及限制 NCCL 的 SM 使用。

我们感谢欧盟委员会授予的 Frontier AI Grand Challenge 奖项。我们希望这项工作能推动对公共计算基础设施的进一步投资，以及欧洲 AI 生态系统的持续增长。

了解更多：

• 在 Jupiter 上训练 1T 参数 MoE（完整文章）
• PyTorch
• TorchTitan
• FlexShard
• Frontier AI Grand Challenge 公告

文章 Training a 1T-Parameter MoE on EU public infrastructure with PyTorch and TorchTitan 首次发表于 PyTorch。

来源：PyTorch：Blog · pytorch.org