""" 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