北大河圖發布分散式訓練神器Galvatron, 一鍵實現大模型高效自動並行

最近一段時間，「大模型」在AI 領域的各種應用場景都大放異彩，其中基於Transformer 的大規模預訓練模型是最典型的大模型之一，已經成為了當前基礎模型（Foundation Model）的核心架構。例如 NLP 領域的 BERT、GPT 系列，CV 領域的 ViT、Swin Transformer 系列，以及最近非常火熱的多專家混合模型 MoE 以及多模態模型 CLIP 等等，都是以 Transformer 作為核心的基礎架構。與之相應的，這類稠密大模型擁有著動輒數十億、百億甚至萬億規模的參數量，面臨高昂的運算、儲存、以及通訊開銷，也為 AI 基礎設施帶來了巨大的挑戰。

為了支援大模型的訓練，人們研發了許多工具（如NVIDIA 提出的「威震天」Megatron，Microsoft 提出的DeepSpeed，Meta 提出的FairSeq 等等）來實現各種各樣的並行方式，資料並行、張量模型並行、流水並行、分片資料並行等等。這些系統對上述並行方式提供了良好的封裝，對外屏蔽相應的實作細節，使得使用者能夠透過添加配置的方式來實現混合並行策略。

基於上述思想，已經有很多工作圍繞如何在張量或算子層面表達各種並行方式進行探究，這類工作的「自動化」 主要體現在並行API 到執行層的轉化過程。但如果只是侷限在設計平行 API 或中間表達式上，這種工程上的封裝並沒有從根本上解決分散式訓練的問題。最直觀的結果就是，使用者仍然無法從分散式部署的難題中解放出來。實際上，大模型的分散式部署是一個非常複雜的問題，目前的絕大多數的分散式訓練系統，都依賴使用者人工反覆嘗試以及系統專家經驗來進行部署，造成嚴重的資源利用效率低下的問題，與真正的「自動並行」 存在相當大的差距。

基於此，北大河圖團隊提出了一款分散式訓練神器 Galvatron，以實現大模型的高效自動並行，研究論文入選國際頂會 VLDB 2023。

• 論文網址：https://arxiv.org/abs/2211.13878
• #專案程式碼連結：https://github.com/PKU-DAIR/Hetu/tree/main/tools/Galvatron

#大模型自動並行到底難在哪裡

研究者認為，大模型的自動並行之難主要體現在以下三個方面：

（1）多樣性：首先，在平行方式方面，目前大模型的平行方式呈現出百花齊放的態勢，即使是對於同一個算子，不考慮混合併行方式，不同的基礎並行方式也會存在顯著的差異，導致不同的記憶體開銷、通訊代價以及運算效率。下圖展示了四種最主要的基礎並行方式，即資料並行（Data Parallelism）、張量模型並行（Tensor Parallelism）、流水模型並行（Pipeline Parallelism）、分片資料並行(Sharded Data Parallelism)，在兩張GPU 上對於簡單的矩陣乘法算子分佈式執行的過程。

#並行方法比較示意圖

其次，在模型方面，各種各樣的模型架構最近也是層出不窮，這往往也伴隨著不同的模型配置（例如不同輸入序列長度，模型層數，模型隱層寬度等），從而造成計算負載上的差異。另外，在硬體方面，使用者往往面臨非常差異化的叢集環境，可能會面臨不同的記憶體容量、通訊頻寬、運算能力等等。整體來看，由於上述多樣性的存在，沒有哪種平行技術總是能夠獲得最佳訓練效率，「自動並行」 也就成為了分散式訓練的核心挑戰。

（2）複雜性：上述分析也相對比較單一，實際上即使是對於同一個算子也可以同時應用多種不同的基礎並行方式，如果考慮到由這些基礎並行方式複合所構成的混合並行方式，則會導致問題變得非常複雜。更重要的是，大模型的計算圖往往結構非常龐大，對應的也需要更大規模的集群，如果對每個算子都進行探索（包括選取集群中合適的計算資源以及設計相應的混合併行方式），會帶來組合空間爆炸的問題，尋找整個模型的最優分散式執行方案變得難以求解。

（3）實用性：除此之外，實用性也是非常重要的問題。一方面，在進行自動平行搜尋的過程中，對於各種分散式執行方案，必須提供比較精確的記憶體、通訊、計算開銷，否則會導致結果與實際執行偏差過大，產生次優解或根本無法使用。為此，就需要非常精準的代價模型，對不同的模型結構和硬體條件進行建模。另一方面，系統提供自動並行能力所帶來的額外時間開銷必須在一個可以接受的範圍內，過於高昂的搜尋代價同樣也無法接受。

分散式訓練神器Galvatron，一鍵實現大模型高效自動並行

系統特性：

為了解決上述問題，研究者們提出了一些系列工作來探索混合併行的自動搜索：一類工作主要討論了同時考慮數據並行和模型並行的搜索空間，代表性工作包括FlexFlow，Tofu，另一類工作則產生於流水並行場景，將其與資料並行結合，代表性工作包括PipeDream，DAPPLE。在此基礎上還有一些衍生工作，如 Unity、Alpa，進一步擴展了自動並行的探索範圍。北大河圖團隊提出的系統「驚破天」Galvatron 同樣屬於自動並行搜尋的研究領域，但相較於現有工作，該系統主要擁有以下三方面優勢：

（1）在多樣性方面，現有工作能夠支持的平行維度仍然比較有限，而 Galvatron 不僅可以支持更多的平行維度，並且面對更加差異化的Transformer 模型結構也能夠精準建模，以及在不同的叢集硬體條件下驗證了其自適應調優的能力。

大模型分散式訓練系統比較示意圖

#（2）在複雜性方面，由於具備在多樣性上的優勢，Galvatron 面臨了前所未有的龐大搜尋空間，為此，研究者探究了幾條目前大規模分佈式訓練過程中的一些經過實驗性或理論性驗證的重要觀察，作為搜尋空間的剪枝準則，從而實現高效率的分散式執行計畫最佳化。

（3）在實用性方面，研究結合了理論建模和實驗測量兩者的優勢，實現對記憶體、通訊、計算開銷的精確估計，甚至考慮到了計算與通訊重疊所導致的GPU 執行效率下降問題，從而確保能夠得到足夠且準確的自動並行最佳化結果。

另外，Galvatron 底層選擇PyTorch 作為執行引擎，相容Huggingface 等常見的主流Transformer 模型實現，所以完全不會對PyTorch 用戶帶來額外的使用負擔；同時也不需要使用者付出額外的系統安裝或偵錯代價，使用時只需要增加幾行程式碼，就可以輕鬆完成自動並行的整個流程。

Galvatron 工作流程及使用者使用介面展示

#關鍵技術

1. 基於決策樹的搜尋空間分解

#

Galvatron 的設計目標是有效地在複雜而龐大的平行策略空間內自動搜索，並為給定的 Transformer 模型和分散式環境產生最佳的平行執行計劃。在搜尋空間上，Galvatron 是業界首個考慮4 種主流平行方法的自動平行訓練系統，包括包含資料並行（data parallelism, DP）、分片資料並行（sharded data parallelism, SDP）、張量並行（tensor parallelism, TP）和流水並行（pipeline parallelism, PP）。由於混合平行策略會包含以上四種平行演算法的任意組合，在多 GPU 的場景下這種組合帶來的搜尋空間十分龐大。例如，對於雙機四卡場景，一個可行的策略是在機內使用2-way TP，機間使用2-way PP，另一個可行的策略是在機內使用2-way PP，機間使用2-way DP。當節點內 GPU 數量擴展到 8 卡時，模型每一層的候選策略都多達上百種。隨著模型層數的增加，其搜尋空間規模指數成長，難以有效探索。

為了有效率地搜尋如此龐大的搜尋空間，研究首先提出了以下觀察值作為指導：

• Takeway#1 ：PP 傾向於被跨設備島放置。此處 「設備島」 指具有高內部頻寬的一組設備，在絕大多數 Transformer 模型中，PP 的通訊量相比於其它並行方式，顯著較少。因此，人們通常優先對模型進行 PP 切分並放置於設備島之間。
• Takeway#2：在同構裝置的前提下，平行策略傾向於將裝置均勻切分。例如，對於 4 卡 GPU 的 2-way DP 傾向於將設備切分兩組 2 卡的設備，而不是一組 1 卡和一組 3 卡的設備。在這種情況下，一個設備組內的最優混合並行策略與其他組內的最優策略保持一致。
• Takeway#3：一般來說，在能夠混合使用 DP 和 SDP 的情況下，只使用 SDP 在理論上表現更優。根據分析結果，N-way SDP 的通訊開銷和記憶體開銷都優於和##的組合，其中。

基於上述重要觀察，研究提出了一種基於決策樹的搜尋空間建構方法：

（1）給定一個Transformer 模型，基於Takeway#1 和Takeway#2，Galvatron 首先用PP 將模型切分成多個階段，同時將設備均勻且連續地切分為多個設備組。例如 8 卡場景下，模型被切割成 1/2/4/8-way PP，分別對應設備組大小為 8/4/2/1。

（2）每種PP 切分對應一棵決策樹及一個子搜尋空間，決策樹葉結點總數為裝置群組大小，決策樹高度為可用的平行方法數量，即決策樹每一層可應用一種平行策略。

（3）並行策略在決策樹不同層間不可重複使用。

（4）非葉結點的度數預設在 2 的指數次方 {2,4,8,…} 中選擇。

基於以上決策樹建構規則，Galvatron 建構的決策樹能表示以上並行的任意組合。 Takeway#1 和 Takeway#2 幫助 Galvatron 規避了低效率的平行組合，縮小了搜尋空間。對於 8 卡 GPU 訓練一層模型的場景，上述規則將產出 34 種候選的混合平行策略。進一步，使用 Takeway#3 將 DP 和 SDP 同時出現在一棵決策樹的情形剪枝後，8 卡候選策略數降至 22 種。

下圖給出了 8 卡 GPU 場景不同 PP 並行度（8/4/2/1）下的決策樹示意圖。

#

8 卡GPU 場景不同PP 並行度（8/4/2/1）下的決策樹示意圖

2 . 基於動態規劃的平行最佳化演算法

現有系統如Megatron 或DeepSpeed 通常由使用者指定全域並行方案及其對應的並行度，嚴重限制了對分佈式執行計劃的表達能力。 Galvatron 的最佳化目標是在使用者給定模型定義和分散式環境的情況下，無需使用者指定任何平行配置，自動產生最優的分散式執行計劃。具體來說，給定一個L 層的模型M 和記憶體容量為E 的N 個GPU 設備，Galvatron 的最佳化目標是搜尋得到最高的系統吞吐T_pt 並傳回對應的並行方案，這裡的平行方案指以層（或算子）為基礎單位的細粒度混合平行策略。

#演算法1: Galvatron 最佳化過程

最佳化過程：Galvatron 的最佳化過程如演算法1 所示。 Galvatron 最外層循環逐步增加搜尋的batch size 直到超出裝置記憶體；給定每個候選batch size B，Galvatron 首先根據Takeaway#1 對模型進行PP 切分並蒐索不同的並行度P（第4 行），選定P- way PP 後模型被切分為P 個階段（第6 行），與之對應的所有設備被切分為P 組，每組含有N/P 個設備；接著Galvatron 構建對應的決策樹，其能不重不漏地表示出DP、SDP、TP 的任意組合，從而獲得策略集合S；然後對於每個模型階段M_i，在設備內存限制E 下，Galvatron 使用動態規劃搜索得到其中每一層的最優混合並行策略並返回最小時間開銷（第9 行）；最後，Galvatron 在所有可能的PP 並行度和batch size 中選擇吞吐率最高的策略並返回（第15 行）。

動態規劃搜尋：以下介紹 Galvatron 並行最佳化工作流程中的動態規劃搜尋演算法。對於一個給定的包含L 層的模型階段，以開銷函數C(L,E) 表示L 層模型在裝置記憶體限制E 下的總執行時間，表示第L 層使用策略S_j 的執行時間，其中策略S_j是並行策略候選集合S 中的策略。設定初始值

，Galvatron 的動態規劃搜尋遵循以下狀態轉移方程式（公式1）：

其中，是第L 層使用策略S_j 的記憶體開銷，是第L 層使用策略S_j、其前一層使用策略S_i 所帶來的轉換開銷。在狀態轉移過程中，當記憶體開銷超過裝置記憶體限制裝置記憶體限制 E，開銷函數 C 傳回無限大。

複雜度分析：Galvatron 使用的動態規劃搜尋（公式 1）計算複雜度為 O(LE|S|)。可見，每一層的搜尋空間 S 的大小對於整體搜尋複雜度非常重要，Galvatron 提出的基於決策樹的搜尋空間分解能夠顯著減小搜尋空間，控制搜尋開銷在合理範圍內。

3. 基於混合建模的執行代價估計方法

Galvatron 使用一個策略開銷估計模組來對混合平行策略的計算、通訊、記憶體開銷進行估計。現有的開銷估計方法主要包括測量（profiling）和模擬（simulating）兩種，Galvatron 吸取了兩者的長處並設計了一種代價低廉、高效且準確的開銷估計方法。具體來說，對於記憶體開銷，Galvatron 使用張量的形狀和資料類型來完成估計；對於計算時間，Galvatron 透過在單一裝置上透過profiling 測量逐樣本計算時間，結合batch size 和擬合函數從而估計總體計算時間；對於通訊時間，Galvatron 透過通訊量除以設備通訊頻寬得到預估的通訊時間，其中通訊量由理論計算得到，通訊頻寬由profiling 測量得到。

基於上述估計結果，Galvatron 透過 simulating 執行過程來計算給定層使用給定策略的開銷 c(l,s)。和現有分散式訓練系統的代價模型不同，Galvatron 在建模中首次考慮了計算和通訊的重疊對 GPU 效能下降的影響。該研究透過實驗發現，由於重疊導致的 GPU 效能下降會顯著影響系統的執行效率，而這在以往的工作中都是被忽略的。因此，Galvatron 的開銷估計更加準確，並行最佳化效果更佳。

實驗結果

實驗設定：在實驗中，研究者將Galvatron 和四種使用單一平行策略的基準系統（DP、SDP、TP、PP）以及專家設定的DeepSpeed 3D Parallelism 進行對比，同時額外設定了兩個弱化版本的Galvatron 作為輔助基線，在有限的平行策略組合空間（即TP DP，PP DP），展開自動並行搜尋。研究選用了 NLP Transformer 模型 BERT 和 T5，CV Transformer 模型 ViT 和 Swin Transformer 作為實驗對象。

Galvatron 與基準系統在8 GPUs 20G 顯存下的吞吐對比情況

實驗比較效果：研究首先在八卡Nvidia RTX TITAN 24GB 環境下進行了實驗。實驗表明，在不同模型規模和不同內存限制下，Galvatron 相比於總是取得最優的吞吐率，與現有的最先進的單一並行方法和混合併行方法相比，訓練吞吐量均有顯著提升。具體而言，在ViT 模型上，Galvatron 相比於單一策略的吞吐率加速比最高可達338%，相比於其他混合並行策略的吞吐量從加速比最高可達55%；在其他三個模型上，Galvatron 相比於單一策略和現存的混合策略加速比最高可達200%-334% 和28%-52%。

Galvatron 搜尋得到的部分最優平行策略示意圖

可解釋性實驗：該研究選取了一些由Galvatron 搜尋得到的最優平行策略進行展示。對於BERT 模型在8GB 情況下(Case A)，Galvatron 選擇了兩種混合平行策略PP-TP-DP 以及PP-TP-SDP，而當可用顯存增大到12GB 時，Galvatron 放棄了PP，選擇使用更多的DP，同時引進SDP 來節省顯存空間。在Swin Transformer 上情況則稍有不同，該模型的不同層呈現出明顯的異質性，當顯存比較緊缺時（Case C），淺層SDP 的並行度較高，隨著層數加深，每層的活化變小，參數變多，所以TP 逐漸取代SDP。當顯存增多時（Case D）, 不僅重新啟用了 PP 取代了一部分低效的 SDP，而且淺層傾向於使用 DP 的趨勢更為明顯。

可擴展性實驗：該研究進一步在更大的叢集上測試了Galvatron，包括16 卡Nvidia RTX TITAN GPUs 的環境和64 卡Nvidia A100 GPUs 的環境。在16 卡環境下，Galvatron 相比於其他策略依然擁有最高吞吐率，與相同內存限制的8 卡實驗結果相比，由於更加多樣化的混合併行策略，Galvatron 在16 卡上能獲得超過2 倍的加速比。在 64 卡實驗中，Galvatron 的吞吐率同樣是其他策略最高的。這說明了 Galvatron 具有良好的可擴展性，詳細結果可以參考論文原文。

北大河圖團隊簡介

河圖（Hetu）開發團隊來自北京大學數據與智慧實驗室（Data and Intelligence Research Lab at Peking Univeristy，以下簡稱：實驗室），實驗室由北京大學電腦學院崔斌教授帶領，多年來主要在人工智慧、大數據等領域進行前沿研究，在理論和技術創新以及系統研發上取得多項成果，已在國際頂尖學術會議和期刊發表學術論文100 餘篇。

河圖系統是一個以超大模型為導向的分散式深度學習系統，相較於現有的老牌分散式深度學習框架，在系統功能性、系統複雜性和系統易用性上有許多創新貢獻，如自動分散式並行策略、一致性協定與通訊架構、GPU 算子最佳化等面向。 Hetu 團隊目前已在多種分散式機器學習或深度學習場景下進行了學術創新，相關成果被SIGMOD、VLDB、ICML、KDD 等國際頂級會議收錄，其中稀疏大模型分佈式訓練系統HET 榮獲VLDB 2022 最佳可擴展數據科學論文獎。這次被 VLDB 2023 錄用的論文 Galvatron 則是河圖團隊在稠密大模型分散式訓練場景上的又一突破成果，已經整合到 Hetu 系統並對外開源。目前，Hetu 團隊已與騰訊、阿里巴巴、快手、位元組跳動等多家知名企業展開科學研究合作與應用落地。

以上是北大河圖發布分散式訓練神器Galvatron, 一鍵實現大模型高效自動並行的詳細內容。更多資訊請關注PHP中文網其他相關文章！