“對硬件特性的深度挖掘和算法創新,是華為引領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 訓練策略。
技術報告標題: Pangu Ultra MoE 模型架構與訓練方法
技術報告地址: https://gitcode.com/ascend-tribe/pangu-ultra-moe/blob/main/Pangu_Ultra_MoE_CN_Report.pdf
01
破解準萬億MoE模型性能瓶頸
打造芯片協同的先進架構
近期,盤古團隊在 MoE 模型訓練領域再進一步,重磅推出參數規模高達 718B 的準萬億全新模型 —— Pangu Ultra MoE 。該模型旨在實現超大規模 MoE 架構在模型效果與效率之間的最佳平衡。
為了達到這個目標, 研究團隊 在設計 Pangu Ultra MoE 架構的時候,充分考慮昇騰硬件特性, 在 昇騰 NPU 平臺上,融合計算、通信和內存等多維度指標,構建 了大規模系統模擬器,并系統性地 探索約 一萬個不同的 MoE 結構組合,最終搜索出一套在訓練與推理吞吐上均達最優的架構方案 。
Pangu Ultra MoE 是一個超大規模、高稀疏比的架構,同時也 包含 MLA 和 MTP 等 先進架構 和特有的 DSSN 穩定性架構 和 EP loss 負載優化 。下面是 Pangu Ultra MoE 的主要的架構和訓練 特性:
超大規模和超高稀疏比 : 采用 256 個路由專家,每個 token 激活 8 個專家 ,模型總參數量 7 18 B ,激活量 39B 。
MLA 注意力機制: 引入 MLA ( Multi-head Latent Attention ),有效壓縮 KV Cache 空間,緩解推理階段的內存帶寬瓶頸,優于傳統 GQA 方案。
MTP 多頭 擴展 :采用單頭 MTP 進行訓練,后續復用 MTP 參數擴展至多頭結構,實現多 Token 投機推理,加速整體推理過程。
Droples s 訓練: 采用 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+Ttokens數據長穩訓練
隨著參數規模和數據體量的激增,大模型訓練面臨前所未有的穩定性挑戰。頻繁的梯度范數突刺已成為阻礙收斂效率與模型性能提升的主要瓶頸。如何在確保訓練深度和寬度擴展的同時,維持梯度信號的穩定傳遞,成為構建高可靠性大模型架構的關鍵課題。在 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 進行擴展來達到類似的效果。為驗證這一策略的有效性,團隊使用 20B MoE 為主干模型,訓練 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 Mod els on Ascend NPUs
華為盤古首次露出,昇騰原生72B MoE架構,SuperCLUE千億內模型并列國內第一
算芯合一!華為披露昇騰體系大模型核心算子設計細節
華為發布OmniPlacement技術,實現超大規模MoE專家最優動態部署,提升昇騰推理系統吞吐10%
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.