DeepSeek開源優化并行策略,提升訓練和通信效率
今早10點,DeepSeek開啟了第四天技術分享,開源了三個優化并行策略。
分別是DualPipe,一種用于 V3/R1 訓練中計算與通信重疊的雙向流水線并行算法;EPLB,針對 V3/R1 的專家并行負載平衡器;用于分析 V3/R1中的計算-通信重疊。
開源地址:https://github.com/deepseek-ai/DualPipe
??https://github.com/deepseek-ai/eplb??
??https://github.com/deepseek-ai/profile-data??
DualPipe是一種創新的雙向流水線并行算法,曾首次在V3版本中使用過。與傳統的1F1B和ZB1P方法相比,DualPipe大幅減少了流水線氣泡,同時僅增加了1倍的激活內存峰值。
DualPipe核心思想是在一個正向和反向傳播的塊對中重疊計算和通信。每個塊被細分為四個部分:注意力機制、全網通信的分發(dispatch)、多層感知機(MLP)以及全網通信的合并(combine)。對于反向傳播塊,注意力機制和MLP進一步被拆分為輸入的反向傳播和權重的反向傳播兩部分。
還有一個流水線通信組件,通過重新排列這些組件,并手動調整GPU流處理器(SM)在通信與計算之間的分配比例,DualPipe能夠在執行過程中完全隱藏全網通信和流水線通信。
這種重疊策略不僅適用于計算密集型的場景,即使在通信負擔較輕的情況下,DualPipe依然展現出顯著的效率優勢。
此外,DualPipe采用了雙向流水線調度策略,從流水線的兩端同時喂入微批次(micro-batch),并通過計算與通信的重疊,確保了在模型進一步擴展時,只要保持恒定的計算與通信比例,就可以在跨節點的情況下使用細粒度的專家,同時實現近乎零的全網通信開銷。
為了進一步優化DualPipe的性能,研究者們還開發了高效的跨節點全網通信內核,以充分利用InfiniBand(IB)和NVLink的帶寬。通過限制每個token最多只能發送到4個節點,減少了IB的流量,并確保token在到達目標節點后能夠通過NVLink快速轉發到特定的GPU上,從而實現了IB和NVLink通信的完全重疊。
在內存優化方面,DualPipe通過重新計算RMSNorm和MLA上投影操作來減少內存占用,避免了持續存儲這些操作的輸出激活值。同時,將模型參數的指數移動平均(EMA)存儲在CPU內存中,并在每次訓練步驟后異步更新,從而在不增加額外內存或時間開銷的情況下保持EMA參數。
在傳統的MoE模型中,專家負載的不平衡會導致計算資源的浪費,尤其是在專家并行(EP)的場景下,這種不平衡會進一步加劇通信開銷。而EPLB通過動態調整每個專家的負載,確保在訓練過程中專家之間的負載保持平衡。
EPLB為每個專家引入了一個偏置項(bias term),這個偏置項被添加到專家的親和力分數(affinity score)中,用于決定每個token的路由(routing)。在訓練過程中,EPLB模塊會監控整個批次的專家負載,并根據負載情況動態調整偏置項。如果某個專家的負載過高,偏置項會減少,反之則增加。這種動態調整機制使得模型能夠在訓練過程中保持專家負載的平衡,同時避免了因強制負載均衡而導致的性能下降。
此外,EPLB模塊還引入了一個互補的序列級負載均衡損失,以防止在單個序列中出現極端的負載不平衡。這個損失函數的權重被設置為一個非常小的值,以確保它不會對模型的整體性能產生顯著影響。
通過這種無輔助損失的負載均衡策略,可幫助模型在訓練過程中能夠有效地利用計算資源,同時保持模型性能的穩定。
本文轉自 AIGC開放社區 ,作者:AIGC開放社區
原文鏈接:??https://mp.weixin.qq.com/s/wgGljTRULzhIGHu0735oFQ??
