0
編譯 | 琰琰
大規(guī)模深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練仍是一項(xiàng)艱巨的挑戰(zhàn)，因?yàn)閯?dòng)輒百億、千億參數(shù)量的語言模型，需要更多的 GPU 內(nèi)存和時(shí)間周期。這篇文章從如何多GPU訓(xùn)練大模型的角度，回顧了現(xiàn)有的并行訓(xùn)練范式，以及主流的模型架構(gòu)和內(nèi)存優(yōu)化設(shè)計(jì)方法。
本文作者Lilian Weng現(xiàn)為OpenAI應(yīng)用人工智能研究負(fù)責(zé)人，主要從事機(jī)器學(xué)習(xí)、深度學(xué)習(xí)和網(wǎng)絡(luò)科學(xué)研究 。她本科畢業(yè)于香港大學(xué)，碩士就讀于北京大學(xué)信息系統(tǒng)與計(jì)算機(jī)科學(xué)系，之后前往印度安納大學(xué)布魯頓分校攻讀博士。
Lilian Weng經(jīng)常在個(gè)人博客分享學(xué)習(xí)和工作筆記，感興趣的可以戳這里：
https://lilianweng.github.io/lil-log/。
“煉大模型”已成為人工智能領(lǐng)域的主流研發(fā)趨勢(shì)。從GPT-3的1750億，到如今悟道2.0的1.75萬億，超大語言模型在 NLP 基準(zhǔn)任務(wù)中不斷刷新SOTA。
然而，參數(shù)和數(shù)據(jù)集的快速增長讓 GPU 算力開始捉襟見肘。單個(gè)GPU內(nèi)存已經(jīng)遠(yuǎn)遠(yuǎn)不能滿足大模型的需求。如，阿里用480塊GPU訓(xùn)練千億模型；英偉達(dá)用3072塊GPU訓(xùn)練萬億模型；谷歌用2048塊TPU訓(xùn)練1.6萬億模型（1 TPU約等于2～3 GPU）。
如何利用上百塊GPU上訓(xùn)練大規(guī)模語言模型？并行計(jì)算是一種行之有效的方法。
近日，OpenAI 研究員Lilian Weng分享干貨文章，從并行訓(xùn)練（數(shù)據(jù)并行、模型并行、管道并行、張量并行）、混合專家、內(nèi)存節(jié)省設(shè)計(jì)（CPU卸載、重新激活計(jì)算、混合精度訓(xùn)練、高效存儲(chǔ)優(yōu)化器）三個(gè)方面回顧了現(xiàn)階段多GPU訓(xùn)練大模型的主流方法。
AI科技評(píng)論這篇文章編譯如下，供想成為“煉丹師”的朋友參考。
并行訓(xùn)練
大規(guī)模神經(jīng)網(wǎng)絡(luò)模型對(duì)存儲(chǔ)空間有強(qiáng)烈的需求，單個(gè)GPU的內(nèi)存已經(jīng)遠(yuǎn)遠(yuǎn)不夠。究其原因，一方面模型權(quán)重有數(shù)百億個(gè)浮點(diǎn)數(shù)，隨機(jī)梯度下降和Adam優(yōu)化需要極高的計(jì)算成本；另一方面在預(yù)訓(xùn)練階段，大模型與大規(guī)模語料庫的配對(duì)需要很長的時(shí)間周期。綜合來看，跨GPU并行計(jì)算顯得尤為重要。
并行計(jì)算在數(shù)據(jù)、模型架構(gòu)和張量等不同維度上都可以操作，接下來本文將具體介紹一些主流方法：
數(shù)據(jù)并行
數(shù)據(jù)并行（ Data parallelism ，DP）最簡單的方法是將相同的模型權(quán)重復(fù)制到worker節(jié)點(diǎn)，并分配一部分?jǐn)?shù)據(jù)以同時(shí)進(jìn)行處理。我們知道，如果模型的參數(shù)量大于單個(gè)GPU節(jié)點(diǎn)的內(nèi)存，DP無法正常工作，GeePS架構(gòu)（Cui等人，2016）的解決思路是使用有限的GPU內(nèi)存。也就是，如果模型太大無法嵌入到一臺(tái)機(jī)器，就將暫時(shí)未使用的參數(shù)卸載回CPU。
數(shù)據(jù)交換傳輸通常在后端進(jìn)行（不干擾訓(xùn)練計(jì)算），在每個(gè)Mini-batch計(jì)算結(jié)束后，worker需要同步梯度或權(quán)重，以保證學(xué)習(xí)效率。現(xiàn)有的同步方法有兩種，各自優(yōu)缺點(diǎn)如下：
1、批量同步并行（BSP）：worker在每個(gè)Mini-batch結(jié)束時(shí)同步數(shù)據(jù)，這種方法保證了模型權(quán)重傳遞的及時(shí)性，但每臺(tái)機(jī)器都必須排隊(duì)等待其他機(jī)器發(fā)送梯度。
2、異步并行（ASP）：每個(gè)GPU采用異步方式處理數(shù)據(jù)，這種方法避免了不同機(jī)器之間的相互等待或暫停，但影響了權(quán)重傳遞的時(shí)效，降低了統(tǒng)計(jì)學(xué)習(xí)效率。而且即使增加計(jì)算時(shí)長，也不會(huì)加快訓(xùn)練的收斂速度。
在中間某些地方的每一次迭代（＞1）都需要同步全局梯度。自Pytorch v1.5版（Li等人，2021年）提出后，該特征在分布式數(shù)據(jù)并行（Distribution Data Parallel，DDP）中被稱為“梯度累積（gradient accumulation）”。分桶梯度(bucketing gradients)避免立即執(zhí)行AllReduce操作，而是將多個(gè)梯度存儲(chǔ)到一個(gè)AllReduce中以提高吞吐量，并基于計(jì)算圖優(yōu)化計(jì)算和通信調(diào)度。
圖1：Pytorch DDP的偽代碼（來源：Li等人，2021年）
模型并行
模型并行（Model parallelism，MP）用于解決模型權(quán)重不能適應(yīng)單個(gè)節(jié)點(diǎn)的情況，在這里，計(jì)算和模型參數(shù)都需要跨多臺(tái)機(jī)器進(jìn)行處理。在數(shù)據(jù)并行中，每個(gè)worker承載著整個(gè)模型的完整副本，而MP只在一個(gè)worker上分配部分模型參數(shù)，因此對(duì)內(nèi)存和計(jì)算的需求要小很多。
深度神經(jīng)網(wǎng)絡(luò)包含一堆垂直層，如果逐層拆分將連續(xù)的小層分配到工作層分區(qū)，操作起來并不難，但通過大量具有順序依賴性的Workers來運(yùn)行每個(gè)數(shù)據(jù)batch會(huì)花費(fèi)大量的時(shí)間，計(jì)算資源的利用率也嚴(yán)重不足。
圖2：一個(gè)包含4個(gè)垂直層的模型并行設(shè)置，由于順序的依賴性，每個(gè)數(shù)據(jù)依次由一個(gè)worker處理，這個(gè)過程會(huì)出現(xiàn)大量多余時(shí)間“氣泡”（來源：Huang等人，2019年）
管道并行
管道并行（Pipeline parallelism，PP）是將模型并行與數(shù)據(jù)并行結(jié)合起來，以減少低效時(shí)間“氣泡”的過程。主要思想是將Mini-batch拆分為更多個(gè)微批次（microbatch），并使每個(gè)階段worker能夠同時(shí)處理。需要注意的是，每個(gè)微批次需要兩次傳遞，一次向前，一次向后。worker分區(qū)的數(shù)量稱為管道深度，不同worker分區(qū)之間的通信僅傳輸激活（向前）和梯度（向后）。這些通道的調(diào)度方式以及梯度的聚合方式在不同的方法中有所不同。
在GPipe（Huang et al.2019）方法中，多個(gè)微批次處理結(jié)束時(shí)會(huì)同時(shí)聚合梯度和應(yīng)用。同步梯度下降保證了學(xué)習(xí)的一致性和效率，與worker數(shù)量無關(guān)。如圖3所示，“氣泡”仍然存在，但比圖2少了很多。給定m個(gè)均勻分割的微批次和d個(gè)分區(qū)，假設(shè)每個(gè)微批次向前和向后都需要一個(gè)時(shí)間單位，則氣泡的分?jǐn)?shù)為：
圖3：帶有4個(gè)微批次 和4個(gè)分區(qū)的GPipe的并行管道（來源：Huang等人，2019年）
GPipe在吞吐量方面與設(shè)備數(shù)量成線性關(guān)系，設(shè)備數(shù)量越多，吞吐量越大。不過，如果模型參數(shù)在worker中分布不均勻，這種線性關(guān)系不會(huì)穩(wěn)定出現(xiàn)。
PipeDream（Narayanan等人，2019年）方法要求每個(gè)worker交替處理向前和向后傳遞的消息（1F1B）。它將每個(gè)模型分區(qū)命名為“stage”，每個(gè)stage worker可以有多個(gè)副本來并行運(yùn)行數(shù)據(jù)。這個(gè)過程使用循環(huán)負(fù)載平衡策略在多個(gè)副本之間分配工作，以確保相同minibatch 向前和向后的傳遞發(fā)生在同一副本上。
圖4：PipeDream中1F1B微批次調(diào)度的圖示（來源：Harlap等人，2018年）
由于PipeDream沒有在所有worker batch結(jié)束時(shí)同步全局梯度，1F1B 很容易導(dǎo)致不同版本的模型權(quán)重的微批次向前和向后傳遞，降低學(xué)習(xí)效率。對(duì)此，PipeDream提供了一些解決的思路：
權(quán)重存儲(chǔ)：每個(gè)worker跟蹤多個(gè)模型版本，給定數(shù)據(jù) batch 的向前和向后傳遞相同版本的權(quán)重。
垂直同步：不同模型權(quán)重版本與激活和梯度一起在全局worker之間傳遞，計(jì)算采用上一個(gè)worker傳播的相對(duì)應(yīng)的隱藏版本。這個(gè)過程確保了worker之間的版本一致性（不同于GPipe，采用異步計(jì)算）。
在訓(xùn)練開始時(shí)，PipeDream會(huì)先分析模型每一層的計(jì)算內(nèi)存和時(shí)間成本，然后將層劃分為不同的stage進(jìn)行優(yōu)化。
圖5：上圖為VGG16在ILSVRC12上的運(yùn)行結(jié)果，ASP=異步并行，BSP=批量同步并行；下圖為不同并行配置的訓(xùn)練時(shí)間加速度（來源：Harlap等人，2018年）
后來有學(xué)者提出了PipeDream兩種變體，主要思路是通過減少模型版本來緩解內(nèi)存占用（Narayanan等人，2021年）。其中，PipeDream-flush增加了定期刷新全局同步管道的功能，就像GPipe一樣，這種方式雖然犧牲了一點(diǎn)吞吐量，但顯著減少了內(nèi)存占用（例如僅需要維護(hù)單一版本的模型權(quán)重）。
圖6：PipeDream flush中管道調(diào)度圖示（來源：Narayanan等人，2021年）
PipeDream-2BW維護(hù)兩個(gè)版本的模型權(quán)重，“2BW”代表“雙緩沖權(quán)重（double-buffered weights）”，它會(huì)在每個(gè)微批次生成一個(gè)新的模型版本K（K>d）。由于一些剩余的向后傳遞仍然依賴于舊版本，新的模型版本無法立即取代舊版本，但因?yàn)橹槐４媪藘蓚€(gè)版本，內(nèi)存占用的也被大大降低了。
圖7：PipeDream-2BW 中的流水線調(diào)度示意圖（來源：Narayanan et al. 2021
張量并行
模型并行和管道并行都會(huì)垂直拆分模型，而張量并行（Tensor Parallelism，TP）是將張量運(yùn)算的計(jì)算水平劃分到多個(gè)設(shè)備上。
以Transformer為例。Transformer架構(gòu)主要由多層MLP和自注意力塊組成。Megatron-LM（Shoeybi et al.2020）采用了一種簡單的方法來并行計(jì)算層內(nèi)MLP和自注意力。
MLP層包含GEMM（通用矩陣乘法）和非線性GeLU傳輸。如果按列拆分權(quán)重矩陣A，可以得到：
注意力塊根據(jù)上述分區(qū)并行運(yùn)行GEMM的 查詢（Q）、鍵（K）和 權(quán)重（V），然后與另一個(gè)GEMM組合以生成頭注意力結(jié)果。
圖8：Megatron-LM中提出的關(guān)鍵Transformer組件的張量平行性說明。（來源：Shoeybi等人，2020年）
今年Narayanan等人將管道、張量和數(shù)據(jù)并行與新的管道調(diào)度策略相結(jié)合，提出了一種名為PTD-P的新方法。該方法不僅在設(shè)備上能夠定位一組連續(xù)的層（“模型塊”），該可以為每個(gè)wokers分配多個(gè)較小的連續(xù)層子集塊（例如，設(shè)備1具有第1、2、9、10層；設(shè)備2具有第3、4、11、12層；每個(gè)具有兩個(gè)模型塊）
每個(gè)batch中，微批次的數(shù)量應(yīng)精確除以wokers數(shù)量（mm）。如果每個(gè)worker有v個(gè)模型塊，那么與GPipe調(diào)度相比，管道的“氣泡”時(shí)間可以減少 v 倍。
圖9：上圖與PipeDream flush中的默認(rèn)1F1B管道明細(xì)表相同；下圖為交錯(cuò)的1F1B管線一覽表（來源：Narayanan等人，202）
混合專家（MoE）
為了突破模型大小的限制，谷歌后來提出一種混合專家（MoE）方法，其核心理念是：綜合學(xué)習(xí)，它假設(shè)多個(gè)弱學(xué)習(xí)者組合起來就會(huì)擁有一個(gè)強(qiáng)學(xué)習(xí)者。
在深度神經(jīng)網(wǎng)絡(luò)中，混合專家（MoE）通過連接多個(gè)專家的門機(jī)制（gating mechanism）實(shí)現(xiàn)集成（Shazeer等人，2017）。門機(jī)制激活不同網(wǎng)絡(luò)的專家以產(chǎn)生不同的輸出。作者在論文將其命名為“稀疏門控專家混合層（sparsely gated MoE）”。
僅一個(gè)MoE層包含：（1）前饋網(wǎng)絡(luò)專家n；（2）可訓(xùn)練的門控網(wǎng)絡(luò)G，通過學(xué)習(xí)n個(gè)專家的概率分布，將流量路由到幾個(gè)特定的專家。
根據(jù)門控輸出，并非每個(gè)專家都必須進(jìn)行評(píng)估。當(dāng)專家的數(shù)量太大時(shí)，可以考慮使用兩層MoE。
圖10：專家混合（MoE）層的圖示，門控網(wǎng)絡(luò)只選擇并激活了n個(gè)專家中的2個(gè)（來源：Shazeer等人，2017年）G將輸入與可訓(xùn)練權(quán)重矩陣Gg相乘，然后執(zhí)行softmax：
由于這個(gè)過程會(huì)產(chǎn)生密集的門控制向量，不利于節(jié)省計(jì)算資源，而且
時(shí)也不需要評(píng)估專家。所以，MoE層僅保留了頂部k值，并通過向G中添加高斯噪聲改進(jìn)負(fù)載平衡，這種機(jī)制被稱為噪聲top-k門。
其中，CV是變異系數(shù)，失重的waux是可調(diào)節(jié)的超參數(shù)。由于每個(gè)專家網(wǎng)絡(luò)只能獲得小部分訓(xùn)練樣本（“收縮批次問題”），所以在MoE中應(yīng)該盡可能使用大batch，但這又會(huì)受到GPU內(nèi)存的限制。數(shù)據(jù)并行和模型并行的應(yīng)用可以提高模型的吞吐量。
圖11：10億單詞的語言建?；鶞?zhǔn)（左）模型從左到右包含4、32、256、256、1024和4096名專家（右）40億參數(shù)的MoE模型在不同計(jì)算預(yù)算下的性能（來源：Shazeer等人，2017年）
GShard（Lepikhin等人，2020年）通過自動(dòng)分片將MoE transformer 模型的參數(shù)擴(kuò)展到了6000億。MoE transformer 用MoE層取代其他每一個(gè)前饋網(wǎng)絡(luò)層。需要說明的是，在多臺(tái)機(jī)器上MoE transformer 僅在MoE層分片，其他層只是復(fù)制。
GShard中的門控功能G有幾種改進(jìn)設(shè)計(jì)：
專家容量：通過一位專家的令牌數(shù)量不應(yīng)超過“專家容量”的閾值。如果令牌被路由到已達(dá)到容量的專家，則令牌將被標(biāo)記為“溢出”，并且門輸出將更改為零向量。
本地組調(diào)度：令牌被均勻地劃分為多個(gè)本地組，專家能力在組水平上得到加強(qiáng)。
輔助損失：與原始MoE aux損失相似，添加輔助損失可以最小化路由到每個(gè)專家的數(shù)據(jù)的均方。
隨機(jī)路由：以與其權(quán)重成比例的概率選擇第二位最佳專家；否則GShard遵循隨機(jī)路由，增加隨機(jī)性。
圖12：GShard中帶輔助損失的組水平top-2門機(jī)制的偽代碼（來源：Lepikhin等人，2020年）
Switch Transformer（Fedus et al.2021）用稀疏開關(guān)FFN層取代了密集前饋層（每個(gè)輸入僅路由到一個(gè)專家網(wǎng)絡(luò)），將模型規(guī)模擴(kuò)展到數(shù)萬億個(gè)參數(shù)。負(fù)載平衡的輔助損失是
，給定n個(gè)專家，fi是路由到第i個(gè)專家的令牌分?jǐn)?shù)，pi是門控網(wǎng)絡(luò)預(yù)測的專家i的路由概率。
圖13：Switch transformer，稀疏Switch FFN層位于藍(lán)色框（來源：Fedus等人，2021年）
為提高訓(xùn)練穩(wěn)定性，switch transformer采用以下設(shè)計(jì)：
選擇精度：使用FP32精度以提高模型局部的穩(wěn)定性，并降低FP32張量的通信成本。FP32精度僅在路由器功能主體內(nèi)使用，結(jié)果將還原到FP16。 較小的初始化：權(quán)重矩陣的初始化從平均μ=0且標(biāo)準(zhǔn)偏差
的正態(tài)分布中采樣，同時(shí)將Transformer初始化參數(shù)從s=1減小到s=0.1 。
使用更高的專家輟學(xué)率：微調(diào)通常適用于小數(shù)據(jù)集，增加每個(gè)專家的輟學(xué)率以避免過度擬合。他們發(fā)現(xiàn)，所有層中的輟學(xué)率增加會(huì)導(dǎo)致性能下降。在論文中，他們?cè)诜菍＜覍又惺褂昧?.1的輟學(xué)率，但在專家FF層中使用了0.4的輟學(xué)率。
switch transformer論文總結(jié)了用于訓(xùn)練大型模型的不同數(shù)據(jù)和模型并行策略，并給出了一個(gè)很好的示例：
圖14：第一行為如何在多個(gè)GPU內(nèi)核拆分模型權(quán)重（頂部），每種顏色代表一個(gè)權(quán)重矩陣；第二行為各種數(shù)據(jù)并行策略的說明，不同顏色表示不同的標(biāo)記集（來源：Fedus等人，2021年）
其他節(jié)省內(nèi)存的設(shè)計(jì)
CPU卸載
如果GPU內(nèi)存已滿，可以將暫時(shí)未使用的數(shù)據(jù)卸載到CPU，并在以后需要時(shí)將其讀回（Rhu等人，2016）。不過，這種方法近年來并不太流行，因?yàn)?/span>它會(huì)延長模型訓(xùn)練的時(shí)間。
激活重新計(jì)算
激活重新計(jì)算，也稱“激活檢查點(diǎn)”或“梯度檢查點(diǎn)”（Chen et al，2016），其核心思路是犧牲計(jì)算時(shí)間來換取內(nèi)存空間。它減少了訓(xùn)練 ? 層深層神經(jīng)網(wǎng)絡(luò)到
的內(nèi)存開銷，每個(gè)batch只消耗額外的前向傳遞計(jì)算。
具體來說，該方法將?層網(wǎng)絡(luò)平均劃分為d個(gè)分區(qū)，僅保存分區(qū)邊界的激活，并在workers之間進(jìn)行通信。計(jì)算梯度仍然需要在分區(qū)內(nèi)層進(jìn)行中間激活，以便在向后過程中重新計(jì)算梯度。在激活重新計(jì)算的情況下，用于訓(xùn)練M(?) 是：
它的最低成本是：
激活重新計(jì)算的方法可以得出與模型大小有關(guān)次線性內(nèi)存開銷，如下圖：
圖15：不同節(jié)省內(nèi)存算法的內(nèi)存開銷。sharing：中間結(jié)果使用的內(nèi)存在不再需要時(shí)被回收。inplace：將輸出直接保存到輸入值的內(nèi)存中（來源：Chen等人，2016）混合精度訓(xùn)練
此前，Narang（Narang&Micikevicius等人，2018年）介紹了一種使用半精度浮點(diǎn)（FP16）數(shù)訓(xùn)練模型而不損失模型精度的方法。
圖16：一層混合精度訓(xùn)練程序（來源：Narang&Micikevicius等人，2018年）
其中涉及三種關(guān)鍵技術(shù)：
全精度權(quán)重復(fù)制：保持累積梯度的模型權(quán)重的全精度（FP32）復(fù)制。對(duì)于向前和向后的傳遞的信息做四舍五入至半精度處理，因?yàn)槊看翁荻雀拢刺荻萖學(xué)習(xí)率）太小，可能無法完全包含在FP16范圍內(nèi)。
縮放損失：放大損失以更好地處理小幅度的梯度（見圖16），放大梯度以使其向可表示范圍的右側(cè)部分（包含較大的值）移動(dòng)，從而保留可能丟失的值。
算術(shù)精度：對(duì)于常見的網(wǎng)絡(luò)算法（如矢量點(diǎn)積、矢量元素求和歸約），將部分結(jié)果累加到FP32中，然后輸出保存為FP16。逐點(diǎn)操作可以在FP16或FP32中執(zhí)行。
圖17：:全精確的梯度直方圖
在這項(xiàng)實(shí)驗(yàn)中，圖像分類、更快的R-CNN等不需要損失縮放，但其他網(wǎng)絡(luò)，如多盒SSD、大LSTM語言模型是需要損失縮放的。
壓縮（Compression）
模型權(quán)重在向前和向后傳遞的過程中會(huì)消耗大量內(nèi)存。考慮到這兩種傳遞方式會(huì)花費(fèi)大量時(shí)間，2018年Jain （Jain et al，2018）提出了一種數(shù)據(jù)編碼策略，即在第一次傳遞后壓縮中間結(jié)果，然后將其解碼用于反向傳播。
Jain和團(tuán)隊(duì)研發(fā)的Gist系統(tǒng)包含兩種編碼方案：一是特定于層的無損編碼，包括 ReLU-Pool和 ReLU-Conv模式；二是有攻擊性的有損編碼，主要使用延遲精度縮減（DPR）。需要注意的是，第一次使用特征圖時(shí)應(yīng)保持高精度，第二次使用時(shí)要適度降低精度。這項(xiàng)實(shí)驗(yàn)表明，Gist可以在5個(gè)最佳圖像分類DNN上減少2倍的內(nèi)存開銷，平均減少1.8倍，性能開銷僅為4%。
內(nèi)存高效優(yōu)化器
優(yōu)化器也會(huì)消耗內(nèi)存。以主流的Adam優(yōu)化器為例，其內(nèi)部需要維護(hù)動(dòng)量和方差，這兩者與梯度和模型參數(shù)比例基本相同。這意味著，我們需要節(jié)省4倍模型權(quán)重的內(nèi)存。
為了減少內(nèi)存消耗，學(xué)術(shù)界已經(jīng)提出了幾款主流優(yōu)化器。與Adam相比，Adafactor（Shazeer et al.2018）優(yōu)化器沒有存儲(chǔ)全部動(dòng)量和變化，只跟蹤移動(dòng)平均數(shù)的每行和每列總和，然后根據(jù)這些總和估計(jì)二階矩。
SM3（Anil et al.2019）優(yōu)化器采用了一種不同的自適應(yīng)優(yōu)化方法。
ZeRO（Rajbhandari et al.2019）零冗余優(yōu)化器節(jié)省了大型模型訓(xùn)練在兩方面的內(nèi)存消耗：
大多數(shù)內(nèi)存由模型狀態(tài)消耗，包括優(yōu)化器狀態(tài)（例如Adam動(dòng)量和方差）、梯度和參數(shù)。混合精度訓(xùn)練也需要大量內(nèi)存，因?yàn)槌薋P16版本之外，優(yōu)化器還需要保存FP32參數(shù)和其他優(yōu)化器狀態(tài)的副本。
未被激活、臨時(shí)緩沖區(qū)以及不可用的碎片內(nèi)存消耗（論文中稱為剩余狀態(tài)）。
ZeRO結(jié)合了ZeRO-DP和ZeRO-R兩種方法。ZeRO-DP是一種增強(qiáng)的數(shù)據(jù)并行，避免了模型狀態(tài)的簡單冗余。它以動(dòng)態(tài)的方式跨多個(gè)并行數(shù)據(jù)劃分優(yōu)化器狀態(tài)、梯度和參數(shù)，以最小化通信量。ZeRO-R使用分區(qū)激活二次計(jì)算、恒定緩沖區(qū)大小和動(dòng)態(tài)內(nèi)存碎片，以優(yōu)化剩余狀態(tài)的內(nèi)存消耗。
參考資料：
[1] Li et al. “PyTorch Distributed: Experiences on Accelerating Data Parallel Training” VLDB 2020.
[2] Cui et al. “GeePS: Scalable deep learning on distributed GPUs with a GPU-specialized parameter server” EuroSys 2016
[3] Shoeybi et al. “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.” arXiv preprint arXiv:1909.08053 (2019).
[4] Narayanan et al. “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM.” arXiv preprint arXiv:2104.04473 (2021).
[5] Huang et al. “GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.” arXiv preprint arXiv:1811.06965 (2018).
[6] Narayanan et al. “PipeDream: Generalized Pipeline Parallelism for DNN Training.” SOSP 2019.
[7] Narayanan et al. “Memory-Efficient Pipeline-Parallel DNN Training.” ICML 2021.
[8] Shazeer et al. “The Sparsely-Gated Mixture-of-Experts Layer Noam.” arXiv preprint arXiv:1701.06538 (2017).
[9] Lepikhin et al. “GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.” arXiv preprint arXiv:2006.16668 (2020).
[10] Fedus et al. “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” arXiv preprint arXiv:2101.03961 (2021).
[11] Narang & Micikevicius, et al. “Mixed precision training.” ICLR 2018.
[12] Chen et al. 2016 “Training Deep Nets with Sublinear Memory Cost.” arXiv preprint arXiv:1604.06174 (2016).
[13] Jain et al. “Gist: Efficient data encoding for deep neural network training.” ISCA 2018.
[14] Shazeer & Stern. “Adafactor: Adaptive learning rates with sublinear memory cost.” arXiv preprint arXiv:1804.04235 (2018).
[15] Anil et al. “Memory-Efficient Adaptive Optimization.” arXiv preprint arXiv:1901.11150 (2019).
[16] Rajbhandari et al. “ZeRO: Memory Optimization Towards Training A Trillion Parameter Models Samyam.” arXiv preprint arXiv:1910.02054 (2019).
編譯鏈接：https://lilianweng.github.io/lil-log/2021/09/24/train-large-neural-networks.html
雷鋒網(wǎng)雷鋒網(wǎng)雷鋒網(wǎng)
雷峰網(wǎng)版權(quán)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。
