近年來，以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%。

超大集群調優案例

一

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%。

技術背景和挑戰

一

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

圖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 自動調優流程

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

一

建模與算法

從整體角度來看，主流的大語言模型都是由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的內存占用（下圖右下角）。

圖12 負載均衡工具尋優流程

獨家干貨&第一手資訊

都在這里

昇騰AI開發者

昇騰官網

【昇騰社區】

https://www.hiascend.com/

【昇騰論壇】

https://bbs.huaweicloud.com/forum/forum-726-1.html

【CANN訓練營】

https://www.hiascend.com/developer/canncamp