#!/usr/bin/env python
import torch
from torch.optim.optimizer import Optimizer
[docs]
class SophiaG(Optimizer):
"""
Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training.
Code from: https://github.com/Liuhong99/Sophia/
"""
def __init__(
self,
params,
lr=1e-4,
betas=(0.965, 0.99),
rho=0.04,
weight_decay=1e-1,
*,
maximize: bool = False,
capturable: bool = False,
):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= betas[0] < 1.0:
raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
if not 0.0 <= betas[1] < 1.0:
raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
if not 0.0 <= rho:
raise ValueError("Invalid rho parameter at index 1: {}".format(rho))
if not 0.0 <= weight_decay:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
defaults = dict(
lr=lr, betas=betas, rho=rho, weight_decay=weight_decay, maximize=maximize, capturable=capturable
)
super().__init__(params, defaults)
[docs]
def __setstate__(self, state):
super().__setstate__(state)
for group in self.param_groups:
group.setdefault("maximize", False)
group.setdefault("capturable", False)
state_values = list(self.state.values())
step_is_tensor = (len(state_values) != 0) and torch.is_tensor(state_values[0]["step"])
if not step_is_tensor:
for s in state_values:
s["step"] = torch.tensor(float(s["step"]))
@torch.no_grad()
[docs]
def update_hessian(self):
for group in self.param_groups:
beta1, beta2 = group["betas"]
for p in group["params"]:
if p.grad is None:
continue
state = self.state[p]
if len(state) == 0:
state["step"] = (
torch.zeros((1,), dtype=torch.float, device=p.device)
if self.defaults["capturable"]
else torch.tensor(0.0)
)
state["exp_avg"] = torch.zeros_like(p, memory_format=torch.preserve_format)
state["hessian"] = torch.zeros_like(p, memory_format=torch.preserve_format)
if "hessian" not in state.keys():
state["hessian"] = torch.zeros_like(p, memory_format=torch.preserve_format)
state["hessian"].mul_(beta2).addcmul_(p.grad, p.grad, value=1 - beta2)
@torch.no_grad()
[docs]
def step(self, closure=None, bs=5120):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
params_with_grad = []
grads = []
exp_avgs = []
state_steps = []
hessian = []
beta1, beta2 = group["betas"]
for p in group["params"]:
if p.grad is None:
continue
params_with_grad.append(p)
if p.grad.is_sparse:
raise RuntimeError("SophiaG does not support sparse gradients")
grads.append(p.grad)
state = self.state[p]
# State initialization
if len(state) == 0:
state["step"] = (
torch.zeros((1,), dtype=torch.float, device=p.device)
if self.defaults["capturable"]
else torch.tensor(0.0)
)
state["exp_avg"] = torch.zeros_like(p, memory_format=torch.preserve_format)
state["hessian"] = torch.zeros_like(p, memory_format=torch.preserve_format)
if "hessian" not in state.keys():
state["hessian"] = torch.zeros_like(p, memory_format=torch.preserve_format)
exp_avgs.append(state["exp_avg"])
state_steps.append(state["step"])
hessian.append(state["hessian"])
if self.defaults["capturable"]:
bs = torch.ones((1,), dtype=torch.float, device=p.device) * bs
# Perform the actual update step here instead of calling SophiaG again
for p, grad, exp_avg, h, step in zip(params_with_grad, grads, exp_avgs, hessian, state_steps):
if group["weight_decay"] != 0:
grad = grad.add(p, alpha=group["weight_decay"])
# Decay the first and second moment running average coefficient
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
step.add_(1)
# Compute the update using the hessian information
update = exp_avg.div(1 - beta1 ** step.item())
h_sqrt = h.sqrt().add_(group["rho"])
p.addcdiv_(update, h_sqrt, value=-group["lr"])
return loss
