in modules/SwissArmyTransformer/sat/ops/fused_ema_adam.py [0:0]
def step(self, closure=None, grads=None, output_params=None, scale=None, grad_norms=None, grad_scaler=None):
"""Performs a single optimization step.
Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
The remaining arguments are deprecated, and are only retained (for the moment) for error-checking purposes.
"""
if any(p is not None for p in [grads, output_params, scale, grad_norms]):
raise RuntimeError(
'FusedAdam has been updated. Simply initialize it identically to torch.optim.Adam, and call step() with no arguments.'
)
loss = None
if closure is not None:
loss = closure()
ema_decay = self.ema_decay
if self.num_updates >= 0:
self.num_updates += 1
ema_decay = min(self.ema_decay,(1 + self.num_updates) / (10 + self.num_updates))
for group in self.param_groups:
if len(group['params']) == 0:
continue
bias_correction = 1 if group['bias_correction'] else 0
beta1, beta2 = group['betas']
# 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' not in group:
group['step'] = 0
# create lists for multi-tensor apply
g_16, p_16, m_16, v_16, s_16 = [], [], [], [], []
g_bf, p_bf, m_bf, v_bf, s_bf = [], [], [], [], []
g_32, p_32, m_32, v_32, s_32 = [], [], [], [], []
for p in group['params']:
if p.grad is None:
continue
if p.grad.data.is_sparse:
raise RuntimeError(
'FusedEmaAdam does not support sparse gradients, please consider SparseAdam instead')
state = self.state[p]
# State initialization
if len(state) == 0:
# DeepSpeed ZeRO 3 processes each subgroup a time, so we need to keep tracking step count for each tensor separately.
# While this is not an issue for ZeRO 1 & 2, since they apply a single optimization step to the whole param group at the same time.
# In order to keep backward compatibility for the existing checkpoints, we use group['state'] to initialize state['step'] if it exists.
state['step'] = group.get('step', 0)
# Exponential moving average of gradient values
state['exp_avg'] = torch.zeros_like(p.data)
# Exponential moving average of squared gradient values
state['exp_avg_sq'] = torch.zeros_like(p.data)
# Exponential moving average of shadow weights
state['shadow'] = p.data.clone()
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'])
s_16.append(state['shadow'])
elif p.dtype == torch.bfloat16:
g_bf.append(p.grad)
p_bf.append(p)
m_bf.append(state['exp_avg'])
v_bf.append(state['exp_avg_sq'])
s_bf.append(state['shadow'])
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'])
s_32.append(state['shadow'])
else:
raise RuntimeError('FusedEmaAdam only support fp16, bf16 and fp32.')
if len(g_16) > 0:
state['step'] += 1
multi_tensor_applier(self.multi_tensor_ema_adam, self._dummy_overflow_buf, [g_16, p_16, m_16, v_16, s_16],
group['lr'], ema_decay, beta1, beta2, group['eps'], state['step'], self.adam_w_mode,
bias_correction, group['weight_decay'])
if len(g_bf) > 0:
state['step'] += 1
multi_tensor_applier(self.multi_tensor_ema_adam, self._dummy_overflow_buf, [g_bf, p_bf, m_bf, v_bf, s_bf],
group['lr'], ema_decay, beta1, beta2, group['eps'], state['step'], self.adam_w_mode,
bias_correction, group['weight_decay'])
if len(g_32) > 0:
state['step'] += 1
multi_tensor_applier(self.multi_tensor_ema_adam, self._dummy_overflow_buf, [g_32, p_32, m_32, v_32, s_32],
group['lr'], ema_decay, beta1, beta2, group['eps'], state['step'], self.adam_w_mode,
bias_correction, group['weight_decay'])
return loss