File size: 5,404 Bytes
ad93bc0 | 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 | # activation_gradient_preconditioned_adam.py
import torch
import math
from torch.optim.optimizer import Optimizer
class SOAP(Optimizer):
"""
Implements Activation-Gradient-Preconditioned Adam.
This optimizer uses second-order statistics from both input activations (X_t)
and output gradients (G_Yt) to form two preconditioners. This is closely
related to the Shampoo optimizer.
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999, 0.95), weight_decay=0.1, eps=1e-6):
# betas[0]: momentum (β1)
# betas[1]: adam second moment (β2 in pseudocode, for V')
# betas[2]: preconditioners (β3 in pseudocode, for L and R)
# if not (isinstance(betas, (list, tuple)) and len(betas) == 3):
# raise ValueError("Betas must be a tuple/list of 3 elements.")
defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay, eps=eps)
super(SOAP, self).__init__(params, defaults)
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
lr = group['lr']
if len(group['betas']) == 2:
beta1, beta2_adam = group['betas']
beta3_precond = beta2_adam
else:
beta1, beta2_adam, beta3_precond = group['betas']
weight_decay = group['weight_decay']
eps = group['eps']
for p in group['params']:
if p.grad is None:
continue
grad = p.grad
state = self.state[p]
if len(state) == 0:
state['step'] = 0
state['momentum_buffer'] = torch.zeros_like(p, memory_format=torch.preserve_format)
# L_t: Gradient preconditioner
state['left_preconditioner'] = torch.zeros((p.shape[0], p.shape[0]), device=p.device)
# R_t: Activation preconditioner
state['right_preconditioner'] = torch.zeros((p.shape[1], p.shape[1]), device=p.device)
# V'_t: Adam's 2nd moment buffer in the *projected* space
state['adam_v_prime'] = torch.zeros_like(p, memory_format=torch.preserve_format)
state['step'] += 1
step = state['step']
momentum_buffer = state['momentum_buffer']
left_preconditioner = state['left_preconditioner']
right_preconditioner = state['right_preconditioner']
adam_v_prime = state['adam_v_prime']
# xxt = p.activation_second_moment
# gygyt = p.output_gradient_second_moment
GGT = torch.mm(grad, grad.t())
GTG = torch.mm(grad.t(), grad)
# Step 5: Update momentum
momentum_buffer.mul_(beta1).add_(grad, alpha=1 - beta1)
# Step 6 & 7: Update preconditioners
left_preconditioner.mul_(beta3_precond).add_(GGT, alpha=1 - beta3_precond)
right_preconditioner.mul_(beta3_precond).add_(GTG, alpha=1 - beta3_precond)
# Step 8 & 9: Eigendecomposition of preconditioners
# we use float32 for eigendecomposition to ensure numerical stability
eps_L = torch.eye(left_preconditioner.size(0), device=p.device) * eps
# S_L, U = torch.linalg.eigh(left_preconditioner.float() + eps_L)
# use svd instead of eigh for better numerical stability
U, S_L, _ = torch.linalg.svd(left_preconditioner.float() + eps_L)
eps_R = torch.eye(right_preconditioner.size(0), device=p.device) * eps
# S_R, V = torch.linalg.eigh(right_preconditioner.float() + eps_R)
# use svd instead of eigh for better numerical stability
V, S_R, _ = torch.linalg.svd(right_preconditioner.float() + eps_R)
# Step 10: Project momentum and gradient
M_prime = torch.mm(U.t(), torch.mm(momentum_buffer, V))
G_prime = torch.mm(U.t(), torch.mm(grad, V))
# +++ ADDED: Store G_prime in state for external logging +++
# will cost more memory, only use it for logging
state['G_prime'] = G_prime
# Step 11: Update Adam's second moment in projected space
adam_v_prime.mul_(beta2_adam).addcmul_(G_prime, G_prime, value=1 - beta2_adam)
# Bias Correction for Adam
bias_correction1 = 1 - beta1 ** step
bias_correction2 = 1 - beta2_adam ** step
M_prime_hat = M_prime / bias_correction1
adam_v_prime_hat = adam_v_prime / bias_correction2
# Step 12: Run Adam in the preconditioned space
denom = adam_v_prime_hat.sqrt().add_(eps)
Q_t = M_prime_hat / denom
if weight_decay > 0.0:
p.add_(p, alpha=-lr * weight_decay)
# Step 13: Project back to do the update
update = torch.mm(U, torch.mm(Q_t, V.t()))
p.add_(update, alpha=-lr)
return loss |