本文摘自我的個人 github
寫在文前:本文所寫其實不僅用於大模型,隻是模型越大,溢出的可能性增加瞭,而大模型訓練的收斂到目前為止仍然是一件不輕松的事情:
fp32/fp16/bf16
- fp32/fp16 絕大多數硬件都支持,所以可以用混合精度訓練提高吞吐;但 bf16/tf32 隻有新的硬件才支持,V100/昇騰910等不支持
- bf16 具有和 fp32 相同的 range,但精度(也就是兩個最小單位之間的間隔)降低
- bf16/fp32 進行混合精度訓練,可以減少溢出幾率
- 對於大型 transformer,bf16 損失的精度被證明不怎麼影響收斂
- tf32 是 A100 中引入的新格式,用於替代 fp32,也即可以全程 tf32 訓練或 bf16/tf32 混合訓練
混合精度訓練的要點
- bf16/fp32 混合訓練因為兩種格式在 range 對齊瞭,並且 bf16 比 fp16 range 更大,所以比 fp16/fp32 混合訓練穩定性更高。但 fp16/fp32 混合訓練 GPT-3 大模型也是完全可行的,隻要解決可溢出問題,有以下幾個要點:
- fp32權重備份 + loss scaling 解決下溢出問題
- 對 loss 進行 scale:見左圖
- 對 gradient 進行 scale:見右圖 由於鏈式法則的存在,對梯度做直接做 scale 也是可以的,反而更劃算。這樣,所有前向後向都可以全用 fp16 (activation、weight、gradient 全用 fp16),隻在進行更新的時候,才用 fp32 和 master weight 更新
7788f050d705ad798077470a212b39f4
- 跳過 NaN 的 batch
- dynamic loss scale (PyTorch 中采用的這種方案):在一個 window 內,如果沒有 NaN,loss scale 變大(例如乘2);如果有 NaN,則降低 loss scale,並跳過 NaN 的 batch,不更新梯度
- 或將 INF 值改為 FP16 的最大值(需要實際驗證)
- 基於 Tensorcore 進行矩陣乘加:在某些模型中,fp16 矩陣乘法的過程中,需要利用 fp32 來進行矩陣乘法中間結果的累加,以減少加法過程中的舍入誤差,保證精度不損失 這也是為什麼有些人說,隻有 Volta 之後有 TensorCore 的 GPU (例如V100),才能利用 fp16+混合精度來進行加速。其他 GPU 如果硬要用 fp16,隻能做 fp16 的 multiply-add operation,會有較大精度損失
3dd3281a18d86a5c2421d1fb87d9e52c
溢出案例分析
- 有時,當 loss scale 降為 1,然後不再上升;同時訓練 loss 要麼跳漲,要麼停滯不下降(相當於沒在訓練瞭)
- 解決方案:可以改變超參,**降低lr**,或** 增大 epsilon**(但這樣做會損失一部分優化器自適應學習率的能力)
- 原因如下:
- loss scale 不再上升:因為每個 window 中都有 Nan
- loss 跳漲:主要來源於 Adam 的缺陷
- Adam 相比 SGD 的優勢是會根據梯度自適應學習率,但當梯度的二階矩過小,導致分母 `v_t+ε` 在 fp16 中下溢出變為0(fp16 最小值為 5.96e-8,但 epsilon 一般設為 1e-8),反而會上溢,模型的改變量變為很大,從而導致訓練不收斂
c347a1cc0735ae26b5e994f10599d835
- 當出現一個難度較大的 batch,在這個 batch 中出現瞭 NaN,這些 NaN 的梯度被跳過瞭。但通過這組難度較大的情況後,因為有 momentum,之後 batch 的梯度可能就一直下溢瞭loss 停滯:可能仍然發生瞭很多 NaN,但 Adam 中分母 `v_t+ε` 這一項沒有嚴重到下溢,隻是因為時不時有 NaN 發生,導致 loss scale 保持為1,使得模型幾乎不再更新
- 對 softmax 改造 > [如何防止softmax函數上溢出 (overflow) 和下溢出 (underflow)](http://www.codelast.com/%E5%8E%9F%E5%88%9B-%E5%A6%82%E4%BD%95%E9%98%B2%E6%AD%A2softmax%E5%87%BD%E6%95%B0%E4%B8%8A%E6%BA%A2%E5%87%BAoverflow%E5%92%8C%E4%B8%8B%E6%BA%A2%E5%87%BAunderflow/) 用 f(x-max) 代替 f(x);再加上 logsoftmax;分子避免瞭下溢、分母避免瞭上溢
如何盡可能早地發現溢出的可能? 如何解決混合精度訓練大模型的局限性問題 - overfit.cn如上所述,當梯度下溢加上 adam,反而可能變成上溢。一般經驗是,在訓練前期,可以用一個 Tensor Collection Hook (TCH),收集溢出比率,比率在 1% 以下並保持穩定而非上升趨勢,後期發生溢出可能性較小