source/gum_optim.py · deeprcurs/MBG-1.0 at main

Upload source/gum_optim.py with huggingface_hub

02da691 verified about 2 months ago

8.38 kB

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

	"""
	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 ~2rn
	instead of 2mn. 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