本文摘自我的個人 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% 以下並保持穩定而非上升趨勢，後期發生溢出可能性較小