File size: 8,380 Bytes
02da691 | 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 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | """
MBG 1.0 — src/gum_optim.py
GUM optimizer (GaLore-Unbiased + Muon) — Innovation I1.
Memory-lean optimizer for the MBG 1.0 ladder. Combines:
* GaLore (Zhao et al. 2024; arXiv:2403.03507): project gradients of weight
matrices into a low-rank subspace, optimize there, project back -> big
optimizer-state savings.
* Muon (Liu et al. 2025, Moonlight; arXiv:2502.16982): orthogonalize momentum
via Newton-Schulz -> better conditioning / ~2x FLOP efficiency.
Dispatch rule (per parameter):
* 2D weight matrices listed in `matrix_groups` -> GUM (low-rank subspace).
* All other params (embeddings not listed, biases, norms, 1D) -> AdamW fallback.
NOTE on memory: GUM stores momentum (mom) and second moment (mom2) in the LOW-RANK
subspace of shape (r, n) instead of full (m, n), so state per matrix is ~2*r*n
instead of 2*m*n. With r = 0.25*min(m,n), savings are large for big matrices.
English-only (project rule). Pure PyTorch.
"""
from __future__ import annotations
import torch
# Newton-Schulz coefficients optimized for Muon (Su et al. 2024).
NS_ALPHA = 3.4445
NS_BETA = -4.7750
NS_GAMMA = 2.0315
def newton_schulz_ortho(x: torch.Tensor, iters: int = 5) -> torch.Tensor:
"""Approximate the polar (orthogonal) factor of x via Newton-Schulz.
Muon's recipe normalizes the input to unit Frobenius norm before iterating
(the Newton-Schulz coefficients are tuned for matrices with spectral norm ~1).
We apply that normalization so the polynomial does not diverge to inf/nan.
"""
a = x / (x.norm() + 1e-8) # unit Frobenius norm
a = a.to(torch.float32)
for _ in range(iters):
a2 = a @ a.transpose(-1, -2)
a = (1 + NS_ALPHA) * a \
- NS_BETA * (a2 @ a) \
+ NS_GAMMA * (a2 @ a2 @ a)
# stability guard: renormalize if it blows up
n = a.norm()
if not torch.isfinite(n):
return x / (x.norm() + 1e-8)
if n > 10.0:
a = a / n
return a
class GUM:
"""Low-rank GaLore + Muon update for a single matrix parameter.
Weight W: (m, n). Gradient G: (m, n).
* P = left singular vectors of G, first r cols -> (m, r).
* project: G_lr = P^T G -> (r, n).
* maintain low-rank momentum mom, mom2 of shape (r, n).
* orthogonalize the momentum direction (Muon) in the subspace.
* update_full = P @ step_lr -> (m, n).
"""
def __init__(self, lr: float = 3e-3, beta1: float = 0.9, beta2: float = 0.95,
eps: float = 1e-8, rank_ratio: float = 0.25,
ns_iters: int = 5, update_proj_gap: int = 20,
weight_decay: float = 0.05) -> None:
self.lr = lr
self.beta1 = beta1
self.beta2 = beta2
self.eps = eps
self.rank_ratio = rank_ratio
self.ns_iters = ns_iters
self.update_proj_gap = update_proj_gap
self.weight_decay = weight_decay
def _rank_for(self, m: int, n: int) -> int:
return max(1, int(min(m, n) * self.rank_ratio))
def step(self, param: torch.nn.Parameter, grad: torch.Tensor,
state: dict, step: int) -> None:
m, n = grad.shape
r = self._rank_for(m, n)
device = grad.device
# ---- maintain low-rank projection (GaLore) ----
# P = top-r LEFT singular vectors of grad -> (m, r)
if "P" not in state or (step % self.update_proj_gap == 0):
try:
U, _, _ = torch.linalg.svd(grad.float(), full_matrices=False)
P = U[:, :r] # (m, r)
except Exception:
P = torch.randn(m, r, device=device)
P = P / P.norm(dim=0, keepdim=True).clamp_min(1e-8)
state["P"] = P.detach()
P = state["P"] # (m, r)
grad_lr = P.t() @ grad.float() # (r, n)
if "mom" not in state:
state["mom"] = torch.zeros_like(grad_lr)
state["mom2"] = torch.zeros_like(grad_lr)
mom = state["mom"]
mom2 = state["mom2"]
mom.mul_(self.beta1).add_(grad_lr, alpha=1 - self.beta1)
mom2.mul_(self.beta2).addcmul_(grad_lr, grad_lr, value=1 - self.beta2)
bc1 = 1 - self.beta1 ** step
bc2 = 1 - self.beta2 ** step
m_hat = mom / bc1
v_hat = mom2 / bc2
# orthogonalize the momentum direction (Muon)
m_ortho = newton_schulz_ortho(m_hat, self.ns_iters)
# adaptive 2nd-moment scaling (Adam-like) in the subspace, then
# normalize the orthogonal step to keep it stable.
denom = v_hat.sqrt().clamp_min(self.eps)
step_lr = m_ortho / denom
if step_lr.numel() > 0 and step_lr.norm() > 0:
step_lr = step_lr * (grad_lr.norm().clamp_max(1.0)
/ (step_lr.norm() + 1e-8))
# ---- project update back to full rank ----
update = P @ step_lr # (m, n)
if self.weight_decay != 0:
param.data.mul_(1 - self.lr * self.weight_decay)
param.data.add_(update, alpha=-self.lr)
state["step"] = step
class AdamWGroup:
"""Vanilla AdamW for 1D / embedding / bias params (stable fallback)."""
def __init__(self, lr: float = 3e-3, beta1: float = 0.9, beta2: float = 0.999,
eps: float = 1e-8, weight_decay: float = 0.05) -> None:
self.lr = lr
self.beta1 = beta1
self.beta2 = beta2
self.eps = eps
self.weight_decay = weight_decay
def step(self, param: torch.nn.Parameter, grad: torch.Tensor,
state: dict, step: int) -> None:
if "exp_avg" not in state:
state["exp_avg"] = torch.zeros_like(param.data)
state["exp_avg_sq"] = torch.zeros_like(param.data)
exp_avg = state["exp_avg"]
exp_avg_sq = state["exp_avg_sq"]
g = grad.float()
exp_avg.mul_(self.beta1).add_(g, alpha=1 - self.beta1)
exp_avg_sq.mul_(self.beta2).addcmul_(g, g, value=1 - self.beta2)
bc1 = 1 - self.beta1 ** step
bc2 = 1 - self.beta2 ** step
step_size = self.lr / bc1
denom = (exp_avg_sq / bc2).sqrt().clamp_min(self.eps)
if self.weight_decay != 0:
param.data.mul_(1 - self.lr * self.weight_decay)
param.data.addcdiv_(exp_avg, denom, value=-step_size)
class GUMOptimizer(torch.optim.Optimizer):
"""Dispatches GUM (low-rank) vs AdamW per parameter group."""
def __init__(self, params, matrix_groups=None, lr: float = 3e-3,
beta1: float = 0.9, beta2: float = 0.95, eps: float = 1e-8,
rank_ratio: float = 0.25, ns_iters: int = 5,
update_proj_gap: int = 20, weight_decay: float = 0.05) -> None:
defaults = dict(lr=lr, weight_decay=weight_decay)
super().__init__(params, defaults)
self._matrix_groups = set(matrix_groups or [])
self._name_by_ptr = {}
self._gum = GUM(lr=lr, beta1=beta1, beta2=beta2, eps=eps,
rank_ratio=rank_ratio, ns_iters=ns_iters,
update_proj_gap=update_proj_gap, weight_decay=weight_decay)
self._adam = AdamWGroup(lr=lr, beta1=beta1, beta2=0.999, eps=eps,
weight_decay=weight_decay)
def set_param_names(self, name_to_param: dict) -> None:
"""Map param object -> name (keys are the param objects)."""
# name_to_param is {name: param_object}
self._name_by_ptr = {id(p): name for name, p in name_to_param.items()}
def _use_gum(self, param) -> bool:
name = self._name_by_ptr.get(id(param), "")
return (param.dim() == 2) and (name in self._matrix_groups)
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
step_count = getattr(self, "_step", 0) + 1
setattr(self, "_step", step_count)
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
state = self.state[p]
if self._use_gum(p):
self._gum.step(p, p.grad, state, step_count)
else:
self._adam.step(p, p.grad, state, step_count)
return loss
|