File size: 3,208 Bytes
554a35b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """GrokAdamW: AdamW c уклоном в обобщение.
- decoupled weight decay (только на матрицах, не на нормах/эмбеддингах);
- cautious-маскинг (C-AdamW): применяем компоненты шага, согласованные по
знаку с градиентом -> стабильнее и чуть быстрее сходимость;
- Grokfast (опц., gf_lambda>0): усиление медленной компоненты градиента.
Для большого претрейна обычно держим gf_lambda=0 (нужен для grokking на
маленьких алгоритмических задачах, не для масштабного обучения).
"""
import math
import torch
class GrokAdamW(torch.optim.Optimizer):
def __init__(self, params, lr=3e-4, betas=(0.9, 0.95), eps=1e-8,
weight_decay=0.1, grokfast_lambda=0.0, grokfast_alpha=0.98,
cautious=True):
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
grokfast_lambda=grokfast_lambda, grokfast_alpha=grokfast_alpha,
cautious=cautious)
super().__init__(params, defaults)
@torch.no_grad()
def step(self, closure=None):
loss = closure() if closure is not None else None
for grp in self.param_groups:
lr, (b1, b2), eps = grp["lr"], grp["betas"], grp["eps"]
wd, gfl, gfa, caut = (grp["weight_decay"], grp["grokfast_lambda"],
grp["grokfast_alpha"], grp["cautious"])
for p in grp["params"]:
if p.grad is None:
continue
g = p.grad
st = self.state[p]
if not st:
st["step"] = 0
st["m"] = torch.zeros_like(p)
st["v"] = torch.zeros_like(p)
if gfl > 0:
st["ema"] = g.clone()
if gfl > 0:
ema = st["ema"]; ema.mul_(gfa).add_(g, alpha=1 - gfa)
g = g.add(ema, alpha=gfl)
m, v = st["m"], st["v"]
st["step"] += 1; t = st["step"]
m.mul_(b1).add_(g, alpha=1 - b1)
v.mul_(b2).addcmul_(g, g, value=1 - b2)
denom = (v.sqrt() / math.sqrt(1 - b2 ** t)).add_(eps)
step_size = lr / (1 - b1 ** t)
if wd != 0:
p.mul_(1 - lr * wd)
upd = m / denom
if caut:
mask = (upd * g > 0).to(upd.dtype)
mask.mul_(mask.numel() / (mask.sum() + 1))
upd.mul_(mask)
p.add_(upd, alpha=-step_size)
return loss
def param_groups(model, weight_decay):
decay, no_decay = [], []
for name, p in model.named_parameters():
if not p.requires_grad:
continue
(no_decay if p.ndim < 2 or name.endswith(".bias") else decay).append(p)
return [{"params": decay, "weight_decay": weight_decay},
{"params": no_decay, "weight_decay": 0.0}]
|