PyTorch 1.6 之前，大家都是用 NVIDIA 的 apex 庫來實現(xiàn) AMP 訓練。1.6 版本之后，PyTorch 出廠自帶 AMP，僅需幾行代碼，就能讓顯存占用減半，訓練速度加倍

混合精度訓練 (Mixed Precision Training)

單精度浮點數(shù) (FP32) 和半精度浮點數(shù) (FP16)

• PyTorch 默認使用單精度浮點數(shù) (FP32) 來進行網(wǎng)絡(luò)模型的計算和權(quán)重存儲，表示范圍為 $\left[?3e^{38},?1e^{?38}\right]\cup \left[1e^{?38},3e^{38}\right]$. 而半精度浮點數(shù) (FP16) 表示范圍只有 $\left[?6.5e^4,?5.9e^{?8}\right]\cup \left[5.9e^{?8},6.5e^4\right]$，可以看到 FP32 能夠表示的范圍要比 FP16 大的多得多
  其中sign 位表示正負，exponent 位表示指數(shù)，fraction 位表示分數(shù)
• 此外浮點數(shù)還存在舍入誤差，當兩個數(shù)字相差太大時，相加是無效的。例如 $2^{?3}+2^{?14}$ 在 FP32 中就不會有問題，但在 FP16 中，由于 FP16 表示的固定間隔為 $2^{?13}$，因此 $2^{?14}$ 加了跟沒加一樣

# FP32
>>> torch.tensor(2**-3) + torch.tensor(2**-14)
tensor(0.1251)# FP16
>>> torch.tensor(2**-3).half() + torch.tensor(2**-14).half()
tensor(0.1250, dtype=torch.float16)

---

對于 float16：

• 如果 Exponent 位全部為 0：
  • 如果 fraction 位全部為 0，則表示數(shù)字 0
  • 如果 fraction 位不為 0，則表示一個非常小的數(shù)字 (subnormal numbers)，其計算方式為 $(?1{)}^{signbit}\times 2^{?14}\times (0+\frac{fraction}{1024})$
• 如果 Exponent 位全部為 1:
  • 如果 fraction 位全部為 0，則表示 $\pm inf$
  • 如果 fraction 位不為0，則表示 NAN
• Exponent 位的其他情況：$(?1{)}^{signbit}\times (exponent\times 2^{?15})\times (1+\frac{fraction}{1024})$

為什么要用 FP16

• 如果我們在訓練過程中將 FP32 替代為 FP16，有以下兩個好處：(1) 減少顯存占用: FP16 的顯存占用只有 FP32 的一半，這使得我們可以用更大的 batch size；(2) 加速訓練: 使用 FP16，模型的訓練速度幾乎可以提升 1 倍

為什么只用 FP16 會有問題

如果我們簡單地把模型權(quán)重和輸入從 FP32 轉(zhuǎn)化成 FP16，雖然速度可以翻倍，但是模型的精度會被嚴重影響。原因如下:

• 上/下溢出: FP16 的表示范圍不大，超過 $6.5e^4$ 的數(shù)字會上溢出變成 inf，小于
  $5.9e^{?8}$ 的數(shù)字會下溢出變成 0。下溢出更加常見，因為在網(wǎng)絡(luò)訓練的后期，模型的梯度往往很小，甚至會小于 FP16 的下限，此時梯度值就會變成 0，模型參數(shù)無法更新。下圖為 SSD 網(wǎng)絡(luò)在訓練過程中的梯度統(tǒng)計，有 67% 的值下溢出變成 0
  
• 舍入誤差: 就算梯度不會上/下溢出，如果梯度值和模型的參數(shù)值相差太遠，也會發(fā)生舍入誤差的問題。假設(shè)模型參數(shù) $w=2^{?3}$，學習率 $\eta =2^{?2}$，梯度 $g=2^{?12}$，則 $w^{\prime }=w+\eta \times g=2^{?3}+2^{?2}\times 2^{?12}=2^{?3}$

解決方案

損失縮放 (Loss Scaling)

• 為了解決下溢出的問題，論文中對計算出來的 loss 值進行縮放 (scale)，由于鏈式法則的存在，對 loss 的縮放會作用在每個梯度上。縮放后的梯度，就會平移到 FP16 的有效范圍內(nèi)。這樣就可以用 FP16 存儲梯度而又不會溢出了。此外，在進行更新之前，需要先將縮放后的梯度轉(zhuǎn)化為 FP32，再將梯度反縮放 (unscale) 回去以便進行參數(shù)的梯度下降 (注意這里一定要先轉(zhuǎn)成 FP32，不然 unscale 的時候還是會下溢出)
• 縮放因子 (loss_scale) 一般都是框架自動確定的，只要沒有發(fā)生 inf 或者 nan，loss_scale 越大越好。因為隨著訓練的進行，網(wǎng)絡(luò)的梯度會越來越小，更大的 loss_scale 可以更加充分地利用 FP16 的表示范圍

FP32 權(quán)重備份

• 為了實現(xiàn) FP16 的訓練，我們需要把模型權(quán)重和輸入數(shù)據(jù)都轉(zhuǎn)成 FP16，反向傳播的時候就會得到 FP16 的梯度。如果此時直接進行更新，因為梯度 $\times$ 學習率的值往往較小，和模型權(quán)重的差距會很大，可能會出現(xiàn)舍入誤差的問題
• 解決思路是: 將模型權(quán)重、激活值、梯度等數(shù)據(jù)用 FP16 來存儲，同時維護一份 FP32 的模型權(quán)重副本用于更新。在反向傳播得到 FP16 的梯度以后，將其轉(zhuǎn)化成 FP32 并 unscale，最后更新 FP32 的模型權(quán)重。因為整個更新過程是在 FP32 的環(huán)境中進行的，所以不會出現(xiàn)舍入誤差

黑名單

• 對于那些在 FP16 環(huán)境中運行不穩(wěn)定的模塊，我們會將其添加到黑名單中，強制它在 FP32 的精度下運行。比如需要計算 batch 均值的 BN 層就應(yīng)該在 FP32 下運行，否則會發(fā)生舍入誤差。還有一些函數(shù)對于算法精度要求很高，比如 torch.acos()，也應(yīng)該在 FP32 下運行
• 如何保證黑名單模塊在 FP32 環(huán)境中運行: 以 BN 層為例，將其權(quán)重轉(zhuǎn)為 FP32，并且將輸入從 FP16 轉(zhuǎn)成 FP32，這樣就可以保證整個模塊是在 FP32 下運行的

Tensor Core

• Tensor Core 可以讓 FP16 做矩陣相乘，然后把結(jié)果累加到 FP32 的矩陣中。這樣既可以享受 FP16 高速的矩陣乘法，又可以利用 FP32 來消除舍入誤差
  

NVIDIA apex 庫代碼解讀

opt-level (o1, o2, o3, o4)

• 首先介紹下 apex 提供的幾種 opt-level: o1, o2, o3, o4
• o0 是純 FP32，用來當精度的基準。o3 是純 FP16，用來當速度的基準
• 重點講 o1 和 o2 。我們之前講的 AMP 策略其實就是 o2: 除了 BN 層的權(quán)重和輸入使用 FP32，模型的其余權(quán)重和輸入都會轉(zhuǎn)化為 FP16。此外還會創(chuàng)建一個 FP32 的權(quán)重副本來執(zhí)行更新操作
• 和 o2 不同， o1 不再需要 FP32 權(quán)重備份，因為 o1 的模型一直都是 FP32。 可能有些讀者會好奇，既然模型參數(shù)是 FP32，那怎么在訓練過程中使用 FP16 呢？答案是 o1 建立了一個 PyTorch 函數(shù)的黑白名單，對于白名單上的函數(shù)，強制要求其用 FP16，即會將函數(shù)的參數(shù)先轉(zhuǎn)化為 FP16，再執(zhí)行函數(shù)本身。黑名單則強制要求 FP32。以 nn.Linear 為例， 這個模塊有兩個權(quán)重參數(shù) weight 和 bias，輸入為 input，前向傳播就是調(diào)用了 torch.nn.functional.linear(input, weight, bias)。 o1 模式會將 input、weight、bias 先轉(zhuǎn)化為 FP16 格式 input_fp16、weight_fp16、bias_fp16，再調(diào)用函數(shù) torch.nn.functional.linear(input_fp16, weight_fp16, bias_fp16)。這樣一來就實現(xiàn)了模型參數(shù)是 FP32，但是仍然可以使用 FP16 來加速訓練。o1 還有一個細節(jié): 雖然白名單上的 PyTorch 函數(shù)是以 FP16 運行的，但是產(chǎn)生的梯度是 FP32，所以不需要手動將其轉(zhuǎn)成 FP32 再 unscale，直接 unscale 即可。通常來說 o1 比 o2 更穩(wěn)，一般先選擇 o1，再嘗試 o2 看是否掉點，如果不掉點就用 o2

apex 的 o1 實現(xiàn)

• (1) 根據(jù)黑白名單對 PyTorch 內(nèi)置的函數(shù)進行包裝。白名單函數(shù)強制 FP16，黑名單函數(shù)強制 FP32。其余函數(shù)則根據(jù)參數(shù)類型自動判斷，如果參數(shù)都是 FP16，則以 FP16 運行，如果有一個參數(shù)為 FP32，則以 FP32 運行
• (2) 將 loss_scale 初始化為一個很大的值
• (3) 對于每次迭代
  • (a). 前向傳播: 模型權(quán)重是 FP32，按照黑白名單自動選擇算子精度
  • (b). 將 loss 乘以 loss_scale
  • ($c$). 反向傳播: 因為模型權(quán)重是 FP32，所以即使函數(shù)以 FP16 運行，也會得到 FP32 的梯度
  • (d). 將梯度 unscale，即除以 loss_scale
  • (e). 如果檢測到 inf 或 nan.
    • i. loss_scale /= 2
    • ii. 跳過此次更新
  • (f). optimizer.step()，執(zhí)行此次更新
  • (g). 如果連續(xù) 2000 次迭代都沒有出現(xiàn) inf 或 nan，則 loss_scale *= 2

apex 的 o2 實現(xiàn)

• (1) 將除了 BN 層以外的模型權(quán)重轉(zhuǎn)化為 FP16，并且包裝了 forward 函數(shù)，將其參數(shù)也轉(zhuǎn)化為 FP16
• (2) 維護一個 FP32 的模型權(quán)重副本用于更新
• (3) 將 loss_scale 初始化為一個很大的值
• (4) 對于每次迭代
  • (a). 前向傳播: 除了 BN 層是 FP32，模型其它部分都是 FP16
  • (b). 將 loss 乘以 loss_scale
  • ($c$). 反向傳播，得到 FP16 的梯度
  • (d). 將 FP16 梯度轉(zhuǎn)化為 FP32，并 unscale
  • (e). 如果檢測到 inf 或 nan
    • i. loss_scale /= 2
    • ii. 跳過此次更新
  • (f). optimizer.step()，執(zhí)行此次更新
  • (g). 如果連續(xù) 2000 次迭代都沒有出現(xiàn) inf 或 nan，則 loss_scale *= 2

在 PyTorch 中使用混合精度訓練

Automatic Mixed Precision (AMP)

from torch.cuda.amp import autocast, GradScaler

• 通常 AMP 需要同時使用 autocast 和 GradScaler，其中 autocast 的實例對象是作為上下文管理器 (context manger) 或裝飾器 (decorator) 來允許用戶代碼的某些區(qū)域在混合精度下運行，自動為 CUDA 算子選擇(單/半)精度來提升性能并保持精度 (See the Autocast Op Reference for details on what precision autocast chooses for each op, and under what circumstances.)，并且 autocast 區(qū)域是可以嵌套的，這可以強制讓 FP16 下可能溢出的模型部分以 FP32 運行；而 GradScaler 則是用來進行 loss scale
• autocast 應(yīng)該只封裝網(wǎng)絡(luò)的前向傳播 (forward pass(es))，以及損失計算 (loss computation(s))。反向傳播不推薦在 autocast 區(qū)域內(nèi)執(zhí)行，反向傳播的操作會自動以對應(yīng)的前向傳播的操作的數(shù)據(jù)類型運行

Typical Mixed Precision Training

# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)# Creates a GradScaler once at the beginning of training.
scaler = GradScaler(enabled=True)for epoch in epochs:for input, target in data:optimizer.zero_grad()# Runs the forward pass with autocasting.with autocast(enabled=True, dtype=torch.float16):output = model(input)loss = loss_fn(output, target)# Scales loss.  Calls backward() on scaled loss to create scaled gradients.scaler.scale(loss).backward()# scaler.step() first unscales the gradients of the optimizer's assigned params.# If these gradients do not contain infs or NaNs, optimizer.step() is then called,# otherwise, optimizer.step() is skipped.scaler.step(optimizer)# Updates the loss scale value for next iteration.scaler.update()

Saving/Resuming

checkpoint = {"model": net.state_dict(),"optimizer": opt.state_dict(),"scaler": scaler.state_dict()}

net.load_state_dict(checkpoint["model"])
opt.load_state_dict(checkpoint["optimizer"])
scaler.load_state_dict(checkpoint["scaler"])

Working with Unscaled Gradients (Gradient Clipping)

• 經(jīng)過 scaler.scale(loss).backward() 得到的梯度是 scaled gradient，如果想要在 scaler.step(optimizer) 前進行梯度裁剪等操作，就必須先用 scaler.unscale_(optimizer) 得到 unscaled gradient

scaler = GradScaler()for epoch in epochs:for input, target in data:optimizer.zero_grad()with autocast(dtype=torch.float16):output = model(input)loss = loss_fn(output, target)scaler.scale(loss).backward()# Unscales the gradients of optimizer's assigned params in-placescaler.unscale_(optimizer)# Since the gradients of optimizer's assigned params are unscaled, clips as usual:torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)# optimizer's gradients are already unscaled, so scaler.step does not unscale them,# although it still skips optimizer.step() if the gradients contain infs or NaNs.scaler.step(optimizer)# Updates the scale for next iteration.scaler.update()

Working with Scaled Gradients

Gradient accumulation

• Gradient accumulation 基于 effective batch of size batch_per_iter * iters_to_accumulate (* num_procs if distributed) 進行梯度累加，因此屬于同一個 effective batch 的多個迭代 batch 內(nèi)，scale factor 應(yīng)該保持不變 (scale updates should occur at effective-batch granularity)，并且累加的梯度應(yīng)該是 Scaled Gradients。因為如果在梯度累加結(jié)束前的某一個迭代中 unscale gradient (或改變 scale factor)，那么下一個迭代的梯度回傳就會把 scaled grads 加到 unscaled grads (或乘上了不同 scale factor 的 scaled grads) 上，這會使得在最后進行梯度更新時，我們無法恢復出 accumulated unscaled grads. 如果想要 unscaled grads，應(yīng)該在梯度累加結(jié)束后調(diào)用 scaler.unscale_(optimizer)

scaler = GradScaler()for epoch in epochs:for i, (input, target) in enumerate(data):with autocast(dtype=torch.float16):output = model(input)loss = loss_fn(output, target)loss = loss / iters_to_accumulate# Accumulates scaled gradients.scaler.scale(loss).backward()if (i + 1) % iters_to_accumulate == 0:# may unscale_ here if desired (e.g., to allow clipping unscaled gradients)scaler.step(optimizer)scaler.update()optimizer.zero_grad()

Gradient penalty

• https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-penalty

Working with Multiple GPUs

• 目前的版本中 (v1.13)，不管是 DP (one GPU per thread) (多線程) 還是 DDP (one GPU per process) (多進程)，上述代碼都無需改動。只有當使用 DDP (multiple GPUs per process) 時，才需要給 model 的 forwad 方法添加 autocast 裝飾器或上下文管理器
• 當然，如果使用老版本的 pytorch，是否需要改動代碼請參考官方文檔

其他注意事項

• 常數(shù)的范圍：為了保證計算不溢出，首先要保證人為設(shè)定的常數(shù)不溢出，如各種 epsilon，INF (改成 -float('inf') 就可以啦)

References

• paper: Micikevicius, Paulius, et al. “Mixed precision training.” (ICLR, 2018).
• AUTOMATIC MIXED PRECISION PACKAGE - TORCH.AMP
• CUDA AUTOMATIC MIXED PRECISION EXAMPLES
• Automatic Mixed Precision Recipe
• 由淺入深的混合精度訓練教程
• 【PyTorch】唯快不破：基于 Apex 的混合精度加速
• 淺談混合精度訓練
• 【Trick2】torch.cuda.amp自動混合精度訓練 —— 節(jié)省顯存并加快推理速度
• 自動混合精度訓練 (AMP) – PyTorch