diff --git a/thorn/thorn.py b/thorn/thorn.py index a5e4232..5a99950 100644 --- a/thorn/thorn.py +++ b/thorn/thorn.py @@ -25,6 +25,7 @@ import torch.distributed as dist from torch.distributed.tensor import DTensor, Replicate, Shard import torch.nn as nn +import torch.nn.functional as F from torch.nn import Parameter from torch.optim import Optimizer from torch.optim.optimizer import _get_value @@ -320,7 +321,7 @@ class _THORNParameterGroup: betas: tuple[float, float] = field(default_factory=lambda: (0.95, 0.95)) iters: int = field(default=5) rectify: bool = field(default=True) - target_rms: float = field(default=0.2) + gain_lr: float = field(default=1e-3) lower_bound: float = field(default=1e-3) safety_factor: float = field(default=0.05) cushion: float = field(default=0.02) @@ -404,14 +405,8 @@ def _weight_decay( update.addcmul_(p, mask.mul_(weight_decay)) return update -def _lr_scale_ortho(p: torch.Tensor, target_rms: float = 0.2): - if target_rms != 0.0: - # Scale LR to match RMS update of AdamW so AdamW's LR can be reused - # per formula 4 of https://arxiv.org/pdf/2502.16982 - return target_rms * (max(1, *p.shape[-2:]) ** 0.5) - else: - # Match original behavior of Jordan et al - return max(1, p.size(-2) / p.size(-1)) ** 0.5 +def _lr_scale_ortho(p: torch.Tensor): + return (max(*p.shape[-2:]) / min(*p.shape[-2:])) ** 0.5 @torch.no_grad() def _compute_rect(group: _THORNParameterGroup, step: float | int): @@ -446,6 +441,26 @@ def _momentum_aligned_mask(g: torch.Tensor, state: dict, group: _THORNParameterG state['random_state'], mask = _w1rand(state['random_state']) return state['s'] * (1.0 if mask % 2 == 0 else 0.0) +def _adam_step( + g: torch.Tensor, + momentum: torch.Tensor, + variance: torch.Tensor, + step: int | torch.Tensor, + beta1: float = 0.9, + beta2: float = 0.95, + degenerate = False +): + momentum.lerp_(g, weight=1 - beta1) + variance.mul_(beta2).addcmul_(g, g, value=1 - beta2) + if not degenerate: + denom = variance.div(1 - beta2 ** step).sqrt_() + # atan2 instead of div per https://arxiv.org/pdf/2407.05872 + u = momentum.div(1 - beta1 ** step).atan2_(denom) + else: + # clone because we might later call _weight_decay which modifies in place + u = momentum.clone() + return u + @dataclass class _DistributedTHORNState: worker_rank: int @@ -605,7 +620,7 @@ def update_param( self.computed_u = None u = _weight_decay(p, self.scattered_u, group.weight_decay) - p.data.sub_(u, alpha=group.lr * _lr_scale_ortho(u, target_rms=group.target_rms) * scale) + p.data.sub_(u, alpha=group.lr * _lr_scale_ortho(u) * scale) self.scattered_u = None u_dtensor = None @@ -622,7 +637,7 @@ class THORNOrthogonalizedParameterGroup(TypedDict, total=True): lower_bound: NotRequired[float] safety_factor: NotRequired[float] cushion: NotRequired[float] - target_rms: NotRequired[float] + gain_lr: NotRequired[float] rectify: NotRequired[bool] momentum_align: NotRequired[bool] coeffs: NotRequired[list[tuple[float, float, float]]] @@ -782,12 +797,30 @@ def _base_ortho_step(self, p: Parameter, group: _THORNParameterGroup): g = p.grad magma_scale = _momentum_aligned_mask(g, state, group) + + # recover direction + row_gain = F.softplus(state['row_gain']) + col_gain = F.softplus(state['col_gain']) + gain = row_gain * col_gain + p.div_(gain) + p_g = p * g + g.mul_(gain) + u = self._update_momentum(p, g, group) if magma_scale != 0.0 and (state['step'] + 1) % self._update_rate == 0: u = _polar_decomp(u, group).to(dtype=p.dtype) - u = _per_neuron_norm(u, state['moment2'], group) - u = _weight_decay(p, u, group.weight_decay) - p.data.sub_(u, alpha=group.lr * _lr_scale_ortho(u, target_rms=group.target_rms) * magma_scale) + + p.sub_(u, alpha=group.lr * _lr_scale_ortho(u) * magma_scale) + p.mul_(state['target_norm'] / (p.norm(dim=(-2, -1), keepdim=True) + 1e-8)) + + grad_row = (p_g * col_gain).sum(dim=-1).mul_(F.sigmoid(state['row_gain']).squeeze(-1)) + grad_col = (p_g * row_gain).sum(dim=-2).mul_(F.sigmoid(state['col_gain']).squeeze(-2)) + state['row_gain'].sub_(_adam_step(grad_row.unsqueeze(-1), state['row_gain_moment'], state['row_gain_variance'], state['step']), alpha=group.gain_lr) + state['col_gain'].sub_(_adam_step(grad_col.unsqueeze(-2), state['col_gain_moment'], state['col_gain_variance'], state['step']), alpha=group.gain_lr) + + # reassemble + p.mul_(F.softplus(state['row_gain'])) + p.mul_(F.softplus(state['col_gain'])) if group.none_grad: del g @@ -883,10 +916,21 @@ def _step_params(self, params: list[torch.nn.Parameter], group: _THORNParameterG state['s'] = 1.0 if group.orthogonalize: state['moment'] = torch.zeros_like(g) - state['moment2'] = torch.zeros((*g.shape[:-1], 1), dtype=g.dtype, device=g.device) else: state['moment'] = torch.zeros_like(g) state['variance'] = torch.zeros_like(g) + # softplus(ln(e - 1)) = 1 + if p.ndim > 1: + state['row_gain'] = torch.full((*g.shape[:-1], 1), 0.5413248546, dtype=g.dtype, device=g.device) + state['col_gain'] = torch.full((*g.shape[:-2], 1, *g.shape[-1:]), 0.5413248546, dtype=g.dtype, device=g.device) + for k in ['row_gain', 'col_gain']: + state[f'{k}_moment'] = torch.zeros_like(state[k]) + state[f'{k}_variance'] = torch.zeros_like(state[k]) + else: + state['gain'] = torch.full((), 0.5413248546, dtype=g.dtype, device=g.device) + state['gain_moment'] = torch.zeros((), dtype=g.dtype, device=g.device) + state['gain_variance'] = torch.zeros((), dtype=g.dtype, device=g.device) + state['target_norm'] = p.norm(dim=(-2, -1) if g.ndim > 1 else -1, keepdim=True) if isinstance(p.data, DTensor): if all(isinstance(placement, Replicate) for placement in cast(DTensor, p).placements) or not group.orthogonalize: @@ -912,24 +956,39 @@ def _step_params(self, params: list[torch.nn.Parameter], group: _THORNParameterG magma_scale = _momentum_aligned_mask(g, state, group) rect = _compute_rect(group, step) - momentum = state['moment'] - variance = state['variance'] - - momentum.lerp_(g, weight=1 - beta1) - variance.mul_(beta2).addcmul_(g, g, value=1 - beta2) - if rect > 0.0: - denom = variance.div(1 - beta2 ** step).sqrt_() - # atan2 instead of div per https://arxiv.org/pdf/2407.05872 - u = momentum.div(1 - beta1 ** step).atan2_(denom) + # recover direction + if p.ndim > 1: + row_gain = F.softplus(state['row_gain']) + col_gain = F.softplus(state['col_gain']) + gain = row_gain * col_gain else: - # clone because _weight_decay modifies in place - u = momentum.clone() + gain = F.softplus(state['gain']) + p.div_(gain) + p_g = p * g + g.mul_(gain) + + u = _adam_step(g, state['moment'], state['variance'], step, beta1, beta2, degenerate=rect == 0.0) should_update = magma_scale != 0.0 and (state['step'] + 1) % self._update_rate == 0 if should_update: u = _weight_decay(p, u, group.weight_decay) - - p.sub_(u, alpha=group.lr * rect * magma_scale) + p.sub_(u, alpha=group.lr * rect * magma_scale) + p.mul_(state['target_norm'] / (p.norm(dim=(-2, -1) if p.ndim > 1 else -1, keepdim=True) + 1e-8)) + + if p.ndim > 1: + grad_row = (p_g * col_gain).sum(dim=-1).mul_(F.sigmoid(state['row_gain']).squeeze(-1)) + grad_col = (p_g * row_gain).sum(dim=-2).mul_(F.sigmoid(state['col_gain']).squeeze(-2)) + state['row_gain'].sub_(_adam_step(grad_row.unsqueeze(-1), state['row_gain_moment'], state['row_gain_variance'], state['step']), alpha=group.gain_lr) + state['col_gain'].sub_(_adam_step(grad_col.unsqueeze(-2), state['col_gain_moment'], state['col_gain_variance'], state['step']), alpha=group.gain_lr) + else: + grad_gain = (p_g * gain).sum().mul_(F.sigmoid(state['gain'])) + state['gain'].sub_(_adam_step(grad_gain, state['gain_moment'], state['gain_variance'], state['step']), alpha=group.gain_lr) + + if p.ndim > 1: + p.mul_(F.softplus(state['row_gain'])) + p.mul_(F.softplus(state['col_gain'])) + else: + p.mul_(F.softplus(state['gain'])) if group.none_grad: del g