File size: 2,651 Bytes
11e2c7c | 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 | """Muon (Newton-Schulz orthogonalised momentum) for 2-D hidden matrices (and 3-D stacks of MoE expert matrices); scale matched to AdamW RMS (Moonlight)."""
import torch
@torch.no_grad()
def newton_schulz(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
a, b, c = 3.4445, -4.7750, 2.0315
X = G.bfloat16()
transposed = X.shape[0] > X.shape[1]
if transposed:
X = X.T
X = X / (X.norm() + 1e-7)
for _ in range(steps):
A = X @ X.T
X = a * X + (b * A + c * (A @ A)) @ X
return X.T if transposed else X
class Muon(torch.optim.Optimizer):
def __init__(self, params, lr=1e-3, momentum=0.95, weight_decay=0.1, split=None):
# split: dict param-id -> number of equal row-chunks to orthogonalise independently (fused gate/up)
super().__init__(params, dict(lr=lr, momentum=momentum, weight_decay=weight_decay))
self.split = split or {}
self.bad = None # device counter of updates dropped for being non-finite (read at log time, no sync here)
@torch.no_grad()
def step(self):
for g in self.param_groups:
for p in g["params"]:
if p.grad is None:
continue
st = self.state[p]
if "buf" not in st:
st["buf"] = torch.zeros_like(p.grad)
buf = st["buf"]
buf.lerp_(p.grad, 1 - g["momentum"])
upd = p.grad.lerp(buf, g["momentum"]) # nesterov
# 3-D = stacked MoE experts: each expert's matrix is orthogonalised on its own
mats = list(upd.unbind(0)) if upd.ndim == 3 else [upd]
done = []
for m in mats:
outs = []
for ch in m.chunk(self.split.get(id(p), 1), dim=0):
o = newton_schulz(ch)
outs.append(o * (0.2 * max(ch.shape) ** 0.5))
done.append(torch.cat(outs, 0) if len(outs) > 1 else outs[0])
o = torch.stack(done) if upd.ndim == 3 else done[0]
# a rare non-finite Newton-Schulz result (seen on the RX 7600 with finite grads) would poison the
# weight for good: drop this parameter's update for the step instead
ok = torch.isfinite(o).all()
if self.bad is None: self.bad = torch.zeros((), device=o.device, dtype=torch.int64)
self.bad += (~ok).long()
o = torch.where(ok, o, torch.zeros((), device=o.device, dtype=o.dtype))
p.mul_(1 - g["lr"] * g["weight_decay"])
p.add_(o.to(p.dtype), alpha=-g["lr"])
|