1.Sophia优化器介绍
斯坦福2023.5月发表的最新研究成果，他们提出了「一种叫Sophia的优化器，相比Adam，它在LLM上能够快2倍，可以大幅降低训练成本」。
论文：https://arxiv.org/pdf/2305.14342.pdf
本文介绍了一种新的模型预训练优化器：Sophia（Second-order Clipped Stochastic Optimization），这是一种轻量级二阶优化器，它使用Hessian对角线的廉价随机估计作为预调节器，并通过限幅机制来控制最坏情况下的更新大小。在GPT-2等预训练语言模型上，Sophia以比Adam少了50%的步骤，且实现了相同的预训练损失。
作者表示 Adam 对于异构曲率（heterogeneous curvatures）的适应性不足。另一方面，vanilla Newton 方法在凸函数中具有最优的 pre-conditioner，但对于负曲率和 Hessian 的快速变化容易受到影响。基于这些见解，该研究设计了一种新的优化器 Sophia，它比 Adam 更适应异构曲率，比 Newton 方法更能抵抗非凸性和 Hessian 的快速变化，并且还使用了成本较低的 pre-conditioner。
研究引入了两个对角 Hessian 估计器，它们的内存和运行时间成本都与计算梯度相似。估计器分别为 Hutchinson 无偏估计器以及 GNB（ Gauss-Newton-Bartlett ） 估计器。伪代码如下所示：
比较 wall-clock 时间与计算量。表 1 比较了每一个 step 的总计算量 (TFLOPs) 和 A100 GPU 上的 wall-clock 时间。本文报告了每个 step 的平均时间，Hessian 计算花费的时间的总计算。较小的批量大小，即每 10 个 step 以计算对角 Hessian 估计，Hessian 计算占总计算量的 6%，与 AdamW 相比，整体 wall-clock 时间开销小于 5%。在内存使用方面，优化器 m 和 h 两个状态，这导致了与 AdamW 相同的内存开销。
2.Sophia引入到yolov8
2.1 修改ultralytics/yolo/engine/trainer.py
核心代码：
代码语言：javascript复制
import torch
from torch.optim.optimizer import Optimizer
import math
from torch import Tensor
from typing import List, Optional
class Sophia(Optimizer):
def __init__(self, model, input_data, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, k=10,
estimator="Hutchinson", rho=1):
self.model = model
self.input_data = input_data
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, k=k, estimator=estimator, rho=rho)
super(Sophia, self).__init__(params, defaults)
def step(self, closure=None):
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
grad = p.grad.data
if grad.is_sparse:
raise RuntimeError("Sophia does not support sparse gradients")
state = self.state[p]
# state init
if len(state) == 0:
state['step'] = 0
state['m'] = torch.zeros_like(p.data)
state['h'] = torch.zeros_like(p.data)
m, h = state['m'], state['h']
beta1, beta2 = group['betas']
state['step'] = 1
if group['weight_decay'] != 0:
grad = grad.add(group["weight_decay"], p.data)
# update biased first moment estimate
m.mul_(beta1).add_(1 - beta1, grad)
# update hessian estimate
if state['step'] % group['k'] == 1:
if group['estimator'] == "Hutchinson":
hessian_estimate = self.hutchinson(p, grad)
elif group['estimator'] == "Gauss-Newton-Bartlett":
hessian_estimate = self.gauss_newton_bartlett(p, grad)
else:
raise ValueError("Invalid estimator choice")
h.mul_(beta2).add_(1 - beta2, hessian_estimate)
# update params
p.data.add_(-group['lr'] * group['weight_decay'], p.data)
p.data.addcdiv_(-group['lr'], m, h.add(group['eps']).clamp(max=group['rho']))
return loss
def hutchinson(self, p, grad):
u = torch.randn_like(grad)
grad_dot_u = torch.sum(grad * u)
hessian_vector_product = torch.autograd.grad(grad_dot_u, p, retain_graph=True)[0]
return u * hessian_vector_product
def gauss_newton_bartlett(self, p, grad):
B = len(self.input_data)
logits = [self.model(xb) for xb in self.input_data]
y_hats = [torch.softmax(logit, dim=0) for logit in logits]
g_hat =
torch.autograd.grad(sum([self.loss_function(logit, y_hat) for logit, y_hat in zip(logits, y_hats)]) / B, p,
retain_graph=True)[0]
return B * g_hat * g_hat
class SophiaG(Optimizer):
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(SophiaG, self).__init__(params, defaults)
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()
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.)
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()
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('Hero 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.)
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
sophiag(params_with_grad,
grads,
exp_avgs,
hessian,
state_steps,
bs=bs,
beta1=beta1,
beta2=beta2,
rho=group['rho'],
lr=group['lr'],
weight_decay=group['weight_decay'],
maximize=group['maximize'],
capturable=group['capturable'])
return loss
def sophiag(params: List[Tensor],
grads: List[Tensor],
exp_avgs: List[Tensor],
hessian: List[Tensor],
state_steps: List[Tensor],
capturable: bool = False,
*,
bs: int,
beta1: float,
beta2: float,
rho: float,
lr: float,
weight_decay: float,
maximize: bool):
if not all(isinstance(t, torch.Tensor) for t in state_steps):
raise RuntimeError("API has changed, `state_steps` argument must contain a list of singleton tensors")
func = _single_tensor_sophiag
#
func(params,
grads,
exp_avgs,
hessian,
state_steps,
bs=bs,
beta1=beta1,
beta2=beta2,
rho=rho,
lr=lr,
weight_decay=weight_decay,
maximize=maximize,
capturable=capturable)
def _single_tensor_sophiag(params: List[Tensor],
grads: List[Tensor],
exp_avgs: List[Tensor],
hessian: List[Tensor],
state_steps: List[Tensor],
*,
bs: int,
beta1: float,
beta2: float,
rho: float,
lr: float,
weight_decay: float,
maximize: bool,
capturable: bool):
for i, param in enumerate(params):
grad = grads[i] if not maximize else -grads[i]
exp_avg = exp_avgs[i]
hess = hessian[i]
step_t = state_steps[i]
if capturable:
assert param.is_cuda and step_t.is_cuda and bs.is_cuda
if torch.is_complex(param):
grad = torch.view_as_real(grad)
exp_avg = torch.view_as_real(exp_avg)
hess = torch.view_as_real(hess)
param = torch.view_as_real(param)
# update step
step_t = 1
# Perform stepweight decay
param.mul_(1 - lr * weight_decay)
# Decay the first and second moment running average coefficient
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
if capturable:
step = step_t
step_size = lr
step_size_neg = step_size.neg()
ratio = (exp_avg.abs() / (rho * bs * hess 1e-15)).clamp(None, 1)
param.addcmul_(exp_avg.sign(), ratio, value=step_size_neg)
else:
step = step_t.item()
step_size_neg = - lr
ratio = (exp_avg.abs() / (rho * bs * hess 1e-15)).clamp(None, 1)
param.addcmul_(exp_avg.sign(), ratio, value=step_size_neg)3.总结
训练稳定性。与 AdamW 和 Lion 相比，Sophia-H 在预训练中具有更好的稳定性。梯度裁剪 (by norm) 是语言模型预训练中的一项重要技术。在实践中，梯度裁剪触发的频率与训练的稳定性有关 —— 如果梯度被频繁裁剪，迭代可能处于非常不稳定的状态。图 7 (a) 比较了 GPT-2 (125M) 触发梯度裁剪的 step 比例。尽管所有方法都使用相同的裁剪阈值 1.0，但 Sophia-H 很少触发梯度裁剪，而 AdamW 和 Lion 在超过 10% 的 step 中触发梯度裁剪。
详见：
https://blog.csdn.net/m0_63774211/article/details/130912702
我正在参与2023腾讯技术创作特训营第三期有奖征文，组队打卡瓜分大奖！
