2604.00785

在 Aurora 超级计算机上实现大规模 MoE 语言模型的可扩展预训练

一句话定位

本文展示了在 Aurora 超算（127,488 个 Intel PVC GPU tile）上使用自研训练库 Optimus 对高达 220B 参数的 MoE 模型进行预训练，并通过 FastSparseMoE 和 EP-Aware 分片优化器等技术实现高效扩展。

小学生也能听懂

这篇论文就像用超级大的乐高积木（Aurora超算）搭一个聪明的机器人（MoE语言模型），他们发明了新工具（Optimus库）和更快的拼装方法（FastSparseMoE和EPSO），让机器人学得又快又好，还能在坏掉几块积木时自动修好。

为什么值得记录

• 首次在 Aurora 超算上完成万亿级 token 的 MoE 模型预训练，验证了 ExaScale 系统的 LLM 训练能力
• 提出 FastSparseMoE 模块，显著优化 MoE 前向-反向计算效率，最高达 2.83x 加速
• 设计 EP-Aware 分片优化器（EPSO），解决专家并行下非专家参数状态冗余问题，提升训练吞吐
• 实现从 384 到 12,288 GPU tiles 的强扩展，Mula-220B-A10B 模型保持约 90% 扩展效率
• 构建完整容错体系（双检查点、持久化模型检查点、软硬节点故障处理），保障千卡级训练稳定性

核心方法

• 开发 Optimus 训练库，支持张量并行（TP）、专家并行（EP）、流水线并行（PP）及混合并行策略
• 实现选择性激活检查点（SAC），对 norm、attention 和 SparseMoE 模块灵活启用重计算以节省显存
• 设计 FastSparseMoE 模块：通过 allgather 替代 all2all 通信、GPU 核函数优化 token 计数与索引生成、Grouped GEMM 合并专家计算，提升 MoE 计算效率
• 提出 EP-Aware 分片优化器（EPSO）：将参数按是否被 EP 复制分为两组，分别沿 DP 和 DP×EP 维度分片优化器状态，减少内存冗余
• 采用双检查点 + 持久化模型检查点机制，结合 DP-Scattered 写入策略，分散 I/O 压力并确保训练可恢复性
• 实现软硬节点故障自动检测与恢复机制，利用缓冲节点无缝替换故障节点

关键结果

• Mula-7B-A1B（6.9B 总参数，1.3B 激活参数）在 4 万亿 token 上预训练，平均 benchmark 准确率达 67.0%，优于同计算量密集模型 Mula-1B（60.3%）
• Mula-7B-A1B 性能接近 AllenAI 的 OLMoE-1B-7B-0924（68.5%），验证软件栈正确性
• FastSparseMoE 在 Mula-7B-A1B 上实现 2.83x 前向-反向加速，端到端训练加速 1.71x
• EPSO 在 Mula-20B-A2B 上实现 1.36x 优化器级加速，端到端加速 1.19x
• Mula-220B-A10B 模型在 12,288 GPU tiles 上训练，扩展效率维持在约 90%（相比 384 tiles）

局限与风险

• 未报告完整 4 万亿 token 下超大模型（如 Mula-220B-A10B）的最终 benchmark 性能，仅训练至 100B token
• 扩展效率在超过 1,000 GPU tiles 后下降约 10%，可能与 MoE 路由非均匀性或通信开销有关，但未深入分析根本原因
• FUR（强制均匀路由）实验表明路由非均匀性不是扩展效率下降主因，暗示系统级通信或负载均衡存在瓶颈
• 所有实验基于 Intel PVC GPU 和 OneCCL 通信库，结果可能不直接适用于其他硬件平台

适合沉淀的概念

mixture-of-experts, expert-parallelism, sharded-optimizer, pipeline-parallelism, activation-checkpointing, model-parallelism, fault-tolerance, gpu-kernel-optimization

阅读注意

• 关注 FastSparseMoE 的 Stage 1 为何选择 allgather 而非理论更优的 all2all——因 OneCCL 对规则通信模式优化更好
• 注意 EPSO 如何区分 expert 和 non-expert 参数的分片策略，这是其性能优势的关键
• 图 4(b) 显示即使使用 FUR，扩展效率仍下降，说明问题不在路由本身，而在系统层面
• Table 3 中 FSMOE+ 列表示同时使用 FSMOE 和 EPSO 的叠加效果，非简单相加，体现优化间的耦合性
• 可靠性设计部分（第4节）是千卡训练成功的关键，值得在工程实践中借鉴

质量说明

• 采用来源：raw，raw/papers/2026/04/2604.00785.md
• 生成模型：LongCat-Flash-Chat
• 源材料判断：论文提供了完整的实验设置、模型配置、优化细节和量化结果，包含训练曲线、benchmark 对比和扩展效率数据，信息充分且自洽。