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