File size: 7,828 Bytes
08af434 | 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 | """Muon optimizer matching `optax.contrib.muon`, with patches for our use."""
import torch
import torch.distributed as dist
import torch.nn as nn
from utils.logging_utils import log_for_0
def muon_with_aux_adam(model: nn.Module, lr: float):
"""Muon optimizer matching `optax.contrib.muon(learning_rate=lr_schedule)`.
Params are partitioned by `ndim == 2 -> Muon`, else Nesterov-Adam (b1=0.9,
b2=0.999, eps=1e-8, wd=0). A safety wrapper zero-fills missing grads
because the decoder/denoiser branches alternate and leave one head's
params at `None` per step.
"""
import muon as _muon_module
from muon import MuonWithAuxAdam, SingleDeviceMuonWithAuxAdam
# Patch muon's Adam update to Nesterov-Adam:
# mu_hat = b1 * (mu / (1 - b1**(t+1))) + (1 - b1) * (g / (1 - b1**t))
def _nesterov_adam_update(grad, mu, nu, step, betas, eps):
b1, b2 = betas
mu.lerp_(grad, 1 - b1)
nu.lerp_(grad.square(), 1 - b2)
mu_hat = b1 * (mu / (1 - b1 ** (step + 1))) + (1 - b1) * (grad / (1 - b1 ** step))
nu_hat = nu / (1 - b2 ** step)
return mu_hat / (nu_hat.sqrt() + eps)
_muon_module.adam_update = _nesterov_adam_update
# Patch NS5 to run in fp32 with eps=1e-8. Upstream muon.py downcasts to
# bf16 and uses eps=1e-7, both of which lose precision through the 5
# squarings.
def _zeropower_via_newtonschulz5_fp32(G, steps):
a, b, c = (3.4445, -4.7750, 2.0315)
X = G.to(torch.float32)
if G.size(-2) > G.size(-1):
X = X.mT
X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-8)
for _ in range(steps):
A = X @ X.mT
B = b * A + c * A @ A
X = a * X + B @ X
if G.size(-2) > G.size(-1):
X = X.mT
return X
_muon_module.zeropower_via_newtonschulz5 = _zeropower_via_newtonschulz5_fp32
# Two patches on the upstream muon_update:
# 1. Add Nesterov bias correction (upstream skips it):
# mu_hat = beta*(mu/(1-beta**(t+1))) + (1-beta)*(g/(1-beta**t))
# Without it, early steps differ by ~20x at t=1.
# 2. Use the `sqrt(max(1, fan_out/fan_in))` shape scaling.
# Bare Parameter tensors like `proj_kernel`/`unembed_kernel` are
# stored (in, out) instead of nn.Linear's (out, in); `flax_layout`
# flags those so we flip the ratio. The ratio can be up to
# sqrt(vocab/bn) ≈ sqrt(63) ≈ 7.9 for unembed_kernel.
def _muon_update_optax(grad, momentum, step, beta=0.95, ns_steps=5,
nesterov=True, flax_layout=False):
momentum.lerp_(grad, 1 - beta)
if nesterov:
mu_corr = momentum / (1 - beta ** (step + 1))
g_corr = grad / (1 - beta ** step)
update = beta * mu_corr + (1 - beta) * g_corr
else:
update = momentum / (1 - beta ** step)
if update.ndim == 4:
update = update.view(len(update), -1)
update = _muon_module.zeropower_via_newtonschulz5(update, steps=ns_steps)
# `sqrt(max(1, fan_out/fan_in))` scaling, accounting for storage layout.
m, n = grad.size(-2), grad.size(-1)
if flax_layout:
update *= max(1, n / m) ** 0.5 # (in, out) layout: fan_out = n
else:
update *= max(1, m / n) ** 0.5 # (out, in) layout: fan_out = m
return update
# Identify which 2D params are nn.Linear.weight (stored (out, in)) vs
# bare Parameter tensors that follow the (in, out) convention. The shape
# scaling formula assumes (in, out); for (out, in) params we flip the
# ratio.
linear_weight_ids = set()
for m in model.modules():
if isinstance(m, nn.Linear) and m.weight is not None:
linear_weight_ids.add(id(m.weight))
muon_params, adam_params = [], []
muon_flax_layout = {} # id(p) -> bool (True = (in, out), False = (out, in))
for _name, p in model.named_parameters():
if not p.requires_grad:
continue
if p.ndim == 2:
muon_params.append(p)
muon_flax_layout[id(p)] = id(p) not in linear_weight_ids
else:
adam_params.append(p)
distributed = dist.is_available() and dist.is_initialized()
base_cls = MuonWithAuxAdam if distributed else SingleDeviceMuonWithAuxAdam
class _SafeMuonAuxAdam(base_cls):
"""Fixes layered on top of upstream `MuonWithAuxAdam.step`:
1. Zero-fill missing grads (decoder/denoiser branches alternate and
leave one head's params at `None` per step).
2. Correct distributed all_gather padding when Muon param count is
not a multiple of world_size: `(-len) % ws`, not `len % ws`.
3. Reimplement the Muon path to call our bias-correcting
`_muon_update_optax`.
"""
@torch.no_grad()
def step(self):
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
p.grad = torch.zeros_like(p)
for group in self.param_groups:
if group["use_muon"]:
params = group["params"]
if distributed:
ws = dist.get_world_size()
pad_n = (-len(params)) % ws
params_pad = params + [torch.empty_like(params[-1])] * pad_n
for base_i in range(0, len(params), ws):
if base_i + dist.get_rank() < len(params):
p = params[base_i + dist.get_rank()]
self._muon_step_one(p, group)
dist.all_gather(
params_pad[base_i:base_i + ws],
params_pad[base_i + dist.get_rank()],
)
else:
for p in params:
self._muon_step_one(p, group)
else:
for p in group["params"]:
state = self.state[p]
if len(state) == 0:
state["exp_avg"] = torch.zeros_like(p)
state["exp_avg_sq"] = torch.zeros_like(p)
state["step"] = 0
state["step"] += 1
update = _muon_module.adam_update(
p.grad, state["exp_avg"], state["exp_avg_sq"],
state["step"], group["betas"], group["eps"],
)
p.mul_(1 - group["lr"] * group["weight_decay"])
p.add_(update, alpha=-group["lr"])
def _muon_step_one(self, p, group):
state = self.state[p]
if len(state) == 0:
state["momentum_buffer"] = torch.zeros_like(p)
state["step"] = 0
state["step"] += 1
update = _muon_update_optax(
p.grad, state["momentum_buffer"], state["step"],
beta=group["momentum"],
flax_layout=muon_flax_layout.get(id(p), False),
)
p.mul_(1 - group["lr"] * group["weight_decay"])
p.add_(update, alpha=-group["lr"])
# Hyperparams hardcoded to `optax.contrib.muon` defaults.
param_groups = [
dict(params=muon_params, lr=lr, momentum=0.95,
weight_decay=0.0, use_muon=True),
dict(params=adam_params, lr=lr, betas=(0.9, 0.999),
eps=1e-8, weight_decay=0.0, use_muon=False),
]
log_for_0(
f"Muon: {len(muon_params)} 2D params; "
f"Nesterov-AdamW: {len(adam_params)} other params"
)
return _SafeMuonAuxAdam(param_groups)
|