| |
|
|
| import torch |
| import math |
| from torch.optim.optimizer import Optimizer |
|
|
| class ActivationPreconditionedOptimizer(Optimizer): |
| """ |
| Implements an Activation-Preconditioned Optimizer with optional normalization, |
| scaling, Nesterov momentum, and grafting, inspired by NormalizedSteepestDescent. |
| """ |
| def __init__(self, params, lr=1e-3, betas=(0.9, 0.95), weight_decay=0.1, |
| scale=0.2, eps=1e-8, use_nesterov=False, use_grafting=False): |
| if not 0.0 <= lr: |
| raise ValueError(f"Invalid learning rate: {lr}") |
| if not 0.0 <= betas[0] < 1.0: |
| raise ValueError(f"Invalid beta1 parameter: {betas[0]}") |
| if not 0.0 <= betas[1] < 1.0: |
| raise ValueError(f"Invalid beta2 parameter: {betas[1]}") |
| if not 0.0 <= weight_decay: |
| raise ValueError(f"Invalid weight_decay value: {weight_decay}") |
|
|
| defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay, scale=scale, |
| eps=eps, use_nesterov=use_nesterov, use_grafting=use_grafting) |
| super(ActivationPreconditionedOptimizer, 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'] |
| beta1, beta2_precond = group['betas'] |
| |
| beta2_grafting = group['betas'][1] |
| weight_decay = group['weight_decay'] |
| eps = group['eps'] |
| scale = group['scale'] |
| use_nesterov = group['use_nesterov'] |
| use_grafting = group['use_grafting'] |
| |
| for p in group['params']: |
| if p.grad is None: |
| continue |
| |
| if not hasattr(p, 'activation_second_moment'): |
| continue |
|
|
| grad = p.grad |
| if grad.is_sparse: |
| raise RuntimeError('This optimizer does not support sparse gradients.') |
|
|
| state = self.state[p] |
|
|
| |
| if len(state) == 0: |
| state['step'] = 0 |
| state['momentum_buffer'] = torch.zeros_like(p, memory_format=torch.preserve_format) |
| state['preconditioner'] = torch.zeros((p.shape[1], p.shape[1]), device=p.device) |
| if use_grafting: |
| state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format) |
|
|
| momentum_buffer = state['momentum_buffer'] |
| preconditioner = state['preconditioner'] |
| |
| state['step'] += 1 |
|
|
| |
| momentum_buffer.mul_(beta1).add_(grad, alpha=1 - beta1) |
| |
| xxt = p.activation_second_moment |
| preconditioner.mul_(beta2_precond).add_(xxt, alpha=1 - beta2_precond) |
|
|
| |
| if use_nesterov: |
| Gt = torch.lerp(grad, momentum_buffer, beta1) |
| else: |
| Gt = momentum_buffer |
|
|
| |
| eps_I = torch.eye(preconditioner.size(0), device=preconditioner.device) * eps |
| precond_inv = torch.linalg.inv(preconditioner + eps_I) |
|
|
| print(f"Preconditioner shape: {preconditioner.shape}, Gt shape: {Gt.shape}") |
| |
| |
| update_direction = torch.matmul(Gt, precond_inv) |
| update_norm = torch.linalg.norm(update_direction) |
| update_direction.div_(update_norm + eps) |
| |
| |
| alpha = 0.0 |
| if use_grafting: |
| exp_avg_sq = state['exp_avg_sq'] |
| exp_avg_sq.mul_(beta2_grafting).addcmul_(grad, grad, value=1 - beta2_grafting) |
| |
| step = state['step'] |
| bias_correction1 = 1 - beta1 ** step |
| bias_correction2 = 1 - beta2_grafting ** step |
| |
| denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps) |
| adam_update_magnitude = torch.linalg.norm((Gt / bias_correction1) / denom) |
| alpha = -lr * adam_update_magnitude |
| else: |
| mn = p.numel() |
| alpha = -lr * scale * math.sqrt(mn) |
|
|
| |
| if weight_decay > 0.0: |
| p.add_(p, alpha=-lr * weight_decay) |
| |
| p.add_(update_direction, alpha=alpha) |
|
|
| delattr(p, 'activation_second_moment') |
|
|
| return loss |
