| """Optimizers for Task 2. Rules: subclass torch.optim.Optimizer directly; no other torch.optim.* anywhere.
|
|
|
| Paper: "Fantastic Pretraining Optimizers and Where to Find Them" (Wen et al., 2025), Appendix A has one algorithm per
|
| optimizer. Every step() below follows the same pattern:
|
| for each param group -> for each parameter p with a gradient:
|
| g = p.grad (same shape as p), state = self.state[p] (dict that persists across steps, empty on step 1)
|
| -> initialise state on the first step -> update the running statistics -> apply the update to p in place.
|
| Hyperparameters come from `group` (the dict built from `defaults`, overridable per group). The training loop sets
|
| group["lr"] every step (learning-rate schedule), so always read the lr from the group, never cache it.
|
| """
|
| import torch
|
| from torch.optim import Optimizer
|
|
|
|
|
| def _adamw_update(p, g, state, group):
|
| lr, (beta1, beta2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"]
|
| if len(state) == 0:
|
| state["step"] = 0
|
| state["exp_avg"] = torch.zeros_like(p)
|
| state["exp_avg_sq"] = torch.zeros_like(p)
|
| state["step"] += 1
|
| t = state["step"]
|
| exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
|
| exp_avg.mul_(beta1).add_(g, alpha=1 - beta1)
|
| exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1 - beta2)
|
| bc1 = 1 - beta1 ** t
|
| bc2 = 1 - beta2 ** t
|
| p.mul_(1 - lr * wd)
|
| denom = (exp_avg_sq / bc2).sqrt_().add_(eps)
|
| p.addcdiv_(exp_avg, denom, value=-lr / bc1)
|
|
|
|
|
| class AdamW(Optimizer):
|
| """Category 1 (baseline). Appendix A, Algorithm 1."""
|
|
|
| def __init__(self, params, lr=6e-4, 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 = None
|
| if closure is not None:
|
| with torch.enable_grad():
|
| loss = closure()
|
| for group in self.param_groups:
|
| for p in group["params"]:
|
| if p.grad is None:
|
| continue
|
| g, state = p.grad, self.state[p]
|
| _adamw_update(p, g, state, group)
|
| return loss
|
|
|
|
|
| class NAdamW(Optimizer):
|
| """Category 2 (variance-reduced AdamW variants). Appendix A, Algorithm 2.
|
| Same state as AdamW; the difference is how the momentum enters the update (Nesterov look-ahead)."""
|
|
|
| def __init__(self, params, lr=6e-4, 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 = None
|
| if closure is not None:
|
| with torch.enable_grad():
|
| loss = closure()
|
| for group in self.param_groups:
|
| lr, (beta1, beta2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"]
|
| for p in group["params"]:
|
| if p.grad is None:
|
| continue
|
| g, state = p.grad, self.state[p]
|
| if len(state) == 0:
|
| state["step"] = 0
|
| state["exp_avg"] = torch.zeros_like(p)
|
| state["exp_avg_sq"] = torch.zeros_like(p)
|
| state["step"] += 1
|
| t = state["step"]
|
| exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
|
| exp_avg.mul_(beta1).add_(g, alpha=1 - beta1)
|
| exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1 - beta2)
|
| bc1 = 1 - beta1 ** t
|
| bc2 = 1 - beta2 ** t
|
| nesterov_m = exp_avg.mul(beta1).add_(g, alpha=1 - beta1)
|
| p.mul_(1 - lr * wd)
|
| denom = (exp_avg_sq / bc2).sqrt_().add_(eps)
|
| p.addcdiv_(nesterov_m, denom, value=-lr / bc1)
|
| return loss
|
|
|
|
|
| class Lion(Optimizer):
|
| """Category 3 (memory-efficient). Appendix A, Algorithm 3.
|
| One state buffer per parameter; the update is a sign, so every coordinate moves by exactly lr
|
| -> needs a much smaller lr (and usually larger weight decay) than AdamW."""
|
|
|
| def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.1):
|
| super().__init__(params, dict(lr=lr, betas=betas, weight_decay=weight_decay))
|
|
|
| @torch.no_grad()
|
| def step(self, closure=None):
|
| loss = None
|
| if closure is not None:
|
| with torch.enable_grad():
|
| loss = closure()
|
| for group in self.param_groups:
|
| lr, (beta1, beta2), wd = group["lr"], group["betas"], group["weight_decay"]
|
| for p in group["params"]:
|
| if p.grad is None:
|
| continue
|
| g, state = p.grad, self.state[p]
|
| if len(state) == 0:
|
| state["exp_avg"] = torch.zeros_like(p)
|
| exp_avg = state["exp_avg"]
|
| update = exp_avg.mul(beta1).add_(g, alpha=1 - beta1).sign_()
|
| p.mul_(1 - lr * wd)
|
| p.add_(update, alpha=-lr)
|
| exp_avg.mul_(beta2).add_(g, alpha=1 - beta2)
|
| return loss
|
|
|
|
|
| def zeropower_via_newtonschulz(G, steps):
|
| """G: (m, n) matrix -> matrix of the same shape with the same singular vectors and singular values pushed to ~1
|
| (an approximate orthogonalisation), computed with `steps` Newton-Schulz iterations. Used by Muon."""
|
| a, b, c = 3.4445, -4.7750, 2.0315
|
| X = G.float()
|
| X = X / (X.norm() + 1e-7)
|
| transposed = G.size(0) > G.size(1)
|
| if transposed:
|
| X = X.T
|
| 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):
|
| """Category 4 (matrix-based). Appendix A, Algorithm 8.
|
| Two kinds of param groups (built in train.py):
|
| use_muon=True -> hidden 2-D weight matrices (attention q/k/v/out, FFN linears): Muon update
|
| use_muon=False -> embedding (= tied output layer), norm gains, biases: AdamW update (reuse your AdamW logic)
|
| The two kinds use different learning rates (Muon's is much larger), set per group in train.py."""
|
|
|
| def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5, weight_decay=0.0,
|
| betas=(0.9, 0.95), eps=1e-8):
|
| super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps,
|
| weight_decay=weight_decay, betas=betas, eps=eps, use_muon=True))
|
|
|
| @torch.no_grad()
|
| def step(self, closure=None):
|
| loss = None
|
| if closure is not None:
|
| with torch.enable_grad():
|
| loss = closure()
|
| for group in self.param_groups:
|
| for p in group["params"]:
|
| if p.grad is None:
|
| continue
|
| g, state = p.grad, self.state[p]
|
| if group["use_muon"]:
|
| lr, momentum, wd = group["lr"], group["momentum"], group["weight_decay"]
|
| if len(state) == 0:
|
| state["momentum_buffer"] = torch.zeros_like(p)
|
| buf = state["momentum_buffer"]
|
| buf.lerp_(g, 1 - momentum)
|
| update = g.lerp(buf, momentum) if group["nesterov"] else buf
|
| update = zeropower_via_newtonschulz(update, group["ns_steps"])
|
| scale = max(1.0, p.size(0) / p.size(1)) ** 0.5
|
| p.mul_(1 - lr * wd)
|
| p.add_(update.to(p.dtype), alpha=-lr * scale)
|
| else:
|
| _adamw_update(p, g, state, group)
|
| return loss
|
|
|
|
|
| OPTIMIZERS = {"adamw": AdamW, "nadamw": NAdamW, "lion": Lion, "muon": Muon}
|
|
|
