機(jī)器之心發(fā)布

機(jī)器之心編輯部

Pangu Ultra MoE 是一個(gè)全流程在昇騰 NPU 上訓(xùn)練的準(zhǔn)萬(wàn)億 MoE 模型，此前發(fā)布了英文技術(shù)報(bào)告[1]。最近華為盤(pán)古團(tuán)隊(duì)發(fā)布了 Pangu Ultra MoE 模型架構(gòu)與訓(xùn)練方法的中文技術(shù)報(bào)告，進(jìn)一步披露了這個(gè)模型的細(xì)節(jié)。

訓(xùn)練超大規(guī)模和極高稀疏性的 MoE 模型極具挑戰(zhàn)，訓(xùn)練過(guò)程中的穩(wěn)定性往往難以保障。針對(duì)這一難題，盤(pán)古團(tuán)隊(duì)在模型架構(gòu)和訓(xùn)練方法上進(jìn)行了創(chuàng)新性設(shè)計(jì)，成功地在昇騰 NPU 上實(shí)現(xiàn)了準(zhǔn)萬(wàn)億 MoE 模型的全流程訓(xùn)練。

盤(pán)古團(tuán)隊(duì)提出Depth-Scaled Sandwich-Norm（DSSN）穩(wěn)定架構(gòu)和TinyInit 小初始化的方法，在昇騰 NPU 上實(shí)現(xiàn)了 10+ T tokens 數(shù)據(jù)的長(zhǎng)期穩(wěn)定訓(xùn)練。此外，他們還提出了EP group loss 負(fù)載優(yōu)化方法，這一設(shè)計(jì)不僅保證了各個(gè)專(zhuān)家之間能保持較好的負(fù)載均衡，也提升了專(zhuān)家的領(lǐng)域特化能力。同時(shí)，Pangu Ultra MoE 使用了業(yè)界先進(jìn)的 MLA 和 MTP 架構(gòu)，在訓(xùn)練時(shí)使用了 Dropless 訓(xùn)練策略。

• 技術(shù)報(bào)告標(biāo)題：Pangu Ultra MoE 模型架構(gòu)與訓(xùn)練方法
• 技術(shù)報(bào)告地址：https://raw.gitcode.com/ascend-tribe/pangu-ultra-moe/raw/main/Pangu_Ultra_MoE_CN_Report.pdf

破解準(zhǔn)萬(wàn)億 MoE 模型性能瓶頸

打造芯片協(xié)同的先進(jìn)架構(gòu)

近期，盤(pán)古團(tuán)隊(duì)在 MoE 模型訓(xùn)練領(lǐng)域再進(jìn)一步，重磅推出參數(shù)規(guī)模高達(dá) 718B 的準(zhǔn)萬(wàn)億全新模型 ——Pangu Ultra MoE。該模型旨在實(shí)現(xiàn)超大規(guī)模 MoE 架構(gòu)在模型效果與效率之間的最佳平衡。

為了達(dá)到這個(gè)目標(biāo)，研究團(tuán)隊(duì)在設(shè)計(jì) Pangu Ultra MoE 架構(gòu)的時(shí)候，充分考慮昇騰硬件特性，在昇騰 NPU 平臺(tái)上，融合計(jì)算、通信和內(nèi)存等多維度指標(biāo)，構(gòu)建了大規(guī)模系統(tǒng)模擬器，并系統(tǒng)性地探索約一萬(wàn)個(gè)不同的 MoE 結(jié)構(gòu)組合，最終搜索出一套在訓(xùn)練與推理吞吐上均達(dá)最優(yōu)的架構(gòu)方案。

Pangu Ultra MoE 是一個(gè)超大規(guī)模、高稀疏比的架構(gòu)，同時(shí)也包含 MLA 和 MTP 等先進(jìn)架構(gòu)和特有的 DSSN 穩(wěn)定性架構(gòu)和 EP group loss 負(fù)載優(yōu)化。下面是 Pangu Ultra MoE 的主要的架構(gòu)和訓(xùn)練特性：

• 超大規(guī)模和超高稀疏比：采用 256 個(gè)路由專(zhuān)家，每個(gè) token 激活 8 個(gè)專(zhuān)家，模型總參數(shù)量 718B，激活量 39B。
• MLA 注意力機(jī)制：引入 MLA（Multi-head Latent Attention），有效壓縮 KV Cache 空間，緩解推理階段的內(nèi)存帶寬瓶頸，優(yōu)于傳統(tǒng) GQA 方案。
• MTP 多頭擴(kuò)展：采用單頭 MTP 進(jìn)行訓(xùn)練，后續(xù)復(fù)用 MTP 參數(shù)擴(kuò)展至多頭結(jié)構(gòu)，實(shí)現(xiàn)多 Token 投機(jī)推理，加速整體推理過(guò)程。
• Dropless 訓(xùn)練：采用 Dropless 訓(xùn)練可以避免 Drop&Pad 訓(xùn)推不一致問(wèn)題，并且提升訓(xùn)練的數(shù)據(jù)效率。
• RL 訓(xùn)練：采用迭代難例挖掘與多能力項(xiàng)均衡的獎(jiǎng)勵(lì)函數(shù)，并參考 GRPO 算法，提升了模型的訓(xùn)練效率與最終推理性能。

以下是 Pangu Ultra MoE 昇騰親和設(shè)計(jì)考慮：

• 隱藏維度貼合硬件：設(shè)置 7680 維隱藏層，精準(zhǔn)匹配昇騰芯片的 16×16 MatMul 單元，充分發(fā)揮 Cube 核心的計(jì)算潛力。
• 層數(shù)親和流水線(xiàn)并行：設(shè)置 61 層 Transformer 結(jié)構(gòu)，并預(yù)留額外 MTP 層空間，保障計(jì)算負(fù)載均衡的 PP/VPP 流水線(xiàn)調(diào)度，減少 pipeline 氣泡，提升整體并行效率。
• 專(zhuān)家規(guī)模符合冪次規(guī)律：路由專(zhuān)家數(shù)量設(shè)為2?=256，在 TP×EP 并行下提升 All-to-All 通信效率，有效加速分布式訓(xùn)練。

Pangu Ultra MoE 的預(yù)訓(xùn)練階段在 6k 到 10k 張 NPU 上進(jìn)行，全流程采用 dropless 訓(xùn)練模式。預(yù)訓(xùn)練階段進(jìn)行了長(zhǎng)序列擴(kuò)展，最終模型具備 128k 長(zhǎng)序列能力。在后訓(xùn)練階段，Pangu Ultra MoE 移除了負(fù)載均衡輔助損失，保留專(zhuān)家間已有的特化能力，從而進(jìn)一步提升模型對(duì)目標(biāo)數(shù)據(jù)的學(xué)習(xí)效率。如表1所示，最終模型在多個(gè)權(quán)威開(kāi)源評(píng)測(cè)集上展現(xiàn)出一流的效果。

表 1: Pangu Ultra MoE 與目前主流模型效果對(duì)比

面向超大MoE模型穩(wěn)定訓(xùn)練新范式：

DSSN結(jié)構(gòu)和TinyInit加持

梯度突刺率下降 51%

支撐 10+T tokens 數(shù)據(jù)長(zhǎng)穩(wěn)訓(xùn)練

隨著參數(shù)規(guī)模和數(shù)據(jù)體量的激增，大模型訓(xùn)練面臨前所未有的穩(wěn)定性挑戰(zhàn)。頻繁的梯度范數(shù)突刺已成為阻礙收斂效率與模型性能提升的主要瓶頸。如何在確保訓(xùn)練深度和寬度擴(kuò)展的同時(shí)，維持梯度信號(hào)的穩(wěn)定傳遞，成為構(gòu)建高可靠性大模型架構(gòu)的關(guān)鍵課題。在 Pangu Ultra 稠密模型 [2] 的訓(xùn)練中，Depth-Scaled Sandwich-Norm和TinyInit方法在保障訓(xùn)練穩(wěn)定性上起到了關(guān)鍵性的作用，所以 Pangu Ultra MoE 依舊采用這個(gè)方案來(lái)控制訓(xùn)練穩(wěn)定性。經(jīng)過(guò)實(shí)驗(yàn)證明，此設(shè)計(jì)在 Pangu Ultra MoE 的訓(xùn)練中同樣能起到增強(qiáng)穩(wěn)定性、加快收斂速度的作用。

Depth-Scaled Sandwich-Norm（DSSN）：傳統(tǒng)的 Pre-LN 結(jié)構(gòu)存在因?yàn)樽訉虞敵鲆?guī)模波動(dòng)而導(dǎo)致訓(xùn)練不穩(wěn)定的現(xiàn)象，DSSN 是為了解決這一問(wèn)題而提出的。通過(guò)在每個(gè)子層輸出后加入額外的層歸一化，并引入深度縮放的初始化方式，從而穩(wěn)定網(wǎng)絡(luò)各層的輸出尺度，達(dá)到抑制梯度異常、降低范數(shù)波動(dòng)的目的。Depth-Scaled Sandwich-Norm + TinyInit 的方案減少了 51% 的突刺量（見(jiàn)圖 1），緩解了梯度范數(shù)頻繁突刺的問(wèn)題，能夠有效降低大模型訓(xùn)練過(guò)程中的不穩(wěn)定性，加快模型收斂，提升模型性能。同時(shí) DSSN+TinyInit 被應(yīng)用到 Pangu Ultra MoE 中實(shí)現(xiàn)了 10+T tokens 數(shù)據(jù)的長(zhǎng)穩(wěn)訓(xùn)練。

圖 1: 訓(xùn)練過(guò)程的梯度范數(shù)對(duì)比圖（黑色實(shí)線(xiàn)為突刺分界線(xiàn)）。DSSN+TinyInit 使梯度突刺率從 1.54% 下降到 0.76%，相對(duì)下降 51%。

基于 EP group 的負(fù)載均衡：

讓計(jì)算效率和路由表達(dá)能力可以兼得

在訓(xùn)練混合專(zhuān)家模型（MoE）時(shí)，容易出現(xiàn)專(zhuān)家負(fù)載不均衡的情況。負(fù)載不均衡指的是不同專(zhuān)家被分配的 token 數(shù)量存在顯著的差距。當(dāng)采用專(zhuān)家并行策略（EP，expert parallelism）時(shí)，負(fù)載不均衡會(huì)影響計(jì)算效率，被分配過(guò)多 token 的專(zhuān)家會(huì)成為計(jì)算瓶頸，而其他專(zhuān)家則處于低利用率狀態(tài)。同時(shí)負(fù)載過(guò)低的專(zhuān)家可能存在訓(xùn)練不充分的問(wèn)題，影響最終的模型效果。因此如何使 token 更均衡地分布至不同專(zhuān)家，對(duì)提高混合專(zhuān)家模型的訓(xùn)練效率和效果非常重要。

為了保證負(fù)載均衡，一般通過(guò)增加輔助的負(fù)載均衡 loss（auxiliary loss）來(lái)約束 tokens 在專(zhuān)家之間均衡分布。然而，如果負(fù)載均衡 loss 過(guò)度地約束 tokens 分配的均衡性，也會(huì)影響模型路由的表達(dá)能力。之前主流的負(fù)載均衡 loss 一般是約束單個(gè)序列或者單個(gè) micro batch 內(nèi)的 token 分配均衡性，而單個(gè)序列往往是來(lái)自同一領(lǐng)域的數(shù)據(jù)，過(guò)度的均衡可能影響專(zhuān)家特化（expert specialization）。

盤(pán)古團(tuán)隊(duì)發(fā)現(xiàn)對(duì)于采用專(zhuān)家并行策略訓(xùn)練的模型，可以設(shè)計(jì)一種對(duì)模型路由約束更小，同時(shí)不影響計(jì)算均衡性的 EP-Group 負(fù)載均衡 loss。當(dāng)采用了專(zhuān)家并行，專(zhuān)家會(huì)被分配到不同卡上進(jìn)行并行計(jì)算。每塊卡上的專(zhuān)家會(huì)接收來(lái)自 EP 組內(nèi)所有卡上的 micro batch 路由給自己的 token。所以可以設(shè)計(jì)一個(gè)負(fù)載均衡 loss，來(lái)約束 EP 組內(nèi)所有 micro batch 路由到組內(nèi)專(zhuān)家之后的均衡性。這相當(dāng)于把 EP 組內(nèi)部的所有 micro batch 聯(lián)合起來(lái)計(jì)算負(fù)載均衡的 loss, 這樣訓(xùn)練時(shí)可以容忍單個(gè) micro batch 的不均衡，只要多個(gè) micro batch 的 token 路由到專(zhuān)家之后是均衡的即可。

為了驗(yàn)證 EP-Group 均衡損失函數(shù)的效果，盤(pán)古團(tuán)隊(duì)使用一個(gè) 20B 參數(shù)量的 MoE 模型進(jìn)行了 100B 數(shù)據(jù)量的對(duì)比實(shí)驗(yàn)。結(jié)果如表 2 所示，可以看到 EP-Group 均衡損失函數(shù)在大部分任務(wù)相比主流的 Micro-batch 上都有顯著的優(yōu)勢(shì)，平均提升了 1.5 個(gè)點(diǎn)。

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

同時(shí)盤(pán)古團(tuán)隊(duì)對(duì) Pangu Ultra MoE 的專(zhuān)家特化進(jìn)行了分析，結(jié)果如圖 2 所示，可以看到不同領(lǐng)域的數(shù)據(jù)對(duì)專(zhuān)家的選擇存在顯著的差異，這表明 EP-Group 均衡損失函數(shù)給模型提供了靈活的路由選擇空間，促進(jìn)了專(zhuān)家特化。

圖 2: Pangu Ultra MoE 的專(zhuān)家特化。其中 ar，de，fr，ru 分別代表阿拉伯語(yǔ)，德語(yǔ)，法語(yǔ)，以及俄語(yǔ)。

多 Token 投機(jī)推理新路徑：

MTP 頭延遲擴(kuò)展策略

投機(jī)接受長(zhǎng)度預(yù)期提升 38%

投機(jī)推理是一種提升大模型生成效率的有效方法，其核心思想是在主模型生成 token 之前，由一個(gè)輕量輔助模塊預(yù)先預(yù)測(cè)多個(gè)候選 token，并通過(guò)快速校驗(yàn)機(jī)制決定是否接納，從而實(shí)現(xiàn)推理過(guò)程的并行化與加速。在當(dāng)前大模型推理中，Multi-token Prediction（MTP）技術(shù)已成為實(shí)現(xiàn)多 token 級(jí)別投機(jī)生成的重要手段。

盤(pán)古團(tuán)隊(duì)在實(shí)踐中發(fā)現(xiàn)，獲取多 token 的投機(jī)推理能力并不需要從訓(xùn)練開(kāi)始便配置多個(gè) MTP 頭，而是可以在訓(xùn)練后期對(duì)單頭 MTP 進(jìn)行擴(kuò)展來(lái)達(dá)到類(lèi)似的效果。為驗(yàn)證這一策略的有效性，團(tuán)隊(duì)使用 20B MoE 為主干模型，訓(xùn)練 185B 數(shù)據(jù)。具體對(duì)比設(shè)置為：以?xún)蓚€(gè) token 的投機(jī)推理為目標(biāo)，分別訓(xùn)練了從頭開(kāi)始配置單 / 兩個(gè) MTP 頭的模型（即單頭從頭訓(xùn)練和雙頭從頭訓(xùn)練），以及在單頭 MTP 模型訓(xùn)練至收斂后，通過(guò)復(fù)制已有頭的參數(shù)再增訓(xùn)出第二個(gè) MTP 頭的模型。對(duì)于擴(kuò)增的模型，對(duì)比全參續(xù)訓(xùn)以及凍結(jié)主干和一頭的續(xù)訓(xùn)的效果，即雙頭擴(kuò)增全參訓(xùn)練和雙頭擴(kuò)增凍結(jié)訓(xùn)練。下游使用 LAMBADA 續(xù)寫(xiě)作為評(píng)測(cè)任務(wù)。

結(jié)果如圖 3 所示。雙頭擴(kuò)增模型的接受長(zhǎng)度和延遲基本和雙頭從頭訓(xùn)練一致，而雙頭的接受長(zhǎng)度約 2.30，單頭的接受長(zhǎng)度約 1.67，雙頭相對(duì)單頭提升約 38%。在模型效果方面，雙頭擴(kuò)增模型全參訓(xùn)練和從零訓(xùn)練相當(dāng)，而由于凍住了主干和一頭，雙頭擴(kuò)增凍結(jié)訓(xùn)練的精度在擴(kuò)增的位置基本保持不變。這表明后期的 MTP 擴(kuò)展可以達(dá)到多頭的從頭訓(xùn)練的投機(jī)推理效果，可以在模型訓(xùn)練早期保持較小的 MTP 配置并在后期再進(jìn)行擴(kuò)展，兼顧計(jì)算成本和推理能力。

圖 3: 20B MoE 的 MTP 在 LAMBADA 續(xù)寫(xiě)上的投機(jī)推理結(jié)果。在接受長(zhǎng)度上，雙頭相對(duì)單頭提升約 38%，而雙頭可以基本無(wú)損地通過(guò)后期擴(kuò)增單頭得到。

迭代難例挖掘與多能力協(xié)同：

后訓(xùn)練強(qiáng)化學(xué)習(xí)持續(xù)提升的關(guān)鍵

模型后訓(xùn)練的過(guò)程中，團(tuán)隊(duì)參考了業(yè)界常規(guī)的 GRPO 算法提升模型的推理性能。然而，在超大參數(shù)規(guī)模情況下，直接應(yīng)用 GRPO 會(huì)帶來(lái)兩方面的問(wèn)題：1. 算法訓(xùn)練需要依賴(lài)多回復(fù)通過(guò)率在 (0,1) 內(nèi)的數(shù)據(jù)，隨著模型性能的提升，相同 prompt 的推理結(jié)果準(zhǔn)確率越來(lái)越高，導(dǎo)致訓(xùn)練過(guò)程中被 “浪費(fèi)” 的數(shù)據(jù)不斷增加，降低推理效率；2. 模型訓(xùn)練需要兼顧多能力協(xié)同提升，包括數(shù)學(xué)、代碼和通用能力等，不同能力項(xiàng)的獎(jiǎng)勵(lì)函數(shù)設(shè)計(jì)會(huì)導(dǎo)致模型能力增長(zhǎng)上的不匹配，出現(xiàn) “蹺蹺板” 問(wèn)題。

圖 4: Pangu Ultra MoE 的強(qiáng)化學(xué)習(xí)訓(xùn)練系統(tǒng)

為了解決上述兩個(gè)實(shí)踐難題，盤(pán)古團(tuán)隊(duì)設(shè)計(jì)了 Pangu Ultra MoE 的強(qiáng)化學(xué)習(xí)訓(xùn)練系統(tǒng)，如圖 4 所示，提升了大 MoE 模型的訓(xùn)練穩(wěn)定性與推理性能。系統(tǒng)設(shè)計(jì)的關(guān)鍵在于兩個(gè)部分：（1）迭代難例挖掘：模型階段性更新后，從初始的數(shù)據(jù)池中進(jìn)行多回復(fù)推理，選取回復(fù)通過(guò)率在 (0,1) 的數(shù)據(jù)組成 RL 訓(xùn)練數(shù)據(jù)池，以保持推理效率最大化；（2）多能力項(xiàng)獎(jiǎng)勵(lì)系統(tǒng)：為了確保模型多能力項(xiàng)協(xié)同提升，數(shù)學(xué)和代碼均采用了基于規(guī)則的獎(jiǎng)勵(lì)，通用獎(jiǎng)勵(lì)模型則使用 LLM-as-a-judge 的方法對(duì)生成的回復(fù)質(zhì)量進(jìn)行評(píng)分，并對(duì)最終的 reward 進(jìn)行歸一化處理，保證了模型在多個(gè)能力項(xiàng)的綜合表現(xiàn)。

[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