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

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

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

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

破解準萬億 MoE 模型性能瓶頸

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

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

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

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

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

以下是 Pangu Ultra MoE 昇騰親和設計考慮：

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

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

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

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

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

梯度突刺率下降 51%

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

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

Depth-Scaled Sandwich-Norm（DSSN）：傳統(tǒng)的 Pre-LN 結(jié)構(gòu)存在因為子層輸出規(guī)模波動而導致訓練不穩(wěn)定的現(xiàn)象，DSSN 是為了解決這一問題而提出的。通過在每個子層輸出后加入額外的層歸一化，并引入深度縮放的初始化方式，從而穩(wěn)定網(wǎng)絡各層的輸出尺度，達到抑制梯度異常、降低范數(shù)波動的目的。

TinyInit：Transformer 模型普遍采用較小的初始化尺度，TinyInit 提出一種標準差為的初始化方案，能夠同時兼顧模型深度與寬度，其中d表示隱藏維度，L表示模型層數(shù)。同時，對詞嵌入層采用標準差為 0.5 的初始化。實驗表明，這樣的初始化策略有助于提升模型性能和訓練穩(wěn)定性。

Depth-Scaled Sandwich-Norm + TinyInit 的方案減少了 51% 的突刺量（見圖 1），緩解了梯度范數(shù)頻繁突刺的問題，能夠有效降低大模型訓練過程中的不穩(wěn)定性，加快模型收斂，提升模型性能。同時 DSSN+TinyInit 被應用到 Pangu Ultra MoE 中實現(xiàn)了 10+T tokens 數(shù)據(jù)的長穩(wěn)訓練。

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

基于 EP group 的負載均衡：

讓計算效率和路由表達能力可以兼得

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

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

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

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

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

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

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

多 Token 投機推理新路徑：

MTP 頭延遲擴展策略

投機接受長度預期提升 38%

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

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

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

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

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

后訓練強化學習持續(xù)提升的關(guān)鍵

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

圖 4:  Pangu Ultra MoE 的強化學習訓練系統(tǒng)

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