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

总览

Pangu Ultra MoE是一个全流程在昇腾NPU上训练的准万亿MoE模型。作为拥有超大规模和极高稀疏性的 MoE 模型，它的稳定训练极具挑战。针对这一难题，盘古团队在模型架构和训练方法上进行了创新性设计，成功地在昇腾 NPU 上实现了准万亿 MoE 模型的全流程训练。本文将对Pangu Ultra MoE的模型架构与训练策略进行简要介绍。

基础架构与训练方法

在针对 Ascend NPU 平台的大模型系统设计中，我们结合计算、通信与内存等多方面因素，构建了完整的系统级模拟器，并基于约一万个不同的 MoE 架构组合，系统性地进行了架构搜索。最终得到一个约718B参数量的结构，它有61层，隐层大小为7680，每一层的注意力模块有128个头，每层MoE结构包含1个共享专家与256个路由专家，每个token激活8个专家，每个路由专家FFN中间维度为2048。

该模型架构在 Ascend NPU 上的模拟训练吞吐和推理吞吐均达到当前搜索空间中的最优，其特性分析如下：

• 隐藏维度贴合硬件：设置7680维隐藏层，精准匹配DaVinci芯片的16×16 MatMul单元，充分发挥Cube核心的计算潜力。

• 层数亲和流水线并行：设置61层Transformer结构，并预留额外MTP层空间，保障计算负载均衡的PP/VPP流水线调度，减少pipeline气泡，提升整体并行效率。

• 专家规模符合幂次规律：路由专家数量设为 $2^8=256$，在TP×EP并行下提升All-to-All通信效率，有效加速分布式训练。

在模型架构设计与训练流程中，我们亦引入了多项关键优化：

• 超大规模和超高稀疏比：采用256个路由专家，每个token激活8个专家，模型总参数量718B，激活量39B。

• MLA注意力机制：引入MLA（Multi-head Latent Attention），有效压缩KV Cache空间，缓解推理阶段的内存带宽瓶颈，优于传统GQA方案。

• MTP多头扩展：采用单头MTP进行训练，后续复用MTP参数扩展至多头结构，实现多Token投机推理，加速整体推理过程。

• Dropless训练：采用Dropless训练可以避免Drop&Pad训推不一致问题，并且提升训练的数据效率。

• RL训练：采用迭代难例挖掘与多能力项均衡的奖励函数，并参考GRPO算法，提升了模型的训练效率与最终推理性能。

Pangu Ultra MoE 的预训练阶段在6k到10k张NPU上进行，全流程采用 dropless 训练模式。预训练阶段进行了长序列扩展，最终模型具备128k长序列能力。在后训练阶段，Pangu Ultra MoE移除了负载均衡辅助损失，保留专家间已有的特化能力，从而进一步提升模型对目标数据的学习效率。最终模型在多个权威开源评测集上展现出稳健性能，如下表所示，整体效果具备一定优势。

范数稳定的模型架构

在模型训练的稳定性方面，Pangu Ultra提出了范数稳定的模型架构，应用于稠密模型的训练。不过，相应技术在稀疏模型中同样适用。Pangu Ultra MoE延续了Pangu Ultra的设计理念，继续应用了Depth-Scaled Sandwich-Norm和TinyInit方法，以有效地稳定梯度并提高训练过程的稳健性。

• Depth-Scaled Sandwich-Norm：传统的Pre-LN结构容易因为子层输出规模波动而导致训练不稳定。为解决这一问题，Depth-Scaled Sandwich-Norm在每个子层输出后加入额外的层归一化，并引入深度缩放的初始化方式。具体地，通过将层归一化的可训练gamma参数初始化为与网络深度的平方根成反比的数值，从而稳定网络各层的输出尺度。这种结构能够有效抑制梯度爆炸或消失，降低梯度范数波动，提升模型的收敛性。

• TinyInit方法：模型初始化对训练稳定性与性能至关重要。Transformer模型通常采用较小的初始化尺度，而TinyInit方法提出了一种同时考虑模型深度与宽度的初始化方案，即采用标准差为$\sqrt{\frac{1}{2dL}}$的初始化方式，其中$d$为隐藏维度，$L$为模型层数。这种方法能够在模型中获得更一致的参数尺度，有助于优化和收敛。

这两种方法的结合显著提高了大规模训练的稳定性，这一点在Pangu Ultra MoE上也得到了验证，减少了51%的突刺量（见下图），缓解了梯度范数频繁突刺的问题，能够有效降低大模型训练过程中的不稳定性，加快模型收敛，提升模型性能。同时DSSN+TinyInit被应用到Pangu Ultra MoE中实现了10+T tokens数据的长稳训练。

专家负载均衡

在混合专家模型（MoE）中，由于每个token只会被分配给一部分专家，因此容易出现负载不均衡的情况。负载不均衡指的是不同专家被分配的token数量存在显著差距。当采用专家并行策略（EP，expert parallelism）时，负载不均衡会影响计算效率，被分配过多token的专家会成为计算瓶颈，而其他专家则处于低利用率状态。同时负载过低的专家可能存在训练不充分的问题，影响最终的模型效果。因此如何使token更均衡地分布至不同专家，对提高混合专家模型的训练效率和效果非常重要。

为了保证负载均衡，一般通过增加辅助的负载均衡loss（auxiliary loss）来鼓励tokens在专家之间均衡分布。然而，如果负载均衡loss过度约束tokens分配的均衡性，也会影响模型路由的表达能力。之前主流的负载均衡loss一般是约束单个序列或者单个micro batch内的token分配均衡性, 由于单个序列往往是来自同一领域的数据，过度的均衡可能影响专家特化（expert specialization）。而我们发现对于采用专家并行策略训练的模型，可以设计一种对模型路由约束更小，同时不影响计算均衡性的EP-Group负载均衡loss。

该方法的基本思想是让每块卡上的专家接收来自EP组内所有卡上的micro batch路由给自己的token。由此我们可以设计一个负载均衡loss，来约束EP组内所有micro batch路由到组内专家之后的均衡性。这相当于把EP组内部的所有micro batch合起来计算负载均衡的loss。具体而言，我们在每个专家并行进程组内部执行all-reduce操作，聚合token选择频率，从而得到每个专家的负载比例，由此来定义负载均衡loss。

为了验证EP-Group均衡损失函数的效果，我们使用一个20B参数量激活1.6B的MoE模型进行了100B数据量的对比实验，结果如下表所示。可见，EP-Group均衡损失函数相比在单个micro batch内计算均衡损失在大部分任务上都有显著的优势，评测指标平均提升了1.5个点。

多token投机推理的MTP扩展策略

投机推理是一种常见的推理加速方法，其基本思想是在主模型生成结果之前，利用一个轻量级辅助模型对未来的 token 进行预测，并通过校验机制决定是否接纳预测结果。该方法通过并行化 token 生成和验证流程，有效降低了推理过程中的延迟，近年来在大语言模型的部署实践中得到广泛应用。

在这一背景下，Multi-Token Prediction（MTP）被提出用于进一步增强投机推理的性能。MTP 在原始模型基础上额外增加一层 Transformer，用于模拟输出层的预测行为，从而提升对原始模型 logit 的拟合能力。在推理过程中，MTP 可以利用该机制进行投机性预测，即在主模型尚未正式生成 token 之前，先行预测下一个甚至多个 token，并通过校验机制决定是否采纳，进而显著加速推理过程。

在原始设计中，一个 MTP 头可以预测一个 token，而引入多个 MTP 头则可以实现多个 token 的并行预测，从而进一步提升推理吞吐量。我们在实践中观察到一个关键现象：多 token 投机推理的能力，并不需要从训练开始便配置多个 MTP 头，而是可以在训练后期对单头 MTP 进行扩展，从而达到类似的推理加速效果。

实验结果表明，MTP 的多 token 投机能力具有良好的可扩展性，训练策略上也可采用“后期扩展”的方式，即先以单头训练，后期再增头精调，以实现高效的训练与推理折中。

后训练强化学习中的迭代难例挖掘与多能力协同

模型后训练的过程中，团队参考了业界常规的GRPO算法提升模型的推理性能。然而，在超大参数规模情况下，直接应用GRPO会带来两方面的问题：1. 算法训练需要依赖多回复通过率在(0, 1)内的数据，随着模型性能的提升，相同输入的推理结果准确率越来越高，导致训练过程中被“浪费”的数据不断增加，降低推理效率；2. 模型训练需要兼顾多能力协同提升，包括数学、代码和通用能力等，不同能力项的奖励函数设计会导致模型能力增长上的不匹配，出现“跷跷板”问题。

为了解决上述两个实践难题，盘古团队设计了Pangu Ultra MoE的强化学习训练系统，如上图所示，提升了大规模MoE模型的训练稳定性与推理性能。系统设计的关键在于两个部分：

• 迭代难例挖掘：系统从初始的数据池中进行多回复推理，选取回复通过率在(0, 1)的数据组成RL训练数据池，避免训练过程中优势函数为0的情况。模型更新多步后，重复上述过程，更新RL训练数据池，以保持推理效率最大化。
• 多能力项奖励系统：为了确保模型多能力项协同提升，数学和代码均采用了基于规则的奖励（数学任务根据最终答案正确与否来分配0/1的奖励，代码任务根据生成代码所能通过的测试用例比例来给予奖励），通用奖励模型则使用LLM-as-a-judge的方法，使用外部模型对生成的回复质量进行评分。所有能力项的奖励最终均进行了归一化处理，保证了模型在多个能力项的综合表现。

许可证

本项目基于Common Creative Attribution 4.0 International许可证。

Credits

感谢所有参与技术报告撰写的作者。