機器之心釋出
機器之心編輯部
最近，飛槳提出了4D混合並行策略，以訓練千億級稠密引數模型。
近幾年，深度學習領域的開發者們對模型效果的追求愈演愈烈，各大榜單紀錄不斷重新整理，而這個現象的背後都有著 “大規模訓練” 的身影。簡單來說，就是使用大規模的資料或大規模引數量的模型來做訓練。大規模的資料可以讓模型有足夠的 “教材” 用於 “學習”，而大規模的引數量則可以讓模型“學習能力” 更強，更容易 “學習” 到“教材”中的“知識”。在資料和引數規模增長的過程中，常規的單機訓練由於硬體資源的限制漸漸顯得捉襟見肘，而分散式訓練則成為了廣大開發者的必然選擇。
所謂分散式訓練，就是使用多臺機器共同完成訓練任務，這其中涉及多機任務拆分、叢集訓練資源配置、平衡訓練速度和收斂速度、彈性訓練與容錯等多項重要技術，同時也是各大深度學習框架彰顯技術實力的重要 “戰略高地”。
飛槳是我國首個開源開放、自主研發、功能完備的產業級深度學習框架，其英文名“PaddlePaddle” 正是 “Parallel Distributed Deep Learning” 並行分散式深度學習的字母縮寫組合。飛槳不僅在業內最早支援了萬億級稀疏引數模型的訓練能力，而且近期又創新性的提出了 4D 混合並行策略，以訓練千億級稠密引數模型，可以說分散式訓練是飛槳最具特色的技術之一。那麼飛槳是如何做到的呢？這與實際業務的錘鍊密不可分。
圖 1 百度豐富業務場景
飛槳的分散式訓練技術在對外提供之前就已經在百度內部廣泛應用，如搜尋引擎、資訊流推薦、百度翻譯、百度地圖、好看影片、文心 ERNIE 等等，既包含網路複雜、稠密引數特點的計算機視覺（CV）\ 自然語言處理（NLP）模型訓練場景，又覆蓋了有著龐大的 Embedding 層模型和超大資料量的推薦搜尋訓練場景，可謂是分散式訓練技術得天獨厚的“練功房”。
圖 2 大規模訓練場景
歷經搜尋推薦業務磨鍊，最成熟萬億稀疏引數模型訓練技術一騎絕塵
搜尋推薦場景經常面臨資料量大、特徵維度高且稀疏化的問題。而分散式訓練的引數伺服器模式採用了一種將模型引數中心化管理的方式來實現模型引數的分散式儲存和更新，該模式有兩個角色 Server 與 Worker：Worker 用於執行模型的前向與反向計算；Server 負責從各個 Worker 收集彙總梯度並更新引數，因此對於儲存超大規模模型引數的訓練場景十分友好，常被用於訓練擁有海量稀疏引數的搜尋推薦領域模型。
圖 3 傳統引數伺服器
百度搜尋作為全球最大的中文搜尋引擎，對模型的規模、效能等要求非常高。為了應對嚴苛的實際業務挑戰，早在 2018 年，飛槳的純 CPU 引數伺服器模式就可以支援萬億規模稀疏引數的模型訓練。之後隨著模型中網路結構更加複雜，以及對訓練效率和價效比的進一步追求，飛槳引數伺服器技術也在更新換代：從早期 Worker 節點的硬體型號必須一致的純 CPU 引數伺服器到純 GPU 引數伺服器，再到 CPU、GPU、其它 AI 硬體混布排程訓練的異構引數伺服器，始終引領引數伺服器技術的發展；同時也得到了更多的應用落地，如 OPPO 應用商店推薦、網易雲音樂主播推薦等等。
從傳統純 CPU 引數伺服器到純 GPU 引數伺服器
傳統的純 CPU 引數伺服器，由高效能非同步訓練 Worker、高效通訊策略和高效能 Server 組成。通常可以使用的 CPU 數量較多，訓練中能夠充分展示 CPU 多核的吞吐量優勢。在非同步訓練模式下訓練簡單模型可以極大提升資料吞吐量，整體訓練速度非常出色。
圖 4 傳統引數伺服器工作流程
但是隨著模型網路越來越複雜，對算力要求越來越高，在資料量不變的情況下，CPU 計算效能差的弱勢就會顯現，雖然可以透過增加 CPU 機器數量來解決，甚至可以增加上百臺，但是這種方法不僅成本大幅提高，而且叢集的穩定性和擴充套件性也存在較大的問題。因此飛槳引入了純 GPU 引數伺服器來提升計算效能，之前 100 臺 CPU 機器才能訓練的模型，僅需 1 臺多卡 GPU 機器即可完成訓練。當然，同時也要解決因為硬體更替所帶來的問題。
GPU 強大的算力毋庸置疑可以提升叢集的計算效能，但隨之而來的是，不僅模型規模會受到機器視訊記憶體和記憶體的制約，而且通訊頻寬也會由於叢集網路卡數量降低而成為瓶頸。為了解決這兩個問題，飛槳引入了兩大亮點技術 SSD-MEM-HBM 三級儲存和 RPC&NCCL 混合通訊，形成了飛槳特有的純 GPU 引數伺服器（PaddleBox）【1】：
SSD-MEM-HBM 三級儲存允許全量引數使用 SSD 硬碟儲存，高頻引數儲存於記憶體，當前 Batch 訓練所用引數使用視訊記憶體，並且同時支援 SSD 的引數在硬碟、記憶體、視訊記憶體之間快速複製。這樣透過非同步流水線執行機制，隱蔽了 IO 帶來的額外效能開銷，在保證訓練速度的同時，使訓練的模型大小不再受制於視訊記憶體和記憶體，極大提升模型的規模。
RPC&NCCL 混合通訊可以將部分稀疏引數採用 RPC 協議跨節點通訊，其餘引數採用卡間 NCCL 方式完成通訊，充分利用頻寬資源。
圖 5 純 GPU 引數伺服器工作流程
飛槳純 GPU 引數伺服器雖然解決了之前純 CPU 模式所面臨的問題，但新的問題又出現了——如何提高訓練資源的利用率？
從傳統純 GPU 引數伺服器到異構引數伺服器
在純 GPU 的引數伺服器下，所有的訓練都在 GPU 中，當模型中部分網路層比較複雜的時候，GPU 利用率很難被打滿，而 GPU 機器中 CPU 與 GPU 的硬體配比是固定的，無法靈活調整。針對這種情況，有兩種解決方案：
定製化 GPU 機型，調整機器內 CPU 與 GPU 的硬體配比。
混布 CPU 和 GPU 機器節點，來調整機器間的硬體配比。
基於這兩種解決方案，飛槳框架 2.0 版本創新性地推出了通用異構引數伺服器功能。一舉解除了傳統引數伺服器模式下，Worker 節點必須嚴格使用同一種硬體型號的枷鎖，使訓練任務對硬體型號不敏感，即可以同時使用不同的硬體混合異構訓練，如 CPU、AI 專用晶片（如百度崑崙 XPU）以及不同型號的 GPU 如 v100、P40、K40 等。同時還可以解決大規模稀疏特徵模型訓練場景下 IO 佔比過高導致的晶片資源利用率過低的問題。透過異構引數伺服器訓練模式，使用者可以在硬體異構叢集中部署分散式訓練任務，例如雲伺服器叢集，高效利用不同算力晶片，為使用者提供更高吞吐、更低資源消耗的訓練能力。
圖 6 異構引數伺服器示意圖
異構引數伺服器的最大亮點是硬體感知的任務切分。如圖 6 所示，針對類似 ERNIE+CTR 這樣計算密集型與 IO 密集型兼有的訓練任務，可以被切分成多個子任務。其中的 IO 密集型任務（如資料讀取、Embedding 查詢）切分給 CPU 機器，計算密集型任務切分給 GPU 機器；使用者可以根據子任務的計算複雜度來靈活決定機器配比，並且還可以相容傳統純 CPU 引數伺服器和純 GPU 引數伺服器所支援的訓練任務。
助力文心 ERNIE 快速迭代，首創 4D 混合並行引領超大規模預訓練潮流
在 NLP 領域中，依託飛槳打造的 “語義理解技術與平臺文心 ERNIE” 曾獲得過無數殊榮：去年 3 月一舉拿下 SemEval 2020 的 5 項冠軍；5 月釋出語言生成預訓練模型 ERNIE-GEN，重新整理語言生成 SOTA；6 月釋出多模態模型 ERNIE-ViL，重新整理 5 項任務紀錄，登頂權威榜單 VCR；7 月亮相 2020 世界人工智慧大會，摘取最高榮譽 SAIL 獎；11 月獲得中國人工智慧學會優秀科技成果獎。在文心 ERNIE 這些閃耀成績的背後，也有飛槳的分散式訓練技術的貢獻。
首先對於 NLP 和 CV 這類擁有複雜網路、稠密引數特點的模型，飛槳分散式訓練技術的集合通訊模式可以很好的支援這類模型的訓練。該模式沒有管理模型引數的中心節點，每個節點都是 Worker，每個 Worker 負責模型訓練的同時還需要掌握當前最新的全域性梯度資訊。集合通訊模式對計算晶片的算力和晶片之間的網路互聯要求較高，如高效能運算的 GPU、晶片之間的高速網路互聯 NVLINK 和 InfiniBand 等，因此非常適合 CV 和 NLP 領域計算密集型訓練任務。
但是在早期的集合通訊架構中，多節點間的引數資訊的傳輸通常是各個 Worker 之間多次點對點通訊完成的，通訊效率較低。百度在 2016 年突破性地提出並使用 Ring-AllReduce 多 GPU 訓練，透過較少的點對點通訊輪數就完成了全域性節點的模型引數傳輸，讓同步並行訓練的多 GPU 擴充套件能力得到極大突破，大幅提升集合通訊模式的訓練速度，使這種模式在 NLP 和 CV 領域得到更廣泛的應用。
圖 7 集合通訊訓練
4D 混合並行策略支援文心 ERNIE 千億語言模型訓練
當前飛槳集合通訊模式已經可以支援文心 ERNIE 千億語言模型的訓練能力，其 Sharding-DP 策略更是在近期助力文心 ERNIE 的多項任務分數重新整理 GLUE 榜單。而這個 Sharding-DP 策略正是飛槳集合通訊模式為了訓練 ERNIE 這樣的大規模複雜模型所支援的多種並行策略中的一種。那麼飛槳是使用哪些策略成功支援文心 ERNIE 千億語言模型訓練的呢？這些策略是如何工作的呢？接下來將為大家詳細介紹。
ERNIE 千億級模型採用 100 多層 Transformer 網路結構，計算複雜，訓練需要佔用 T 級視訊記憶體資源，如果想用更少的機器高效訓練，必須採取一系列效能最佳化和視訊記憶體最佳化措施。
首先看如何效能最佳化。我們透過一個公式來看哪些因素可以影響訓練速度，在固定的硬體環境下：
總訓練速度 ∝ 單卡速度 * 卡數 * 多卡加速比
其中單卡速度由資料讀取和計算速度決定；多卡加速比由計算 / 通訊效率決定。顯而易見，這三個是關鍵因素。除了單卡可以使用的運算元融合、混合精度之類的基礎效能最佳化策略之外，分散式訓練還引入一系列並行策略。並行策略的核心思想是將資料和計算有關的圖 / 運算元切分到不同裝置上，同時儘可能降低裝置間通訊所需的代價，合理使用多臺裝置資源，實現高效的併發排程訓練，最大化提升訓練速度。常見並行策略有資料並行 DP（Data Parallel）、Layer 間並行（流水線並行 PP，Pipeline Parallel）、Layer 內並行（模型並行 MP，Model Parallel）。如下表所示。我們從裝置資源和計算 / 通訊效率來分析三種策略的優缺點：
資料並行訓練加速比最高，但要求每個裝置上都備份一份模型，視訊記憶體佔用比較高。為此我們的改進方案是分組引數切片資料並行策略（具體原理後文介紹），相容了 MP+DP 的優勢，但缺點是通訊量大。
模型並行，通訊佔比高，適合在機器內做模型並行且支援的模型型別有限。
流水線並行，訓練裝置容易出現空閒狀態，加速效率沒有 DP 高；但能減少通訊邊界支援更多的層數，適合在機器間使用。
其次看視訊記憶體問題，透過下表分析的視訊記憶體佔用來源可以看出，上述的並行策略同樣可以很好地應對不同來源的視訊記憶體佔用，更多的層數可以透過流水線並行和分組引數切分策略來解決；某層引數很大可以透過模型並行來解決；其次飛槳還提供一些其它靈活的最佳化方式，例如每層輸出佔用的視訊記憶體，可以透過重計算和 Offload 來解決。
綜上所述，針對效能最佳化和視訊記憶體最佳化，幾種並行策略都有用武之地，但是同時也有各自的侷限性，所以如果想高效訓練千億模型，需要這幾種策略相互組合，取長補短，發揮各自的優勢。
那麼如何組合呢？飛槳研發人員首先在單機內使用模型並行和分組引數切片組合的 2D 策略，這麼選擇的原因是這兩個策略通訊量較大，適合使用機器內的卡間通訊；然後為了承載千億規模模型，再疊加流水線並行策略，使用多臺機器共同分擔；最後為了做到高效，在外層又疊加了資料並行來增加併發數量，提升整體訓練速度。這樣業內首個 4D 混合並行策略就誕生了。
圖 8 4D 混合並行策略示意圖
下面我們們再來簡單介紹下幾個並行策略的原理。
模型並行策略指的是將某一層網路切成多份，並分給不同的卡平行計算，每張卡僅需要計算部分結果。對於 ERNIE 中的 Transformer 網路結構，模型並行就可以對全連線層 FC 切分，然後透過通訊操作合併計算結果【2】。
流水線並行策略支援將模型的不同層放置到不同的裝置上，透過多個裝置來共同分擔視訊記憶體消耗，實現超大規模模型訓練。相鄰裝置間透過通訊鏈路傳輸資料。由於各個裝置間傳輸的僅是相鄰裝置間的輸出張量，因此通訊量較小，相對來說較為適合機器間通訊的場景。
值得注意的是，流水線並行可以說是廣義模型並行的一種特例，本文中的模型並行僅指 Tensor 切分，也就是會出現同一層網路交由不同卡來計算的情況，而流水線並行則是按照網路層的粒度切分。
圖 9 流水線並行策略示意圖
流水線並行策略本身也有很大的最佳化空間。如圖 10（a）所示，最佳化前，任意時刻只有單個計算裝置處於計算狀態，其它計算裝置則處於空閒狀態，這個空閒時間我們稱之為 Bubble 時間【3】。為了減少 Bubble 時間，如圖 10（b）所示，飛槳進一步將 mini-batch 切分成若干更小粒度的 micro-batch，每個裝置依次計算單個 micro-batch 的結果，從而增加了裝置間的併發度，降低了流水線並行 Bubble 時間比例。
此外飛槳研發人員經過對流水線並行訓練過程更加深入的剖析，發現還可以進一步最佳化視訊記憶體的利用率。採用如圖 10（c）中所示的方式，在一個 micro-batch 完成前向計算後，提前排程完成相應後向計算，這樣就能釋放部分視訊記憶體，用以接納新的資料，提升整體訓練效能。使用 ERNIE 模型實測，從 10(b)到 10(c)，總 BatchSize 可以提升 32 倍，效能可以提升 9 倍。
圖 10 流水線並行時序圖
最後再來看下飛槳的分組引數切片策略，其特色是在使用引數切片方式節省視訊記憶體的基礎上，與資料並行策略組合成更加強大的 Sharding-DP 策略。簡而言之，這種組合後的策略擁有很強的靈活性，使用者可以根據現有硬體環境情況，自由設定模型引數切分的數量（sharding_degree）和資料並行的路數（dp_degree），僅需要確保 sharding_degree × dp_degree = 總卡數即可。
舉個例子，假設使用者有 4 臺單機四卡的機器（共 16 張卡），訓一個 16 層網路的模型。如果模型引數規模可以使用一臺機器承載，則推薦使用 dp_degree=4 & sharding_degree=4 的方式，如圖 11 所示。這種方式的優勢在於只有機器內卡間通訊，但是模型最大不能超過單臺機器所能承受儲存範圍。
圖 11 dp_degree=4 & sharding_degree=4 的 Sharding-DP 示意圖
如果模型大小超過了單臺機器，問題也不大，使用者可以靈活地選擇 dp_degree=2 & sharding_degree=8 方式，如圖 12 所示。與上一種方式相比，這種方式支援的模型引數規模翻倍。
圖 12 dp_degree=2 & sharding_degree=8 的 Sharding-DP 示意圖
但是在一些特殊的情況下，如果模型引數規模非常大，半數機器都無法承載，則可以進一步使用 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 機器的環境來驗證不同策略組合的訓練效果，而測試的物件，正是擁有 2300 億引數規模的“重量級”ERNIE 模型。測試後我們可以看到 4D 混合並行策略訓練速度高於其它兩種 3D 混合並行策略，達到了8698 tokens/s，至少可以提速 23.7%。
寫在最後
自飛槳設計之初就開始潛心研究分散式訓練技術以應對大規模引數模型的訓練任務。在豐富的搜尋推薦業務的驅動下，飛槳分散式訓練引數伺服器模式歷經三代。最早的純 CPU 引數伺服器就已經可以訓練萬億級規模的稀疏引數模型。其後隨著業務的需要以及前沿技術的發展，產生了計算能力更強的純 GPU 引數伺服器模式。最近新推出的業內首創的異構引數伺服器模式，支援場景更多而且可以極大地提升硬體資源利用效率。對於大規模稠密引數模型，飛槳分散式訓練技術同樣與業務緊密結合，其集合通訊模式透過最新的 4D 混合並行策略，支援了 2300 億引數規模的文心 ERNIE 模型的分散式訓練。如今飛槳已經開始研究下一代分散式技術，來同時相容超大規模稠密引數和稀疏引數模型的訓練。相信在實際產業應用這個核心驅動力推動下，飛槳分散式訓練必將成為星辰大海上的那顆北極星，為廣大開發者們指引航向。