最近一段時間，「大模型」在 AI 領域的各種應用場景都大放異彩，其中基于 Transformer 的大規(guī)模預訓練模型是最典型的大模型之一，已經成為了當前基礎模型（Foundation Model）的核心架構。例如 NLP 領域的 BERT、GPT 系列，CV 領域的 ViT、Swin Transformer 系列，以及最近非?；馃岬亩鄬＜一旌夏Ｐ?MoE 以及多模態(tài)模型 CLIP 等等，都是以 Transformer 作為核心的基礎架構。與之相應的，這類稠密大模型擁有著動輒數十億、百億甚至萬億規(guī)模的參數量，面臨高昂的計算、存儲、以及通信開銷，也為 AI 基礎設施帶來了巨大的挑戰(zhàn)。

為了支持大模型的訓練，人們研發(fā)了很多工具（如 NVIDIA 提出的「威震天」Megatron，Microsoft 提出的 DeepSpeed，Meta 提出的 FairSeq 等等）來實現各種各樣的并行方式，數據并行、張量模型并行、流水并行、分片數據并行等等。這些系統(tǒng)對上述并行方式提供了良好的封裝，對外屏蔽相應的實現細節(jié)，使得用戶能夠通過添加配置的方式來實現混合并行策略。

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

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

• 論文地址：https://arxiv.org/abs/2211.13878
• 項目代碼鏈接：https://github.com/PKU-DAIR/Hetu/tree/main/tools/Galvatron

大模型自動并行到底難在哪里

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

（1）多樣性：首先，在并行方式方面，目前大模型的并行方式呈現出百花齊放的態(tài)勢，即使是對于同一個算子，不考慮混合并行方式，不同的基礎并行方式也會存在顯著的差異，從而導致不同的內存開銷、通信代價以及計算效率。下圖展示了四種最主要的基礎并行方式，即數據并行（Data Parallelism）、張量模型并行（Tensor Parallelism）、流水模型并行（Pipeline Parallelism）、分片數據并行 (Sharded Data Parallelism)，在兩張 GPU 上對于簡單的矩陣乘法算子分布式執(zhí)行的過程。

?

并行方法對比示意圖

其次，在模型方面，各種各樣的模型架構最近也是層出不窮，這往往也伴隨著不同的模型配置（例如不同輸入序列長度，模型層數，模型隱層寬度等），從而造成計算負載上的差異。另外，在硬件方面，用戶往往面臨著非常差異化的集群環(huán)境，可能會面臨不同的內存容量、通信帶寬、計算能力等等。總體上來看，由于上述多樣性的存在，沒有哪種并行技術總是能夠獲得最佳訓練效率，“自動并行” 也就成為了分布式訓練的核心挑戰(zhàn)。

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

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

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

系統(tǒng)特性：

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

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

大模型分布式訓練系統(tǒng)對比示意圖?

（2）在復雜性方面，由于具備在多樣性上的優(yōu)勢，Galvatron 面臨了前所未有的龐大搜索空間，為此，研究者探究了幾條目前大規(guī)模分布式訓練過程中的一些經過實驗性或理論性驗證的重要觀察，作為搜索空間的剪枝準則，從而實現高效的分布式執(zhí)行計劃優(yōu)化。

（3）在實用性方面，該研究結合了理論建模和實驗測量兩者的優(yōu)勢，實現對內存、通信、計算開銷的精確估計，甚至考慮到了計算與通信重疊所導致的 GPU 執(zhí)行效率下降問題，從而保證能夠得到足夠準確的自動并行優(yōu)化結果。

另外，Galvatron 底層選擇 PyTorch 作為執(zhí)行引擎，兼容 Huggingface 等常見的主流 Transformer 模型實現，所以完全不會對 PyTorch 用戶帶來額外的使用負擔；同時也不需要用戶付出額外的系統(tǒng)安裝或者調試代價，使用時只需要添加幾行代碼，就可以輕松完成自動并行的整個流程。

Galvatron 工作流程及用戶使用接口展示?

關鍵技術

1. 基于決策樹的搜索空間分解

Galvatron 的設計目標是高效地在復雜而龐大的并行策略空間內自動搜索，并為給定的 Transformer 模型和分布式環(huán)境生成最佳的并行執(zhí)行計劃。在搜索空間上，Galvatron 是業(yè)界首個考慮 4 種主流并行方法的自動并行訓練系統(tǒng)，包括包括數據并行（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。當節(jié)點內 GPU 數量擴展到 8 卡時，模型每一層的候選策略都多達上百種。隨著模型層數的增加，其搜索空間規(guī)模指數增長，難以有效探索。

為了高效地搜索如此龐大的搜索空間，該研究首先提出了以下觀察作為指導：

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

基于以上重要觀察，該研究提出了一種基于決策樹的搜索空間構建方法：

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

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

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

（4）非葉結點的度數默認在 2 的指數次冪 {2,4,8,…} 中選擇。

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

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

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

2. 基于動態(tài)規(guī)劃的并行優(yōu)化算法

現有系統(tǒng)如 Megatron 或者 DeepSpeed 通常由用戶指定全局并行方案及其對應的并行度，嚴重限制了對分布式執(zhí)行計劃的表達能力。Galvatron 的優(yōu)化目標是在用戶給定模型定義和分布式環(huán)境的情況下，無需用戶指定任何并行配置，自動生成最優(yōu)的分布式執(zhí)行計劃。具體來說，給定一個 L 層的模型 M 和內存容量為 E 的 N 個 GPU 設備，Galvatron 的優(yōu)化目標是搜索得到最高的系統(tǒng)吞吐 T_pt 并返回對應的并行方案，這里的并行方案指以層（或算子）為基礎單位的細粒度混合并行策略。

?

算法 1: Galvatron 優(yōu)化過程

優(yōu)化過程：Galvatron 的優(yōu)化過程如算法 1 所示。Galvatron 最外層循環(huán)逐步增加搜索的 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 使用動態(tài)規(guī)劃搜索得到其中每一層的最優(yōu)混合并行策略并返回最小時間開銷（第 9 行）；最后，Galvatron 在所有可能的 PP 并行度和 batch size 中選擇吞吐率最高的策略并返回（第 15 行）。

動態(tài)規(guī)劃搜索：下面介紹 Galvatron 并行優(yōu)化工作流中的動態(tài)規(guī)劃搜索算法。對于一個給定的包含 L 層的模型階段，用開銷函數 C(L,E) 表示 L 層模型在設備內存限制 E 下的總執(zhí)行時間，表示第 L 層使用策略 S_j 的執(zhí)行時間，其中策略 S_j 是并行策略候選集合 S 中的策略。設定初始值

，Galvatron 的動態(tài)規(guī)劃搜索遵循以下狀態(tài)轉移方程（公式 1）：

其中，?是第 L 層使用策略 S_j 的內存開銷，?是第 L 層使用策略 S_j、其前一層使用策略 S_i 帶來的轉換開銷。在狀態(tài)轉移過程中，當內存開銷超過設備內存限制設備內存限制 E，開銷函數 C 返回無窮大。

復雜度分析：Galvatron 使用的動態(tài)規(guī)劃搜索（公式 1）計算復雜度為 O(LE|S|)?？梢?，每一層的搜索空間 S 的大小對于整體搜索復雜度非常重要，Galvatron 提出的基于決策樹的搜索空間分解能夠顯著減小搜索空間，控制搜索開銷在合理范圍內。

3. 基于混合建模的執(zhí)行代價估計方法

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

基于以上估計結果，Galvatron 通過 simulating 執(zhí)行過程來計算給定層使用給定策略的開銷 c(l,s)。和現有分布式訓練系統(tǒng)的代價模型不同，Galvatron 在建模中首次考慮了計算和通信的重疊對 GPU 性能下降帶來的影響。該研究通過實驗發(fā)現，由于重疊導致的 GPU 性能下降會顯著影響系統(tǒng)的執(zhí)行效率，而這在以往的工作中都是被忽視的。因此，Galvatron 的開銷估計更加準確，并行優(yōu)化效果更佳。

實驗結果

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

Galvatron 和基線系統(tǒng)在 8 GPUs 20G 顯存下的吞吐對比情況

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

Galvatron 搜索得到的部分最優(yōu)并行策略示意圖

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

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

北大河圖團隊簡介

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

河圖系統(tǒng)是一個面向超大模型的分布式深度學習系統(tǒng)，相比現有的老牌分布式深度學習框架，在系統(tǒng)功能性、系統(tǒng)復雜性和系統(tǒng)易用性上有諸多創(chuàng)新貢獻，如自動分布式并行策略、一致性協議和通信架構、GPU 算子優(yōu)化等方面。Hetu 團隊目前已在多種分布式機器學習或深度學習場景下進行了學術創(chuàng)新，相關成果被 SIGMOD、VLDB、ICML、KDD 等國際頂級會議收錄，其中稀疏大模型分布式訓練系統(tǒng) HET 榮獲 VLDB 2022 最佳可擴展數據科學論文獎。本次被 VLDB 2023 錄用的論文 Galvatron 則是河圖團隊在稠密大模型分布式訓練場景上的又一突破成果，已經集成到 Hetu 系統(tǒng)并對外開源。目前，Hetu 團隊已經與騰訊、阿里巴巴、快手、字節(jié)跳動等多家知名企業(yè)展開科研合作和應用落地。