新智元報道

編輯：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