一只小狐貍帶你解鎖?煉丹術&NLP?秘籍

前陣子微軟開源了DeepSpeed訓練框架，從測試效果來看有10倍的速度提升，而且對內存進行了各種優化，最大可以訓練100B(illion)參數的模型。同時發布了這個框架訓練出的17B模型 Turing-NLG，處于目前壕賽事的頂端。

訓100B的模型就先別想了(狗頭)，先把110M的BERT-base訓好上線吧。本文主要介紹模型訓練中速度和內存的優化策略，針對以下幾種情況：

我明天就要答辯了，今天必須把這十個實驗跑完

我的模型有些大，好不容易放到一張卡上，訓完一億樣本之前我就可以領N+1了

我想出了一個絕妙的T6模型，卻加載不進12GB的卡里，又拿不到今年的best paper了

(以上純屬虛構，如有雷同請趕緊看下文)現實總是殘酷的，其實限制大模型訓練只有兩個因素：時間和空間(=GPU=錢)，根據不同情況可以使用的方案大致如下：

1. 梯度累加 Gradient Accumulation

如果只有單卡，且可以加載模型，但batch受限的話可以使用梯度累加，進行N次前向后反向更新一次參數，相當于擴大了N倍的batch size。正常的訓練代碼是這樣的：for i, (inputs, labels) in enumerate(training_set):
loss = model(inputs, labels) # 計算loss
optimizer.zero_grad() # 清空梯度
loss.backward() # 反向計算梯度
optimizer.step() # 更新參數加入梯度累加后：for i, (inputs, labels) in enumerate(training_set):
loss = model(inputs, labels) # 計算loss
loss = loss / accumulation_steps # Normalize our loss (if averaged)
loss.backward() # 反向計算梯度，累加到之前梯度上
if (i+1) % accumulation_steps == 0:
optimizer.step() # 更新參數
model.zero_grad() # 清空梯度要注意的是，batch擴大后，如果想保持樣本權重相等，學習率也要線性擴大或者適當調整。另外batchnorm也會受到影響，小batch下的均值和方差肯定不如大batch的精準，可以調整BN中的momentum參數解決[2]。

2. 梯度檢查點 Gradient Checkpointing

如果只有一張卡，又想訓大模型，可以嘗試壓縮模型所占顯存。梯度檢查點是一種以時間換空間的方法，通過減少保存的激活值壓縮模型占用空間，但是在計算梯度時必須從新計算沒有存儲的激活值。細節可以參考陳天奇的Training Deep Nets with Sublinear Memory Cost[3]。注：第一行節點是前向，第二行是反向

3. 混合精度訓練 Mixed Precision Training

混合精度訓練在單卡和多卡情況下都可以使用，通過cuda計算中的half2類型提升運算效率。一個half2類型中會存儲兩個FP16的浮點數，在進行基本運算時可以同時進行，因此FP16的期望速度是FP32的兩倍。舉個Gelu的FP16優化栗子：//FP32的gelu運算float gelu(float x){float cdf = 0.5f * (1.0f + tanhf((0.7978845608028654f * (x + 0.044715f * x * x * x))));return x * cdf;
}//FP16的gelu運算half2 gelu(half2 val){
half2 val_pow3 = __hmul2(val, __hmul2(val, val)); //同時計算兩個x*x*x
float2 tmp_pow = __half22float2(val_pow3);
float2 cdf = __half22float2(val); //由于tanhf不支持half2類型，只能分開算
cdf.x = 0.5f * (1.0f + tanhf((0.7978845608028654f * (cdf.x + 0.044715f * tmp_pow.x))));
cdf.y = 0.5f * (1.0f + tanhf((0.7978845608028654f * (cdf.y + 0.044715f * tmp_pow.y))));//同時計算兩個x * cdf;return __hmul2(val, __float22half2_rn(cdf));
}混合精度訓練[5]不是很難理解，但要注意以下幾點：

混合精度訓練不是單純地把FP32轉成FP16去計算就可以了，只用FP16會造成80%的精度損失

Loss scaling：由于梯度值都很小，用FP16會下溢，因此先用FP32存儲loss并放大，使得梯度也得到放大，可以用FP16存儲，更新時變成FP32再縮放

在涉及到累加操作時，比如BatchNorm、Softmax，FP16會上溢，需要用FP32保存，一般使用GPU中TensorCore的FP16*FP16+FP32=FP32運算

整體流程：FP32權重 -> FP16權重 -> FP16計算前向 -> FP32的loss，擴大 -> 轉為FP16 -> FP16反向計算梯度 -> 縮放為FP32的梯度更新權重！！手工分割線：接下來就是壕賽道了！！

4. 分布式訓練 Distributed Training

分布式訓練就是多張卡并行訓練，一般有以下兩種情況：

• Multi-GPU：單機多卡，通過PCIE、NVlink、GPU Direct P2P來通信
• Multi-Node：多機多卡，通過Sockets (Ethernet) 或者InfiniBand with GPU Direct RDMA通信

實踐中可以使用英偉達的NCCL通信框架，多機通過IB(InfiniBand)可以接近機內的通信速度[6]。底層的東西就不多說了(我也不太懂)，實際上對于煉丹師來說就是找運維爸爸提供幫助，并借助開源框架配置上服務器地址就行了。并行訓練有多種優化策略，主要目的就是減少計算中的參數同步(Sync)和數據傳輸。目前32GB的卡最多能放1.3B參數的模型，塞得下的話可以使用數據并行的方式，否則可以把不同層放在不同機器上進行訓練。兩種方式的區別看下圖[7]就明白啦：

4.1 數據并行 Data Parallelism

數據并行有兩種方式[9]：Parameter Server集群中有一個master和多個worker，master需要等待所有節點計算完畢統一計算梯度，在master上更新參數，之后把新的參數廣播給worker。這種方式的主要瓶頸在master，因此也可以異步訓練，即不等待其他節點，收到一個worker的梯度后就更新參數，但這樣其他worker在舊參數上算完后的梯度會作用到新參數上，導致模型優化過頭，陷入次優解。Ring All-Reduce集群中所有worker形成一個閉環，把數據分成K份，計算完一份就把累加好的梯度傳給下家，同時接受上家的梯度，迭代到最后所有worker的梯度都是相等的，可以同步更新參數，比PS架構要高效，是目前的主流方式。下圖[10]展示了Scatter Reduce和All Gather兩個階段：preview

4.2 模型并行 Model Parallelism

模型并行目前并不常見，一是因為大部分模型單卡都放得下，二是因為通訊開銷比數據并行多，因為反向傳播需要把loss對每層激活值的梯度都傳回去，樣本數量大的話激活值也有很多。Pipelined ParallelismPipeline的并行方式就是把模型的不同層放到不同機器上，順序地進行前向和反向計算。19年谷歌和微軟先后放出了GPipe[11]和PipeDream[12]的論文和源碼，給大家梳理一下他們的心路歷程：首先來看最naive的模型并行方式，實在是有些浪費生命：注：反向需要計算對參數和激活值的偏導，所以耗時更長。所以谷歌GPipe提出了一個改進，其實就是把數據分片，像allreduce一樣計算完一些就傳給下個節點，最后同步更新參數，但這樣看還是不能挽救我們的青春：于是微軟提出了PipeDream，其實就是把同步變為了小數據上的異步，計算完一個數據分片就立刻反向，反向完了就更新梯度，誰也別等誰，大家一起瘋狂干起來：但這樣就有一個問題，就是大家越干越亂，比如worker1在計算5的前向時用的是1反向后的參數，但之后計算5反向的梯度時參數早就被2/3/4更新了。于是作者加入了Weight stashing機制，把每個數據對應的參數都存起來！這樣worker1在5反向的時候就可以從百寶箱里拿出之前的參數，進行更新：那問題又來了：worker1上5的前向是用1的參數，但worker3上是用3的，最后匯總的時候不就又亂了？于是作者又加入了Vertical Sync機制，強制所有worker在計算5的時候都用1的參數。這樣在最后匯總模型的時候，就能拿到一致的參數了。但這樣同步會導致很多計算作廢，比如5更新時用的1的權重，但2/3/4的權重都白計算了，所以默認是不用Vertical Sync的，這樣每層雖然不完全一致，但由于weight stashing，所有的參數都是有效的。Tensor Slicing神經網絡可以看作一個復合函數，本質就是各個tensor之間的計算，我們定義好的CNN、RNN其實就是計算函數的集合。從這個角度來思考，模型并行其實就是把各個tensor計算分散到不同的機器上。這方面的研究有18年的FlexFLow和Mesh-TensorFlow，英偉達的威震天[13]也是使用這個策略。下面以Transformer為例說明一下如何拆分。Transformer主要有self-attention和FFN組成，對于FFN中的第一層Y=GLUE(XA)可以有兩種拆分方式：可以看到，第一種需要在計算GLUE時同步，因此威震天通過第二種方式進行tensor切片，self-attention也采用類似的策略，這樣只需要在前向時通過g聚合，反向時通過f聚合就可以了：剩下的Layernorm和dropout還是需要同步后計算：同時，作者也在vocab的維度對embedding進行了切分，并把最后的MLM預測和cross-entropy融合到一起，減少網絡通信量(否則需要傳輸batch_size*seq_len *vocab_size個prob，改過后只傳batch_size *seq_len個loss值)。隨著模型越來越大，分布式訓練甚至推理肯定是一個趨勢，在工程上還有很多可以優化的點，不僅是上面介紹的分布式策略，還有網絡通信優化、內存優化等。

5. 加速優化器 LAMB

上文提到的數據并行雖然可以接近線性地提升訓練速度，但過大的Batch會降低模型精度和收斂速度(對數據的擬合變差)。因此谷歌在19年推出了LAMB[14]優化器，全稱為Layer-wise Adaptive Moments optimizer for Batch training，針對大batch做了優化，在分布式訓練的場景下可訓65536/32768的樣本，減少迭代次數，從而縮短訓練時間，感受一下金錢的味道：LAMB主要是綜合了Adam和LARS(Layerwise Adaptive Rate Scaling)，對學習率進行調整。上文提到當batch變大時學習率也需要變大，這樣會導致收斂不穩定，LARS通過給LR乘上權重與梯度的norm比值來解決這個問題[15]：這里的norm都是取一層的權重計算，所以是layerwise。可以這樣理解上面的公式：剛開始訓練時，權重比較小，而loss和梯度比較大，所以學習率開始較小，但隨著權重變大&梯度變小會慢慢warmup。當對一些樣本擬合很好，loss接近0時，梯度變小，學習率又會增大，跳出局部最優，防止過擬合。LAMB融合了這種layerwise的自適應思想：圖中的公式稍稍有改動，一個是給權重norm加了映射，本質都是起scale的作用；另一個是梯度公式中加了weight decay，也就是目標函數中的L2正則化。

總結

本文介紹了從速度和內存去優化模型訓練的幾種方式，實踐中各種都是可以混合起來的，比如混合精度+數據并行、數據并行+模型并行、數據并行+梯度檢查點等。DeepSpeed里基本涵蓋了本文所講的策略，用pytorch的同學可以安排起來了～最后，在介紹各種策略的時候，由于篇幅原因也有省略一些假設和最終效果，感興趣的同學們可以深入研讀參考資料里的內容～如果路過的大佬們發現哪里有錯誤煩請指出～

可

能

喜

歡

• 如何讓BERT擁有視覺感知能力？兩種方式將視頻信息注入BERT

• 模型訓練太慢？顯存不夠用？這個算法讓你的GPU老樹開新花

• 訓練效率低？GPU利用率上不去？快來看看別人家的tricks吧～

• 如何打造高質量的NLP數據集

• 萬萬沒想到，我的煉丹爐玩壞了

參考文獻[1]?微軟Turing-NLG：https://www.microsoft.com/en-us/research/blog/turing-nlg-a-17-billion-parameter-language-model-by-microsoft/[2]?梯度累加：https://www.zhihu.com/question/303070254/answer/573037166[3]?陳天奇 Training Deep Nets with Sublinear Memory Cost:?https://www.zhihu.com/question/274635237/answer/755102181[4]?高開遠 Reformer解讀：https://zhuanlan.zhihu.com/p/104935987[5]?混合精度訓練：https://zhuanlan.zhihu.com/p/84219777[6]?英偉達NCCL：https://www.zhihu.com/question/63219175/answer/206697974[7]?數據并行與模型并行：https://www.zhihu.com/question/53851014/answer/158794752[8]?分布式之數據并行：https://zhuanlan.zhihu.com/p/68615246[9] AllReduce：https://zhuanlan.zhihu.com/p/100012827[10] AllReduce細節：https://zhuanlan.zhihu.com/p/56991108[11] GPipe：https://arxiv.org/pdf/1811.06965.pdf[12] PipeDream：https://arxiv.org/pdf/1806.03377.pdf[13] Megatron-LM：https://arxiv.org/abs/1909.08053[14] LAMB：https://arxiv.org/abs/1904.00962v3[15] LAMB解讀：https://towardsdatascience.com/an-intuitive-understanding-of-the-lamb-optimizer-46f8c0ae4866

夕小瑤的賣萌屋

_

關注&星標小夕，帶你解鎖煉丹秘籍

訂閱號主頁下方「撩一下」有驚喜哦

總結

以上是生活随笔為你收集整理的训练好的神经网络 如何预测_显存不够，如何训练大型神经网络？的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。