wehub-resource-sync
249 lines
11 KiB
Python
249 lines
11 KiB
Python
"""Callback and utility functions to allow mixed precision training
|
|
|
|
Docs: https://docs.fast.ai/callback.fp16.html.md"""
|
|
|
|
# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/18_callback.fp16.ipynb.
|
|
|
|
# %% auto #0
|
|
__all__ = ['AMPMode', 'MixedPrecision', 'get_master', 'to_master_grads', 'to_model_params', 'test_overflow', 'grad_overflow',
|
|
'copy_clone', 'ModelToHalf', 'NonNativeMixedPrecision']
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #8d2f192e
|
|
from ..basics import *
|
|
from .progress import *
|
|
|
|
from torch.amp import GradScaler,autocast
|
|
from torch.amp.grad_scaler import OptState
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #2739b62c
|
|
class AMPMode(Enum):
|
|
"Automatic mixed precision modes for ease of completion"
|
|
FP16 = 'fp16'
|
|
BF16 = 'bf16'
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #3d30daa9
|
|
@delegates(GradScaler)
|
|
class MixedPrecision(Callback):
|
|
"Mixed precision training using Pytorch's Automatic Mixed Precision (AMP)"
|
|
order = 10
|
|
def __init__(self,
|
|
amp_mode:str|AMPMode=AMPMode.FP16, # Mixed Precision training mode. Supports fp16 and bf16.
|
|
**kwargs
|
|
):
|
|
amp_mode = AMPMode(amp_mode)
|
|
store_attr(names='amp_mode')
|
|
self.kwargs = kwargs
|
|
|
|
def before_fit(self):
|
|
if self.amp_mode == AMPMode.BF16:
|
|
if torch.cuda.is_available() and not torch.cuda.is_bf16_supported():
|
|
raise ValueError("Unsupported GPU for bfloat16 mixed precision training")
|
|
dtype = torch.bfloat16
|
|
elif self.amp_mode == AMPMode.FP16:
|
|
dtype = torch.float16
|
|
else:
|
|
raise ValueError(f"Unrecognized precision: {self.amp_mode}")
|
|
# `GradScaler` is not needed for bfloat16 as fp32 and bf16 have the same range
|
|
self.kwargs['enabled'] = dtype == torch.float16
|
|
self.autocast,self.learn.scaler,self.scales = autocast('cuda', dtype=dtype),GradScaler('cuda', **self.kwargs),L()
|
|
|
|
def before_batch(self): self.autocast.__enter__()
|
|
def after_pred(self):
|
|
self.learn.pred = to_float(self.pred)
|
|
def after_loss(self): self.autocast.__exit__(None, None, None)
|
|
def before_backward(self): self.learn.loss_grad = self.scaler.scale(self.loss_grad)
|
|
def before_step(self):
|
|
"Use `self` as a fake optimizer. `self.skipped` will be set to True `after_step` if gradients overflow."
|
|
self.skipped=True
|
|
self.scaler.step(self)
|
|
if self.skipped: raise CancelStepException()
|
|
self.scales.append(self.scaler.get_scale())
|
|
def after_step(self): self.learn.scaler.update()
|
|
def after_fit(self): self.autocast,self.learn.scaler,self.scales = None,None,None
|
|
|
|
@property
|
|
def param_groups(self):
|
|
"Pretend to be an optimizer for `GradScaler`"
|
|
return self.opt.param_groups
|
|
def step(self, *args, **kwargs):
|
|
"Fake optimizer step to detect whether this batch was skipped from `GradScaler`"
|
|
self.skipped=False
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #d56db335
|
|
@patch
|
|
@delegates(GradScaler)
|
|
def to_fp16(self:Learner, **kwargs):
|
|
"Set `Learner` to float16 mixed precision using PyTorch AMP"
|
|
return self.add_cb(MixedPrecision(**kwargs))
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #6c6c0e6d
|
|
@patch
|
|
def to_bf16(self:Learner):
|
|
"Set `Learner` to bfloat16 mixed precision using PyTorch AMP"
|
|
return self.add_cb(MixedPrecision(amp_mode=AMPMode.BF16))
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #3d78320f
|
|
@patch
|
|
def to_fp32(self:Learner):
|
|
"Set `Learner` to float32 precision"
|
|
return self.remove_cb(MixedPrecision)
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #3a800e16
|
|
from ..fp16_utils import convert_network, model_grads_to_master_grads, master_params_to_model_params
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #73fd0b34
|
|
from torch.nn.utils import parameters_to_vector
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #1e983fc9
|
|
def get_master(
|
|
opt:Optimizer, # Optimizer from which to retrieve model params
|
|
flat_master:bool=False, # Flatten fp32 params into a vector for better performance
|
|
) -> list: # List of fp16 params, and list of fp32 params
|
|
"Creates fp16 model params given an initialized `Optimizer`, also returning fp32 model params. "
|
|
model_params = [[param for param in pg if getattr(param, 'requires_grad', False) and hasattr(param, 'data')] for pg in opt.param_lists]
|
|
if flat_master:
|
|
master_params = []
|
|
for pg in model_params:
|
|
mp = parameters_to_vector([param.data.float() for param in pg])
|
|
mp = nn.Parameter(mp, requires_grad=True)
|
|
if mp.grad is None: mp.grad = mp.new(*mp.size())
|
|
master_params.append([mp])
|
|
else:
|
|
master_params = [[nn.Parameter(param.data.clone().float().detach(), requires_grad=True) for param in pg] for pg in model_params]
|
|
return model_params, master_params
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #84d2c17a
|
|
def to_master_grads(
|
|
model_pgs:list, # Fp16 model parameters to copy gradients from
|
|
master_pgs:list, # Fp32 model parameters to copy gradients to
|
|
flat_master:bool=False, # Whether or not fp32 parameters were previously flattened
|
|
):
|
|
"Move fp16 model gradients to fp32 master gradients"
|
|
for (model_params,master_params) in zip(model_pgs,master_pgs):
|
|
model_grads_to_master_grads(model_params, master_params, flat_master=flat_master)
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #d11afa3d
|
|
def to_model_params(
|
|
model_pgs:list, # Fp16 model params to copy to
|
|
master_pgs:list, # Fp32 master params to copy from
|
|
flat_master:bool=False # Whether master_pgs was previously flattened
|
|
)->None:
|
|
"Copy updated fp32 master params to fp16 model params after gradient step. "
|
|
for (model_params,master_params) in zip(model_pgs,master_pgs):
|
|
master_params_to_model_params(model_params, master_params, flat_master=flat_master)
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #4301268d
|
|
def test_overflow(x:torch.Tensor):
|
|
"Tests whether fp16 gradients have overflown."
|
|
s = float(x.float().sum())
|
|
return (s == float('inf') or s == float('-inf') or s != s)
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #f2e64c61
|
|
def grad_overflow(pgs:list)->bool:
|
|
"Tests all fp16 parameters in pgs for gradient overflow"
|
|
for pg in pgs:
|
|
for p in pg:
|
|
if p.grad is not None and test_overflow(p.grad.data): return True
|
|
return False
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #65bc0f15
|
|
def copy_clone(d):
|
|
return {k:(v.detach().clone().float() if isinstance(v,Tensor) else v) for k,v in d.items()}
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #a35f40af
|
|
def _copy_state(opt, pgs1, pgs2):
|
|
opt.param_lists = pgs2
|
|
for pg1,pg2 in zip(pgs1, pgs2):
|
|
for p1,p2 in zip(pg1, pg2): opt.state[p2] = copy_clone(opt.state.pop(p1, {}))
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #ead99755
|
|
class ModelToHalf(Callback):
|
|
"Use with NonNativeMixedPrecision callback (but it needs to run at the very beginning)"
|
|
order=-50
|
|
def before_fit(self): self.learn.model = convert_network(self.model, dtype=torch.float16)
|
|
def after_fit (self): self.learn.model = convert_network(self.model, dtype=torch.float32)
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #86746ab6
|
|
@docs
|
|
class NonNativeMixedPrecision(Callback):
|
|
"Run training in mixed precision"
|
|
order=10
|
|
def __init__(self,
|
|
loss_scale:int=512, # Non-dynamic loss scale, used to avoid underflow of gradients.
|
|
flat_master:bool=False, # Whether to flatten fp32 parameters for performance
|
|
dynamic:bool=True, # Whether to automatically determine loss scaling
|
|
max_loss_scale:float=2.**24, # Starting value for dynamic loss scaling
|
|
div_factor:float=2., # Divide by this on overflow, multiply by this after scale_wait batches
|
|
scale_wait:int=500, # Number of batches to wait for increasing loss scale
|
|
clip:float=None, # Value to clip gradients at, max_norm, as in `nn.utils.clip_grad_norm_`
|
|
):
|
|
assert torch.backends.cudnn.enabled, "Mixed precision training requires cudnn."
|
|
self.flat_master,self.dynamic,self.max_loss_scale = flat_master,dynamic,max_loss_scale
|
|
self.div_factor,self.scale_wait,self.clip = div_factor,scale_wait,clip
|
|
self.loss_scale = max_loss_scale if dynamic else loss_scale
|
|
|
|
def before_fit(self):
|
|
assert self.dls.device.type == 'cuda', "Mixed-precision training requires a GPU, remove the call `to_fp16`"
|
|
if self.learn.opt is None: self.learn.create_opt()
|
|
self.model_pgs,self.master_pgs = get_master(self.opt, self.flat_master)
|
|
self.old_pgs = self.opt.param_lists
|
|
#Changes the optimizer so that the optimization step is done in FP32.
|
|
_copy_state(self.learn.opt, self.model_pgs, self.master_pgs)
|
|
if self.dynamic: self.count = 0
|
|
|
|
def before_batch(self): self.learn.xb = to_half(self.xb)
|
|
def after_pred(self): self.learn.pred = to_float(self.pred)
|
|
def before_backward(self): self.learn.loss_grad *= self.loss_scale
|
|
|
|
def before_step(self):
|
|
#First, check for an overflow
|
|
if self.dynamic and grad_overflow(self.model_pgs):
|
|
self.loss_scale /= self.div_factor
|
|
self.learn.loss_grad /= self.div_factor #to record correct loss
|
|
self.model.zero_grad()
|
|
raise CancelBatchException() #skip step and zero_grad
|
|
to_master_grads(self.model_pgs, self.master_pgs, self.flat_master)
|
|
for master_params in self.master_pgs:
|
|
for param in master_params:
|
|
if param.grad is not None: param.grad.div_(self.loss_scale)
|
|
if self.clip is not None:
|
|
for group in self.master_pgs: nn.utils.clip_grad_norm_(group, self.clip)
|
|
# Check if it's been long enough without overflow
|
|
if self.dynamic:
|
|
self.count += 1
|
|
if self.count == self.scale_wait:
|
|
self.count = 0
|
|
self.loss_scale *= self.div_factor
|
|
|
|
def after_step(self):
|
|
self.model.zero_grad() #Zero the gradients of the model manually (optimizer disconnected)
|
|
to_model_params(self.model_pgs, self.master_pgs, self.flat_master)
|
|
|
|
def after_batch(self):
|
|
if self.training: self.learn.loss_grad /= self.loss_scale #Log correct loss
|
|
def after_fit(self):
|
|
if not hasattr(self,'master_pgs'): return
|
|
_copy_state(self.learn.opt, self.master_pgs, self.model_pgs)
|
|
self.learn.opt.param_lists = self.old_pgs
|
|
delattr(self, "master_pgs")
|
|
delattr(self, "model_pgs")
|
|
delattr(self, "old_pgs")
|
|
|
|
_docs = dict(before_fit="Put the model in FP16 and prepare the two copies of the parameters",
|
|
before_batch="Put the input in FP16",
|
|
after_pred="Put the output back to FP32 so that the loss is computed in FP32",
|
|
before_backward="Apply loss scaling to avoid gradient underflow",
|
|
before_step="Update and apply dynamic loss scaling, move gradients to fp32, apply gradient clipping",
|
|
after_step="Zero fp16 grads and update fp16 params with fp32 params. ",
|
|
after_batch="Ensure loss is logged correctly",
|
|
after_fit="Put the model back in FP32")
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #571f71c1
|
|
@patch
|
|
@delegates(NonNativeMixedPrecision.__init__)
|
|
def to_non_native_fp16(self:Learner, **kwargs): return self.add_cbs([ModelToHalf(), NonNativeMixedPrecision(**kwargs)])
|
|
|
|
# %% ../../nbs/18_callback.fp16.ipynb #f202da47
|
|
@patch
|
|
def to_non_native_fp32(self: Learner): return self.remove_cbs([ModelToHalf, NonNativeMixedPrecision])
|