华为盘古 Ultra MoE:昇腾 NPU 上准万亿 MoE 大模型训练方案详解

华为盘古发布 Ultra MoE 模型:在昇腾 NPU 上训练的准万亿 MoE 模型,通过 DSSN、TinyInit、EP group loss 等技术,实现了稳定训练和性能优化。

原文标题:还得是华为!Pangu Ultra MoE架构:不用GPU,你也可以这样训练准万亿MoE大模型

原文作者:机器之心

冷月清谈:

本文详细介绍了华为盘古团队发布的 Pangu Ultra MoE 模型架构与训练方法,该模型是一个全流程在昇腾 NPU 上训练的准万亿 MoE 模型。文章重点介绍了以下几个关键技术点:

1. **DSSN 稳定架构和 TinyInit 小初始化:** 这两种方法保证了模型在昇腾 NPU 上 10+ T tokens 数据的长期稳定训练,通过减少梯度突刺率,提升训练的稳定性。
2. **EP group loss 负载优化:** 这种方法保证了各个专家之间能保持较好的负载均衡,提升了专家的领域特化能力,同时约束更小,不影响计算均衡性。
3. **MLA 和 MTP 架构:** 模型使用了业界先进的 MLA 和 MTP 架构,MLA 有效压缩 KV Cache 空间,MTP 加速整体推理过程,同时使用了 Dropless 训练策略。
4. **硬件亲和设计:** 盘古 Ultra MoE 在设计时充分考虑了昇腾硬件特性,包括隐藏维度、层数和专家规模,以充分发挥硬件的计算潜力。
5. **多 Token 投机推理:** 通过 MTP 头延迟扩展策略,在训练后期扩展单头 MTP,获取多 token 的投机推理能力,提升推理效率。
6. **迭代难例挖掘与多能力协同:** 在后训练阶段,通过迭代难例挖掘和多能力项均衡的奖励函数,借鉴 GRPO 算法,提升了模型的训练效率与最终推理性能。

总而言之,Pangu Ultra MoE 模型在模型架构、训练方法和硬件协同方面都进行了创新性的设计,成功地在昇腾 NPU 上实现了准万亿 MoE 模型的全流程训练,并在多个权威开源评测集上展现出一流的效果

怜星夜思:

1、Pangu Ultra MoE 模型在昇腾 NPU 上的训练,对国产 AI 芯片的发展有哪些启示?除了算力之外,昇腾 NPU 还需要在哪些方面进行提升才能更好地支持大模型的训练和推理?
2、文章中提到的 EP group loss 负载优化方法,如何在保证负载均衡的同时,避免过度约束 tokens 分配的均衡性,从而影响模型路由的表达能力?这种方法对其他 MoE 模型的训练有哪些借鉴意义?
3、Pangu Ultra MoE 模型在后训练阶段采用了强化学习方法,并针对超大参数规模带来的问题进行了优化。在实际应用中,如何选择合适的奖励函数,并避免不同能力项之间的“跷跷板”问题?

原文内容

机器之心发布

机器之心编辑部


Pangu Ultra MoE 是一个全流程在昇腾 NPU 上训练的准万亿 MoE 模型,此前发布了英文技术报告[1]。最近华为盘古团队发布了 Pangu Ultra MoE 模型架构与训练方法的中文技术报告,进一步披露了这个模型的细节。


超大规模和极高稀疏性的 MoE 模型极具挑战,训练过程中的稳定性往往难以保障。针对这一难题,盘古团队在模型架构和训练方法上进行了创新性设计,成功地在昇腾 NPU 上实现了准万亿 MoE 模型的全流程训练。


盘古团队提出 Depth-Scaled Sandwich-Norm(DSSN)稳定架构和 TinyInit 小初始化的方法,在昇腾 NPU 上实现了 10+ T tokens 数据的长期稳定训练。此外,他们还提出了 EP group loss 负载优化方法,这一设计不仅保证了各个专家之间能保持较好的负载均衡,也提升了专家的领域特化能力。同时,Pangu Ultra MoE 使用了业界先进的 MLA 和 MTP 架构,在训练时使用了 Dropless 训练策略。



  • 技术报告标题:Pangu Ultra MoE 模型架构与训练方法

  • 技术报告地址:https://raw.gitcode.com/ascend-tribe/pangu-ultra-moe/raw/main/Pangu_Ultra_MoE_CN_Report.pdf


破解准万亿 MoE 模型性能瓶颈

打造芯片协同的先进架构


近期,盘古团队在 MoE 模型训练领域再进一步,重磅推出参数规模高达 718B 的准万亿全新模型 ——Pangu Ultra MoE。该模型旨在实现超大规模 MoE 架构在模型效果与效率之间的最佳平衡。


为了达到这个目标,研究团队在设计 Pangu Ultra MoE 架构的时候,充分考虑昇腾硬件特性,在昇腾 NPU 平台上,融合计算、通信和内存等多维度指标,构建了大规模系统模拟器,并系统性地探索约一万个不同的 MoE 结构组合,最终搜索出一套在训练与推理吞吐上均达最优的架构方案。


Pangu Ultra MoE 是一个超大规模、高稀疏比的架构,同时也包含 MLA 和 MTP 等先进架构和特有的 DSSN 稳定性架构和 EP group loss 负载优化。下面是 Pangu Ultra MoE 的主要的架构和训练特性:


  • 超大规模和超高稀疏比:采用 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 昇腾亲和设计考虑:


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

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

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


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


表 1: Pangu Ultra MoE 与目前主流模型效果对比


面向超大MoE模型稳定训练新范式:

DSSN结构和TinyInit加持

梯度突刺率下降 51%

支撑 10+T tokens 数据长稳训练


随着参数规模和数据体量的激增,大模型训练面临前所未有的稳定性挑战。频繁的梯度范数突刺已成为阻碍收敛效率与模型性能提升的主要瓶颈。如何在确保训练深度和宽度扩展的同时,维持梯度信号的稳定传递,成为构建高可靠性大模型架构的关键课题。在 Pangu Ultra 稠密模型 [2] 的训练中,Depth-Scaled Sandwich-Norm 和 TinyInit 方法在保障训练稳定性上起到了关键性的作用,所以 Pangu Ultra MoE 依旧采用这个方案来控制训练稳定性。经过实验证明,此设计在 Pangu Ultra MoE 的训练中同样能起到增强稳定性、加快收敛速度的作用。


Depth-Scaled Sandwich-Norm(DSSN):传统的 Pre-LN 结构存在因为子层输出规模波动而导致训练不稳定的现象,DSSN 是为了解决这一问题而提出的。通过在每个子层输出后加入额外的层归一化,并引入深度缩放的初始化方式,从而稳定网络各层的输出尺度,达到抑制梯度异常、降低范数波动的目的。


TinyInit:Transformer 模型普遍采用较小的初始化尺度,TinyInit 提出一种标准差为图片的初始化方案,能够同时兼顾模型深度与宽度,其中d表示隐藏维度,L表示模型层数。同时,对词嵌入层采用标准差为 0.5 的初始化。实验表明,这样的初始化策略有助于提升模型性能和训练稳定性。


Depth-Scaled Sandwich-Norm + TinyInit 的方案减少了 51% 的突刺量(见图 1),缓解了梯度范数频繁突刺的问题,能够有效降低大模型训练过程中的不稳定性,加快模型收敛,提升模型性能。同时 DSSN+TinyInit 被应用到 Pangu Ultra MoE 中实现了 10+T tokens 数据的长稳训练。


图 1: 训练过程的梯度范数对比图(黑色实线为突刺分界线)。DSSN+TinyInit 使梯度突刺率从 1.54% 下降到 0.76%,相对下降 51%。


基于 EP group 的负载均衡:

让计算效率和路由表达能力可以兼得


在训练混合专家模型(MoE)时,容易出现专家负载不均衡的情况。负载不均衡指的是不同专家被分配的 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, 这样训练时可以容忍单个 micro batch 的不均衡,只要多个 micro batch 的 token 路由到专家之后是均衡的即可。


为了验证 EP-Group 均衡损失函数的效果,盘古团队使用一个 20B 参数量的 MoE 模型进行了 100B 数据量的对比实验。结果如表 2 所示,可以看到 EP-Group 均衡损失函数在大部分任务相比主流的 Micro-batch 上都有显著的优势,平均提升了 1.5 个点。


表 2:  Micro-batch 和 EP-Group 的 auxiliary loss 效果比较


同时盘古团队对 Pangu Ultra MoE 的专家特化进行了分析,结果如图 2 所示,可以看到不同领域的数据对专家的选择存在显著的差异,这表明 EP-Group 均衡损失函数给模型提供了灵活的路由选择空间,促进了专家特化。


图 2:  Pangu Ultra MoE 的专家特化。其中 ar,de,fr,ru 分别代表阿拉伯语,德语,法语,以及俄语。


多 Token 投机推理新路径:

MTP 头延迟扩展策略

投机接受长度预期提升 38%


投机推理是一种提升大模型生成效率的有效方法,其核心思想是在主模型生成 token 之前,由一个轻量辅助模块预先预测多个候选 token,并通过快速校验机制决定是否接纳,从而实现推理过程的并行化与加速。在当前大模型推理中,Multi-token Prediction(MTP)技术已成为实现多 token 级别投机生成的重要手段。


盘古团队在实践中发现,获取多 token 的投机推理能力并不需要从训练开始便配置多个 MTP 头,而是可以在训练后期对单头 MTP 进行扩展来达到类似的效果。为验证这一策略的有效性,团队使用 20B MoE 为主干模型,训练 185B 数据。具体对比设置为:以两个 token 的投机推理为目标,分别训练了从头开始配置单 / 两个 MTP 头的模型(即单头从头训练和双头从头训练),以及在单头 MTP 模型训练至收敛后,通过复制已有头的参数再增训出第二个 MTP 头的模型。对于扩增的模型,对比全参续训以及冻结主干和一头的续训的效果,即双头扩增全参训练和双头扩增冻结训练。下游使用 LAMBADA 续写作为评测任务。


结果如图 3 所示。双头扩增模型的接受长度和延迟基本和双头从头训练一致,而双头的接受长度约 2.30,单头的接受长度约 1.67,双头相对单头提升约 38%。在模型效果方面,双头扩增模型全参训练和从零训练相当,而由于冻住了主干和一头,双头扩增冻结训练的精度在扩增的位置基本保持不变。这表明后期的 MTP 扩展可以达到多头的从头训练的投机推理效果,可以在模型训练早期保持较小的 MTP 配置并在后期再进行扩展,兼顾计算成本和推理能力。


图 3:  20B MoE 的 MTP 在 LAMBADA 续写上的投机推理结果。在接受长度上,双头相对单头提升约 38%,而双头可以基本无损地通过后期扩增单头得到。


迭代难例挖掘与多能力协同:

后训练强化学习持续提升的关键


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


图 4:  Pangu Ultra MoE 的强化学习训练系统


为了解决上述两个实践难题,盘古团队设计了 Pangu Ultra MoE 的强化学习训练系统,如图 4 所示,提升了大 MoE 模型的训练稳定性与推理性能。系统设计的关键在于两个部分:(1)迭代难例挖掘:模型阶段性更新后,从初始的数据池中进行多回复推理,选取回复通过率在 (0,1) 的数据组成 RL 训练数据池,以保持推理效率最大化;(2)多能力项奖励系统:为了确保模型多能力项协同提升,数学和代码均采用了基于规则的奖励,通用奖励模型则使用 LLM-as-a-judge 的方法对生成的回复质量进行评分,并对最终的 reward 进行归一化处理,保证了模型在多个能力项的综合表现。


[1] Pangu Ultra MoE: How to Train Your Big MoE on Ascend NPUs

https://arxiv.org/abs/2505.04519

[2] Pangu Ultra: Pushing the Limits of Dense Large Language Models on Ascend NPUs

https://arxiv.org/abs/2504.07866


© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

这个问题有点深奥啊,不过我试着理解了一下。感觉 EP group loss 的巧妙之处在于,它不是一刀切地追求每个 micro batch 的均衡,而是关注整个 EP 组内的均衡。这就给了模型一定的自由度,允许它在单个 micro batch 内根据数据特点进行灵活的路由选择。 这样既能保证整体的负载均衡,又能避免过度干预模型的学习。

这种方法对其他 MoE 模型的借鉴意义在于,提供了一种更加精细化的负载均衡控制思路。在实际应用中,可以根据模型的具体情况和数据特点,灵活调整 EP group 的大小和负载均衡的约束强度,以达到最佳的训练效果。

奖励函数的设计,简直就是玄学!我觉得关键在于,要对模型的能力进行拆解,然后针对每个能力设计合理的指标。比如,对于通用能力,可以采用 LLM-as-a-judge 的方法,让另一个更强大的模型来评价生成结果的质量。

至于“跷跷板”问题,我的理解是,要尽量平衡不同能力项的奖励力度。可以对奖励进行归一化处理,避免某个能力项的奖励过大,从而抑制了其他能力项的提升。另外,也要根据实际情况,动态调整奖励策略,比如,可以优先提升短板能力。

EP group loss 的核心在于将负载均衡的粒度从单个 micro-batch 扩展到 EP 组,其关键在于利用了专家并行(EP)的特性。 在 EP 中,每个 GPU/NPU 负责一部分专家,因此可以将一个 EP 组内的 micro-batch 视为一个整体来做负载均衡,允许单个 micro-batch 的不均衡,只要组内整体均衡即可。 这样的设计有几个优点:

* 降低了对专家特化的约束: 允许专家在不同领域的数据上有所偏好,避免了过度均衡导致的“一刀切”现象。
* 提高了计算效率: 减少了因负载不均衡导致的 GPU/NPU 空闲时间。

对其他 MoE 模型的借鉴意义在于,提供了一种新的负载均衡思路。 这种思路可以推广到其他并行训练框架中,例如数据并行、模型并行等。 通过将负载均衡的粒度扩展到更大的范围,可以更好地平衡负载均衡和模型表达能力之间的关系。

从学术角度来看,Pangu Ultra MoE 的成功验证了国产 AI 芯片在特定架构下进行大规模 MoE 模型训练的可行性。 它启示我们,国产 AI 芯片的发展不应仅仅着眼于通用算力的提升,更应该结合自身的特点,发展差异化竞争优势。 对于昇腾 NPU 而言,除了算力之外,以下几个方面也至关重要:

* 互联带宽: 大规模 MoE 模型训练涉及大量的数据传输和通信,因此需要提升 NPU 之间的互联带宽,降低通信延迟。
* 内存容量和带宽: MoE 模型参数量巨大,需要更大的内存容量和带宽来支持模型的存储和访问。
* 软件生态: 完善的软件生态能够降低开发门槛,吸引更多的开发者,促进模型的应用和创新。
* 异构计算支持: 考虑支持 CPU、GPU 等异构计算资源,实现更灵活的资源调度和优化。

在强化学习中,奖励函数的设计至关重要,它直接影响着模型的学习方向和最终性能。 针对超大参数规模的 MoE 模型,选择合适的奖励函数并避免“跷跷板”问题需要考虑以下几个方面:

* 多维度评估: 奖励函数应该综合考虑多个维度,例如通用能力(流畅度、相关性、逻辑性等)、专业能力(数学、代码等)和安全性(避免生成有害信息)。
* 基于规则与模型结合: 对于容易量化的能力(如数学、代码),可以采用基于规则的奖励函数;对于难以量化的能力(如通用能力),可以采用 LLM-as-a-judge 的方法,利用大型语言模型进行评估。
* 奖励归一化: 为了避免不同能力项之间的奖励尺度差异过大,需要对奖励进行归一化处理,保证每个能力项的贡献度相对均衡。
* Curriculum Learning: 可以采用 Curriculum Learning 的策略,先训练模型的基础能力,再逐步引入更复杂的能力,避免模型在训练初期就被复杂的奖励信号所淹没。

总而言之,奖励函数的设计是一个迭代优化的过程,需要根据模型的实际表现不断调整和完善。

EP group loss,我感觉就像是“抓大放小”的策略。单个token可能跑偏,但整体上要保证均衡。这让模型有更大的空间去学习,不至于被强制“平均主义”。

借鉴意义嘛,我觉得是提醒我们,不能为了均衡而均衡,还是要考虑模型的实际需求。不同的MoE模型,数据特点和硬件环境都不一样,要灵活应用这种思想。

这个问题问得好!这次Pangu Ultra MoE在昇腾上的成功,绝对是给国产AI芯片打了一剂强心针。这说明了国产芯片在特定场景下,完全有能力支撑起超大规模模型的训练。感觉以后可以多关注一下国产芯片在特定领域的突破,说不定能给我们带来意想不到的惊喜。

至于昇腾NPU,我觉得除了继续提升算力之外,更重要的是构建一个完善的生态系统。要吸引更多的开发者参与进来,提供易用的开发工具和丰富的模型库。只有这样,才能真正发挥出国产芯片的潜力。

国产芯片要崛起,除了硬实力,软实力也很重要!这次盘古Ultra MoE在昇腾上的成功,我觉得更多的是一种“软实力”的体现——也就是对硬件特性的理解和优化。只有真正了解硬件,才能把它的潜力发挥到极致。所以,昇腾NPU需要加强与算法团队的合作,共同优化模型和训练方法,才能更好地支持大模型的训练和推理。

奖励函数这东西,就像是给孩子设定目标。目标太单一,孩子就可能只关注这一方面,忽略了其他重要的东西;目标太多,孩子又可能顾此失彼,啥都做不好。

要避免“跷跷板”问题,我的建议是:1. 奖励要多样化,不能只看重一个指标;2. 要根据模型的实际情况,动态调整奖励的权重。比如,如果发现模型在数学方面比较弱,就可以适当提高数学奖励的权重。