<code>關注微信公眾號：人工智能前沿講習，
重磅乾貨，第一時間送達/<code>
前陣子微軟開源了DeepSpeed訓練框架，從測試效果來看有10倍的速度提升，而且對內存進行了各種優化，最大可以訓練100B(illion)參數的模型。同時發佈了這個框架訓練出的17B模型 Turing-NLG，處於目前壕賽事的頂端。
訓100B的模型就先別想了（狗頭），先把110M的BERT-base訓好上線吧。本文主要介紹模型訓練中速度和內存的優化策略，針對以下幾種情況：
- 我明天就要答辯了，今天必須把這十個實驗跑完
- 我的模型有些大，好不容易放到一張卡上，訓完一億樣本之前我就可以領N+1了
- 我想出了一個絕妙的T6模型，卻加載不進12GB的卡里，又拿不到今年的best paper了
（以上純屬虛構，如有雷同請趕緊看下文）
現實總是殘酷的，其實限制大模型訓練只有兩個因素：時間和空間(=GPU=錢)，根據不同情況可以使用的方案大致如下：
一、梯度累加 Gradient Accumulation
如果只有單卡，且可以加載模型，但batch受限的話可以使用梯度累加，進行N次前向後反向更新一次參數，相當於擴大了N倍的batch size。
正常的訓練代碼是這樣的：
<code>for i, (inputs, labels) in enumerate(training_set): loss = model(inputs, labels)
# 計算loss optimizer.zero_grad() # 清空梯度 loss.backward() # 反向計算梯度 optimizer.step() # 更新參數/<code>
加入梯度累加後：
<code>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() # 清空梯度/<code>
要注意的是，batch擴大後，如果想保持樣本權重相等，學習率也要線性擴大或者適當調整。另外batchnorm也會受到影響，小batch下的均值和方差肯定不如大batch的精準，可以調整BN中的momentum參數解決[2]。
二、梯度檢查點 Gradient Checkpointing
如果只有一張卡，又想訓大模型，可以嘗試壓縮模型所佔顯存。
梯度檢查點是一種以時間換空間的方法，通過減少保存的激活值壓縮模型佔用空間，但是在計算梯度時必須從新計算沒有存儲的激活值。
細節可以參考陳天奇的Training Deep Nets with Sublinear Memory Cost[3]，最近的Reformer（參考RevNet）也是這種思想[4]。
注：第一行節點是前向，第二行是反向
三、混合精度訓練 Mixed Precision Training
混合精度訓練在單卡和多卡情況下都可以使用，通過cuda計算中的half2類型提升運算效率。一個half2類型中會存儲兩個FP16的浮點數，在進行基本運算時可以同時進行，因此FP16的期望速度是FP32的兩倍。舉個Gelu的FP16優化栗子：
<code>//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));
}/<code>
混合精度訓練[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的梯度更新權重
！！手工分割線：接下來就是壕賽道了！！
四、分佈式訓練 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值）。
隨著模型越來越大，分佈式訓練甚至推理肯定是一個趨勢，在工程上還有很多可以優化的點，不僅是上面介紹的分佈式策略，還有網絡通信優化、內存優化等。
五、加速優化器 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的同學可以安排起來了～
最後，在介紹各種策略的時候，由於篇幅原因也有省略一些假設和最終效果，感興趣的同學們可以深入研讀參考資料裡的內容～如果路過的大佬們發現哪裡有錯誤煩請指出～
接下來的功課：
一文搞懂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
閱讀更多 人工智能前沿學生論壇 的文章
