Spaces:
Running on Zero
Running on Zero
| """AdamAtan2 (Everett et al., 2024; used by TRM) and a warmup + cosine LR schedule.""" | |
| import math | |
| import torch | |
| class AdamAtan2(torch.optim.Optimizer): | |
| """Adam with the division m / sqrt(v) replaced by a * atan2(m, b * sqrt(v)); decoupled weight decay.""" | |
| def __init__(self, params, lr=1e-4, betas=(0.9, 0.95), weight_decay=0.0, a=1.27, b=1.0): | |
| super().__init__(params, dict(lr=lr, betas=betas, weight_decay=weight_decay, a=a, b=b)) | |
| def step(self): | |
| for g in self.param_groups: | |
| b1, b2 = g["betas"] | |
| for p in g["params"]: | |
| if p.grad is None: | |
| continue | |
| s = self.state[p] | |
| if not s: | |
| s["step"], s["m"], s["v"] = 0, torch.zeros_like(p), torch.zeros_like(p) | |
| s["step"] += 1 | |
| p.mul_(1 - g["lr"] * g["weight_decay"]) | |
| s["m"].lerp_(p.grad, 1 - b1) | |
| s["v"].lerp_(p.grad * p.grad, 1 - b2) | |
| m_hat = s["m"] / (1 - b1 ** s["step"]) | |
| v_hat = s["v"] / (1 - b2 ** s["step"]) | |
| p.add_(torch.atan2(m_hat, g["b"] * v_hat.sqrt()), alpha=-g["lr"] * g["a"]) | |
| def lr_at(step, base_lr, warmup, total, min_ratio=1.0): | |
| if step < warmup: | |
| return base_lr * (step + 1) / warmup | |
| progress = min(1.0, (step - warmup) / max(1, total - warmup)) | |
| return base_lr * (min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * progress))) | |
| def set_lr(opt, lr): | |
| for g in opt.param_groups: | |
| g["lr"] = lr | |