mirror of https://github.com/hpcaitech/ColossalAI
226 lines
8.8 KiB
Python
226 lines
8.8 KiB
Python
# modified from https://github.com/NVIDIA/apex/blob/master/apex/optimizers/fused_lamb.py
|
|
import torch
|
|
|
|
from colossalai.utils import multi_tensor_applier
|
|
|
|
|
|
class FusedLAMB(torch.optim.Optimizer):
|
|
"""Implements LAMB algorithm.
|
|
|
|
`FusedLAMB` requires CUDA extensions which can be built during installation or runtime.
|
|
|
|
This version of fused LAMB implements 2 fusions.
|
|
|
|
* Fusion of the LAMB update's elementwise operations
|
|
* A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
|
|
|
|
:class:`colossalai.nn.optimizer.FusedLAMB`'s usage is identical to any ordinary Pytorch optimizer
|
|
|
|
:class:`colossalai.nn.optimizer.FusedLAMB` may be used with or without Amp.
|
|
|
|
LAMB was proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
|
|
|
|
Arguments:
|
|
params (iterable): iterable of parameters to optimize or dicts defining
|
|
parameter groups.
|
|
lr (float, optional): learning rate. (default: 1e-3)
|
|
betas (Tuple[float, float], optional): coefficients used for computing
|
|
running averages of gradient and its norm. (default: (0.9, 0.999))
|
|
eps (float, optional): term added to the denominator to improve
|
|
numerical stability. (default: 1e-6)
|
|
weight_decay (float, optional): weight decay (L2 penalty) (default: 0.01)
|
|
amsgrad (boolean, optional): whether to use the AMSGrad variant of this
|
|
algorithm from the paper `On the Convergence of Adam and Beyond`_
|
|
NOT SUPPORTED now! (default: False)
|
|
adam_w_mode (boolean, optional): Apply L2 regularization or weight decay
|
|
True for decoupled weight decay(also known as AdamW) (default: True)
|
|
grad_averaging (bool, optional): whether apply (1-beta2) to grad when
|
|
calculating running averages of gradient. (default: True)
|
|
set_grad_none (bool, optional): whether set grad to None when zero_grad()
|
|
method is called. (default: True)
|
|
max_grad_norm (float, optional): value used to clip global grad norm
|
|
(default: 1.0)
|
|
use_nvlamb (boolean, optional): Apply adaptive learning rate to 0.0
|
|
weight decay parameter (default: False)
|
|
|
|
.. _Large Batch Optimization for Deep Learning: Training BERT in 76 minutes:
|
|
https://arxiv.org/abs/1904.00962
|
|
.. _On the Convergence of Adam and Beyond:
|
|
https://openreview.net/forum?id=ryQu7f-RZ
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
params,
|
|
lr=1e-3,
|
|
bias_correction=True,
|
|
betas=(0.9, 0.999),
|
|
eps=1e-6,
|
|
weight_decay=0.01,
|
|
amsgrad=False,
|
|
adam_w_mode=True,
|
|
grad_averaging=True,
|
|
set_grad_none=True,
|
|
max_grad_norm=1.0,
|
|
use_nvlamb=False,
|
|
):
|
|
if amsgrad:
|
|
raise RuntimeError("FusedLAMB does not support the AMSGrad variant.")
|
|
defaults = dict(
|
|
lr=lr,
|
|
bias_correction=bias_correction,
|
|
betas=betas,
|
|
eps=eps,
|
|
weight_decay=weight_decay,
|
|
grad_averaging=grad_averaging,
|
|
max_grad_norm=max_grad_norm,
|
|
)
|
|
super(FusedLAMB, self).__init__(params, defaults)
|
|
if multi_tensor_applier.available:
|
|
from colossalai.kernel.kernel_loader import FusedOptimizerLoader
|
|
|
|
fused_optim = FusedOptimizerLoader().load()
|
|
|
|
self.multi_tensor_l2norm = fused_optim.multi_tensor_l2norm
|
|
# Skip buffer
|
|
self._dummy_overflow_buf = torch.tensor(
|
|
[0], dtype=torch.int, device=self.param_groups[0]["params"][0].device
|
|
)
|
|
self.multi_tensor_lamb = fused_optim.multi_tensor_lamb
|
|
else:
|
|
raise RuntimeError("FusedLAMB requires cuda extensions")
|
|
|
|
self.adam_w_mode = 1 if adam_w_mode else 0
|
|
self.set_grad_none = set_grad_none
|
|
self.use_nvlamb = use_nvlamb
|
|
|
|
def zero_grad(self):
|
|
if self.set_grad_none:
|
|
for group in self.param_groups:
|
|
for p in group["params"]:
|
|
p.grad = None
|
|
else:
|
|
super(FusedLAMB, self).zero_grad()
|
|
|
|
def step(self, closure=None):
|
|
"""Performs a single optimization step.
|
|
|
|
Arguments:
|
|
closure (callable, optional): A closure that reevaluates the model
|
|
and returns the loss.
|
|
"""
|
|
loss = None
|
|
if closure is not None:
|
|
loss = closure()
|
|
|
|
# create separate grad lists for fp32 and fp16 params
|
|
g_all_32, g_all_16 = [], []
|
|
for group in self.param_groups:
|
|
for p in group["params"]:
|
|
if p.grad is None:
|
|
continue
|
|
if p.dtype == torch.float32:
|
|
g_all_32.append(p.grad.data)
|
|
elif p.dtype == torch.float16:
|
|
g_all_16.append(p.grad.data)
|
|
else:
|
|
raise RuntimeError("FusedLAMB only support fp16 and fp32.")
|
|
|
|
device = self.param_groups[0]["params"][0].device
|
|
g_norm_32, g_norm_16 = torch.zeros(1, device=device), torch.zeros(1, device=device)
|
|
# compute grad norm for two lists
|
|
if len(g_all_32) > 0:
|
|
g_norm_32 = multi_tensor_applier(self.multi_tensor_l2norm, self._dummy_overflow_buf, [g_all_32], False)[0]
|
|
if len(g_all_16) > 0:
|
|
g_norm_16 = multi_tensor_applier(self.multi_tensor_l2norm, self._dummy_overflow_buf, [g_all_16], False)[0]
|
|
|
|
# blend two grad norms to get global grad norm
|
|
global_grad_norm = multi_tensor_applier(
|
|
self.multi_tensor_l2norm, self._dummy_overflow_buf, [[g_norm_32, g_norm_16]], False
|
|
)[0]
|
|
max_grad_norm = self.defaults["max_grad_norm"]
|
|
|
|
for group in self.param_groups:
|
|
bias_correction = 1 if group["bias_correction"] else 0
|
|
beta1, beta2 = group["betas"]
|
|
grad_averaging = 1 if group["grad_averaging"] else 0
|
|
|
|
# assume same step across group now to simplify things
|
|
# per parameter step can be easily support by making it tensor, or pass list into kernel
|
|
if "step" in group:
|
|
group["step"] += 1
|
|
else:
|
|
group["step"] = 1
|
|
|
|
# create lists for multi-tensor apply
|
|
g_16, p_16, m_16, v_16 = [], [], [], []
|
|
g_32, p_32, m_32, v_32 = [], [], [], []
|
|
|
|
for p in group["params"]:
|
|
if p.grad is None:
|
|
continue
|
|
if p.grad.data.is_sparse:
|
|
raise RuntimeError(
|
|
"FusedLAMB does not support sparse gradients, please consider SparseAdam instead"
|
|
)
|
|
|
|
state = self.state[p]
|
|
# State initialization
|
|
if len(state) == 0:
|
|
# Exponential moving average of gradient values
|
|
state["exp_avg"] = torch.zeros_like(p)
|
|
# Exponential moving average of gradient values
|
|
state["exp_avg_sq"] = torch.zeros_like(p)
|
|
|
|
if p.dtype == torch.float16:
|
|
g_16.append(p.grad.data)
|
|
p_16.append(p.data)
|
|
m_16.append(state["exp_avg"])
|
|
v_16.append(state["exp_avg_sq"])
|
|
elif p.dtype == torch.float32:
|
|
g_32.append(p.grad.data)
|
|
p_32.append(p.data)
|
|
m_32.append(state["exp_avg"])
|
|
v_32.append(state["exp_avg_sq"])
|
|
else:
|
|
raise RuntimeError("FusedLAMB only support fp16 and fp32.")
|
|
|
|
if len(g_16) > 0:
|
|
multi_tensor_applier(
|
|
self.multi_tensor_lamb,
|
|
self._dummy_overflow_buf,
|
|
[g_16, p_16, m_16, v_16],
|
|
group["lr"],
|
|
beta1,
|
|
beta2,
|
|
group["eps"],
|
|
group["step"],
|
|
bias_correction,
|
|
group["weight_decay"],
|
|
grad_averaging,
|
|
self.adam_w_mode,
|
|
global_grad_norm,
|
|
max_grad_norm,
|
|
self.use_nvlamb,
|
|
)
|
|
if len(g_32) > 0:
|
|
multi_tensor_applier(
|
|
self.multi_tensor_lamb,
|
|
self._dummy_overflow_buf,
|
|
[g_32, p_32, m_32, v_32],
|
|
group["lr"],
|
|
beta1,
|
|
beta2,
|
|
group["eps"],
|
|
group["step"],
|
|
bias_correction,
|
|
group["weight_decay"],
|
|
grad_averaging,
|
|
self.adam_w_mode,
|
|
global_grad_norm,
|
|
max_grad_norm,
|
|
self.use_nvlamb,
|
|
)
|
|
|
|
return loss
|