點藍色字關注「機器學習算法工程師」
設為星標，乾貨直達！
AI編輯：我是小將
本文作者：OpenMMLab @202011
https://zhuanlan.zhihu.com/p/348554267
本文已由原作者授權轉載
Nvidia 在 Volta 架構中引入 Tensor Core 單元，來支持 FP32 和 FP16 混合精度計算。也在 2018 年提出一個 PyTorch 拓展 apex，來支持模型參數自動混合精度訓練。自動混合精度（Automatic Mixed Precision, AMP)訓練，是在訓練一個數值精度 FP32 的模型，一部分算子的操作時，數值精度為 FP16，其餘算子的操作精度是 FP32，而具體哪些算子用 FP16，哪些用 FP32，不需要用戶關心，amp 自動給它們都安排好了。這樣在不改變模型、不降低模型訓練精度的前提下，可以縮短訓練時間，降低存儲需求，因而能支持更多的 batch size、更大模型和尺寸更大的輸入進行訓練。PyTorch 從 1.6 以後（在此之前 OpenMMLab 已經支持混合精度訓練，即 Fp16OptimizerHook），開始原生支持 amp，即torch.cuda.amp module。2020 ECCV，英偉達官方做了一個 tutorial 推廣 amp。從官方各種文檔網頁 claim 的結果來看，amp 在分類、檢測、圖像生成、3D CNNs、LSTM，以及 NLP 中機器翻譯、語義識別等應用中，都在沒有降低模型訓練精度都前提下，加速了模型的訓練速度。
本文是對torch.cuda.amp工作機制，和 module 中接口使用方法介紹，以及在算法角度上對 amp 不掉點原因進行分析，最後補充一點對 amp 存儲消耗的解釋。
1. 混合精度訓練機制torch.cuda.amp 給用戶提供了較為方便的混合精度訓練機制，「方便」體現在兩個方面：
以上兩點，分別是通過使用amp.autocast和amp.GradScaler來實現的。
autocast可以作為 Python 上下文管理器和裝飾器來使用，用來指定腳本中某個區域、或者某些函數，按照自動混合精度來運行。混合精度在操作的時候，是先將 FP32 的模型的參數拷貝一份，拷貝的參數轉換成 FP16，而 amp 規定了的 FP16 的算子（例如卷積、全連接），對 FP16 的數值進行操作；FP32 的算子（例如涉及 reduction 的算子，BatchNormalize，softmax...），輸入和輸出是 FP16，計算的精度是 FP32。在反向傳播時，依然是混合精度計算，得到數值精度為 FP16 的梯度。最後，由於 GPU 中的 Tensor Core 天然支持 FP16 乘積的結果與 FP32 的累加（Tensor Core math），優化器的操作是利用 FP16 的梯度對 FP32 的參數進行更新。
對於 FP16 不可避免的問題就是：表示的範圍較窄，如下圖所示，大量非 0 梯度會遇到溢出問題。解決辦法是：對梯度乘一個 的係數，稱為 scale factor，把梯度 shift 到 FP16 的表示範圍。
GradScaler的工作就是在反向傳播前給 loss 乘一個 scale factor，所以之後反向傳播得到的梯度都乘了相同的 scale factor。並且為了不影響學習率，在梯度更新前將梯度unscale。總結amp的基本訓練流程：
維護一個 FP32 數值精度模型的副本
在每個iteration
拷貝並且轉換成 FP16 模型
前向傳播（FP16 的模型參數）
loss 乘 scale factor s
反向傳播（FP16 的模型參數和參數梯度）
參數梯度乘 1/s
利用 FP16 的梯度更新 FP32 的模型參數
但是，這裡會有一個問題，scale factor 應該如何選取？選一個常量顯然是不合適的，因為 loss 和梯度的數值在變，scale factor 需要跟隨 loss 動態變化。健康的 loss 是振蕩中下降，因此GradScaler設計的 scale factor 每隔 個 iteration 乘一個大於 1 的係數，再 scale loss；並且每次更新前檢查溢出問題（檢查梯度中有沒有inf和nan），如果有，scale factor 乘一個小於 1 的係數並跳過該 iteration 的參數更新環節，如果沒有，就正常更新參數。動態更新 scale factor 是 amp 實際操作中的流程。總結 amp 動態 scale factor 的訓練流程：
維護一個 FP32 數值精度模型的副本
初始化 s
在每個 iteration + a 拷貝並且轉換成FP16模型 + b 前向傳播（FP16 的模型參數） + c loss 乘 scale factor s + d 反向傳播（FP16 的模型參數和參數梯度） + e 檢查有沒有inf或者nan的參數梯度 + 如果有：降低 s，回到步驟a + f 參數梯度乘 1/s + g 利用 FP16 的梯度更新 FP32 的模型參數
2. amp模塊的API用戶使用混合精度訓練基本操作：
# amp依賴Tensor core架構，所以model參數必須是cuda tensor類型
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# GradScaler對象用來自動做梯度縮放
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
# 在autocast enable 區域運行forward
with autocast():
# model做一個FP16的副本，forward
output = model(input)
loss = loss_fn(output, target)
# 用scaler，scale loss(FP16)，backward得到scaled的梯度(FP16)
scaler.scale(loss).backward()
# scaler 更新參數，會先自動unscale梯度
# 如果有nan或inf，自動跳過
scaler.step(optimizer)
# scaler factor更新
scaler.update()
autocast(enable=True)`` 可以作為上下文管理器和裝飾器來使用，給算子自動安排按照 FP16 或者 FP32 的數值精度來操作。
2.1.1 autocast算子PyTorch中，只有 CUDA 算子有資格被 autocast，而且只有 「out-of-place」 才可以被 autocast，例如：a.addmm(b, c)是可以被 autocast，但是a.addmm_(b, c)和a.addmm(b, c, out=d)不可以 autocast。amp autocast 成 FP16 的算子有：
autocast 成 FP32 的算子：
剩下沒有列出的算子，像dot,add,cat...都是按數據中較大的數值精度，進行操作，即有 FP32 參與計算，就按 FP32，全是 FP16 參與計算，就是 FP16。
2.1.2 MisMatch error作為上下文管理器使用時，混合精度計算 enable 區域得到的 FP16 數值精度的變量在 enable 區域外需要顯式的轉成 FP32：
# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")
with autocast():
# torch.mm is on autocast's list of ops that should run in float16.
e_float16 = torch.mm(a_float32, b_float32)
# Also handles mixed input types
f_float16 = torch.mm(d_float32, e_float16)
# After exiting autocast, calls f_float16.float() to use with d_float32
g_float32 = torch.mm(d_float32, f_float16.float())
# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")
with autocast():
e_float16 = torch.mm(a_float32, b_float32)
with autocast(enabled=False):
f_float32 = torch.mm(c_float32, e_float16.float())
g_float16 = torch.mm(d_float32, f_float32)
這種情況一般用於 data parallel 的模型的，autocast 設計為 「thread local」 的，所以只在 main thread 上設 autocast 區域是不 work 的:
model = MyModel()
dp_model = nn.DataParallel(model)
with autocast(): # dp_model's internal threads won't autocast.
#The main thread's autocast state has no effect.
output = dp_model(input) # loss_fn still autocasts, but it's too late...
loss = loss_fn(output)
正確姿勢是對 forward 裝飾：
MyModel(nn.Module):
...
@autocast()
def forward(self, input):
...
另一個正確姿勢是在 forward 的裡面設 autocast 區域：
MyModel(nn.Module):
...
def forward(self, input):
with autocast():
...
forward 函數處理之後，在 main thread 裡 autocast
model = MyModel()
dp_model = nn.DataParallel(model)
with autocast():
output = dp_model(input)
loss = loss_fn(output)
對於用戶自定義的 autograd 函數，需要用amp.custom_fwd裝飾 forward 函數，amp.custom_bwd裝飾 backward 函數：
class MyMM(torch.autograd.Function):
@staticmethod
@custom_fwd
def forward(ctx, a, b):
ctx.save_for_backward(a, b)
return a.mm(b)
@staticmethod
@custom_bwd
def backward(ctx, grad):
a, b = ctx.saved_tensors
return grad.mm(b.t()), a.t().mm(grad)
調用時再 autocast
mymm = MyMM.apply
with autocast():
output = mymm(input1, input2)
autocast主要實現接口有：
A. __enter__
def __enter__(self):
self.prev = torch.is_autocast_enabled()
torch.set_autocast_enabled(self._enabled)
torch.autocast_increment_nesting()
B. __exit__
def __exit__(self, *args):
if torch.autocast_decrement_nesting() == 0:
torch.clear_autocast_cache()
torch.set_autocast_enabled(self.prev)
return False
C. __call__
def __call__(self, func):
@functools.wraps(func)
def decorate_autocast(*args, **kwargs):
with self:
return func(*args, **kwargs)
return decorate_autocast
其中torch.*autocast*函數是在 pytorch/aten/src/ATen/autocast_mode.cpp 裡實現。PyTorch ATen 是 A TENsor library for C++11，ATen 部分有大量的代碼是來聲明和定義 Tensor 運算相關的邏輯的。autocast_mode.cpp 實現策略是 「 cache fp16 casts of fp32 model weights」。
2.2 GradScaler 類torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)用於動態 scale 梯度
+. init_scale: scale factor 的初始值 +. growth_factor: 每次 scale factor 的增長係數 +. backoff_factor: scale factor 下降係數 +. growth_interval: 每隔多個 interval 增長 scale factor +. enabled: 是否做 scale
2.2.1 scale(output)方法對outputs乘 scale factor，並返回，如果enabled=False就原樣返回。
2.2.3 step(optimizer, *args, **kwargs)方法step 方法在做兩件事情：
注意：GradScaler的step不支持傳 closure。
2.2.4 update(new_scale=None)方法update方法在每個 iteration 結束前都需要調用，如果參數更新跳過，會給 scale factor 乘backoff_factor，或者到了該增長的 iteration，就給 scale factor 乘growth_factor。也可以用new_scale直接更新 scale factor。
2.3 舉例2.3.1 Gradient clippingscaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
with autocast():
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
# unscale 梯度，可以不影響clip的threshold
scaler.unscale_(optimizer)
# clip梯度
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
# unscale_（）已經被顯式調用了，scaler正常執行step更新參數，有nan/inf也會跳過
scaler.step(optimizer)
scaler.update()
scaler = GradScaler()
for epoch in epochs:
for i, (input, target) in enumerate(data):
with autocast():
output = model(input)
loss = loss_fn(output, target)
# loss 根據 累加的次數歸一一下
loss = loss / iters_to_accumulate
# scale 歸一的loss 並backward
scaler.scale(loss).backward()
if (i + 1) % iters_to_accumulate == 0:
# may unscale_ here if desired
# (e.g., to allow clipping unscaled gradients)
# step() and update() proceed as usual.
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
with autocast():
output = model(input)
loss = loss_fn(output, target)
# 防止溢出，在不是autocast 區域，先用scaled loss 得到 scaled 梯度
scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss),
inputs=model.parameters(),
create_graph=True)
# 梯度unscale
inv_scale = 1./scaler.get_scale()
grad_params = [p * inv_scale for p in scaled_grad_params]
# 在autocast 區域，loss 加上梯度懲罰項
with autocast():
grad_norm = 0
for grad in grad_params:
grad_norm += grad.pow(2).sum()
grad_norm = grad_norm.sqrt()
loss = loss + grad_norm
scaler.scale(loss).backward()
# may unscale_ here if desired
# (e.g., to allow clipping unscaled gradients)
# step() and update() proceed as usual.
scaler.step(optimizer)
scaler.update()
scaler 一個就夠，但 scale(loss) 和 step(optimizer) 要分別執行
scaler = torch.cuda.amp.GradScaler()
for epoch in epochs:
for input, target in data:
optimizer0.zero_grad()
optimizer1.zero_grad()
with autocast():
output0 = model0(input)
output1 = model1(input)
loss0 = loss_fn(2 * output0 + 3 * output1, target)
loss1 = loss_fn(3 * output0 - 5 * output1, target)
# (retain_graph here is unrelated to amp, it's present because in this
# example, both backward() calls share some sections of graph.)
scaler.scale(loss0).backward(retain_graph=True)
scaler.scale(loss1).backward()
# You can choose which optimizers receive explicit unscaling, if you
# want to inspect or modify the gradients of the params they own.
scaler.unscale_(optimizer0)
scaler.step(optimizer0)
scaler.step(optimizer1)
scaler.update()
torch DDP 和 torch DP model 的處理方式一樣
Q1. amp 是如何做到 FP16 和 FP32 混合使用，「還不掉點」模型量化、模型壓縮的算法挺多的，但都做不 amp 這樣，對多數模型訓練不掉點（但是實操中，聽有經驗的大神介紹，完全不到點還是有點難度的）。amp 能做成這樣，因為它對模型沒有壓縮和量化，維護的還是一個 32 位的模型。只是用 16 位去表示原來 32 位的梯度：通常模型訓練依賴 FP32 的精度，因為梯度會有一部分 FP16 表示不了，而 scale factor 把梯度 shift 到 FP16 能表示範圍，使得梯度方面精度的損失較小，可能 forward 時候的直接的精度壓縮是訓練最大的損失。
Q2. 沒有 Tensor Core 架構能否使用 amp沒有 Tensor Core 架構的 GPU 試用 amp，速度反而下降，但顯存會明顯減少。作者在 Turing 架構的 GTX 1660 上試用 amp，運算時間增加了一倍，但顯存不到原來的一半。
Q3. 為什麼 amp 中有兩份參數，存儲消耗反而更小相比與模型參數，對中間層結果的存儲更是 deep learning 的 bottleneck。當對中間結果的存儲砍半，整個存儲消耗就基本上原來的一半。
·················END·················