飛槳分布式訓練又推新品，4D混合并行可訓千億級AI模型

最近，飛槳提出了4D混合并行策略，以訓練千億級稠密參數模型。

近幾年，深度學(xué)習領(lǐng)域的開(kāi)發(fā)者們對模型效果的追求愈演愈烈，各大榜單紀錄不斷刷新，而這個(gè)現象的背后都有著(zhù) “大規模訓練” 的身影。簡(jiǎn)單來(lái)說(shuō)，就是使用大規模的數據或大規模參數量的模型來(lái)做訓練。大規模的數據可以讓模型有足夠的 “教材” 用于 “學(xué)習”，而大規模的參數量則可以讓模型“學(xué)習能力” 更強，更容易 “學(xué)習” 到“教材”中的“知識”。在數據和參數規模增長(cháng)的過(guò)程中，常規的單機訓練由于硬件資源的限制漸漸顯得捉襟見(jiàn)肘，而分布式訓練則成為了廣大開(kāi)發(fā)者的必然選擇。

所謂分布式訓練，就是使用多臺機器共同完成訓練任務(wù)，這其中涉及多機任務(wù)拆分、集群訓練資源配置、平衡訓練速度和收斂速度、彈性訓練與容錯等多項重要技術(shù)，同時(shí)也是各大深度學(xué)習框架彰顯技術(shù)實(shí)力的重要 “戰略高地”。

飛槳是我國首個(gè)開(kāi)源開(kāi)放、自主研發(fā)、功能完備的產(chǎn)業(yè)級深度學(xué)習框架，其英文名“PaddlePaddle” 正是 “Parallel Distributed Deep Learning” 并行分布式深度學(xué)習的字母縮寫(xiě)組合。飛槳不僅在業(yè)內最早支持了萬(wàn)億級稀疏參數模型的訓練能力，而且近期又創(chuàng )新性的提出了 4D 混合并行策略，以訓練千億級稠密參數模型，可以說(shuō)分布式訓練是飛槳最具特色的技術(shù)之一。那么飛槳是如何做到的呢？這與實(shí)際業(yè)務(wù)的錘煉密不可分。

圖 1 百度豐富業(yè)務(wù)場(chǎng)景

飛槳的分布式訓練技術(shù)在對外提供之前就已經(jīng)在百度內部廣泛應用，如搜索引擎、信息流推薦、百度翻譯、百度地圖、好看視頻、文心 ERNIE 等等，既包含網(wǎng)絡(luò )復雜、稠密參數特點(diǎn)的計算機視覺(jué)（CV）\ 自然語(yǔ)言處理（NLP）模型訓練場(chǎng)景，又覆蓋了有著(zhù)龐大的 Embedding 層模型和超大數據量的推薦搜索訓練場(chǎng)景，可謂是分布式訓練技術(shù)得天獨厚的“練功房”。

圖 2 大規模訓練場(chǎng)景

歷經(jīng)搜索推薦業(yè)務(wù)磨煉，最成熟萬(wàn)億稀疏參數模型訓練技術(shù)一騎絕塵

搜索推薦場(chǎng)景經(jīng)常面臨數據量大、特征維度高且稀疏化的問(wèn)題。而分布式訓練的參數服務(wù)器模式采用了一種將模型參數中心化管理的方式來(lái)實(shí)現模型參數的分布式存儲和更新，該模式有兩個(gè)角色 Server 與 Worker：Worker 用于執行模型的前向與反向計算；Server 負責從各個(gè) Worker 收集匯總梯度并更新參數，因此對于存儲超大規模模型參數的訓練場(chǎng)景十分友好，常被用于訓練擁有海量稀疏參數的搜索推薦領(lǐng)域模型。

圖 3 傳統參數服務(wù)器

百度搜索作為全球最大的中文搜索引擎，對模型的規模、性能等要求非常高。為了應對嚴苛的實(shí)際業(yè)務(wù)挑戰，早在 2018 年，飛槳的純 CPU 參數服務(wù)器模式就可以支持萬(wàn)億規模稀疏參數的模型訓練。之后隨著(zhù)模型中網(wǎng)絡(luò )結構更加復雜，以及對訓練效率和性?xún)r(jià)比的進(jìn)一步追求，飛槳參數服務(wù)器技術(shù)也在更新?lián)Q代：從早期 Worker 節點(diǎn)的硬件型號必須一致的純 CPU 參數服務(wù)器到純 GPU 參數服務(wù)器，再到 CPU、GPU、其它 AI 硬件混布調度訓練的異構參數服務(wù)器，始終引領(lǐng)參數服務(wù)器技術(shù)的發(fā)展；同時(shí)也得到了更多的應用落地，如 OPPO 應用商店推薦、網(wǎng)易云音樂(lè )主播推薦等等。

從傳統純 CPU 參數服務(wù)器到純 GPU 參數服務(wù)器

傳統的純 CPU 參數服務(wù)器，由高性能異步訓練 Worker、高效通信策略和高性能 Server 組成。通?？梢允褂玫?CPU 數量較多，訓練中能夠充分展示 CPU 多核的吞吐量?jì)?yōu)勢。在異步訓練模式下訓練簡(jiǎn)單模型可以極大提升數據吞吐量，整體訓練速度非常出色。

圖 4 傳統參數服務(wù)器工作流程

但是隨著(zhù)模型網(wǎng)絡(luò )越來(lái)越復雜，對算力要求越來(lái)越高，在數據量不變的情況下，CPU 計算性能差的弱勢就會(huì )顯現，雖然可以通過(guò)增加 CPU 機器數量來(lái)解決，甚至可以增加上百臺，但是這種方法不僅成本大幅提高，而且集群的穩定性和擴展性也存在較大的問(wèn)題。因此飛槳引入了純 GPU 參數服務(wù)器來(lái)提升計算性能，之前 100 臺 CPU 機器才能訓練的模型，僅需 1 臺多卡 GPU 機器即可完成訓練。當然，同時(shí)也要解決因為硬件更替所帶來(lái)的問(wèn)題。

GPU 強大的算力毋庸置疑可以提升集群的計算性能，但隨之而來(lái)的是，不僅模型規模會(huì )受到機器顯存和內存的制約，而且通信帶寬也會(huì )由于集群網(wǎng)卡數量降低而成為瓶頸。為了解決這兩個(gè)問(wèn)題，飛槳引入了兩大亮點(diǎn)技術(shù) SSD-MEM-HBM 三級存儲和 RPC&NCCL 混合通信，形成了飛槳特有的純 GPU 參數服務(wù)器（PaddleBox）【1】：

SSD-MEM-HBM 三級存儲允許全量參數使用 SSD 硬盤(pán)存儲，高頻參數存儲于內存，當前 Batch 訓練所用參數使用顯存，并且同時(shí)支持 SSD 的參數在硬盤(pán)、內存、顯存之間快速拷貝。這樣通過(guò)異步流水線(xiàn)執行機制，隱蔽了 IO 帶來(lái)的額外性能開(kāi)銷(xiāo)，在保證訓練速度的同時(shí)，使訓練的模型大小不再受制于顯存和內存，極大提升模型的規模。

RPC&NCCL 混合通信可以將部分稀疏參數采用 RPC 協(xié)議跨節點(diǎn)通信，其余參數采用卡間 NCCL 方式完成通信，充分利用帶寬資源。

圖 5 純 GPU 參數服務(wù)器工作流程

飛槳純 GPU 參數服務(wù)器雖然解決了之前純 CPU 模式所面臨的問(wèn)題，但新的問(wèn)題又出現了——如何提高訓練資源的利用率？

從傳統純 GPU 參數服務(wù)器到異構參數服務(wù)器

在純 GPU 的參數服務(wù)器下，所有的訓練都在 GPU 中，當模型中部分網(wǎng)絡(luò )層比較復雜的時(shí)候，GPU 利用率很難被打滿(mǎn)，而 GPU 機器中 CPU 與 GPU 的硬件配比是固定的，無(wú)法靈活調整。針對這種情況，有兩種解決方案：

定制化 GPU 機型，調整機器內 CPU 與 GPU 的硬件配比。

混布 CPU 和 GPU 機器節點(diǎn)，來(lái)調整機器間的硬件配比。

基于這兩種解決方案，飛槳框架 2.0 版本創(chuàng )新性地推出了通用異構參數服務(wù)器功能。一舉解除了傳統參數服務(wù)器模式下，Worker 節點(diǎn)必須嚴格使用同一種硬件型號的枷鎖，使訓練任務(wù)對硬件型號不敏感，即可以同時(shí)使用不同的硬件混合異構訓練，如 CPU、AI 專(zhuān)用芯片（如百度昆侖 XPU）以及不同型號的 GPU 如 v100、P40、K40 等。同時(shí)還可以解決大規模稀疏特征模型訓練場(chǎng)景下 IO 占比過(guò)高導致的芯片資源利用率過(guò)低的問(wèn)題。通過(guò)異構參數服務(wù)器訓練模式，用戶(hù)可以在硬件異構集群中部署分布式訓練任務(wù)，例如云服務(wù)器集群，高效利用不同算力芯片，為用戶(hù)提供更高吞吐、更低資源消耗的訓練能力。

圖 6 異構參數服務(wù)器示意圖

異構參數服務(wù)器的最大亮點(diǎn)是硬件感知的任務(wù)切分。如圖 6 所示，針對類(lèi)似 ERNIE+CTR 這樣計算密集型與 IO 密集型兼有的訓練任務(wù)，可以被切分成多個(gè)子任務(wù)。其中的 IO 密集型任務(wù)（如數據讀取、Embedding 查詢(xún)）切分給 CPU 機器，計算密集型任務(wù)切分給 GPU 機器；用戶(hù)可以根據子任務(wù)的計算復雜度來(lái)靈活決定機器配比，并且還可以兼容傳統純 CPU 參數服務(wù)器和純 GPU 參數服務(wù)器所支持的訓練任務(wù)。

助力文心 ERNIE 快速迭代，首創(chuàng ) 4D 混合并行引領(lǐng)超大規模預訓練潮流

在 NLP 領(lǐng)域中，依托飛槳打造的 “語(yǔ)義理解技術(shù)與平臺文心 ERNIE” 曾獲得過(guò)無(wú)數殊榮：去年 3 月一舉拿下 SemEval 2020 的 5 項冠軍；5 月發(fā)布語(yǔ)言生成預訓練模型 ERNIE-GEN，刷新語(yǔ)言生成 SOTA；6 月發(fā)布多模態(tài)模型 ERNIE-ViL，刷新 5 項任務(wù)紀錄，登頂權威榜單 VCR；7 月亮相 2020 世界人工智能大會(huì )，摘取最高榮譽(yù) SAIL 獎；11 月獲得中國人工智能學(xué)會(huì )優(yōu)秀科技成果獎。在文心 ERNIE 這些閃耀成績(jì)的背后，也有飛槳的分布式訓練技術(shù)的貢獻。

首先對于 NLP 和 CV 這類(lèi)擁有復雜網(wǎng)絡(luò )、稠密參數特點(diǎn)的模型，飛槳分布式訓練技術(shù)的集合通信模式可以很好的支持這類(lèi)模型的訓練。該模式?jīng)]有管理模型參數的中心節點(diǎn)，每個(gè)節點(diǎn)都是 Worker，每個(gè) Worker 負責模型訓練的同時(shí)還需要掌握當前最新的全局梯度信息。集合通信模式對計算芯片的算力和芯片之間的網(wǎng)絡(luò )互聯(lián)要求較高，如高性能計算的 GPU、芯片之間的高速網(wǎng)絡(luò )互聯(lián) NVLINK 和 InfiniBand 等，因此非常適合 CV 和 NLP 領(lǐng)域計算密集型訓練任務(wù)。

但是在早期的集合通信架構中，多節點(diǎn)間的參數信息的傳輸通常是各個(gè) Worker 之間多次點(diǎn)對點(diǎn)通信完成的，通訊效率較低。百度在 2016 年突破性地提出并使用 Ring-AllReduce 多 GPU 訓練，通過(guò)較少的點(diǎn)對點(diǎn)通信輪數就完成了全局節點(diǎn)的模型參數傳輸，讓同步并行訓練的多 GPU 擴展能力得到極大突破，大幅提升集合通信模式的訓練速度，使這種模式在 NLP 和 CV 領(lǐng)域得到更廣泛的應用。

圖 7 集合通信訓練

4D 混合并行策略支持文心 ERNIE 千億語(yǔ)言模型訓練

當前飛槳集合通信模式已經(jīng)可以支持文心 ERNIE 千億語(yǔ)言模型的訓練能力，其 Sharding-DP 策略更是在近期助力文心 ERNIE 的多項任務(wù)分數刷新 GLUE 榜單。而這個(gè) Sharding-DP 策略正是飛槳集合通信模式為了訓練 ERNIE 這樣的大規模復雜模型所支持的多種并行策略中的一種。那么飛槳是使用哪些策略成功支持文心 ERNIE 千億語(yǔ)言模型訓練的呢？這些策略是如何工作的呢？接下來(lái)將為大家詳細介紹。

ERNIE 千億級模型采用 100 多層 Transformer 網(wǎng)絡(luò )結構，計算復雜，訓練需要占用 T 級顯存資源，如果想用更少的機器高效訓練，必須采取一系列性能優(yōu)化和顯存優(yōu)化措施。

首先看如何性能優(yōu)化。我們通過(guò)一個(gè)公式來(lái)看哪些因素可以影響訓練速度，在固定的硬件環(huán)境下：

總訓練速度 ∝ 單卡速度 * 卡數 * 多卡加速比 

其中單卡速度由數據讀取和計算速度決定；多卡加速比由計算 / 通信效率決定。顯而易見(jiàn)，這三個(gè)是關(guān)鍵因素。除了單卡可以使用的算子融合、混合精度之類(lèi)的基礎性能優(yōu)化策略之外，分布式訓練還引入一系列并行策略。并行策略的核心思想是將數據和計算有關(guān)的圖 / 算子切分到不同設備上，同時(shí)盡可能降低設備間通信所需的代價(jià)，合理使用多臺設備資源，實(shí)現高效的并發(fā)調度訓練，最大化提升訓練速度。常見(jiàn)并行策略有數據并行 DP（Data Parallel）、Layer 間并行（流水線(xiàn)并行 PP，Pipeline Parallel）、Layer 內并行（模型并行 MP，Model Parallel）。如下表所示。我們從設備資源和計算 / 通信效率來(lái)分析三種策略的優(yōu)缺點(diǎn)： 

數據并行訓練加速比最高，但要求每個(gè)設備上都備份一份模型，顯存占用比較高。為此我們的改進(jìn)方案是分組參數切片數據并行策略（具體原理后文介紹），兼容了 MP+DP 的優(yōu)勢，但缺點(diǎn)是通信量大。

模型并行，通信占比高，適合在機器內做模型并行且支持的模型類(lèi)型有限。

流水線(xiàn)并行，訓練設備容易出現空閑狀態(tài)，加速效率沒(méi)有 DP 高；但能減少通信邊界支持更多的層數，適合在機器間使用。

其次看顯存問(wèn)題，通過(guò)下表分析的顯存占用來(lái)源可以看出，上述的并行策略同樣可以很好地應對不同來(lái)源的顯存占用，更多的層數可以通過(guò)流水線(xiàn)并行和分組參數切分策略來(lái)解決；某層參數很大可以通過(guò)模型并行來(lái)解決；其次飛槳還提供一些其它靈活的優(yōu)化方式，例如每層輸出占用的顯存，可以通過(guò)重計算和 Offload 來(lái)解決。

綜上所述，針對性能優(yōu)化和顯存優(yōu)化，幾種并行策略都有用武之地，但是同時(shí)也有各自的局限性，所以如果想高效訓練千億模型，需要這幾種策略相互組合，取長(cháng)補短，發(fā)揮各自的優(yōu)勢。

那么如何組合呢？飛槳研發(fā)人員首先在單機內使用模型并行和分組參數切片組合的 2D 策略，這么選擇的原因是這兩個(gè)策略通信量較大，適合使用機器內的卡間通信；然后為了承載千億規模模型，再疊加流水線(xiàn)并行策略，使用多臺機器共同分擔；最后為了做到高效，在外層又疊加了數據并行來(lái)增加并發(fā)數量，提升整體訓練速度。這樣業(yè)內首個(gè) 4D 混合并行策略就誕生了。

圖 8 4D 混合并行策略示意圖

下面咱們再來(lái)簡(jiǎn)單介紹下幾個(gè)并行策略的原理。

模型并行策略指的是將某一層網(wǎng)絡(luò )切成多份，并分給不同的卡并行計算，每張卡僅需要計算部分結果。對于 ERNIE 中的 Transformer 網(wǎng)絡(luò )結構，模型并行就可以對全連接層 FC 切分，然后通過(guò)通信操作合并計算結果【2】。

流水線(xiàn)并行策略支持將模型的不同層放置到不同的設備上，通過(guò)多個(gè)設備來(lái)共同分擔顯存消耗，實(shí)現超大規模模型訓練。相鄰設備間通過(guò)通信鏈路傳輸數據。由于各個(gè)設備間傳輸的僅是相鄰設備間的輸出張量，因此通信量較小，相對來(lái)說(shuō)較為適合機器間通信的場(chǎng)景。

值得注意的是，流水線(xiàn)并行可以說(shuō)是廣義模型并行的一種特例，本文中的模型并行僅指 Tensor 切分，也就是會(huì )出現同一層網(wǎng)絡(luò )交由不同卡來(lái)計算的情況，而流水線(xiàn)并行則是按照網(wǎng)絡(luò )層的粒度切分。

圖 9 流水線(xiàn)并行策略示意圖

流水線(xiàn)并行策略本身也有很大的優(yōu)化空間。如圖 10（a）所示，優(yōu)化前，任意時(shí)刻只有單個(gè)計算設備處于計算狀態(tài)，其它計算設備則處于空閑狀態(tài)，這個(gè)空閑時(shí)間我們稱(chēng)之為 Bubble 時(shí)間【3】。為了減少 Bubble 時(shí)間，如圖 10（b）所示，飛槳進(jìn)一步將 mini-batch 切分成若干更小粒度的 micro-batch，每個(gè)設備依次計算單個(gè) micro-batch 的結果，從而增加了設備間的并發(fā)度，降低了流水線(xiàn)并行 Bubble 時(shí)間比例。

此外飛槳研發(fā)人員經(jīng)過(guò)對流水線(xiàn)并行訓練過(guò)程更加深入的剖析，發(fā)現還可以進(jìn)一步優(yōu)化顯存的利用率。采用如圖 10（c）中所示的方式，在一個(gè) micro-batch 完成前向計算后，提前調度完成相應后向計算，這樣就能釋放部分顯存，用以接納新的數據，提升整體訓練性能。使用 ERNIE 模型實(shí)測，從 10(b)到 10(c)，總 BatchSize 可以提升 32 倍，性能可以提升 9 倍。

圖 10 流水線(xiàn)并行時(shí)序圖

最后再來(lái)看下飛槳的分組參數切片策略，其特色是在使用參數切片方式節省顯存的基礎上，與數據并行策略組合成更加強大的 Sharding-DP 策略。簡(jiǎn)而言之，這種組合后的策略擁有很強的靈活性，用戶(hù)可以根據現有硬件環(huán)境情況，自由設定模型參數切分的數量（sharding_degree）和數據并行的路數（dp_degree），僅需要確保 sharding_degree × dp_degree = 總卡數即可。

舉個(gè)例子，假設用戶(hù)有 4 臺單機四卡的機器（共 16 張卡），訓一個(gè) 16 層網(wǎng)絡(luò )的模型。如果模型參數規?？梢允褂靡慌_機器承載，則推薦使用 dp_degree=4 & sharding_degree=4 的方式，如圖 11 所示。這種方式的優(yōu)勢在于只有機器內卡間通信，但是模型最大不能超過(guò)單臺機器所能承受存儲范圍。

圖 11 dp_degree=4 & sharding_degree=4 的 Sharding-DP 示意圖

如果模型大小超過(guò)了單臺機器，問(wèn)題也不大，用戶(hù)可以靈活地選擇 dp_degree=2 & sharding_degree=8 方式，如圖 12 所示。與上一種方式相比，這種方式支持的模型參數規模翻倍。

圖 12 dp_degree=2 & sharding_degree=8 的 Sharding-DP 示意圖

但是在一些特殊的情況下，如果模型參數規模非常大，半數機器都無(wú)法承載，則可以進(jìn)一步使用 dp_degree=1 & sharding_degree=16 方式，即將整個(gè)模型參數交由全部機器承載，這也是標準的 ZeRO-DP【4】方式，如圖 11 所示。這種方式跨機器通信數非常高，對訓練速度影響很大。其實(shí) Sharding-DP 可以說(shuō)是 ZeRO-DP 的一種升華，讓用戶(hù)可以使用更加高效方式應對特殊場(chǎng)景之外的絕大部分訓練任務(wù)。

圖 13 dp_degree=1 &  sharding_degree=16 的 Sharding-DP 示意圖示意圖

最后我們從理論性能角度對比分析了幾組混合并行策略，即 DP2+PP32+Sharding2+MP4、PP64+Sharding2+MP4 和 DP2+PP32+MP8。如下表所示，與兩種 3D 方式相比，4D 混合并行策略在通信量和 Bubble 時(shí)間上并未明顯增長(cháng)（具體公式推導和示例請參見(jiàn)相關(guān)教程【5】），但是大幅提升了數據并行路數！ 

測試驗證

從上面理論分析上來(lái)看，4D 混合并行策略應該會(huì )具有更好的性能。那么實(shí)際效果如何呢？咱們進(jìn)入實(shí)測階段。我們使用 64 臺 8 卡 GPU v100 機器的環(huán)境來(lái)驗證不同策略組合的訓練效果，而測試的對象，正是擁有 2300 億參數規模的“重量級”ERNIE 模型。測試后我們可以看到 4D 混合并行策略訓練速度高于其它兩種 3D 混合并行策略，達到了8698 tokens/s，至少可以提速 23.7%。

寫(xiě)在最后

自飛槳設計之初就開(kāi)始潛心研究分布式訓練技術(shù)以應對大規模參數模型的訓練任務(wù)。在豐富的搜索推薦業(yè)務(wù)的驅動(dòng)下，飛槳分布式訓練參數服務(wù)器模式歷經(jīng)三代。最早的純 CPU 參數服務(wù)器就已經(jīng)可以訓練萬(wàn)億級規模的稀疏參數模型。其后隨著(zhù)業(yè)務(wù)的需要以及前沿技術(shù)的發(fā)展，產(chǎn)生了計算能力更強的純 GPU 參數服務(wù)器模式。最近新推出的業(yè)內首創(chuàng )的異構參數服務(wù)器模式，支持場(chǎng)景更多而且可以極大地提升硬件資源利用效率。對于大規模稠密參數模型，飛槳分布式訓練技術(shù)同樣與業(yè)務(wù)緊密結合，其集合通信模式通過(guò)最新的 4D 混合并行策略，支持了 2300 億參數規模的文心 ERNIE 模型的分布式訓練。如今飛槳已經(jīng)開(kāi)始研究下一代分布式技術(shù)，來(lái)同時(shí)兼容超大規模稠密參數和稀疏參數模型的訓練。相信在實(shí)際產(chǎn)業(yè)應用這個(gè)核心驅動(dòng)力推動(dòng)下，飛槳分布式訓練必將成為星辰大海上的那顆北極星，為廣大開(kāi)發(fā)者們指引航向。

【1】Zhao W, Xie D, Jia R, et al. Distributed hierarchical gpu parameter server for massive scale deep learning ads systems[J]. arXiv preprint arXiv:2003.05622, 2020

【2】Shoeybi M, Patwary M, Puri R, et al. Megatron-lm: Training multi-billion parameter language models using model parallelism[J]. arXiv preprint arXiv:1909.08053, 2019.

【3】Huang Y, Cheng Y, Bapna A, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism[J]. arXiv preprint arXiv:1811.06965, 2018.

【4】Rajbhandari S, Rasley J, Ruwase O, et al. Zero: Memory optimizations toward training trillion parameter models[C]//SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020: 1-16.

【5】相關(guān)教程：https://fleet-x.readthedocs.io/en/latest/paddle_fleet_rst/collective/collective_mp/hybrid_parallelism.html