Spaces:
Paused
Paused
File size: 7,173 Bytes
6f75f79 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """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)
@torch.no_grad()
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"]
@torch.no_grad()
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
@torch.no_grad()
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
|