加速DeepSeek訓練效能調優，揭秘昇思MindSpore大模型自動流水負載均衡技術

近年來，以DeepSeek、Llama、Qwen、Mixtral為代表的新一代大模型不斷突破引數規模瓶頸，推動模型體量向萬億級躍進，流水線並行(Pipeline Parallelism)已演變為大模型分散式訓練的必備技術之一。

流水線並行透過將模型的不同層放置於不同的計算裝置，降低單個計算裝置的記憶體（視訊記憶體，後文統一稱為記憶體）消耗，從而實現超大規模模型訓練。相較於張量並行中大量的All-Gather等通訊，流水線帶來的通訊代價是每個stage之間的P2P通訊，通訊總量相對來說較小。

然而，流水線並行的特點決定了流水線的不同stage之間必然存在著依賴計算關係，這種依賴關係會產生額外等待時間（又叫做“空泡”或“bubble”），制約著模型的訓練效能。為了儘量壓縮bubble，提升訓練效能，業界提出瞭如GPipe、PipeDream和Megatron-LM等當前應用廣泛的幾種流水線編排技術。

這些技術在拉高了訓練效率的上限的同時也引入了不同裝置上記憶體開銷的不均衡，增大了策略調優的難度和成本，流水線並行策略的配置成為了影響大模型端到端效能的核心因素之一。

洞察到這個挑戰，昇思MindSpore團隊提出了一種同時均衡計算和記憶體的全域性求解方法，面向複雜的流水線編排邏輯，自動生成最優的stage-layer分配和重計算策略。研發人員將演算法打包成為自動流水負載均衡工具，結合Dryrun叢集記憶體模擬工具，實現低成本高效模型調優：

自動流水負載均衡工具是基於符號抽象的全域性線性規劃演算法，實現同時均衡計算和記憶體的負載。透過對模型的效能和記憶體進行分解，構建出線性規劃問題，在分鐘級的時間內求解出一個最優的layer分配和重計算策略，大幅度提升模型調優的效率。
Dryrun是模擬叢集實際執行訓練任務的工具，可以在單卡上完美模擬AI框架的構圖、編譯、記憶體分配的過程，其輸出的記憶體預估和真實執行之間的誤差極小，從而實現單卡模擬大叢集並行策略的調優，降低調優資源消耗、提升調優效率。

MindSpore TransFormers Pipeline_balance地址：

https://gitee.com/mindspore/mindformers/tree/dev/toolkit/pipeline_balance

透過結合昇思MindSpore提供的自動流水負載均衡工具與Dryrun叢集記憶體模擬工具，在大叢集上實測DeepSeek-V3 671B調優，實現訓練效能倍增；實測千億引數Llama模型自動負載均衡調優相比於專家調優提升效能14%。

# 01

超大叢集調優案例

一

DeepSeek-V3 671B調優案例

DeepSeek的架構中不同的MoE層專家數存在著區別，並且還引入了MTP層，這讓Deepseek模型的不同層之間記憶體和計算建模有著顯著差異。手工調優需要工程師多次實驗總結不同層處在不同stage時的記憶體變化規律，相比於Llama系列模型複雜度和工作量都是成倍提升。

面對DeepSeek一類的異構模型，流水負載均衡工具可以有效相容。表3展示了DeepSeek-V3 671B 大叢集調優的例項。

當模型引數量來到了671B這樣的規模，即使配置pp=16，將模型的layer平均切分到每個stage時，裝置記憶體壓力仍然較大，因此只能選擇開啟完全重計算。表中圖示展示了流水執行的模擬示意，藍色方框為正向計算，黃色為反向計算，縱軸表示16個stage，橫軸表示時間，下面的紅色real bubble代表bubble相比於真實計算時間的比值，real bubble越大，預估效能越差。

在平均切分全重計算的情況下，可以觀察到stage0上每層的計算時間相比於stage1-15較小，這是因為Dense層的計算時間小於MoE層；另外stage16的每個小方框較大，是因為MTP層的計算時間較大。從圖中觀察1F1B的流水執行中，可以看到瓶頸是stage16的正反向時間，其餘的stage不得不等待stage16的計算結束，因此可以看出在stable階段產生了不小的bubble，而在cooldown階段，因為前面stage的反向已經結束，產生了非常明顯的巨大的bubble，最終real bubble值為1.45，很明顯該平均切分策略效能較差。

對於DeepSeek-V3 671B使用自動流水負載均衡工具，妥善配置了每個stage的layer數和重計算策略，流水線各stage上的計算、記憶體負載基本均衡。從第二張圖可以看出最後一個stage的計算量變小，是透過挪動部分layer到其他stage來調整的。透過解決瓶頸stage的計算問題，stable階段的bubble被基本消除，cooldown階段的bubble也大幅度降低，最終real bubble降為0.91，大幅提高了訓練效率。

在大叢集實測訓練效能相比開箱效能提升一倍以上。

圖1 DeepSeek-V3 671B 大叢集調優

二

千億引數Llama模型

大模型訓練中為了保障訓練的效果（loss的收斂和最終的評測效果），通常會控制global batch size不能太大。在限定global batch size的情況下，流水並行的micro size數也受到了限制，尤其是在大叢集訓練時，資料並行dp的維度數非常大。根據bubble計算公式（p-1）/m，當m過小時，bubble相對較大，裝置空等時間較多，影響訓練效能。針對interleave場景的調優難題，自動負載工具做到了對記憶體排布和計算時間的建模，使用者可以任意給定chunk數，工具將給出此條件下的最優解。同時工具自帶模擬器，方便使用者對比各種不同的並行策略。

下圖透過流水並行模擬器展示了超大叢集的千億引數Llama模型訓練的調優例項。

圖2 超大叢集千億引數規模Llama調優

圖2中共有三種流水線負載均衡配置，分別是對layer進行平均切分，且不開啟重計算；平均切分開啟全部重計算以及經過自動負載均衡調優後的策略。

最終透過超大叢集實測，自動負載均衡調優相比於專家調優提升效能14%。

# 02

技術背景和挑戰

一

流水線並行的典型編排方式

圖1展示了最簡單的流水並行切分，將模型的層進行了切分，分佈在不同的裝置上，雖然節省了記憶體佔用，但是可以很明顯的看到非常大的bubble，在實際訓練時效能會非常差，因此在實際訓練中基本上不會直接使用native pipeline。谷歌為了提高流水線的效率，提出了Gpipe，如圖2所示，裝置上的訓練資料被分成了更細的micro batch，不同的micro batch可以平行計算，這樣就壓縮了訓練過程中的bubble，提升了訓練效能。

圖3 Native Pipeline Parallelism

圖4 GPipe

在這種正規化下，每個device上的記憶體壓力是一樣的，都需要在反向之前儲存一個step中所有layer的啟用值（activation/Acti.）。因此，GPipe雖然透過流水線的方式把模型分割到了不同的裝置上，但其所採取的流水線編排方式在調大micro batch時，記憶體佔用過高，影響了其擴充套件性。因為通常來說，micro batch越大，可以流水並行的計算部分越多，bubble越小。

圖5 1F1B & interleaved pipeline

PipeDream在GPipe batch-forward-batch-backward的編排基礎上提出了一種新的編排方式，1-forward-1-backward(1F1B)(圖5上半部分），流水線在一個前向處理完成後便可以開始計算反向，這樣啟用值在反向執行後就可以及時釋放，相比於Gpipe提高了記憶體使用效率，增加micro batch的數量不再會對裝置上的峰值記憶體有影響。在這種1F1B的編排方式之下，流水線的最小bubble為（p-1）/m，p和m分別是流水線stage數和micro batch數。效能持平GPipe的同時進一步節省了記憶體壓力。以下表為例，當micro batch數量為8的時候，stage = 4的流水線並行最少也能為裝置節省50%的記憶體。

因為在峰值記憶體方面的優勢，業界的流水線並行普遍採用了1F1B的編排方式。經過深入分析可以發現，1F1B編排模式下，裝置的記憶體壓力相比於GPipe而言減少了很多，然而流水線中每個stage的記憶體壓力是不平衡的（見上表最右列）。每個裝置上的峰值記憶體由進入穩定狀態之前1F1B累積的啟用值決定，第i個stage累計的啟用值帶來的記憶體開銷可以透過公式計算，計算結果顯示在流水線中靠前的裝置（即i數值較小的）峰值記憶體更高。

Megatron團隊在2021年提出了流水線交織（interleaved pipeline）的排布方式（圖5下半部分），這是一種基於1F1B的最佳化，透過把一個裝置上原本連續的layer繼續分割成多個chunk，進一步細化了流水線的顆粒度，進一步壓縮bubble到（p-1）/vm，v是每個stage切分出的chunk數。當stage = 4時，如果micro batch數為16，理論上1F1B下bubble最小佔比為18%，而分出兩個chunk之後，bubble的佔比極限下可以進一步被壓縮到9%左右。不過，流水線交織並不能解決記憶體不平衡，只是引入了新變數chunk，導致調優又新增了一個維度，複雜度幾何級上升。

二

流水線並行調優難點

綜上所述，在當前主流的流水線並行技術之下，stage之間記憶體資源的使用是天然不平衡的，如果模型規模繼續增大，需要切分更多的流水線stage數，這種不平衡會更加明顯。

流水負載均衡的調優需要找到記憶體和計算的trade-off，既要讓每個裝置的記憶體資源都要得到充分利用，也要讓每個裝置的計算負載儘量相等。為了達到這一點，調優中至少要考慮以下幾個維度：

每個stage要分配多少layer；

每個stage中有多少的layer需要進行重計算，如果涉及多種重計算，比如完全重計算、選擇重計算，則每種都需要單獨考慮；

如果採取流水線交織的話，則還需考慮每個chunk中layer的分配；

當各個stage的記憶體得到了充分利用之後，能否真的在短板效應的影響下提高階到端的效能。

三

當前的流水線並行策略尋優方式

並行策略對模型整體訓練效率的PipeDream和Dapple這些工作中都提出過針對性的策略尋優演算法。這些演算法實際應用中還是依賴工程師的手工調優，依靠經驗和直覺多次迭代調整去發現最好的策略。常常需要花費數天至數週的時間去完成策略的微調，期間還伴隨著訓練任務的多次拉起。

圖6 專家手工調優流程

在大多數情況下，人工調優會採取貪心的策略，在1F1B模式且stage數目較少的情形下是可行的。但當stage增多的時候調優難度會逐步提高，尤其是，開啟流水線交織後，記憶體模式會有所不同，且這幾個調優維度之間是互相影響的，調優難度指數級增加。

對每個維度獨立考慮，最終組合起來並一定是最優解。這就使得貪心的調優策略在很多場景下會無能為力，必須用一種全域性的求解思路去尋找最優策略。

面對這個問題，我們提出了基於符號抽象的全域性線性規劃演算法，同時均衡計算和記憶體的負載。透過對模型的效能和記憶體進行分解，構建出線性規劃問題，在分鐘級的時間內求解出一個最優的layer分配和重計算策略，大幅度提升模型調優的效率。

圖7 自動調優流程

# 03

基於符號代價模型和線性規劃的自動負載均衡演算法

一

建模與演算法

從整體角度來看，主流的大語言模型都是由embedding+多層layer重複+lm_head構成的，我們只需要知道layer層面的效能資訊即可。在記憶體方面，只要有一定的stage記憶體開銷的已知資料，就可以分解出模型的記憶體構成，完成記憶體建模。

二

效能模型分解

傳統效能模型的求解能力和測量出來的數值精度繫結，需要多次profiling以識別規律，不能很好預估某一次執行的絕對速度。白盒拆解網路關係則可以更好地修正時間資訊。在負載均衡的建模中，我們把時間分成head、body和tail。以Llama型別的網路為例，模型的輸入經過了embedding之後是重複的decoder，最後以lm_head和loss結束。此時，embedding就是head，而decoder layer是body，tail則是lm_head和loss。這些時間都可以在profiling中找到。如果不考慮流水線並行，那麼一個micro batch的前向時間就是head + n × body + tail。同時這種白盒方案也有利於在反向時間中考慮重計算的影響。此外，這種抽象正規化也能泛化到Deepseek這種“異構模型”上。在Deepseek中，其decode layer既有dense MoE也有sparse MoE，還會接上一層MTP，不過對於我們的演算法來說，這些不同的層抽象之後就變成了不同的body，可以被尋優演算法天然相容。

圖8 效能模型分解：DeepSeek-V3

三

記憶體建模

記憶體的建模是演算法中需要謹慎設計的一環。一方面，透過多次拉起的方式去獲得記憶體，這種純黑盒的方式時間成本過高，違背了工具設計的初衷。另一方面，白盒演算法透過計算模型中各個引數和啟用值的shape來推演理論記憶體的方式準確度欠佳：難以全面考慮AI框架、特性最佳化、記憶體複用、碎片處理、硬體驅動記憶體分配器的參與。

Dryrun是MindSpore提供的一種能夠模擬實際執行的工具，除了真實的在硬體上執行計算和通訊，Dryrun可以完美模擬AI框架的構圖、編譯、記憶體分配的過程，其輸出的記憶體預估和真實執行之間的誤差非常小。在最終的方案中，我們透過提前進行一次Dryrun獲得每個stage的真實記憶體，再把真實記憶體透過線性方程組的方式進行分解。由於演算法最終的輸出是每個stage中的layer分配以及重計算策略，記憶體可以被分解到layer和重計算的層面上。假設有一個stage數為pp的流水線，每個stage的記憶體值為M，stage中的layer數為li，重計算層數為ri，i為stage id。記憶體將被分解為如下幾項：

圖9 記憶體建模：DeepSeek-V3

在上文對1F1B的分析中得知，第i個stage由於activation帶來的記憶體開銷，根據公式每個stage的峰值記憶體可表示為：

由此，只要透過Dryrun等手段得到足夠多的stage峰值記憶體，就可以透過方程組求解完成對記憶體的分解。在程式碼中可以透過numpy方便地實現這一邏輯。

四

線性規劃

按照上面兩節描述的方式分解出了效能和記憶體資訊之後，透過整數線性規劃建模流水線並行的編排模式，就可以衡量不同的配置對效能和記憶體的影響。以裝置的記憶體作為約束條件，以最小化端到端訓練時間為目標，可以自動求出最優的layer分配和重計算策略。這裡需要考慮的是如何對端到端時間建模。原封不動地刻畫layer間的依賴關係是一種自然的思路，但在複雜的調優條件下這種建模方式的時間複雜度很快就會超出可接受的範圍。實際上，如果考慮1F1B自身的性質，就可以簡化問題、分組描述，降低複雜度。

如下圖，在1F1B編排下，流水線一個step的端到端時間可以分為三個部分，分別為warmup、steady和cooldown。warmup階段處於反向開始之前，steady階段處於穩定的前反向交替狀態，到了cooldown階段則已經不再有需要計算前向的micro batch了，只需處理完剩餘的反向過程。

圖10 1F1B編排模式

三個部分的時間表示為：

這種建模思路也能夠直接複用到interleave的場景下，完成更復雜的負載均衡。基於pulp的程式碼邏輯表達如下：

當然實際的建模中需要處理很多的細節問題，比如說開啟interleave之後不同chunk之間記憶體不一致帶來的擾動等等，只有儘可能考慮到實際執行中的各項影響因素才能準確的做出效能和記憶體預估。

五

尋優方法邏輯

對於使用者來說，使用自動負載均衡工具需要提供如下的yaml檔案。pipeline_config用來描述模型的流水線配置。time_config用於描述時間資訊，來源於profiling資料（profiling要知道decoder layer和embedding、lm_head之間的耗時關係）。recofmpute_config和memory_uage是需要的dryrun資料，有了這些資訊之後演算法會自動構建線性規劃問題，利用求解器完成對最佳策略的尋找。

圖11 輸入輸出樣例

負載均衡工具的尋優流程如下圖所示，yaml檔案中包含的時間和記憶體資訊會被工具中的parser模組所分解，構成一套抽象表徵提供給演算法的interface。求解出來的策略以一種相容MindSpore TransFormers配置的格式給出並給出這個配置的bubble比例（上圖右），這個策略可以直接提供給MindSpore TransFormers使用，同時工具也包含了一個流水線模擬器，會根據算出來的策略模擬出效能和各個stage的記憶體佔用（下圖右下角）。