| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| import math |
| import torch |
| from torch.optim import Optimizer |
|
|
|
|
| def clip_grad_norm(params, max_norm): |
| """Global L2 gradient clipping (re-implemented; returns the pre-clip norm).""" |
| grads = [p.grad for p in params if p.grad is not None] |
| if not grads: |
| return torch.zeros(()) |
| total = torch.linalg.vector_norm(torch.stack([torch.linalg.vector_norm(g.float()) for g in grads])) |
| if max_norm and max_norm > 0: |
| coef = (max_norm / (total + 1e-6)).clamp(max=1.0) |
| for g in grads: |
| g.mul_(coef.to(g.dtype)) |
| return total |
|
|
|
|
| def split_decay_groups(model, weight_decay): |
| decay, no_decay = [], [] |
| for n, p in model.named_parameters(): |
| if not p.requires_grad: |
| continue |
| (decay if p.dim() >= 2 else no_decay).append(p) |
| return [{"params": decay, "weight_decay": weight_decay}, |
| {"params": no_decay, "weight_decay": 0.0}] |
|
|
|
|
| |
| |
| |
| class AdamW(Optimizer): |
| """m_t = b1 m + (1-b1) g ; v_t = b2 v + (1-b2) g^2 |
| w <- w - lr * (m_hat / (sqrt(v_hat) + eps)) - lr * wd * w""" |
| def __init__(self, params, lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1): |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)) |
|
|
| @torch.no_grad() |
| def step(self, closure=None): |
| loss = closure() if closure is not None else None |
| for group in self.param_groups: |
| lr, (b1, b2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"] |
| for p in group["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, memory_format=torch.preserve_format) |
| st["v"] = torch.zeros_like(p, memory_format=torch.preserve_format) |
| st["step"] += 1 |
| t = st["step"] |
| m, v = st["m"], st["v"] |
| m.lerp_(g, 1 - b1) |
| v.mul_(b2).addcmul_(g, g, value=1 - b2) |
| bc1, bc2 = 1 - b1 ** t, 1 - b2 ** t |
| if wd: |
| p.mul_(1 - lr * wd) |
| denom = (v.sqrt() / math.sqrt(bc2)).add_(eps) |
| p.addcdiv_(m, denom, value=-lr / bc1) |
| return loss |
|
|
|
|
| |
| |
| |
| class MARS(Optimizer): |
| """c_t = g_t + gamma * b1/(1-b1) * (g_t - g_{t-1}) |
| c_t <- c_t / max(1, ||c_t||) (per-tensor clipping of the corrected gradient) |
| then AdamW moments / update on c_t. |
| Groups with gamma=0 (1-D params, as in the official MARS code) reduce to AdamW.""" |
| def __init__(self, params, lr=3e-3, betas=(0.95, 0.99), eps=1e-8, weight_decay=0.1, |
| gamma=0.025, max_c_norm=1.0): |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, |
| gamma=gamma, max_c_norm=max_c_norm)) |
|
|
| @torch.no_grad() |
| def step(self, closure=None): |
| loss = closure() if closure is not None else None |
| for group in self.param_groups: |
| lr, (b1, b2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"] |
| gamma, max_c = group["gamma"], group["max_c_norm"] |
| for p in group["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 gamma: |
| st["prev_grad"] = torch.zeros_like(p) |
| st["step"] += 1 |
| t = st["step"] |
| if gamma: |
| if t == 1: |
| c = g.clone() |
| else: |
| c = g + (gamma * b1 / (1 - b1)) * (g - st["prev_grad"]) |
| st["prev_grad"].copy_(g) |
| if max_c: |
| n = torch.linalg.vector_norm(c) |
| c.mul_((max_c / (n + 1e-6)).clamp(max=1.0)) |
| else: |
| c = g |
| m, v = st["m"], st["v"] |
| m.lerp_(c, 1 - b1) |
| v.mul_(b2).addcmul_(c, c, value=1 - b2) |
| bc1, bc2 = 1 - b1 ** t, 1 - b2 ** t |
| if wd: |
| p.mul_(1 - lr * wd) |
| denom = (v.sqrt() / math.sqrt(bc2)).add_(eps) |
| p.addcdiv_(m, denom, value=-lr / bc1) |
| return loss |
|
|
|
|
| |
| |
| |
| class Lion(Optimizer): |
| """u = sign(b1 m + (1-b1) g) ; w <- w - lr*(u + wd*w) ; m <- b2 m + (1-b2) g""" |
| def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=1.0): |
| super().__init__(params, dict(lr=lr, betas=betas, weight_decay=weight_decay)) |
|
|
| @torch.no_grad() |
| def step(self, closure=None): |
| loss = closure() if closure is not None else None |
| for group in self.param_groups: |
| lr, (b1, b2), wd = group["lr"], group["betas"], group["weight_decay"] |
| for p in group["params"]: |
| if p.grad is None: |
| continue |
| g = p.grad |
| st = self.state[p] |
| if not st: |
| st["m"] = torch.zeros_like(p) |
| m = st["m"] |
| u = torch.lerp(g, m, b1).sign_() |
| if wd: |
| p.mul_(1 - lr * wd) |
| p.add_(u, alpha=-lr) |
| m.lerp_(g, 1 - b2) |
| return loss |
|
|
|
|
| |
| |
| |
| def zeropower_via_newtonschulz5(G, steps=5, eps=1e-7): |
| """Quintic Newton-Schulz iteration approximating the orthogonal polar factor U V^T |
| of G = U S V^T (coefficients from Jordan et al. 2024). Runs in fp32 (T4/P100 lack bf16).""" |
| assert G.ndim == 2 |
| a, b, c = 3.4445, -4.7750, 2.0315 |
| X = G.float() |
| transposed = X.size(0) > X.size(1) |
| if transposed: |
| X = X.T |
| X = X / (X.norm() + eps) |
| for _ in range(steps): |
| A = X @ X.T |
| B = b * A + c * (A @ A) |
| X = a * X + B @ X |
| if transposed: |
| X = X.T |
| return X |
|
|
|
|
| class Muon(Optimizer): |
| """Hidden 2-D weight matrices: Nesterov momentum -> Newton-Schulz orthogonalisation -> |
| scale by sqrt(max(1, rows/cols)). Everything else (embeddings / tied LM head, |
| LayerNorm gains, biases) is updated with AdamW, as in Alg. 8 of the paper.""" |
| def __init__(self, muon_params, adamw_params, lr=0.02, momentum=0.95, nesterov=True, |
| ns_steps=5, weight_decay=0.01, adamw_lr=3e-3, adamw_betas=(0.9, 0.95), |
| adamw_eps=1e-8, adamw_wd=0.1): |
| groups = [dict(params=list(muon_params), use_muon=True, lr=lr, momentum=momentum, |
| nesterov=nesterov, ns_steps=ns_steps, weight_decay=weight_decay)] |
| adamw_params = list(adamw_params) |
| dec = [p for p in adamw_params if p.dim() >= 2] |
| nodec = [p for p in adamw_params if p.dim() < 2] |
| for ps, wd in ((dec, adamw_wd), (nodec, 0.0)): |
| if ps: |
| groups.append(dict(params=ps, use_muon=False, lr=adamw_lr, betas=adamw_betas, |
| eps=adamw_eps, weight_decay=wd)) |
| super().__init__(groups, dict(lr=lr)) |
|
|
| @torch.no_grad() |
| def step(self, closure=None): |
| loss = closure() if closure is not None else None |
| for group in self.param_groups: |
| lr, wd = group["lr"], group["weight_decay"] |
| if group["use_muon"]: |
| mom = group["momentum"] |
| for p in group["params"]: |
| if p.grad is None: |
| continue |
| g = p.grad |
| st = self.state[p] |
| if not st: |
| st["momentum_buffer"] = torch.zeros_like(p) |
| buf = st["momentum_buffer"] |
| buf.mul_(mom).add_(g) |
| u = g.add(buf, alpha=mom) if group["nesterov"] else buf |
| u = zeropower_via_newtonschulz5(u, group["ns_steps"]) |
| scale = max(1.0, p.size(0) / p.size(1)) ** 0.5 |
| if wd: |
| p.mul_(1 - lr * wd) |
| p.add_(u.to(p.dtype), alpha=-lr * scale) |
| else: |
| (b1, b2), eps = group["betas"], group["eps"] |
| for p in group["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) |
| st["step"] += 1 |
| t = st["step"] |
| st["m"].lerp_(g, 1 - b1) |
| st["v"].mul_(b2).addcmul_(g, g, value=1 - b2) |
| if wd: |
| p.mul_(1 - lr * wd) |
| denom = (st["v"].sqrt() / math.sqrt(1 - b2 ** t)).add_(eps) |
| p.addcdiv_(st["m"], denom, value=-lr / (1 - b1 ** t)) |
| return loss |
|
|
|
|
| def muon_param_split(model): |
| """2-D weights inside transformer blocks -> Muon; embeddings / lm_head / 1-D -> AdamW.""" |
| muon, other, seen = [], [], set() |
| for n, p in model.named_parameters(): |
| if id(p) in seen or not p.requires_grad: |
| continue |
| seen.add(id(p)) |
| if p.dim() == 2 and ".h." in f".{n}" and "wte" not in n and "wpe" not in n and "lm_head" not in n: |
| muon.append(p) |
| else: |
| other.append(p) |
| return muon, other |
|
|
|
|
| |
| |
| |
| class Sophia(Optimizer): |
| """m_t = b1 m + (1-b1) g |
| every k steps: h = b2 h + (1-b2) h_hat (h_hat from Hutchinson u*(Hu) or GNB B*g_hat^2) |
| w <- w - lr*wd*w ; w <- w - lr * sign(m) * min(|m| / max(rho*h, eps), 1)""" |
| def __init__(self, params, lr=6e-4, betas=(0.965, 0.99), rho=0.05, weight_decay=0.2, eps=1e-15): |
| super().__init__(params, dict(lr=lr, betas=betas, rho=rho, weight_decay=weight_decay, eps=eps)) |
| self.last_clip_frac = float("nan") |
|
|
| @torch.no_grad() |
| def update_hessian(self, h_hats): |
| """h_hats: dict param -> diagonal Hessian estimate (same shape).""" |
| for group in self.param_groups: |
| b2 = group["betas"][1] |
| for p in group["params"]: |
| if p not in h_hats: |
| continue |
| st = self.state[p] |
| if "h" not in st: |
| st["h"] = torch.zeros_like(p) |
| st["h"].mul_(b2).add_(h_hats[p], alpha=1 - b2) |
|
|
| @torch.no_grad() |
| def step(self, closure=None): |
| loss = closure() if closure is not None else None |
| clipped, total = 0, 0 |
| for group in self.param_groups: |
| lr, (b1, _), rho, wd, eps = (group["lr"], group["betas"], group["rho"], |
| group["weight_decay"], group["eps"]) |
| for p in group["params"]: |
| if p.grad is None: |
| continue |
| st = self.state[p] |
| if "m" not in st: |
| st["m"] = torch.zeros_like(p) |
| if "h" not in st: |
| st["h"] = torch.zeros_like(p) |
| m, h = st["m"], st["h"] |
| m.lerp_(p.grad, 1 - b1) |
| if wd: |
| p.mul_(1 - lr * wd) |
| ratio = (m.abs() / (rho * h).clamp(min=eps)).clamp(max=1.0) |
| p.addcmul_(m.sign(), ratio, value=-lr) |
| clipped += int((ratio >= 1.0).sum()); total += ratio.numel() |
| self.last_clip_frac = clipped / max(total, 1) |
| return loss |
|
|
|
|
| def sophia_hutchinson_estimate(model, x, y, params): |
| """h_hat = u * (H u), u ~ Rademacher. Needs double backward -> fp32 + math attention.""" |
| ctx = None |
| try: |
| from torch.nn.attention import sdpa_kernel, SDPBackend |
| ctx = sdpa_kernel(SDPBackend.MATH) |
| except Exception: |
| import contextlib |
| ctx = torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, |
| enable_math=True) if torch.cuda.is_available() else contextlib.nullcontext() |
| with ctx: |
| _, loss, _ = model(x, y) |
| grads = torch.autograd.grad(loss, params, create_graph=True) |
| us = [torch.randint_like(p, 0, 2) * 2 - 1 for p in params] |
| hvps = torch.autograd.grad(grads, params, grad_outputs=us) |
| return {p: (u * hv).detach() for p, u, hv in zip(params, us, hvps)} |
|
|
|
|
| def sophia_gnb_estimate(model, x, params, amp_ctx, scaler=None): |
| """Gauss-Newton-Bartlett: sample y_hat ~ softmax(logits), h_hat = N_tokens * grad(CE(y_hat))^2.""" |
| import torch.nn.functional as F |
| model.zero_grad(set_to_none=True) |
| with amp_ctx: |
| logits, _, _ = model(x) |
| with torch.no_grad(): |
| probs = F.softmax(logits.float().view(-1, logits.size(-1)), -1) |
| y_hat = torch.multinomial(probs, 1).view(x.shape) |
| del probs |
| loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), y_hat.view(-1)) |
| scale = scaler.get_scale() if (scaler is not None and scaler.is_enabled()) else 1.0 |
| (loss * scale).backward() |
| n_tok = x.numel() |
| out = {} |
| ok = True |
| for p in params: |
| if p.grad is None: |
| continue |
| g = p.grad.float() / scale |
| if not torch.isfinite(g).all(): |
| ok = False; break |
| out[p] = n_tok * g * g |
| model.zero_grad(set_to_none=True) |
| return out if ok else {} |
|
|
|
|
| |
| |
| |
| OPTIMIZER_DEFAULTS = { |
| |
| |
| "adamw": dict(lr=2e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1, grad_clip=1.0), |
| "mars": dict(lr=2e-3, betas=(0.95, 0.99), eps=1e-8, weight_decay=0.1, gamma=0.025, |
| max_c_norm=1.0, grad_clip=0.0), |
| "lion": dict(lr=2e-4, betas=(0.9, 0.99), weight_decay=1.0, grad_clip=1.0), |
| "muon": dict(lr=0.02, momentum=0.95, nesterov=True, ns_steps=5, weight_decay=0.01, |
| adamw_lr=2e-3, adamw_betas=(0.9, 0.95), adamw_wd=0.1, grad_clip=1.0), |
| "sophia": dict(lr=6e-4, betas=(0.965, 0.99), rho=0.05, weight_decay=0.2, grad_clip=1.0, |
| hess_interval=10, estimator="gnb", hess_batch=16), |
| } |
|
|
| OPTIMIZER_CATEGORY = { |
| "adamw": "AdamW (baseline)", |
| "mars": "Variance-reduced AdamW variant", |
| "lion": "Memory-efficient", |
| "muon": "Matrix-based", |
| "sophia": "Hessian-approximation (bonus)", |
| } |
|
|
| SWEEP_GRIDS = { |
| "adamw": [5e-4, 1e-3, 2e-3, 4e-3], |
| "mars": [5e-4, 1e-3, 2e-3, 4e-3], |
| "lion": [5e-5, 1e-4, 2e-4, 4e-4], |
| "muon": [0.005, 0.01, 0.02, 0.04], |
| "sophia": [2e-4, 4e-4, 8e-4, 1.6e-3], |
| } |
|
|
|
|
| def build_optimizer(name, model, hp): |
| """Returns (optimizer, hp_used). hp overrides OPTIMIZER_DEFAULTS[name].""" |
| h = dict(OPTIMIZER_DEFAULTS[name]); h.update(hp or {}) |
| if name == "adamw": |
| opt = AdamW(split_decay_groups(model, h["weight_decay"]), lr=h["lr"], betas=tuple(h["betas"]), eps=h["eps"]) |
| elif name == "mars": |
| dec, nodec = [], [] |
| for n, p in model.named_parameters(): |
| (dec if p.dim() >= 2 else nodec).append(p) |
| groups = [dict(params=dec, weight_decay=h["weight_decay"], gamma=h["gamma"]), |
| dict(params=nodec, weight_decay=0.0, gamma=0.0)] |
| opt = MARS(groups, lr=h["lr"], betas=tuple(h["betas"]), eps=h["eps"], max_c_norm=h["max_c_norm"]) |
| elif name == "lion": |
| opt = Lion(split_decay_groups(model, h["weight_decay"]), lr=h["lr"], betas=tuple(h["betas"])) |
| elif name == "muon": |
| mp, op = muon_param_split(model) |
| opt = Muon(mp, op, lr=h["lr"], momentum=h["momentum"], nesterov=h["nesterov"], ns_steps=h["ns_steps"], |
| weight_decay=h["weight_decay"], adamw_lr=h["adamw_lr"], adamw_betas=tuple(h["adamw_betas"]), |
| adamw_wd=h["adamw_wd"]) |
| elif name == "sophia": |
| opt = Sophia(split_decay_groups(model, h["weight_decay"]), lr=h["lr"], betas=tuple(h["betas"]), rho=h["rho"]) |
| else: |
| raise ValueError(name) |
| return opt, h |
|
|
|
|
| def optimizer_state_bytes(opt): |
| n = 0 |
| for st in opt.state.values(): |
| for v in st.values(): |
| if torch.is_tensor(v): |
| n += v.numel() * v.element_size() |
| return n |
|
|
