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

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

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

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

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

圖 1 百度豐富業(yè)務場景

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

圖 2 大規(guī)模訓練場景

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

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

圖 3 傳統(tǒng)參數服務器

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

從傳統(tǒng)純 CPU 參數服務器到純 GPU 參數服務器

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

圖 4 傳統(tǒng)參數服務器工作流程

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

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

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

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

圖 5 純 GPU 參數服務器工作流程

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

從傳統(tǒng)純 GPU 參數服務器到異構參數服務器

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

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

混布 CPU 和 GPU 機器節(jié)點，來調整機器間的硬件配比。

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

圖 6 異構參數服務器示意圖

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

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

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

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

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

圖 7 集合通信訓練

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

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

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

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

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

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

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

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

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

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

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

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

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

下面咱們再來簡單介紹下幾個并行策略的原理。

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

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

值得注意的是，流水線并行可以說是廣義模型并行的一種特例，本文中的模型并行僅指 Tensor 切分，也就是會出現同一層網絡交由不同卡來計算的情況，而流水線并行則是按照網絡層的粒度切分。

圖 9 流水線并行策略示意圖

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

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

圖 10 流水線并行時序圖

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

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

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

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

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

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

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

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

測試驗證

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

寫在最后

自飛槳設計之初就開始潛心研究分布式訓練技術以應對大規(guī)模參數模型的訓練任務。在豐富的搜索推薦業(yè)務的驅動下，飛槳分布式訓練參數服務器模式歷經三代。最早的純 CPU 參數服務器就已經可以訓練萬億級規(guī)模的稀疏參數模型。其后隨著業(yè)務的需要以及前沿技術的發(fā)展，產生了計算能力更強的純 GPU 參數服務器模式。最近新推出的業(yè)內首創(chuàng)的異構參數服務器模式，支持場景更多而且可以極大地提升硬件資源利用效率。對于大規(guī)模稠密參數模型，飛槳分布式訓練技術同樣與業(yè)務緊密結合，其集合通信模式通過最新的 4D 混合并行策略，支持了 2300 億參數規(guī)模的文心 ERNIE 模型的分布式訓練。如今飛槳已經開始研究下一代分布式技術，來同時兼容超大規(guī)模稠密參數和稀疏參數模型的訓練。相信在實際產業(yè)應用這個核心驅動力推動下，飛槳分布式訓練必將成為星辰大海上的那顆北極星，為廣大開發(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】相關教程：https://fleet-x.readthedocs.io/en/latest/paddle_fleet_rst/collective/collective_mp/hybrid_parallelism.html