混合精度訓(xùn)練

轉(zhuǎn)自：https://zhuanlan.zhihu.com/p/441591808

通常我們訓(xùn)練神經(jīng)網(wǎng)絡(luò)模型的時候默認(rèn)使用的數(shù)據(jù)類型為單精度FP32。近年來，為了加快訓(xùn)練時間、減少網(wǎng)絡(luò)訓(xùn)練時候所占用的內(nèi)存，并且保存訓(xùn)練出來的模型精度持平的條件下，業(yè)界提出越來越多的混合精度訓(xùn)練的方法。這里的混合精度訓(xùn)練是指在訓(xùn)練的過程中，同時使用單精度（FP32）和半精度（FP16）。

1 浮點數(shù)據(jù)類型

根據(jù)IEEE二進(jìn)制浮點數(shù)算術(shù)標(biāo)準(zhǔn)（IEEE 754）的定義，浮點數(shù)據(jù)類型分為雙精度（Fp64）、單精度（Fp32）、半精度（FP16）三種，其中每一種都有三個不同的位來表示。FP64表示采用8個字節(jié)共64位，來進(jìn)行的編碼存儲的一種數(shù)據(jù)類型；同理，FP32表示采用4個字節(jié)共32位來表示；FP16則是采用2字節(jié)共16位來表示。如圖所示：

從圖中可以看出，與FP32相比，FP16的存儲空間是FP32的一半，FP32則是FP16的一半。主要分為三個部分：

• 最高位表示符號位sign bit。
• 中間表示指數(shù)位exponent bit。
• 低位表示分?jǐn)?shù)位fraction bit。

以FP16為例子，第一位符號位sign表示正負(fù)符號，接著5位表示指數(shù)exponent，最后10位表示分?jǐn)?shù)fraction。公式為：
$$x=(?1{)}^S\times 2^{E?15}\times (1+\frac{fraction}{1024})$$
同理，一個規(guī)則化FP32的真值為：
$$x=(?1{)}^S\times 2^{E?127}\times 1.M$$
一個規(guī)則化的FP64的真值為：
$$x=(?1{)}^S\times 2^{E?1023}\times (1.M)$$
FP16可以表示的最大值為 $0111101111111111$ ，計算方法為：
$$\begin{gathered} (?1{)}^0\times 2^{30?15}\times 1.1111111111 \\ =1.1111111111(b)\times 2^{15} \\ =1.9990234375(d)times{2}^{15} \\ =65504 \end{gathered}$$
FP16可以表示的（絕對值）最小值為 $0000010000000000$，計算方法為：
$$(?1{)}^0\times 2^{1?15}=6.104\times 10^{?5}$$
因此FP16的最大取值范圍是[-65504 - 66504]，能表示的精度范圍 $2^{?24}$，超過這個數(shù)值的數(shù)字會被直接置0。

2 使用FP16訓(xùn)練問題

首先來看看為什么需要混合精度。使用FP16訓(xùn)練神經(jīng)網(wǎng)絡(luò)，相對比使用FP32帶來的優(yōu)點有：

減少內(nèi)存占用：FP16的位寬是FP32的一半，因此權(quán)重等參數(shù)所占用的內(nèi)存也是原來的一半，節(jié)省下來的內(nèi)存可以放更大的網(wǎng)絡(luò)模型或者使用更多的數(shù)據(jù)進(jìn)行訓(xùn)練。

加快通訊效率：針對分布式訓(xùn)練，特別是在大模型訓(xùn)練的過程中，通訊的開銷制約了網(wǎng)絡(luò)模型訓(xùn)練的整體性能，通訊的位寬少了意味著可以提升通訊性能，減少等待時間，加快數(shù)據(jù)的流通。

計算效率更高：在特殊的AI加速芯片如華為Ascend 910和310系列，或者NVIDIA VOTAL架構(gòu)的Titan V and Tesla V100的GPU上，使用FP16的執(zhí)行運算性能比FP32更加快。

但是使用FP16同樣會帶來一些問題，其中最重要的是1）精度溢出和2）舍入誤差。

**數(shù)據(jù)溢出：**數(shù)據(jù)溢出比較好理解，FP16的有效數(shù)據(jù)表示范圍為 $6.10\times 10^{?5}?65504$，FP32的有效數(shù)據(jù)表示范圍為 $1.4\times 10^{?45}?1.7\times 10^{38}$。可見FP16相比FP32的有效范圍要窄很多，使用FP16替換FP32會出現(xiàn)上溢（Overflow）和下溢（Underflow）的情況。而在深度學(xué)習(xí)中，需要計算網(wǎng)絡(luò)模型中權(quán)重的梯度（一階導(dǎo)數(shù)），因此梯度會比權(quán)重值更加小，往往容易出現(xiàn)下溢情況。

**舍入誤差：**Rounding Error指示是當(dāng)網(wǎng)絡(luò)模型的反向梯度很小，一般FP32能夠表示，但是轉(zhuǎn)換到FP16會小于當(dāng)前區(qū)間內(nèi)的最小間隔，會導(dǎo)致數(shù)據(jù)溢出。如0.00006666666在FP32中能正常表示，轉(zhuǎn)換到FP16后會表示成為0.000067，不滿足FP16最小間隔的數(shù)會強(qiáng)制舍入。

3 混合精度相關(guān)技術(shù)

為了想讓深度學(xué)習(xí)訓(xùn)練可以使用FP16的好處，又要避免精度溢出和舍入誤差。于是可以通過FP16和FP32的混合精度訓(xùn)練（Mixed-Precision），混合精度訓(xùn)練過程中可以引入權(quán)重備份（Weight Backup）、損失放大（Loss Scaling）、精度累加（Precision Accumulated）三種相關(guān)的技術(shù)。

3.1 權(quán)重備份（Weight Backup）

權(quán)重備份主要用于解決舍入誤差的問題。其主要思路是把神經(jīng)網(wǎng)絡(luò)訓(xùn)練過程中產(chǎn)生的激活activations、梯度 gradients、中間變量等數(shù)據(jù)，在訓(xùn)練中都利用FP16來存儲，同時復(fù)制一份FP32的權(quán)重參數(shù)weights，用于訓(xùn)練時候的更新。具體如下圖所示。

從圖中可以了解，在計算過程中所產(chǎn)生的權(quán)重weights，激活activations，梯度gradients等均使用 FP16 來進(jìn)行存儲和計算，其中權(quán)重使用FP32額外進(jìn)行備份。由于在更新權(quán)重公式為:
$$weight=weight+\eta ?gradient$$
深度模型中，lr x gradent的參數(shù)值可能會非常小，利用FP16來進(jìn)行相加的話，則很可能會出現(xiàn)舍入誤差問題，導(dǎo)致更新無效。因此通過將權(quán)重weights拷貝成FP32格式，并且確保整個更新過程是在 fp32 格式下進(jìn)行的。即：
$$weigh{t}_{32}=weigh{t}_{32}+\eta ?gradien{t}_{16}$$
權(quán)重用FP32格式備份一次，那豈不是使得內(nèi)存占用反而更高了呢？是的，額外拷貝一份weight的確增加了訓(xùn)練時候內(nèi)存的占用。 但是實際上，在訓(xùn)練過程中內(nèi)存中分為動態(tài)內(nèi)存和靜態(tài)內(nèi)容，其中動態(tài)內(nèi)存是靜態(tài)內(nèi)存的3-4倍，主要是中間變量值和激活activations的值。而這里備份的權(quán)重增加的主要是靜態(tài)內(nèi)存。只要動態(tài)內(nèi)存的值基本都是使用FP16來進(jìn)行存儲，則最終模型與整網(wǎng)使用FP32進(jìn)行訓(xùn)練相比起來， 內(nèi)存占用也基本能夠減半。

3.2 損失縮放（Loss Scaling）

如圖所示，如果僅僅使用FP32訓(xùn)練，模型收斂得比較好，但是如果用了混合精度訓(xùn)練，會存在網(wǎng)絡(luò)模型無法收斂的情況。原因是梯度的值太小，使用FP16表示會造成了數(shù)據(jù)下溢出（Underflow）的問題，導(dǎo)致模型不收斂，如圖中灰色的部分。于是需要引入損失縮放（Loss Scaling）技術(shù)。

下面是在網(wǎng)絡(luò)模型訓(xùn)練階段， 某一層的激活函數(shù)梯度分布式中，其中有68%的網(wǎng)絡(luò)模型激活參數(shù)位0，另外有4%的精度在 $2^{?32}?2^{?20}$這個區(qū)間內(nèi)，直接使用FP16對這里面的數(shù)據(jù)進(jìn)行表示，會截斷下溢的數(shù)據(jù)，所有的梯度值都會變?yōu)?。

為了解決梯度過小數(shù)據(jù)下溢的問題，對前向計算出來的Loss值進(jìn)行放大操作，也就是把FP32的參數(shù)乘以某一個因子系數(shù)后，把可能溢出的小數(shù)位數(shù)據(jù)往前移，平移到FP16能表示的數(shù)據(jù)范圍內(nèi)。根據(jù)鏈?zhǔn)角髮?dǎo)法則，放大Loss后會作用在反向傳播的每一層梯度，這樣比在每一層梯度上進(jìn)行放大更加高效。

損失放大是需要結(jié)合混合精度實現(xiàn)的，其主要的主要思路是：

• Scale up階段，網(wǎng)絡(luò)模型前向計算后在反響傳播前，將得到的損失變化值DLoss增大 $2^K$倍。
• Scale down階段，反向傳播后，將權(quán)重梯度縮 $2^K$ 倍，恢復(fù)FP32值進(jìn)行存儲。

**動態(tài)損失縮放（Dynamic Loss Scaling）：**上面提到的損失縮放都是使用一個默認(rèn)值對損失值進(jìn)行縮放，為了充分利用FP16的動態(tài)范圍，可以更好地緩解舍入誤差，盡量使用比較大的放大倍數(shù)。總結(jié)動態(tài)損失縮放算法，就是每當(dāng)梯度溢出時候減少損失縮放規(guī)模，并且間歇性地嘗試增加損失規(guī)模，從而實現(xiàn)在不引起溢出的情況下使用最高損失縮放因子，更好地恢復(fù)精度。

動態(tài)損失縮放的算法如下：

動態(tài)損失縮放的算法會從比較高的縮放因子開始（如2^24），然后開始進(jìn)行訓(xùn)練迭代中檢查數(shù)是否會溢出（Infs/Nans）；

如果沒有梯度溢出，則不進(jìn)行縮放，繼續(xù)進(jìn)行迭代；如果檢測到梯度溢出，則縮放因子會減半，重新確認(rèn)梯度更新情況，直到數(shù)不產(chǎn)生溢出的范圍內(nèi)；

在訓(xùn)練的后期，loss已經(jīng)趨近收斂穩(wěn)定，梯度更新的幅度往往小了，這個時候可以允許更高的損失縮放因子來再次防止數(shù)據(jù)下溢。

因此，動態(tài)損失縮放算法會嘗試在每N（N=2000）次迭代將損失縮放增加F倍數(shù)，然后執(zhí)行步驟2檢查是否溢出。

3.3 精度累加（Precision Accumulated）

在混合精度的模型訓(xùn)練過程中，使用FP16進(jìn)行矩陣乘法運算，利用FP32來進(jìn)行矩陣乘法中間的累加（accumulated），然后再將FP32的值轉(zhuǎn)化為FP16進(jìn)行存儲。簡單而言，就是利用FP16進(jìn)行矩陣相乘，利用FP32來進(jìn)行加法計算彌補(bǔ)丟失的精度。 這樣可以有效減少計算過程中的舍入誤差，盡量減緩精度損失的問題。

例如在Nvidia Volta 結(jié)構(gòu)中帶有Tensor Core，可以利用FP16混合精度來進(jìn)行加速，還能保持精度。Tensor Core主要用于實現(xiàn)FP16的矩陣相乘，在利用FP16或者FP32進(jìn)行累加和存儲。在累加階段能夠使用FP32大幅減少混合精度訓(xùn)練的精度損失。

4 混合精度訓(xùn)練策略（Automatic Mixed Precision，AMP）

混合精度訓(xùn)練有很多有意思的地方，不僅僅是在深度學(xué)習(xí)，另外在HPC的迭代計算場景下，從迭代的開始、迭代中期和迭代后期，都可以使用不同的混合精度策略來提升訓(xùn)練性能的同時保證計算的精度。以動態(tài)的混合精度達(dá)到計算和內(nèi)存的最高效率比也是一個較為前言的研究方向。

以NVIDIA的APEX混合精度庫為例，里面提供了4種策略，分別是默認(rèn)使用FP32進(jìn)行訓(xùn)練的O0，只優(yōu)化前向計算部分O1、除梯度更新部分以外都使用混合精度的O2和使用FP16進(jìn)行訓(xùn)練的O3。具體如圖所示。

O1策略中，會根據(jù)實際Tensor和Ops之間的關(guān)系建立黑白名單來使用FP16。例如GEMM和CNN卷積操作對于FP16操作特別友好的計算，會把輸入的數(shù)據(jù)和權(quán)重轉(zhuǎn)換成FP16進(jìn)行運算，而softmax、batchnorm等標(biāo)量和向量在FP32操作好的計算，則是繼續(xù)使用FP32進(jìn)行運算，另外還提供了動態(tài)損失縮放（dynamic loss scaling）。

而O2策略中，模型權(quán)重參數(shù)會轉(zhuǎn)化為FP16，輸入的網(wǎng)絡(luò)模型參數(shù)也轉(zhuǎn)換為FP16，Batchnorms使用FP32，另外模型權(quán)重文件復(fù)制一份FP32用于跟優(yōu)化器更新梯度保持一致都是FP32，另外還提供動態(tài)損失縮放（dynamic loss scaling）。使用了權(quán)重備份來減少舍入誤差和使用損失縮放來避免數(shù)據(jù)溢出。

當(dāng)然上面提供的策略是跟硬件有關(guān)系，并不是所有的AI加速芯片都使用，這時候針對自研的AI芯片，需要找到適合得到混合精度策略。

5 實驗結(jié)果

從下圖的Accuracy結(jié)果可以看到，混合精度基本沒有精度損失：

Loss scale的效果：

題外話，前不久去X公司跟X總監(jiān)聊下一代AI芯片架構(gòu)的時候，他認(rèn)為下一代芯片可以不需要加入INT8數(shù)據(jù)類型，因為Transformer結(jié)構(gòu)目前有大一統(tǒng)NLP和CV等領(lǐng)域的趨勢，從設(shè)計、流片到量產(chǎn)，2年后預(yù)計Transformer會取代CNN成為最流行的架構(gòu)。我倒是不同意這個觀點，目前來看神經(jīng)網(wǎng)絡(luò)的4個主要的結(jié)構(gòu)MLP、CNN、RNN、Transformer都有其對應(yīng)的使用場景，并沒有因為某一種結(jié)構(gòu)的出現(xiàn)而推翻以前的結(jié)構(gòu)。只能說根據(jù)使用場景的側(cè)重點比例有所不同，我理解Int8、fp16、fp32的數(shù)據(jù)類型在AI芯片中仍然會長期存在，針對不同的應(yīng)用場景和計算單元會有不同的比例。

參考文獻(xiàn)：

Micikevicius, Paulius, et al. "Mixed precision training."arXiv preprint arXiv:1710.03740(2017).

Ott, Myle, et al. "Scaling neural machine translation."arXiv preprint arXiv:1806.00187(2018).

https://en.wikipedia.org/wiki/Half-precision_floating-point_format

apex.amp - Apex 0.1.0 documentation.

Automatic Mixed Precision for Deep Learning.

Training With Mixed Precision.

Dreaming.O：淺談混合精度訓(xùn)練.

總結(jié)

以上是生活随笔為你收集整理的混合精度训练的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

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