Pangu Ultra MoE 模型架构与训练方法

Pangu Ultra MoE 是一款在 Ascend NPU 平台上全流程开发的准万亿参数 MoE 大模型，其在开源评测集上达到一流水平。本文介绍了 Pangu Ultra MoE 面向超大规模参数与高稀疏比场景下的关键技术方案，涵盖模型架构、训练方法和高效推理系统等方面。

模型架构与训练方法：

• 基础架构：通过系统级模拟器进行架构搜索，最终得到约 718B 参数量的模型，包含 61 层，隐层大小为 7680，每层注意力模块有 128 个头，每层 MoE 结构包含 1 个共享专家与 256 个路由专家，每个 token 激活 8 个专家。
• 范数稳定的模型架构：引入 DSSN 稳定结构与 TinyInit 小初始化策略，有效缓解训练初期的梯度波动，显著提升模型的稳定性与收敛效率。实验表明，DSSN/TinyInit 相比 Pre-LN/Init0.006，突刺比例减少了 51%。
• 专家负载均衡优化：针对高稀疏比 MoE 架构中的负载不均问题，设计了 EP-Group Auxiliary Loss，在专家并行组粒度引入正则项，兼顾通信效率与正则强度，实现专家激活均衡并促进特化。实验表明，EP-Group 负载均衡 loss 在大部分任务上都有显著优势，平均提升了 1.5 个百分点。
• MTP 扩头策略：通过扩展 Multi-Token Prediction（MTP）多头能力，支持多 token 投机推理，后续可用于 Pangu Ultra MoE 以提升模型生成质量和推理速度。实验表明，后期 MTP 扩展能够达到多 token 投机推理的效果。
• 迭代式强化学习：在模型后训练阶段，构建了强化学习训练系统，增强数据利用效率与能力协同，缓解能力增长不均衡。

推理优化：

• 并行优化：采用 PD 分离技术为预填充和解码实例部署不同的并行化策略，实现极致的融合大算子优化降低内存访问开销。在 Atlas 800I A3 384 卡集群，将 EP 规模扩展到 144 卡以获得更低的解码延迟。
• 融合大算子：实现了一系列融合大算子以提高模型推理效率，包括 MLA 预处理融合大算子、加法 & RmsNorm & 量化、门控 & Top-K 等。
• 量化压缩优化：引入 W8A8/W4A4/KV 等多级量化机制，在保持精度的前提下提升了内存与计算效率。W8A8 量化显存需求减少约一半，吞吐能力提升 20%；W4A4 配置在多种任务中均展现出良好的鲁棒性。

结论： Pangu Ultra MoE 在模型稳定性、训练效率与推理性能方面实现了系统性优化，展现出高稀疏比大模型在硬件受限环境下的良好扩展性与工程可行性。通过软硬件协同设计，模型在保障训练稳定的同时有效压缩了推理成本，并提升了整体吞吐性能，体现出稀疏化技术在构建超大规模模型中的应用潜力，为大模型稀疏架构的发展与落地提供了重要参考。