在欧洲公共超算上训练1万亿参数MoE，PyTorch如何突破带宽瓶颈

导语

训练近1万亿参数的混合专家模型（MoE），通常需要专属的大规模AI集群。EUROPA联盟在EuroHPC Frontier AI Grand Challenge中选择了另一条路径：利用面向科研开放的JUPITER超级计算机，结合PyTorch和TorchTitan，将前沿规模预训练搬到共享公共基础设施上。

项目使用的模型总参数约为9519亿，但每个token实际激活约328亿参数。模型包含74层、384个路由专家，每次选择8个专家，隐藏维度为6144，词表规模为25.6万，预训练序列长度达到65536 token。实验在384块NVIDIA GH200 GPU上进行，覆盖96个节点。

核心瓶颈不是参数量本身

JUPITER的每个节点配备4个GH200超芯片。每块GPU拥有96 GB HBM3，并通过高速NVLink-C2C连接72核Grace CPU和120 GB LPDDR5X内存；但节点间通信依赖InfiniBand，带宽远低于GPU与本地CPU之间的链路。

这意味着训练过程不能简单依赖传统的数据并行。团队采用384路FSDP，在节点内安排4路专家并行，同时放弃流水线并行，将MoE的token全互联尽量限制在节点内的NVLink上。长序列则被用来增加计算时间，让跨节点通信有更多机会与计算重叠。

三项关键优化

• FlexShard降低Muon通信成本。 大多数权重使用Muon优化器，其更新需要对完整矩阵执行Newton-Schulz正交化。团队按完整专家切分参数，使每个rank能够在本地处理专家矩阵，避免频繁收集和分散；密集矩阵则在传输期间与专家计算重叠。将Newton-Schulz从Grace CPU移到GPU后，优化器步骤由约41.7秒缩短至2.2秒，吞吐从441提升到597 tokens/s/GPU。
• 激活检查点卸载到Grace。 65536 token序列下，仅保存层边界激活就约需55 GiB显存。TorchTitan的检查点机制被扩展为把部分激活转移到Grace内存，在反向计算前再取回，从而释放HBM空间。32K序列配合完整激活重计算时，吞吐已由90提升至370 tokens/s/GPU；卸载后进一步达到508 tokens/s/GPU。
• 让FSDP通信进入计算路径。 FSDP下一层的all-gather可以与当前层计算并行。项目使用unshard_async_op在计算流上执行相关操作，释放每块GPU约14 GiB HBM，并改善通信与计算的衔接。最终，在384块GH200上，单GPU吞吐达到约1300 tokens/s。

意义与边界

这一案例的价值不只是刷新一个吞吐数字。它证明，公共科研超算即使不是为AI训练专门设计，也可以通过并行策略、内存分层和框架级改造承载前沿模型。PyTorch的可编程性，以及TorchTitan对分片、检查点和优化器流程的支持，使团队能够针对具体硬件逐项调整，并将部分能力回馈到开源项目。

不过，结果高度依赖JUPITER的GH200拓扑、内存配置、序列长度和并行布局，不能直接外推为所有1万亿参数MoE的通用方案。它更重要的启示是：当跨节点带宽成为稀缺资源时，模型结构、优化器状态、激活存储与通信调度必须作为一个整体设计。

来源：PyTorch Blog