Spaces:
Paused
Paused
| """Optimizer factory for the baselines the counter synapse should be compared against. | |
| The point of the method is memory, so the honest comparison is not just FP32 AdamW but the | |
| real memory-efficient-training optimizers. Supported names: | |
| adamw : torch.optim.AdamW (FP32 moments) -- the heavy reference. | |
| bnb8 : bitsandbytes 8-bit AdamW (optional dep, CUDA only) -- 4x smaller optimizer state. | |
| galore : GaLore-style low-rank projected AdamW (moments live in a low-rank subspace). | |
| lomo : LoMo-style fused-backward SGD (update during backward, no grad/moment storage). | |
| GaLore and LoMo are implemented here in plain PyTorch so they run on CPU and need no extra | |
| deps. bnb8 is imported lazily and only works where bitsandbytes + CUDA are available. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| __all__ = ["build_optimizer", "GaLoreAdamW", "LoMo", "available_optimizers"] | |
| def available_optimizers() -> list[str]: | |
| return ["adamw", "bnb8", "galore", "lomo"] | |
| def build_optimizer(name: str, params, lr: float, **kw): | |
| params = list(params) | |
| if name == "adamw": | |
| return torch.optim.AdamW(params, lr=lr) | |
| if name == "bnb8": | |
| try: | |
| import bitsandbytes as bnb | |
| except Exception as exc: # pragma: no cover - depends on env | |
| raise RuntimeError( | |
| "optimizer 'bnb8' needs bitsandbytes (CUDA): pip install bitsandbytes" | |
| ) from exc | |
| return bnb.optim.AdamW8bit(params, lr=lr) | |
| if name == "galore": | |
| return GaLoreAdamW(params, lr=lr, **kw) | |
| if name == "lomo": | |
| return LoMo(params, lr=lr, **kw) | |
| raise ValueError(f"unknown optimizer {name!r}; choices: {available_optimizers()}") | |
| class GaLoreAdamW(torch.optim.Optimizer): | |
| """Gradient Low-Rank Projection AdamW (Zhao et al. 2024), minimal pure-PyTorch version. | |
| For each 2-D parameter, the gradient G [m,n] is projected to a rank-r subspace P^T G (P | |
| from a periodic SVD of G), Adam runs on the small [r,n] moments, and the update is | |
| projected back. Optimizer-state memory for that parameter drops from 2*m*n to ~2*r*n. | |
| 1-D parameters fall back to plain AdamW. This is a faithful-enough reimplementation for | |
| a memory baseline, not the reference package. | |
| """ | |
| def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2, | |
| rank=128, update_proj_gap=200, scale=0.25): | |
| defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, | |
| rank=rank, update_proj_gap=update_proj_gap, scale=scale) | |
| super().__init__(params, defaults) | |
| def step(self, closure=None): | |
| loss = closure() if closure is not None else None | |
| for group in self.param_groups: | |
| b1, b2 = group["betas"] | |
| for p in group["params"]: | |
| if p.grad is None: | |
| continue | |
| g = p.grad | |
| st = self.state[p] | |
| low_rank = g.ndim == 2 and min(g.shape) > group["rank"] | |
| if "step" not in st: | |
| st["step"] = 0 | |
| if low_rank: | |
| st["proj"] = None | |
| else: | |
| st["exp_avg"] = torch.zeros_like(p) | |
| st["exp_avg_sq"] = torch.zeros_like(p) | |
| st["step"] += 1 | |
| t = st["step"] | |
| if group["weight_decay"]: | |
| p.mul_(1 - group["lr"] * group["weight_decay"]) | |
| if not low_rank: | |
| m, v = st["exp_avg"], st["exp_avg_sq"] | |
| m.mul_(b1).add_(g, alpha=1 - b1) | |
| v.mul_(b2).addcmul_(g, g, value=1 - b2) | |
| mhat = m / (1 - b1 ** t) | |
| vhat = v / (1 - b2 ** t) | |
| p.addcdiv_(mhat, vhat.sqrt().add_(group["eps"]), value=-group["lr"]) | |
| continue | |
| # tall vs wide: project the long dimension down to rank r. | |
| m_, n_ = g.shape | |
| r = group["rank"] | |
| if st["proj"] is None or t % group["update_proj_gap"] == 1: | |
| # periodic projection basis from the current gradient (left singular vecs) | |
| U, _, Vh = torch.linalg.svd(g.float(), full_matrices=False) | |
| if m_ >= n_: | |
| st["side"] = "left" | |
| st["proj"] = U[:, :r].contiguous() # [m, r] | |
| else: | |
| st["side"] = "right" | |
| st["proj"] = Vh[:r, :].contiguous() # [r, n] | |
| st["exp_avg"] = torch.zeros( | |
| (r, n_) if st["side"] == "left" else (m_, r), device=p.device, dtype=p.dtype) | |
| st["exp_avg_sq"] = torch.zeros_like(st["exp_avg"]) | |
| P = st["proj"] | |
| if st["side"] == "left": | |
| gr = P.t().to(g.dtype) @ g # [r, n] | |
| else: | |
| gr = g @ P.t().to(g.dtype) # [m, r] | |
| m, v = st["exp_avg"], st["exp_avg_sq"] | |
| m.mul_(b1).add_(gr, alpha=1 - b1) | |
| v.mul_(b2).addcmul_(gr, gr, value=1 - b2) | |
| mhat = m / (1 - b1 ** t) | |
| vhat = v / (1 - b2 ** t) | |
| upd = mhat / vhat.sqrt().add_(group["eps"]) | |
| full = (P.to(g.dtype) @ upd) if st["side"] == "left" else (upd @ P.to(g.dtype)) | |
| p.add_(full, alpha=-group["lr"] * group["scale"]) | |
| return loss | |
| class LoMo(torch.optim.Optimizer): | |
| """LoMo-style fused SGD (Lv et al. 2023): apply the update during backward via a grad | |
| hook and immediately free the gradient, so the full gradient tensor and any optimizer | |
| moments are never simultaneously resident. Pure SGD (no momentum) -> zero optimizer state. | |
| Usage differs from a normal optimizer: register hooks once, then DON'T call .backward via | |
| an outer optimizer.step that needs grads -- the hook does the update. We expose .step() as | |
| a no-op so the training loop stays uniform. | |
| """ | |
| def __init__(self, params, lr=1e-3, weight_decay=0.0): | |
| params = list(params) | |
| defaults = dict(lr=lr, weight_decay=weight_decay) | |
| super().__init__(params, defaults) | |
| self._handles = [] | |
| for group in self.param_groups: | |
| for p in group["params"]: | |
| if p.requires_grad: | |
| self._handles.append(p.register_post_accumulate_grad_hook(self._make_hook(group))) | |
| def _make_hook(self, group): | |
| lr, wd = group["lr"], group["weight_decay"] | |
| def hook(p): | |
| if p.grad is None: | |
| return | |
| if wd: | |
| p.mul_(1 - lr * wd) | |
| p.add_(p.grad, alpha=-lr) | |
| p.grad = None # free immediately: no grad lingers, no moments exist | |
| return hook | |
| def step(self, closure=None): | |
| return closure() if closure is not None else None | |
| def zero_grad(self, set_to_none: bool = True): | |
| # grads are already freed in the hook; nothing to clear. | |
| return | |