混合精度訓(xùn)練原理總結(jié)

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

1、浮點(diǎn)數(shù)據(jù)類型

浮點(diǎn)數(shù)據(jù)類型主要分為雙精度（Fp64）、單精度（Fp32）、半精度（FP16）。在神經(jīng)網(wǎng)絡(luò)模型的訓(xùn)練過(guò)程中，一般默認(rèn)采用單精度（FP32）浮點(diǎn)數(shù)據(jù)類型，來(lái)表示網(wǎng)絡(luò)模型權(quán)重和其他參數(shù)。在了解混合精度訓(xùn)練之前，這里簡(jiǎn)單了解浮點(diǎn)數(shù)據(jù)類型。

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

從圖中可以看出，與FP32相比，F(xiàn)P16的存儲(chǔ)空間是FP32的一半，F(xiàn)P32則是FP16的一半。主要分為三個(gè)部分：

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

以FP16為例子，第一位符號(hào)位sign表示正負(fù)符號(hào)，接著5位表示指數(shù)exponent，最后10位表示分?jǐn)?shù)fraction。公式為：

同理，一個(gè)規(guī)則化的FP32的真值為：

一個(gè)規(guī)格化的FP64的真值為：

FP16可以表示的最大值為 0 11110 1111111111，計(jì)算方法為：

FP16可以表示的最小值為 0 00001 0000000000，計(jì)算方法為：

因此FP16的最大取值范圍是[-65504 - 66504]，能表示的精度范圍是 ，超過(guò)這個(gè)數(shù)值的數(shù)字會(huì)被直接置0。

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

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

1. 減少內(nèi)存占用：FP16的位寬是FP32的一半，因此權(quán)重等參數(shù)所占用的內(nèi)存也是原來(lái)的一半，節(jié)省下來(lái)的內(nèi)存可以放更大的網(wǎng)絡(luò)模型或者使用更多的數(shù)據(jù)進(jìn)行訓(xùn)練。
2. 加快通訊效率：針對(duì)分布式訓(xùn)練，特別是在大模型訓(xùn)練的過(guò)程中，通訊的開(kāi)銷制約了網(wǎng)絡(luò)模型訓(xùn)練的整體性能，通訊的位寬少了意味著可以提升通訊性能，減少等待時(shí)間，加快數(shù)據(jù)的流通。
3. 計(jì)算效率更高：在特殊的AI加速芯片如華為Ascend 910和310系列，或者NVIDIA VOTAL架構(gòu)的Titan V and Tesla V100的GPU上，使用FP16的執(zhí)行運(yùn)算性能比FP32更加快。

但是使用FP16同樣會(huì)帶來(lái)一些問(wèn)題，其中最重要的是1）精度溢出和2）舍入誤差。

1. 數(shù)據(jù)溢出： 數(shù)據(jù)溢出比較好理解，F(xiàn)P16的有效數(shù)據(jù)表示范圍為 ，F(xiàn)P32的有效數(shù)據(jù)表示范圍為 ?？梢?jiàn)FP16相比FP32的有效范圍要窄很多，使用FP16替換FP32會(huì)出現(xiàn)上溢（Overflow）和下溢（Underflow）的情況。而在深度學(xué)習(xí)中，需要計(jì)算網(wǎng)絡(luò)模型中權(quán)重的梯度（一階導(dǎo)數(shù)），因此梯度會(huì)比權(quán)重值更加小，往往容易出現(xiàn)下溢情況。
2. 舍入誤差： Rounding Error指示是當(dāng)網(wǎng)絡(luò)模型的反向梯度很小，一般FP32能夠表示，但是轉(zhuǎn)換到FP16會(huì)小于當(dāng)前區(qū)間內(nèi)的最小間隔，會(huì)導(dǎo)致數(shù)據(jù)溢出。如0.00006666666在FP32中能正常表示，轉(zhuǎn)換到FP16后會(huì)表示成為0.000067，不滿足FP16最小間隔的數(shù)會(huì)強(qiáng)制舍入。

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

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

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

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

從圖中可以了解，在計(jì)算過(guò)程中所產(chǎn)生的權(quán)重weights，激活activations，梯度gradients等均使用 FP16 來(lái)進(jìn)行存儲(chǔ)和計(jì)算，其中權(quán)重使用FP32額外進(jìn)行備份。由于在更新權(quán)重公式為:

深度模型中，lr x gradent的參數(shù)值可能會(huì)非常小，利用FP16來(lái)進(jìn)行相加的話，則很可能會(huì)出現(xiàn)舍入誤差問(wèn)題，導(dǎo)致更新無(wú)效。因此通過(guò)將權(quán)重weights拷貝成FP32格式，并且確保整個(gè)更新過(guò)程是在 fp32 格式下進(jìn)行的。即：

權(quán)重用FP32格式備份一次，那豈不是使得內(nèi)存占用反而更高了呢？是的，額外拷貝一份weight的確增加了訓(xùn)練時(shí)候內(nèi)存的占用。但是實(shí)際上，在訓(xùn)練過(guò)程中內(nèi)存中分為動(dòng)態(tài)內(nèi)存和靜態(tài)內(nèi)容，其中動(dòng)態(tài)內(nèi)存是靜態(tài)內(nèi)存的3-4倍，主要是中間變量值和激活activations的值。而這里備份的權(quán)重增加的主要是靜態(tài)內(nèi)存。只要?jiǎng)討B(tài)內(nèi)存的值基本都是使用FP16來(lái)進(jìn)行存儲(chǔ)，則最終模型與整網(wǎng)使用FP32進(jìn)行訓(xùn)練相比起來(lái)， 內(nèi)存占用也基本能夠減半。

3.2、損失縮放（Loss Scaling）

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

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

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

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

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

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

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

1. 動(dòng)態(tài)損失縮放的算法會(huì)從比較高的縮放因子開(kāi)始（如2^24），然后開(kāi)始進(jìn)行訓(xùn)練迭代中檢查數(shù)是否會(huì)溢出（Infs/Nans）；
2. 如果沒(méi)有梯度溢出，則不進(jìn)行縮放，繼續(xù)進(jìn)行迭代；如果檢測(cè)到梯度溢出，則縮放因子會(huì)減半，重新確認(rèn)梯度更新情況，直到數(shù)不產(chǎn)生溢出的范圍內(nèi)；
3. 在訓(xùn)練的后期，loss已經(jīng)趨近收斂穩(wěn)定，梯度更新的幅度往往小了，這個(gè)時(shí)候可以允許更高的損失縮放因子來(lái)再次防止數(shù)據(jù)下溢。
4. 因此，動(dòng)態(tài)損失縮放算法會(huì)嘗試在每N（N=2000）次迭代將損失縮放增加F倍數(shù)，然后執(zhí)行步驟2檢查是否溢出。

3.3、精度累加（Precision Accumulated）

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

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

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

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

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

這里面比較有意思的是O1和O2策略。

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

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

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

5、實(shí)驗(yàn)結(jié)果

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

Loss scale的效果：

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

參考文獻(xiàn)：

1. Micikevicius, Paulius, et al. "Mixed precision training."arXiv preprint arXiv:1710.03740(2017).
2. Ott, Myle, et al. "Scaling neural machine translation."arXiv preprint arXiv:1806.00187(2018).
3. https://en.wikipedia.org/wiki/Half-precision_floating-point_format
4. apex.amp - Apex 0.1.0 documentation.
5. Automatic Mixed Precision for Deep Learning.
6. Training With Mixed Precision.
7. Dreaming.O：淺談混合精度訓(xùn)練.

如果覺(jué)得有用，就請(qǐng)分享到朋友圈吧！