新智元報道
編輯:KingHZ 桃子
【新智元導讀】最近,華為在MoE訓練系統方面,給出了MoE訓練算子和內存優化新方案:三大核心算子全面提速,系統吞吐再提20%,Selective R/S實現內存節省70%。
在通往更強大的AI路上,MoE已成為科技巨頭另一個首選路徑。
只要Scaling Law沒有失效,大模型的參數規模依舊不斷擴大,由此AI智能水平才能不斷攀升。
憑借獨特的架構設計,MoE正以前所未有的參數規模,成為突破大規模模型訓練的算力瓶頸的關鍵路徑之一。
然而,如何將MoE潛力真正轉化為高效的訓練實踐,一直是業界探索的難題。
此前,華為曾通過Adaptive Pipe&EDPB框架,實現了集群級高效分布式計算,讓通信和計算能完美并行,提高訓練集群效率。
本次,華為通過昇騰與鯤鵬算力的深度協同,進一步實現了訓練算子計算效率和內存利用率大幅提升。
他們從單節點視角出發,深入到NPU和CPU內部,對算子計算、下發、訓練內存使用等進行細粒度拆解。
令人驚喜的是,結果顯示,MoE訓練在之前的基礎上,吞吐又提升了20%,內存占用降低了70%。
這不僅是一次技術突破,更是引領MoE訓練的風向標。
MoE訓練單節點效率提升挑戰
現實中,MoE模型的訓練并非易事。
因其框架的復雜性,除了集群分布式訓練效率優化外,在單節點訓練效率提升方面有兩大核心挑戰。
· 算子計算效率低,存在等待
首先,硬件核心計算單元,如Cube利用率不足,存在冗余操作和可優化的數據流水搬運,進而拖累了整體計算吞吐。
同時,專家路由機制導致算子下發頻繁,且容易中斷。
因為,復雜的專家路由機制增加了算子下發的調度壓力,如同高速閘道入口既窄還有頻繁紅燈,形成了Host-Bound瓶頸。
· 「昂貴的」NPU內存永遠不夠用
為了擴展模型參數量以提高模型智能水平,MoE模型參數量動輒千億甚至萬億。
加之訓練過程中前向傳播累積的海量激活值,讓內存資源顯得捉襟見肘。
如果太過擠壓內存,還容易引發NPU內存溢出(OOM),造成訓練中斷。
因此,NPU內存優化是大規模MoE訓練永恒的主題。
針對這些難題,華為給出了業界最完整的解決方案。
昇騰算子計算加速
訓練吞吐飆升15%
毋庸置疑,只有更高算力的利用率,才能將AI系統的效用最大化。
除了并行策略、通算掩蓋等框架層優化方法,硬件親和算子優化,也是進一步獲得潛在性能優化的新路徑。
在MoE模型中,最「吃時間」的幾個核心算子有:融合算子FlashAttention、基礎算子MatMul,以及負責數據重排與反重排的Vector(矢量)算子。
這些算子,占據了總計算耗時75%以上。
從數學等價和昇騰硬件親和角度出發,華為提出了「瘦身術」、「均衡術」、「搬運術」三大優化策略。
得益于這些方案,MoE模型計算冗余消除,數據流效率提升,同時計算單元間數據移動減少,充分發揮出昇騰的硬件能力。
在Pangu Ultra MoE 718B模型訓練實踐中,三大算子加速實現整體訓練吞吐量提升15%。
FlashAttention「瘦身術」(RECT-FA)
針對FA算子,華為團隊優化了計算順序,進而消除了冗余計算,進一步讓FA內部處理流水排布更緊密。
它能支持原生非對齊計算,直接處理Key/Query維度不匹配場景,省去填充與切片開銷。
利用昇騰片上緩存原位累加技術,可基于數學等價計算消除旋轉位置編碼中復雜的拼接操作(如圖1所示)。
圖1:基于數學等價計算消除拼接開銷原理
通過核間高效同步與緩存智能預搬運技術,實現FA內部計算步驟的高效流水線銜接,消除等待阻塞。
通過這三點優化,實現FA前/反向性能分別提升50%/30%;
MatMul(矩陣乘法)「均衡術」(AscendP-MM)
針對矩陣乘法算子,華為優化了雙級數據流水搬運問題,如下圖2所示。
在數據從通用內存傳輸至L1緩存時,基于更小的L0緩存容量調整搬運量,從而更早啟動從L1至L0的搬運,以及后續計算。
通過矩陣子塊斜向分配計算核,降低并發沖突,提升數據在芯片內部的流水效率。
最終,實現了Cube(核心計算單元)的利用率提升10%。
圖2:MatMul算子數據流水優化前后示意圖
Vector算子「搬運術」(VectorSort)
針對Vector算子,華為充分利用了昇騰芯片Vector指令特性,融合了多個細粒度小算子,以降低內存搬運耗時。
通過等價變換(如圖3所示),減少了重排與反重排操作中數據的反復搬運。
上述優化操作的效果立竿見影,VectorSort類算子性能直接飆升3倍以上。
圖3:VectorSort重排操作優化前后流水示意圖
昇騰鯤鵬協同再加速
吞吐提升至20%,內存節省70%
通過昇騰和鯤鵬的高效協同,華為的研究團隊實現了算子下發幾乎「零」等待,內存節省70%,其中關鍵在于Host-Device協同的算子下發優化和Selective R/S精準內存手術兩項創新。
鯤鵬Host-昇騰Device協同優化算子下發
· 等效計算的同步消除
算子下發中斷通常是因為Host需要等待Device返回結果(即同步),就如同高速有車需要逆行回到閘道入口,會阻塞所有后續算子的下發。
在Host或Device上就地等效計算,避免了數據逆行,實現算子下發無同步。
圖4展示Token分發預處理過程的同步消除結果,同步消除后,算子下發不再中斷,提高Device上算力利用率。
圖4:Token分發結果預處理同步消除前后對比
· 重排下發序規避空閑等待
通過對模塊的細粒度分離并重排序(圖5所示),如同在綠燈后閘道入口優先放行起步快的小車。
這讓Device上盡快接收到計算任務,避免過長時間的空閑等待,實現單次Host-Bound從2.1ms降低至0.6ms,超過70%的降幅。
圖5:重排下發序規避空閑等待
· 昇騰+鯤鵬親和提升下發速度
為了發揮算子下發的極致性能,可以通過任務綁核來控制CPU端算子任務的處理器親和性,將主要任務綁定在最親和的核上,并隔離其余任務對算子下發線程的影響。
這如同加寬閘道,有助于提升算子的下發速度。
采用自定義粗粒度綁核方式(每NPU綁24核),完全消除系統型持續Host-Bound。
通過上述Host-Device協同優化,有效緩解了同步型與持續型Host-Bound瓶頸。
在Pangu Ultra MoE 718B模型訓練實踐中,團隊實現了算子下發幾乎零等待(free時間占比<2%),訓練吞吐量進一步提升4%,充分發揮算子加速優勢,兩者疊加可加速訓練19.6%。
Selective R/S精準內存手術方案
基于昇騰+鯤鵬內存協同架構,研究團隊提出了Selective R/S內存優化技術。
這項創新實現了對Pangu Ultra MoE 718B模型訓練多維度、定制化的「內存解剖」,在訓練實踐中可節省超過70%的激活值內存。
這項技術主要分為兩大部分:
(1)豐富多樣的、通用化、張量級的細粒度重計算以及Swap策略等組成的「顯微手術器械庫」;
(2)可自適應內存系統優化的「手術臺管理機制」。
· 細粒度重計算(R)與Swap(S)
華為實現了多個模塊重計算的細粒度支持(如表1左所示)。
而且系統可以自適應調整重計算算子的執行順序,巧妙地「隱藏」了重計算產生的額外耗時。
對于重計算過程也做了優化,包括:
·MLA重計算創新性地調整了計算順序,將KV的計算與Q解耦;
·RmsNorm重計算同時兼容了Sandwich Norm和Pre Norm;
·Permute重計算支持Token重排操作中計算與通信重計算;
·Activation重計算支持激活函數以及Prob前移乘法重計算。
在Swap方面,昇騰+鯤鵬內存協同化管理,實現了Attention模塊內激活值的張量級卸載和預取,提供了靈活的Swap策略選擇,并為MoE層中重計算代價高昂的模塊設計了Swap方案(如表1右所示)。
表1:已實現的細粒度重計算與Swap方案
如圖6所示,通過對激活值卸載與預取位置的精準調整,有效規避HBM讀寫競爭帶來的性能劣化。
圖6:Swap_attention卸載&預取的時機
· 自適應內存優化管理機制
內存優化管理機制的主要指標是Memory-Runtime(節省內存/額外耗時)性價比。
為了找到最優的自適應內存優化策略組合,需要先計算模型訓練的內存占用。
以「先增后減」的貪心算法作為基礎選擇機制,并結合對Swap帶寬競爭的分析,基于已實現的內存優化策略,最終給出最優的策略組合。
國產AI訓練優化新方案
華為昇騰+鯤鵬深度協同,結合算子計算加速和內存優化技術,為MoE訓練提供了高效、低成本的解決方案。
無論是三大算子加速、鯤鵬昇騰協同的算子下發「零等待」,還是激活內存節省70%,都展現了華為在AI算力領域的深厚技術積累。
這不僅為大規模MoE模型訓練掃清了障礙,也為行業提供了寶貴的參考路徑。
想要了解更多技術細節,請查閱完整技術報告。
報告地址:https://gitcode.com/ascend-tribe/ascend-training-system/tree/main/StandaloneOptimization
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
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.