一只小狐貍帶你解鎖?煉丹術&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) # 計算lossoptimizer.zero_grad() # 清空梯度loss.backward() # 反向計算梯度optimizer.step() # 更新參數

加入梯度累加后：

for i, (inputs, labels) in enumerate(training_set):loss = model(inputs, labels) # 計算lossloss = 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*xfloat2 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 Parallelism

Pipeline的并行方式就是把模型的不同層放到不同機器上，順序地進行前向和反向計算。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

夕小瑤的賣萌屋

_

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

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

總結

以上是生活随笔為你收集整理的训练大型神经网络方法总结的全部內容，希望文章能夠幫你解決所遇到的問題。

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