當前DeepSeek等MoE(Mixture of Experts)混合專家模型已逐步成為LLM(大語言模型)領域的主流架構,但MoE模型訓練面臨幾個關鍵問題,如內存占用大、通信時延高、負載不均衡、算力利用率低等,其中All2All通信耗時占比高和顯存開銷大是顯著瓶頸。MindSpeed基于昇騰硬件特點實現(xiàn)了DualPipe通信掩蓋,創(chuàng)新提出重計算細粒度通信掩蓋和Zero-Memory內存優(yōu)化技術,緩解顯存瓶頸的同時保持計算效率,并支持兼容多種流水線調度算法,供用戶靈活選擇。
MoE模型的All2All通信耗時和顯存占用面臨的挑戰(zhàn)
在MoE模型每一層的前向和反向計算中,各有兩次All2All通信,共計四次All2All通信,稱為Dispatch(F), Combine(F), Combine(B), Dispatch(B)。當模型專家數(shù)量較多時需要進行EP(專家并行)切分,EP切分涉及跨節(jié)點服務器傳輸數(shù)據(jù),通信耗時受到節(jié)點間帶寬限制,大規(guī)模集群中AllToAll通信耗時占單次訓練迭代時間高達50%。
MoE模型雖比稠密模型模型激活參數(shù)更小,但顯存占用更高,首先所有專家參數(shù)都需要被保存,總參數(shù)量增大,單卡顯存占用是稠密模型的4~8倍;其次激活值膨脹,顯存需求增加30%以上,加劇顯存壓力。
昇騰MindSpeed細粒度DualPipe通信掩蓋
MindSpeed基于昇騰硬件特點實現(xiàn)DualPipe方案,達成細粒度的1F1B階段All2All通信全掩蓋、Warmup/Cooldown階段All2All自掩蓋的效果,創(chuàng)新性地提出重計算細粒度通信掩蓋和ZeroMemory內存優(yōu)化技術,掩蓋了重計算階段通信,緩解顯存瓶頸的同時保持計算效率。同時,MindSpeed的DualPipe通信掩蓋特性兼容VPP、Dualpipe、Cut-in-half[1]等三種流水線調度算法,用戶根據(jù)使用場景靈活選擇。
01基于昇騰硬件特點實現(xiàn)DualPipe通信掩蓋,達成1F1B階段All2All通信100%掩蓋的同時,Warmup和Cooldown階段額外掩蓋50% EP通信
在1F1B階段,通過將計算和通信任務分離為兩條獨立的處理流,實現(xiàn)細粒度的調度與編排,通過拆分模型計算圖中的數(shù)據(jù)流與梯度流,完成通信效率與計算資源的雙重優(yōu)化(dw分離技術)提升通信掩蓋比例,實現(xiàn)All2All通信的100%掩蓋。在1F1B基礎上,MindSpeed在Warmup/Cooldown階段,利用共享專家的計算掩蓋了部分的Dispatch(F)/Combine(B) 通信,采用無依賴關系的專家反向計算的dw進一步掩蓋了部分Dispatch(B) 通信,在1F1B掩蓋方案的基礎上額外掩蓋50%的EP通信。
掩蓋計算和流水掩蓋示意圖
02支持重計算細粒度通信掩蓋和Zero-Memory/Gelu重計算技術,實現(xiàn)內存優(yōu)化40%
針對MoE模型訓練中高顯存占用的問題,MindSpeed創(chuàng)造性地提出了重計算細粒度通信編排和ZeroMemory技術,通過細粒度聯(lián)合編排重計算以及前向和反向過程中所有的計算通信隱藏重計算的通信開銷,進而丟棄前向計算中每層的GroupedMatmul融合算子的輸入,從而實現(xiàn)內存優(yōu)化。該特性在DeepSeek V3模型上每層能節(jié)省接近200M的激活值開銷,疊加Gelu重計算等技術,整體內存優(yōu)化40%,且性能幾乎無損。
重計算細粒度通信掩蓋示意圖
03MindSpeed的DualPipe通信掩蓋特性兼容多種流水線調度算法,用戶可靈活選擇
DualPipe的掩蓋收益與流水線調度策略并不是強綁定關系,這種掩蓋的使能前提是在流水線策略中當前可做兩個無依賴關系的MicroBatch(微批次)計算。同時,針對社區(qū)提出的DualPipe的改進算法Cut-In-Half[1]調度(只取DualPipe鏡像對稱結構的一半,呈V字排布,相比DualPipe在同PP并行度下算法啟動規(guī)模和靜態(tài)顯存減半,消除冗余參數(shù)且bubble保持一致)與DualPipe類似。針對VPP(Virtual Pipeline Parallelism,虛擬流水線并行)的1F1B 階段以跨MicroBatch的前反向計算為主,稍改調度也能實現(xiàn)EP通信掩蓋。因此,MindSpeed的通信掩蓋技術可兼容Dualpipe、Cut-in-half、VPP等三種流水線調度算法,用戶可根據(jù)使用場景進行自由選擇。
結語
本期為大家介紹了MindSpeed細粒度DualPipe通信掩蓋特性,在訓練性能加速上,MindSpeed將持續(xù)探索演進,歡迎訪問MindSpeed開源社區(qū)
http://www.picoinsstore.com/uploadfile/pic2020/2025/0401/20250401132002136G>
參考鏈接:
[1] :http://www.picoinsstore.com/uploadfile/pic2020/2025/0401/20250401132002137G>
(免責聲明:本網(wǎng)站內容主要來自原創(chuàng)、合作伙伴供稿和第三方自媒體作者投稿,凡在本網(wǎng)站出現(xiàn)的信息,均僅供參考。本網(wǎng)站將盡力確保所提供信息的準確性及可靠性,但不保證有關資料的準確性及可靠性,讀者在使用前請進一步核實,并對任何自主決定的行為負責。本網(wǎng)站對有關資料所引致的錯誤、不確或遺漏,概不負任何法律責任。
任何單位或個人認為本網(wǎng)站中的網(wǎng)頁或鏈接內容可能涉嫌侵犯其知識產(chǎn)權或存在不實內容時,應及時向本網(wǎng)站提出書面權利通知或不實情況說明,并提供身份證明、權屬證明及詳細侵權或不實情況證明。本網(wǎng)站在收到上述法律文件后,將會依法盡快聯(lián)系相關文章源頭核實,溝通刪除相關內容或斷開相關鏈接。 )