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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

從傳統(tǒng)純 GPU 參數(shù)服務器到異構(gòu)參數(shù)服務器

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

定制化 GPU 機型，調(diào)整機器內(nèi) CPU 與 GPU 的硬件配比。

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

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

圖 6 異構(gòu)參數(shù)服務器示意圖

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

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

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

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

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

圖 7 集合通信訓練

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

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

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

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

總訓練速度 ∝ 單卡速度 * 卡數(shù) * 多卡加速比 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

圖 10 流水線并行時序圖

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

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

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

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

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

但是在一些特殊的情況下，如果模型參數(shù)規(guī)模非常大，半數(shù)機器都無法承載，則可以進一步使用 dp_degree=1 & sharding_degree=16 方式，即將整個模型參數(shù)交由全部機器承載，這也是標準的 ZeRO-DP【4】方式，如圖 11 所示。這種方式跨機器通信數(shù)非常高，對訓練速度影響很大。其實 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 時間上并未明顯增長（具體公式推導和示例請參見相關(guān)教程【5】），但是大幅提升了數(shù)據(jù)并行路數(shù)！ 

測試驗證

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

寫在最后

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