def step(self,
closure
) -> torch.Tensor:
"""
Args:
closure: A closure that reevaluates the model and returns the loss.
Returns: the loss value evaluated on the original point
"""
closure = torch.enable_grad()(closure)
loss = closure().detach()
for group in self.param_groups:
grads = []
params_with_grads = []
rho = group['rho']
# update internal_optim's learning rate
for p in group['params']:
if p.grad is not None:
# without clone().detach(), p.grad will be zeroed by closure()
grads.append(p.grad.clone().detach())
params_with_grads.append(p)
device = grads[0].device
# compute \hat{\epsilon}=\rho/\norm{g}\|g\|
grad_norm = torch.stack([g.detach().norm(2).to(device) for g in grads]).norm(2)
epsilon = grads # alias for readability
torch._foreach_mul_(epsilon, rho / grad_norm)
# virtual step toward \epsilon
torch._foreach_add_(params_with_grads, epsilon)
# compute g=\nabla_w L_B(w)|_{w+\hat{\epsilon}}
closure()
# virtual step back to the original point
torch._foreach_sub_(params_with_grads, epsilon)
super().step()
return loss
def _multi_tensor_adamw(params: List[Tensor], grads: List[Tensor],
exp_avgs: List[Tensor], exp_avg_sqs: List[Tensor],
max_exp_avg_sqs: List[Tensor],
state_steps: List[Tensor], *, amsgrad: bool,
beta1: float, beta2: float, lr: float,
weight_decay: float, eps: float, maximize: bool,
capturable: bool):
if len(params) == 0:
return
if capturable:
assert all(p.is_cuda and step.is_cuda for p, step in zip(params, state_steps)), \
"If capturable=True, params and state_steps must be CUDA tensors."
if maximize:
grads = torch._foreach_neg(tuple(grads)) # type: ignore[assignment]
grads = [
torch.view_as_real(x) if torch.is_complex(x) else x for x in grads
]
exp_avgs = [
torch.view_as_real(x) if torch.is_complex(x) else x for x in exp_avgs
]
exp_avg_sqs = [
torch.view_as_real(x) if torch.is_complex(x) else x
for x in exp_avg_sqs
]
params = [
torch.view_as_real(x) if torch.is_complex(x) else x for x in params
]
# update steps
torch._foreach_add_(state_steps, 1)
# Perform stepweight decay
torch._foreach_mul_(params, 1 - lr * weight_decay)
# Decay the first and second moment running average coefficient
torch._foreach_mul_(exp_avgs, beta1)
torch._foreach_add_(exp_avgs, grads, alpha=1 - beta1)
torch._foreach_mul_(exp_avg_sqs, beta2)
torch._foreach_addcmul_(exp_avg_sqs, grads, grads, 1 - beta2)
if capturable:
# TODO: use foreach_pow if/when foreach_pow is added
bias_correction1 = [torch.pow(beta1, step) for step in state_steps]
bias_correction2 = [torch.pow(beta2, step) for step in state_steps]
# foreach_sub doesn't allow a scalar as the first arg
torch._foreach_sub_(bias_correction1, 1)
torch._foreach_sub_(bias_correction2, 1)
torch._foreach_neg_(bias_correction1)
torch._foreach_neg_(bias_correction2)
# foreach_div doesn't allow a scalar as the first arg
step_size = torch._foreach_div(bias_correction1, lr)
torch._foreach_reciprocal_(step_size)
torch._foreach_neg_(step_size)
bias_correction2_sqrt = torch._foreach_sqrt(bias_correction2)
if amsgrad:
# Maintains the maximum of all 2nd moment running avg. till now
torch._foreach_maximum_(max_exp_avg_sqs, exp_avg_sqs)
# Use the max. for normalizing running avg. of gradient
max_exp_avg_sq_sqrt = torch._foreach_sqrt(max_exp_avg_sqs)
# Folds in (admittedly ugly) 1-elem step_size math here to avoid extra param-set-sized read+write
# (can't fold it into addcdiv_ below because addcdiv_ requires value is a Number, not a Tensor)
torch._foreach_div_(
max_exp_avg_sq_sqrt,
torch._foreach_mul(bias_correction2_sqrt, step_size))
eps_over_step_size = torch._foreach_div(step_size, eps)
torch._foreach_reciprocal_(eps_over_step_size)
denom = torch._foreach_add(max_exp_avg_sq_sqrt, eps_over_step_size)
else:
exp_avg_sq_sqrt = torch._foreach_sqrt(exp_avg_sqs)
torch._foreach_div_(
exp_avg_sq_sqrt,
torch._foreach_mul(bias_correction2_sqrt, step_size))
eps_over_step_size = torch._foreach_div(step_size, eps)
torch._foreach_reciprocal_(eps_over_step_size)
denom = torch._foreach_add(exp_avg_sq_sqrt, eps_over_step_size)
torch._foreach_addcdiv_(params, exp_avgs, denom)
else:
bias_correction1 = [1 - beta1**step.item() for step in state_steps]
bias_correction2 = [1 - beta2**step.item() for step in state_steps]
step_size = [(lr / bc) * -1 for bc in bias_correction1]
bias_correction2_sqrt = [math.sqrt(bc) for bc in bias_correction2]
if amsgrad:
# Maintains the maximum of all 2nd moment running avg. till now
torch._foreach_maximum_(max_exp_avg_sqs, exp_avg_sqs)
# Use the max. for normalizing running avg. of gradient
max_exp_avg_sq_sqrt = torch._foreach_sqrt(max_exp_avg_sqs)
torch._foreach_div_(max_exp_avg_sq_sqrt, bias_correction2_sqrt)
denom = torch._foreach_add(max_exp_avg_sq_sqrt, eps)
else:
exp_avg_sq_sqrt = torch._foreach_sqrt(exp_avg_sqs)
torch._foreach_div_(exp_avg_sq_sqrt, bias_correction2_sqrt)
denom = torch._foreach_add(exp_avg_sq_sqrt, eps)
torch._foreach_addcdiv_(params, exp_avgs, denom, step_size)
