北大校友「煉丹」分享：OpenAI如何訓練千億級模型？

• 2021 年 10 月 9 日
• AI

編譯 | 琰琰

大規模深度神經網絡訓練仍是一項艱巨的挑戰，因為動輒百億、千億參數量的語言模型，需要更多的 GPU 內存和時間周期。這篇文章從如何多GPU訓練大模型的角度，回顧了現有的並行訓練範式，以及主流的模型架構和內存優化設計方法。

本文作者Lilian Weng現為OpenAI應用人工智能研究負責人，主要從事機器學習、深度學習和網絡科學研究 。她本科畢業於香港大學，碩士就讀於北京大學信息系統與計算機科學系，之後前往印度安納大學布魯頓分校攻讀博士。

Lilian Weng經常在個人博客分享學習和工作筆記，感興趣的可以戳這裡：
//lilianweng.github.io/lil-log/。

「煉大模型」已成為人工智能領域的主流研發趨勢。從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上訓練大規模語言模型？並行計算是一種行之有效的方法。

近日，OpenAI 研究員Lilian Weng分享乾貨文章，從並行訓練（數據並行、模型並行、管道並行、張量並行）、混合專家、內存節省設計（CPU卸載、重新激活計算、混合精度訓練、高效存儲優化器）三個方面回顧了現階段多GPU訓練大模型的主流方法。

AI科技評論這篇文章編譯如下，供想成為「煉丹師」的朋友參考。

大規模神經網絡模型對存儲空間有強烈的需求，單個GPU的內存已經遠遠不夠。究其原因，一方面模型權重有數百億個浮點數，隨機梯度下降和Adam優化需要極高的計算成本；另一方面在預訓練階段，大模型與大規模語料庫的配對需要很長的時間周期。綜合來看，跨GPU並行計算顯得尤為重要。

並行計算在數據、模型架構和張量等不同維度上都可以操作，接下來本文將具體介紹一些主流方法：

數據並行 

數據並行（ Data parallelism ，DP）最簡單的方法是將相同的模型權重複制到worker節點，並分配一部分數據以同時進行處理。我們知道，如果模型的參數量大於單個GPU節點的內存，DP無法正常工作，GeePS架構（Cui等人，2016）的解決思路是使用有限的GPU內存。也就是，如果模型太大無法嵌入到一台機器，就將暫時未使用的參數卸載回CPU。

數據交換傳輸通常在後端進行（不干擾訓練計算），在每個Mini-batch計算結束後，worker需要同步梯度或權重，以保證學習效率。現有的同步方法有兩種，各自優缺點如下：

1、批量同步並行（BSP）：worker在每個Mini-batch結束時同步數據，這種方法保證了模型權重傳遞的及時性，但每台機器都必須排隊等待其他機器發送梯度。

2、異步並行（ASP）：每個GPU採用異步方式處理數據，這種方法避免了不同機器之間的相互等待或暫停，但影響了權重傳遞的時效，降低了統計學習效率。而且即使增加計算時長，也不會加快訓練的收斂速度。 

在中間某些地方的每一次迭代（＞1）都需要同步全局梯度。自Pytorch v1.5版（Li等人，2021年）提出後，該特徵在分佈式數據並行（Distribution Data Parallel，DDP）中被稱為「梯度累積（gradient accumulation）」。分桶梯度(bucketing gradients)避免立即執行AllReduce操作，而是將多個梯度存儲到一個AllReduce中以提高吞吐量，並基於計算圖優化計算和通信調度。

圖1：Pytorch DDP的偽代碼（來源：Li等人，2021年）

模型並行 

模型並行（Model parallelism，MP）用於解決模型權重不能適應單個節點的情況，在這裡，計算和模型參數都需要跨多台機器進行處理。在數據並行中，每個worker承載着整個模型的完整副本，而MP只在一個worker上分配部分模型參數，因此對內存和計算的需求要小很多。

 深度神經網絡包含一堆垂直層，如果逐層拆分將連續的小層分配到工作層分區，操作起來並不難，但通過大量具有順序依賴性的Workers來運行每個數據batch會花費大量的時間，計算資源的利用率也嚴重不足。

圖2：一個包含4個垂直層的模型並行設置，由於順序的依賴性，每個數據依次由一個worker處理，這個過程會出現大量多餘時間「氣泡」（來源：Huang等人，2019年）

管道並行

管道並行（Pipeline parallelism，PP）是將模型並行與數據並行結合起來，以減少低效時間「氣泡」的過程。主要思想是將Mini-batch拆分為更多個微批次（microbatch），並使每個階段worker能夠同時處理。需要注意的是，每個微批次需要兩次傳遞，一次向前，一次向後。worker分區的數量稱為管道深度，不同worker分區之間的通信僅傳輸激活（向前）和梯度（向後）。這些通道的調度方式以及梯度的聚合方式在不同的方法中有所不同。 

在GPipe（Huang et al.2019）方法中，多個微批次處理結束時會同時聚合梯度和應用。同步梯度下降保證了學習的一致性和效率，與worker數量無關。如圖3所示，「氣泡」仍然存在，但比圖2少了很多。給定m個均勻分割的微批次和d個分區，假設每個微批次向前和向後都需要一個時間單位，則氣泡的分數為：

 圖3：帶有4個微批次 和4個分區的GPipe的並行管道（來源：Huang等人，2019年）

 GPipe在吞吐量方面與設備數量成線性關係，設備數量越多，吞吐量越大。不過，如果模型參數在worker中分佈不均勻，這種線性關係不會穩定出現。 

PipeDream（Narayanan等人，2019年）方法要求每個worker交替處理向前和向後傳遞的消息（1F1B）。它將每個模型分區命名為「stage」，每個stage worker可以有多個副本來並行運行數據。這個過程使用循環負載平衡策略在多個副本之間分配工作，以確保相同minibatch 向前和向後的傳遞發生在同一副本上。 

圖4：PipeDream中1F1B微批次調度的圖示（來源：Harlap等人，2018年）

由於PipeDream沒有在所有worker batch結束時同步全局梯度，1F1B 很容易導致不同版本的模型權重的微批次向前和向後傳遞，降低學習效率。對此，PipeDream提供了一些解決的思路：

• 權重存儲：每個worker跟蹤多個模型版本，給定數據 batch 的向前和向後傳遞相同版本的權重。

• 垂直同步：不同模型權重版本與激活和梯度一起在全局worker之間傳遞，計算採用上一個worker傳播的相對應的隱藏版本。這個過程確保了worker之間的版本一致性（不同於GPipe，採用異步計算）。 

在訓練開始時，PipeDream會先分析模型每一層的計算內存和時間成本，然後將層劃分為不同的stage進行優化。

圖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，可以得到：

注意力塊根據上述分區並行運行GEMM的 查詢（Q）、鍵（K）和 權重（V），然後與另一個GEMM組合以生成頭注意力結果。

圖8：Megatron-LM中提出的關鍵Transformer組件的張量平行性說明。（來源：Shoeybi等人，2020年）

今年Narayanan等人將管道、張量和數據並行與新的管道調度策略相結合，提出了一種名為PTD-P的新方法。該方法不僅在設備上能夠定位一組連續的層（「模型塊」），該可以為每個wokers分配多個較小的連續層子集塊（例如，設備1具有第1、2、9、10層；設備2具有第3、4、11、12層；每個具有兩個模型塊）

每個batch中，微批次的數量應精確除以wokers數量（mm）。如果每個worker有v個模型塊，那麼與GPipe調度相比，管道的「氣泡」時間可以減少 v 倍。

圖9：上圖與PipeDream flush中的默認1F1B管道明細表相同；下圖為交錯的1F1B管線一覽表（來源：Narayanan等人，202）

為了突破模型大小的限制，谷歌后來提出一種混合專家（MoE）方法，其核心理念是：綜合學習，它假設多個弱學習者組合起來就會擁有一個強學習者。

在深度神經網絡中，混合專家（MoE）通過連接多個專家的門機制（gating mechanism）實現集成（Shazeer等人，2017）。門機制激活不同網絡的專家以產生不同的輸出。作者在論文將其命名為「稀疏門控專家混合層（sparsely gated MoE）」。

僅一個MoE層包含：（1）前饋網絡專家n；（2）可訓練的門控網絡G，通過學習n個專家的概率分佈，將流量路由到幾個特定的專家。 

根據門控輸出，並非每個專家都必須進行評估。當專家的數量太大時，可以考慮使用兩層MoE。

 圖10：專家混合（MoE）層的圖示，門控網絡只選擇並激活了n個專家中的2個（來源：Shazeer等人，2017年）

G將輸入與可訓練權重矩陣Gg相乘，然後執行softmax：

由於這個過程會產生密集的門控制向量，不利於節省計算資源，而且時也不需要評估專家。所以，MoE層僅保留了頂部k值，並通過向G中添加高斯噪聲改進負載平衡，這種機制被稱為噪聲top-k門。

其中，CV是變異係數，失重的waux是可調節的超參數。由於每個專家網絡只能獲得小部分訓練樣本（「收縮批次問題」），所以在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層取代了密集前饋層（每個輸入僅路由到一個專家網絡），將模型規模擴展到數萬億個參數。負載平衡的輔助損失是，給定n個專家，fi是路由到第i個專家的令牌分數，pi是門控網絡預測的專家i的路由概率。

‍‍ 圖13：Switch transformer，稀疏Switch FFN層位於藍色框（來源：Fedus等人，2021年）

為提高訓練穩定性，switch transformer採用以下設計：

• 選擇精度：使用FP32精度以提高模型局部的穩定性，並降低FP32張量的通信成本。FP32精度僅在路由器功能主體內使用，結果將還原到FP16。 較小的初始化：權重矩陣的初始化從平均μ=0且標準偏差的正態分佈中採樣，同時將Transformer初始化參數從s=1減小到s=0.1 。

• 使用更高的專家輟學率：微調通常適用於小數據集，增加每個專家的輟學率以避免過度擬合。他們發現，所有層中的輟學率增加會導致性能下降。在論文中，他們在非專家層中使用了0.1的輟學率，但在專家FF層中使用了0.4的輟學率。

switch transformer論文總結了用於訓練大型模型的不同數據和模型並行策略，並給出了一個很好的示例：

 圖14：第一行為如何在多個GPU內核拆分模型權重（頂部），每種顏色代表一個權重矩陣；第二行為各種數據並行策略的說明，不同顏色表示不同的標記集（來源：Fedus等人，2021年）

CPU卸載

如果GPU內存已滿，可以將暫時未使用的數據卸載到CPU，並在以後需要時將其讀回（Rhu等人，2016）。不過，這種方法近年來並不太流行，因為它會延長模型訓練的時間。 

激活重新計算 

激活重新計算，也稱「激活檢查點」或「梯度檢查點」（Chen et al，2016），其核心思路是犧牲計算時間來換取內存空間。它減少了訓練 ℓ 層深層神經網絡到的內存開銷，每個batch只消耗額外的前向傳遞計算。

具體來說，該方法將ℓ層網絡平均劃分為d個分區，僅保存分區邊界的激活，並在workers之間進行通信。計算梯度仍然需要在分區內層進行中間激活，以便在向後過程中重新計算梯度。在激活重新計算的情況下，用於訓練M(ℓ) 是：

它的最低成本是：

激活重新計算的方法可以得出與模型大小有關次線性內存開銷，如下圖：

 圖15：不同節省內存算法的內存開銷。sharing：中間結果使用的內存在不再需要時被回收。inplace：將輸出直接保存到輸入值的內存中（來源：Chen等人，2016）

混合精度訓練

此前，Narang（Narang&Micikevicius等人，2018年）介紹了一種使用半精度浮點（FP16）數訓練模型而不損失模型精度的方法。

圖16：一層混合精度訓練程序（來源：Narang&Micikevicius等人，2018年） 

其中涉及三種關鍵技術：

• 全精度權重複制：保持累積梯度的模型權重的全精度（FP32）複製。對於向前和向後的傳遞的信息做四捨五入至半精度處理，因為每次梯度更新（即梯度X學習率）太小，可能無法完全包含在FP16範圍內。

• 縮放損失：放大損失以更好地處理小幅度的梯度（見圖16），放大梯度以使其向可表示範圍的右側部分（包含較大的值）移動，從而保留可能丟失的值。

• 算術精度：對於常見的網絡算法（如矢量點積、矢量元素求和歸約），將部分結果累加到FP32中，然後輸出保存為FP16。逐點操作可以在FP16或FP32中執行。

圖17：:全精確的梯度直方圖

在這項實驗中，圖像分類、更快的R-CNN等不需要損失縮放，但其他網絡，如多盒SSD、大LSTM語言模型是需要損失縮放的。

壓縮（Compression） 

模型權重在向前和向後傳遞的過程中會消耗大量內存。考慮到這兩種傳遞方式會花費大量時間，2018年Jain （Jain et al，2018）提出了一種數據編碼策略，即在第一次傳遞後壓縮中間結果，然後將其解碼用於反向傳播。 

Jain和團隊研發的Gist系統包含兩種編碼方案：一是特定於層的無損編碼，包括 ReLU-Pool和 ReLU-Conv模式；二是有攻擊性的有損編碼，主要使用延遲精度縮減（DPR）。需要注意的是，第一次使用特徵圖時應保持高精度，第二次使用時要適度降低精度。這項實驗表明，Gist可以在5個最佳圖像分類DNN上減少2倍的內存開銷，平均減少1.8倍，性能開銷僅為4%。 

內存高效優化器

優化器也會消耗內存。以主流的Adam優化器為例，其內部需要維護動量和方差，這兩者與梯度和模型參數比例基本相同。這意味着，我們需要節省4倍模型權重的內存。 

為了減少內存消耗，學術界已經提出了幾款主流優化器。與Adam相比，Adafactor（Shazeer et al.2018）優化器沒有存儲全部動量和變化，只跟蹤移動平均數的每行和每列總和，然後根據這些總和估計二階矩。

SM3（Anil et al.2019）優化器採用了一種不同的自適應優化方法。

ZeRO（Rajbhandari et al.2019）零冗餘優化器節省了大型模型訓練在兩方面的內存消耗：

• 大多數內存由模型狀態消耗，包括優化器狀態（例如Adam動量和方差）、梯度和參數。混合精度訓練也需要大量內存，因為除了FP16版本之外，優化器還需要保存FP32參數和其他優化器狀態的副本。 

• 未被激活、臨時緩衝區以及不可用的碎片內存消耗（論文中稱為剩餘狀態）。

ZeRO結合了ZeRO-DP和ZeRO-R兩種方法。ZeRO-DP是一種增強的數據並行，避免了模型狀態的簡單冗餘。它以動態的方式跨多個並行數據劃分優化器狀態、梯度和參數，以最小化通信量。ZeRO-R使用分區激活二次計算、恆定緩衝區大小和動態內存碎片，以優化剩餘狀態的內存消耗。

參考資料：

[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).

編譯鏈接：//lilianweng.github.io/lil-log/2021/09/24/train-large-neural-networks.html

由於微信公眾號試行亂序推送，您可能不再能準時收到AI科技評論的推送。為了第一時間收到AI科技評論的報道， 請將「AI科技評論」設為星標賬號，以及常點文末右下角的「在看」。

雷鋒網雷鋒網雷鋒網