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

前陣子微軟開源了DeepSpeed訓(xùn)練框架，從測(cè)試效果來(lái)看有10倍的速度提升，而且對(duì)內(nèi)存進(jìn)行了各種優(yōu)化，最大可以訓(xùn)練100B(illion)參數(shù)的模型。同時(shí)發(fā)布了這個(gè)框架訓(xùn)練出的17B模型 Turing-NLG，處于目前壕賽事的頂端。

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

我明天就要答辯了，今天必須把這十個(gè)實(shí)驗(yàn)跑完

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

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

（以上純屬虛構(gòu)，如有雷同請(qǐng)趕緊看下文）

現(xiàn)實(shí)總是殘酷的，其實(shí)限制大模型訓(xùn)練只有兩個(gè)因素：時(shí)間和空間(=GPU=錢)，根據(jù)不同情況可以使用的方案大致如下：

1. 梯度累加 Gradient Accumulation

如果只有單卡，且可以加載模型，但batch受限的話可以使用梯度累加，進(jìn)行N次前向后反向更新一次參數(shù)，相當(dāng)于擴(kuò)大了N倍的batch size。

正常的訓(xùn)練代碼是這樣的：

for i, (inputs, labels) in enumerate(training_set):loss = model(inputs, labels) # 計(jì)算lossoptimizer.zero_grad() # 清空梯度loss.backward() # 反向計(jì)算梯度optimizer.step() # 更新參數(shù)

加入梯度累加后：

for i, (inputs, labels) in enumerate(training_set):loss = model(inputs, labels) # 計(jì)算lossloss = loss / accumulation_steps # Normalize our loss (if averaged)loss.backward() # 反向計(jì)算梯度，累加到之前梯度上if (i+1) % accumulation_steps == 0:optimizer.step() # 更新參數(shù)model.zero_grad() # 清空梯度

要注意的是，batch擴(kuò)大后，如果想保持樣本權(quán)重相等，學(xué)習(xí)率也要線性擴(kuò)大或者適當(dāng)調(diào)整。另外batchnorm也會(huì)受到影響，小batch下的均值和方差肯定不如大batch的精準(zhǔn)，可以調(diào)整BN中的momentum參數(shù)解決[2]。

2. 梯度檢查點(diǎn) Gradient Checkpointing

如果只有一張卡，又想訓(xùn)大模型，可以嘗試壓縮模型所占顯存。

梯度檢查點(diǎn)是一種以時(shí)間換空間的方法，通過(guò)減少保存的激活值壓縮模型占用空間，但是在計(jì)算梯度時(shí)必須從新計(jì)算沒有存儲(chǔ)的激活值。

細(xì)節(jié)可以參考陳天奇的Training Deep Nets with Sublinear Memory Cost[3]。

注：第一行節(jié)點(diǎn)是前向，第二行是反向

3. 混合精度訓(xùn)練 Mixed Precision Training

混合精度訓(xùn)練在單卡和多卡情況下都可以使用，通過(guò)cuda計(jì)算中的half2類型提升運(yùn)算效率。一個(gè)half2類型中會(huì)存儲(chǔ)兩個(gè)FP16的浮點(diǎn)數(shù)，在進(jìn)行基本運(yùn)算時(shí)可以同時(shí)進(jìn)行，因此FP16的期望速度是FP32的兩倍。舉個(gè)Gelu的FP16優(yōu)化栗子：

//FP32的gelu運(yùn)算float gelu(float x) {float cdf = 0.5f * (1.0f + tanhf((0.7978845608028654f * (x + 0.044715f * x * x * x))));return x * cdf; } //FP16的gelu運(yùn)算half2 gelu(half2 val) {half2 val_pow3 = __hmul2(val, __hmul2(val, val)); //同時(shí)計(jì)算兩個(gè)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))));//同時(shí)計(jì)算兩個(gè)x * cdf;return __hmul2(val, __float22half2_rn(cdf)); }

混合精度訓(xùn)練[5]不是很難理解，但要注意以下幾點(diǎn)：

混合精度訓(xùn)練不是單純地把FP32轉(zhuǎn)成FP16去計(jì)算就可以了，只用FP16會(huì)造成80%的精度損失

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

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

整體流程：FP32權(quán)重 -> FP16權(quán)重 -> FP16計(jì)算前向 -> FP32的loss，擴(kuò)大 -> 轉(zhuǎn)為FP16 -> FP16反向計(jì)算梯度 -> 縮放為FP32的梯度更新權(quán)重

！！手工分割線：接下來(lái)就是壕賽道了！！

4. 分布式訓(xùn)練 Distributed Training

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

• Multi-GPU：單機(jī)多卡，通過(guò)PCIE、NVlink、GPU Direct P2P來(lái)通信

• Multi-Node：多機(jī)多卡，通過(guò)Sockets (Ethernet) 或者InfiniBand with GPU Direct RDMA通信

實(shí)踐中可以使用英偉達(dá)的NCCL通信框架，多機(jī)通過(guò)IB(InfiniBand)可以接近機(jī)內(nèi)的通信速度[6]。底層的東西就不多說(shuō)了（我也不太懂），實(shí)際上對(duì)于煉丹師來(lái)說(shuō)就是找運(yùn)維爸爸提供幫助，并借助開源框架配置上服務(wù)器地址就行了。

并行訓(xùn)練有多種優(yōu)化策略，主要目的就是減少計(jì)算中的參數(shù)同步(Sync)和數(shù)據(jù)傳輸。

目前32GB的卡最多能放1.3B參數(shù)的模型，塞得下的話可以使用數(shù)據(jù)并行的方式，否則可以把不同層放在不同機(jī)器上進(jìn)行訓(xùn)練。兩種方式的區(qū)別看下圖[7]就明白啦：

4.1 數(shù)據(jù)并行 Data Parallelism

數(shù)據(jù)并行有兩種方式[9]：

Parameter Server

集群中有一個(gè)master和多個(gè)worker，master需要等待所有節(jié)點(diǎn)計(jì)算完畢統(tǒng)一計(jì)算梯度，在master上更新參數(shù)，之后把新的參數(shù)廣播給worker。這種方式的主要瓶頸在master，因此也可以異步訓(xùn)練，即不等待其他節(jié)點(diǎn)，收到一個(gè)worker的梯度后就更新參數(shù)，但這樣其他worker在舊參數(shù)上算完后的梯度會(huì)作用到新參數(shù)上，導(dǎo)致模型優(yōu)化過(guò)頭，陷入次優(yōu)解。

Ring All-Reduce

集群中所有worker形成一個(gè)閉環(huán)，把數(shù)據(jù)分成K份，計(jì)算完一份就把累加好的梯度傳給下家，同時(shí)接受上家的梯度，迭代到最后所有worker的梯度都是相等的，可以同步更新參數(shù)，比PS架構(gòu)要高效，是目前的主流方式。下圖[10]展示了Scatter Reduce和All Gather兩個(gè)階段：

preview

4.2 模型并行 Model Parallelism

模型并行目前并不常見，一是因?yàn)榇蟛糠帜Ｐ蛦慰ǘ挤诺孟?#xff0c;二是因?yàn)橥ㄓ嶉_銷比數(shù)據(jù)并行多，因?yàn)榉聪騻鞑バ枰裭oss對(duì)每層激活值的梯度都傳回去，樣本數(shù)量大的話激活值也有很多。

Pipelined Parallelism

Pipeline的并行方式就是把模型的不同層放到不同機(jī)器上，順序地進(jìn)行前向和反向計(jì)算。19年谷歌和微軟先后放出了GPipe[11]和PipeDream[12]的論文和源碼，給大家梳理一下他們的心路歷程：

首先來(lái)看最naive的模型并行方式，實(shí)在是有些浪費(fèi)生命：

注：反向需要計(jì)算對(duì)參數(shù)和激活值的偏導(dǎo)，所以耗時(shí)更長(zhǎng)。

所以谷歌GPipe提出了一個(gè)改進(jìn)，其實(shí)就是把數(shù)據(jù)分片，像allreduce一樣計(jì)算完一些就傳給下個(gè)節(jié)點(diǎn)，最后同步更新參數(shù)，但這樣看還是不能挽救我們的青春：

于是微軟提出了PipeDream，其實(shí)就是把同步變?yōu)榱诵?shù)據(jù)上的異步，計(jì)算完一個(gè)數(shù)據(jù)分片就立刻反向，反向完了就更新梯度，誰(shuí)也別等誰(shuí)，大家一起瘋狂干起來(lái)：

但這樣就有一個(gè)問(wèn)題，就是大家越干越亂，比如worker1在計(jì)算5的前向時(shí)用的是1反向后的參數(shù)，但之后計(jì)算5反向的梯度時(shí)參數(shù)早就被2/3/4更新了。于是作者加入了Weight stashing機(jī)制，把每個(gè)數(shù)據(jù)對(duì)應(yīng)的參數(shù)都存起來(lái)！這樣worker1在5反向的時(shí)候就可以從百寶箱里拿出之前的參數(shù)，進(jìn)行更新：

那問(wèn)題又來(lái)了：worker1上5的前向是用1的參數(shù)，但worker3上是用3的，最后匯總的時(shí)候不就又亂了？于是作者又加入了Vertical Sync機(jī)制，強(qiáng)制所有worker在計(jì)算5的時(shí)候都用1的參數(shù)。這樣在最后匯總模型的時(shí)候，就能拿到一致的參數(shù)了。但這樣同步會(huì)導(dǎo)致很多計(jì)算作廢，比如5更新時(shí)用的1的權(quán)重，但2/3/4的權(quán)重都白計(jì)算了，所以默認(rèn)是不用Vertical Sync的，這樣每層雖然不完全一致，但由于weight stashing，所有的參數(shù)都是有效的。

Tensor Slicing

神經(jīng)網(wǎng)絡(luò)可以看作一個(gè)復(fù)合函數(shù)，本質(zhì)就是各個(gè)tensor之間的計(jì)算，我們定義好的CNN、RNN其實(shí)就是計(jì)算函數(shù)的集合。從這個(gè)角度來(lái)思考，模型并行其實(shí)就是把各個(gè)tensor計(jì)算分散到不同的機(jī)器上。這方面的研究有18年的FlexFLow和Mesh-TensorFlow，英偉達(dá)的威震天[13]也是使用這個(gè)策略。下面以Transformer為例說(shuō)明一下如何拆分。

Transformer主要有self-attention和FFN組成，對(duì)于FFN中的第一層Y=GLUE(XA)可以有兩種拆分方式：

可以看到，第一種需要在計(jì)算GLUE時(shí)同步，因此威震天通過(guò)第二種方式進(jìn)行tensor切片，self-attention也采用類似的策略，這樣只需要在前向時(shí)通過(guò)g聚合，反向時(shí)通過(guò)f聚合就可以了：

剩下的Layernorm和dropout還是需要同步后計(jì)算：

同時(shí)，作者也在vocab的維度對(duì)embedding進(jìn)行了切分，并把最后的MLM預(yù)測(cè)和cross-entropy融合到一起，減少網(wǎng)絡(luò)通信量（否則需要傳輸batch_size*seq_len *vocab_size個(gè)prob，改過(guò)后只傳batch_size *seq_len個(gè)loss值）。

隨著模型越來(lái)越大，分布式訓(xùn)練甚至推理肯定是一個(gè)趨勢(shì)，在工程上還有很多可以優(yōu)化的點(diǎn)，不僅是上面介紹的分布式策略，還有網(wǎng)絡(luò)通信優(yōu)化、內(nèi)存優(yōu)化等。

5. 加速優(yōu)化器 LAMB

上文提到的數(shù)據(jù)并行雖然可以接近線性地提升訓(xùn)練速度，但過(guò)大的Batch會(huì)降低模型精度和收斂速度（對(duì)數(shù)據(jù)的擬合變差）。因此谷歌在19年推出了LAMB[14]優(yōu)化器，全稱為Layer-wise Adaptive Moments optimizer for Batch training，針對(duì)大batch做了優(yōu)化，在分布式訓(xùn)練的場(chǎng)景下可訓(xùn)65536/32768的樣本，減少迭代次數(shù)，從而縮短訓(xùn)練時(shí)間，感受一下金錢的味道：

LAMB主要是綜合了Adam和LARS（Layerwise Adaptive Rate Scaling），對(duì)學(xué)習(xí)率進(jìn)行調(diào)整。上文提到當(dāng)batch變大時(shí)學(xué)習(xí)率也需要變大，這樣會(huì)導(dǎo)致收斂不穩(wěn)定，LARS通過(guò)給LR乘上權(quán)重與梯度的norm比值來(lái)解決這個(gè)問(wèn)題[15]：

這里的norm都是取一層的權(quán)重計(jì)算，所以是layerwise。可以這樣理解上面的公式：剛開始訓(xùn)練時(shí)，權(quán)重比較小，而loss和梯度比較大，所以學(xué)習(xí)率開始較小，但隨著權(quán)重變大&梯度變小會(huì)慢慢warmup。當(dāng)對(duì)一些樣本擬合很好，loss接近0時(shí)，梯度變小，學(xué)習(xí)率又會(huì)增大，跳出局部最優(yōu)，防止過(guò)擬合。

LAMB融合了這種layerwise的自適應(yīng)思想：

圖中的公式稍稍有改動(dòng)，一個(gè)是給權(quán)重norm加了映射，本質(zhì)都是起scale的作用；另一個(gè)是梯度公式中加了weight decay，也就是目標(biāo)函數(shù)中的L2正則化。

總結(jié)

本文介紹了從速度和內(nèi)存去優(yōu)化模型訓(xùn)練的幾種方式，實(shí)踐中各種都是可以混合起來(lái)的，比如混合精度+數(shù)據(jù)并行、數(shù)據(jù)并行+模型并行、數(shù)據(jù)并行+梯度檢查點(diǎn)等。DeepSpeed里基本涵蓋了本文所講的策略，用pytorch的同學(xué)可以安排起來(lái)了～

最后，在介紹各種策略的時(shí)候，由于篇幅原因也有省略一些假設(shè)和最終效果，感興趣的同學(xué)們可以深入研讀參考資料里的內(nèi)容～如果路過(guò)的大佬們發(fā)現(xiàn)哪里有錯(cuò)誤煩請(qǐng)指出～

可

能

喜

歡

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

• 模型訓(xùn)練太慢？顯存不夠用？這個(gè)算法讓你的GPU老樹開新花

• 訓(xùn)練效率低？GPU利用率上不去？快來(lái)看看別人家的tricks吧～
  

• 如何打造高質(zhì)量的NLP數(shù)據(jù)集

• 萬(wàn)萬(wàn)沒想到，我的煉丹爐玩壞了

參考文獻(xiàn)

[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]?高開遠(yuǎn) Reformer解讀：https://zhuanlan.zhihu.com/p/104935987

[5]?混合精度訓(xùn)練：https://zhuanlan.zhihu.com/p/84219777

[6]?英偉達(dá)NCCL：https://www.zhihu.com/question/63219175/answer/206697974

[7]?數(shù)據(jù)并行與模型并行：https://www.zhihu.com/question/53851014/answer/158794752

[8]?分布式之?dāng)?shù)據(jù)并行：https://zhuanlan.zhihu.com/p/68615246

[9] AllReduce：https://zhuanlan.zhihu.com/p/100012827

[10] AllReduce細(xì)節(jié)：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

夕小瑤的賣萌屋

_

關(guān)注&星標(biāo)小夕，帶你解鎖煉丹秘籍

訂閱號(hào)主頁(yè)下方「撩一下」有驚喜哦

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。