編譯 | 琰琰

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

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

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

“煉大模型”已成為人工智能領域的主流研發(fā)趨勢。從GPT-3的1750億，到如今悟道2.0的1.75萬億，超大語言模型在 NLP 基準任務中不斷刷新SOTA。

然而，參數(shù)和數(shù)據(jù)集的快速增長讓 GPU 算力開始捉襟見肘。單個GPU內存已經遠遠不能滿足大模型的需求。如，阿里用480塊GPU訓練千億模型；英偉達用3072塊GPU訓練萬億模型；谷歌用2048塊TPU訓練1.6萬億模型（1 TPU約等于2～3 GPU）。

如何利用上百塊GPU上訓練大規(guī)模語言模型？并行計算是一種行之有效的方法。

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

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

1

并行訓練

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

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

數(shù)據(jù)并行 

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

數(shù)據(jù)交換傳輸通常在后端進行（不干擾訓練計算），在每個Mini-batch計算結束后，worker需要同步梯度或權重，以保證學習效率?，F(xiàn)有的同步方法有兩種，各自優(yōu)缺點如下：

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

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

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

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

模型并行 

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

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

圖2：一個包含4個垂直層的模型并行設置，由于順序的依賴性，每個數(shù)據(jù)依次由一個worker處理，這個過程會出現(xiàn)大量多余時間“氣泡”（來源：Huang等人，2019年）

管道并行

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

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

 GPipe論文表明，如果微批次的數(shù)量超過分區(qū)數(shù)量4倍（m>4d），則“氣泡”開銷幾乎可以忽略不計。

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

 GPipe在吞吐量方面與設備數(shù)量成線性關系，設備數(shù)量越多，吞吐量越大。不過，如果模型參數(shù)在worker中分布不均勻，這種線性關系不會穩(wěn)定出現(xiàn)。 

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

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

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

• 權重存儲：每個worker跟蹤多個模型版本，給定數(shù)據(jù) 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，可以得到：

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

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

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

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

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

2

混合專家（MoE）

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

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

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

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

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

G將輸入與可訓練權重矩陣Gg相乘，然后執(zhí)行softmax：

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

 為了避免門控網絡可能始終偏向少數(shù)強勢專家的自我強化效應，Shazeer等人（2017）提出了通過額外重要損失的軟約束，以鼓勵所有專家擁有相同的權重。其數(shù)值相當于每個專家的分批平均值變異系數(shù)的平方：

 

其中，CV是變異系數(shù)，失重的waux是可調節(jié)的超參數(shù)。由于每個專家網絡只能獲得小部分訓練樣本（“收縮批次問題”），所以在MoE中應該盡可能使用大batch，但這又會受到GPU內存的限制。數(shù)據(jù)并行和模型并行的應用可以提高模型的吞吐量。

圖11：10億單詞的語言建?；鶞剩ㄗ螅┠Ｐ蛷淖蟮接野?、32、256、256、1024和4096名專家（右）40億參數(shù)的MoE模型在不同計算預算下的性能（來源：Shazeer等人，2017年）

GShard（Lepikhin等人，2020年）通過自動分片將MoE transformer 模型的參數(shù)擴展到了6000億。MoE transformer 用MoE層取代其他每一個前饋網絡層。需要說明的是，在多臺機器上MoE  transformer 僅在MoE層分片，其他層只是復制。 

GShard中的門控功能G有幾種改進設計：

• 專家容量：通過一位專家的令牌數(shù)量不應超過“專家容量”的閾值。如果令牌被路由到已達到容量的專家，則令牌將被標記為“溢出”，并且門輸出將更改為零向量。

• 本地組調度：令牌被均勻地劃分為多個本地組，專家能力在組水平上得到加強。

• 輔助損失：與原始MoE aux損失相似，添加輔助損失可以最小化路由到每個專家的數(shù)據(jù)的均方。

• 隨機路由：以與其權重成比例的概率選擇第二位最佳專家；否則GShard遵循隨機路由，增加隨機性。

圖12：GShard中帶輔助損失的組水平top-2門機制的偽代碼（來源：Lepikhin等人，2020年） 

Switch Transformer（Fedus et al.2021）用稀疏開關FFN層取代了密集前饋層（每個輸入僅路由到一個專家網絡），將模型規(guī)模擴展到數(shù)萬億個參數(shù)。負載平衡的輔助損失是，給定n個專家，fi是路由到第i個專家的令牌分數(shù)，pi是門控網絡預測的專家i的路由概率。

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

為提高訓練穩(wěn)定性，switch transformer采用以下設計：

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

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

switch transformer論文總結了用于訓練大型模型的不同數(shù)據(jù)和模型并行策略，并給出了一個很好的示例：

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

 

3

其他節(jié)省內存的設計 

CPU卸載

如果GPU內存已滿，可以將暫時未使用的數(shù)據(jù)卸載到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）數(shù)訓練模型而不損失模型精度的方法。

圖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）提出了一種數(shù)據(jù)編碼策略，即在第一次傳遞后壓縮中間結果，然后將其解碼用于反向傳播。 

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

內存高效優(yōu)化器

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

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

SM3（Anil et al.2019）優(yōu)化器采用了一種不同的自適應優(yōu)化方法。

ZeRO（Rajbhandari et al.2019）零冗余優(yōu)化器節(jié)省了大型模型訓練在兩方面的內存消耗：

• 大多數(shù)內存由模型狀態(tài)消耗，包括優(yōu)化器狀態(tài)（例如Adam動量和方差）、梯度和參數(shù)?；旌暇扔柧氁残枰罅績却?，因為除了FP16版本之外，優(yōu)化器還需要保存FP32參數(shù)和其他優(yōu)化器狀態(tài)的副本。 

• 未被激活、臨時緩沖區(qū)以及不可用的碎片內存消耗（論文中稱為剩余狀態(tài)）。

ZeRO結合了ZeRO-DP和ZeRO-R兩種方法。ZeRO-DP是一種增強的數(shù)據(jù)并行，避免了模型狀態(tài)的簡單冗余。它以動態(tài)的方式跨多個并行數(shù)據(jù)劃分優(yōu)化器狀態(tài)、梯度和參數(shù)，以最小化通信量。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

雷鋒網雷鋒網雷鋒網

雷峰網版權文章，未經授權禁止轉載。詳情見轉載須知。