飛槳分布式訓(xùn)練又推新品，4D混合并行可訓(xùn)千億級(jí)AI模型

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

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

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

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

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

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

圖 2 大規(guī)模訓(xùn)練場(chǎng)景

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

圖 7 集合通信訓(xùn)練

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

測(cè)試驗(yàn)證

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

寫在最后

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