From 3331f7f8540c1fdc68182c106fc8b3d2e8a4ee04 Mon Sep 17 00:00:00 2001 From: "Carson M." Date: Thu, 30 Jul 2026 22:49:30 -0500 Subject: [PATCH 1/3] Magnitude-direction decoupling --- thorn/thorn.py | 54 ++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 39 insertions(+), 15 deletions(-) diff --git a/thorn/thorn.py b/thorn/thorn.py index a5e4232..43f873c 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): @@ -605,7 +600,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 +617,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 +777,35 @@ def _base_ortho_step(self, p: Parameter, group: _THORNParameterGroup): g = p.grad magma_scale = _momentum_aligned_mask(g, state, group) + + gain = F.softplus(state['row_gain']) * F.softplus(state['col_gain']) + p_hat = p / gain + p_g = p_hat * g + g.mul_(gain) + + def _adam_step(g, momentum, variance, step, beta1 = 0.9, beta2 = 0.95): + momentum.lerp_(g, weight=1 - beta1) + variance.mul_(beta2).addcmul_(g, g, value=1 - beta2) + denom = variance.div(1 - beta2 ** step).sqrt_() + return momentum.div(1 - beta1 ** step).atan2_(denom) + 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_hat.sub_(u, alpha=group.lr * _lr_scale_ortho(u) * magma_scale) + p_hat.mul_(state['target_norm'] / (p_hat.norm(dim=(-2, -1), keepdim=True) + 1e-8)) + + row_gain = F.softplus(state['row_gain']) + col_gain = F.softplus(state['col_gain']) + + grad_row = (p_g * col_gain).sum(dim=-1) * F.sigmoid(state['row_gain']).squeeze(-1) + grad_col = (p_g * row_gain).sum(dim=-2) * 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) + + gain = F.softplus(state['row_gain']) * F.softplus(state['col_gain']) + p.data.copy_(p_hat * gain) if group.none_grad: del g @@ -883,7 +901,13 @@ 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) + # softplus(ln(e - 1)) = 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]) + state['target_norm'] = p.norm(dim=(-2, -1), keepdim=True) else: state['moment'] = torch.zeros_like(g) state['variance'] = torch.zeros_like(g) From 4629efabc945fe3808f530f7eb5c30abd1672b9b Mon Sep 17 00:00:00 2001 From: "Carson M." Date: Sat, 1 Aug 2026 12:47:29 -0500 Subject: [PATCH 2/3] Do some stuff inplace --- thorn/thorn.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/thorn/thorn.py b/thorn/thorn.py index 43f873c..73dfa47 100644 --- a/thorn/thorn.py +++ b/thorn/thorn.py @@ -778,9 +778,12 @@ def _base_ortho_step(self, p: Parameter, group: _THORNParameterGroup): magma_scale = _momentum_aligned_mask(g, state, group) - gain = F.softplus(state['row_gain']) * F.softplus(state['col_gain']) - p_hat = p / gain - p_g = p_hat * g + # 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) def _adam_step(g, momentum, variance, step, beta1 = 0.9, beta2 = 0.95): @@ -793,19 +796,17 @@ def _adam_step(g, momentum, variance, step, beta1 = 0.9, beta2 = 0.95): if magma_scale != 0.0 and (state['step'] + 1) % self._update_rate == 0: u = _polar_decomp(u, group).to(dtype=p.dtype) - p_hat.sub_(u, alpha=group.lr * _lr_scale_ortho(u) * magma_scale) - p_hat.mul_(state['target_norm'] / (p_hat.norm(dim=(-2, -1), keepdim=True) + 1e-8)) + 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)) - row_gain = F.softplus(state['row_gain']) - col_gain = F.softplus(state['col_gain']) - - grad_row = (p_g * col_gain).sum(dim=-1) * F.sigmoid(state['row_gain']).squeeze(-1) - grad_col = (p_g * row_gain).sum(dim=-2) * F.sigmoid(state['col_gain']).squeeze(-2) + 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) - gain = F.softplus(state['row_gain']) * F.softplus(state['col_gain']) - p.data.copy_(p_hat * gain) + # reassemble + p.mul_(F.softplus(state['row_gain'])) + p.mul_(F.softplus(state['col_gain'])) if group.none_grad: del g From 23ba21902c073e98585a565b8687e14055bc3db3 Mon Sep 17 00:00:00 2001 From: "Carson M." Date: Mon, 3 Aug 2026 12:51:12 -0500 Subject: [PATCH 3/3] MD for Adam params --- thorn/thorn.py | 80 +++++++++++++++++++++++++++++++++++--------------- 1 file changed, 57 insertions(+), 23 deletions(-) diff --git a/thorn/thorn.py b/thorn/thorn.py index 73dfa47..5a99950 100644 --- a/thorn/thorn.py +++ b/thorn/thorn.py @@ -441,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 @@ -786,12 +806,6 @@ def _base_ortho_step(self, p: Parameter, group: _THORNParameterGroup): p_g = p * g g.mul_(gain) - def _adam_step(g, momentum, variance, step, beta1 = 0.9, beta2 = 0.95): - momentum.lerp_(g, weight=1 - beta1) - variance.mul_(beta2).addcmul_(g, g, value=1 - beta2) - denom = variance.div(1 - beta2 ** step).sqrt_() - return momentum.div(1 - beta1 ** step).atan2_(denom) - 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) @@ -902,16 +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) - # softplus(ln(e - 1)) = 1 + 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]) - state['target_norm'] = p.norm(dim=(-2, -1), keepdim=True) else: - state['moment'] = torch.zeros_like(g) - state['variance'] = torch.zeros_like(g) + 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: @@ -937,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