Add MP-01 (9500) 100M base: weights, config, tokenizer, model code, model card and eval results
11e2c7c verified | """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 | |
| 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) | |
| 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"]) | |