华为盘古Ultra准万亿MoE模型:业界一流,源自昇腾原生的长稳训练
“对硬件特性的深度挖掘和算法创新,是华为引领AI技术的硬实力。”
作者丨李希
本篇关键词:“昇腾亲和的先进架构,准万亿大稀疏长稳训练,实现高MFU高吞吐”
Pangu Ultra MoE是一个全流程在昇腾NPU上训练的准万亿MoE模型。最近华为盘古团队发布了Pangu Ultra MoE模型架构和训练方法的技术报告[1],进一步披露了这个模型的细节。
训练超大规模和极高稀疏性的MoE模型极具挑战,训练过程中的稳定性往往难以保障。针对这一难题,盘古团队在模型架构和训练方法上进行了创新性设计,成功地在昇腾NPU上实现了准万亿MoE模型的全流程训练。
盘古团队提出Depth-Scaled Sandwich-Norm(DSSN)稳定架构和TinyInit小初始化的方法,在昇腾NPU上实现了10+ T tokens数据的长期稳定训练。此外,他们还提出了EP loss负载优化方法,这一设计不仅保证了各个专家之间的能保持较好的负载均衡,也提升了专家的领域特化能力。同时,Pangu Ultra MoE使用了业界先进的MLA和MTP架构,在训练时使用了Dropless训练策略。
技术报告标题:PanguUltraMoE模型架构与训练方法
技术报告地址:https://gitcode.com/ascend-tribe/pangu-ultra-moe/blob/main/Pangu_Ultra_MoE_CN_Report.pdf
01
破解准万亿MoE模型性能瓶颈
打造芯片协同的先进架构
近期,盘古团队在MoE模型训练领域再进一步,重磅推出参数规模高达718B的准万亿全新模型——Pangu Ultra MoE。该模型旨在实现超大规模MoE架构在模型效果与效率之间的最佳平衡。
为了达到这个目标,研究团队在设计PanguUltra MoE架构的时候,充分考虑昇腾硬件特性,在昇腾NPU平台上,融合计算、通信和内存等多维度指标,构建了大规模系统模拟器,并系统性地探索约一万个不同的MoE结构组合,最终搜索出一套在训练与推理吞吐上均达最优的架构方案。
Pangu Ultra MoE是一个超大规模、高稀疏比的架构,同时也包含MLA和MTP等先进架构和特有的DSSN稳定性架构和EP 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维隐藏层,精准匹配DaVinci芯片的16×16 MatMul单元,充分发挥Cube核心的计算潜力。
层数亲和流水线并行:设置61层Transformer结构,并预留额外MTP层空间,保障计算负载均衡的PP/VPP流水线调度,减少pipeline气泡,提升整体并行效率。
专家规模符合幂次规律:路由专家数量设为 2^8 = 256 ,在TP×EP并行下提升All-to-All通信效率,有效加速分布式训练。
Pangu Ultra MoE的预训练阶段在6k到10k张NPU上进行,全流程采用dropless训练模式。预训练阶段进行了长序列扩展,最终模型具备128k长序列能力。在后训练阶段,Pangu Ultra MoE移除了负载均衡辅助损失,保留专家间已有的特化能力,从而进一步提升模型对目标数据的学习效率。最终模型在多个权威开源评测集上展现出稳健性能,如表 1 所示,整体效果具备一定优势。
表1:Pangu Ultra MoE与目前主流模型效果对比
02
面向超大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提出一种标准差为√1/2dL的初始化方案,能够同时兼顾模型深度与宽度,其中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%。
03
基于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分别代表阿拉伯语,德语,法语,以及俄语。
04
多Token投机推理新路径:
MTP头延迟扩展策略
投机接受长度预期提升38%
投机推理是一种提升大模型生成效率的有效方法,其核心思想是在主模型生成token之前,由一个轻量辅助模块预先预测多个候选token,并通过快速校验机制决定是否接纳,从而实现推理过程的并行化与加速。在当前大模型推理中,Multi-token Prediction(MTP)技术已成为实现多token级别投机生成的重要手段。
盘古团队在实践中发现,获取多token的投机推理能力并不需要从训练开始便配置多个MTP头,而是可以在训练后期对单头MTP进行扩展来达到类似的效果。为验证这一策略的有效性,团队使用20BMoE为主干模型,训练185B数据。具体对比设置为:以两个 token 的投机推理为目标,分别训练了从头开始配置单/两个MTP头的模型(即单头从头训练和双头从头训练),以及在单头MTP模型训练至收敛后,通过复制已有头的参数再增训出第二个MTP头的模型。对于扩增的模型,对比全参续训以及冻结主干和一头的续训的效果,即双头扩增全参训练和双头扩增冻结训练。下游使用LAMBADA续写作为评测任务。
结果如图3所示。双头扩增模型的接受长度和延迟基本和双头从头训练一致,而双头的接受长度约2.30,单头的接受长度约1.67,双头相对单头提升约38%。在模型效果方面,双头扩增模型全参训练和从零训练相当,而由于冻住了主干和一头,双头扩增冻结训练的精度在扩增的位置基本保持不变。这表明后期的MTP扩展可以达到多头的从头训练的投机推理效果,可以在模型训练早期保持较小的MTP配置并在后期再进行扩展,兼顾计算成本和推理能力。
图3: 20BMoE的MTP在 LAMBADA 续写上的投机推理结果。在接受长度上,双头相对单头提升约38%,而双头可以基本无损地通过后期扩增单头得到。
05
迭代难例挖掘与多能力协同:
后训练强化学习持续提升的关键
模型后训练的过程中,团队参考了业界常规的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
【2】Pangu Ultra: Pushing the Limits of Dense Large Language Models on Ascend NPUs
(来源:新浪科技)