File size: 9,417 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 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 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | """
Implementation of a hybrid optimizer.
For high-dimensional parameters (tensors with more than 2 dimensions), this optimizer
applies steepest descent with respect to a specified matrix norm (induced 1, 2, or infinity norm).
The update rules are derived from the provided formulas.
For all other parameters (2D or 1D), it defaults to the Adam optimization algorithm.
"""
import torch
import math
from optimizer.numerical_methods import zeropower_via_newtonschulz5
class MatrixNormSteepestDescent(torch.optim.Optimizer):
"""
Implements a hybrid optimization algorithm.
It uses matrix norm steepest descent for tensors with dim > 2,
and Adam for all other parameters.
"""
def __init__(self, params, norm_p=2, lr=1e-3, momentum=0.95, betas=(0.9, 0.999), weight_decay=0.1, scale=0.2, eps=1e-8, use_nesterov=False):
"""
Initializes the optimizer.
Args:
params (iterable): Iterable of parameters to optimize or dicts defining parameter groups.
norm_p (int or float): The matrix norm to use for high-dimensional parameters.
Supported values are 1, 2, and float('inf').
lr (float, optional): Learning rate (default: 1e-3).
betas (Tuple[float, float], optional): Coefficients used for computing running averages
of gradient and its square in Adam (default: (0.9, 0.999)).
beta1 is also used as momentum for the matrix norm updates.
weight_decay (float, optional): Weight decay (L2 penalty) (default: 0.1).
eps (float, optional): Term added to the denominator to improve numerical stability (default: 1e-8).
use_nesterov (bool, optional): Enables Nesterov momentum for matrix norm updates (default: False).
"""
if lr < 0.0:
raise ValueError(f"Invalid learning rate: {lr}")
if not 0.0 <= betas[0] < 1.0:
raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
if not 0.0 <= betas[1] < 1.0:
raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
if weight_decay < 0.0:
raise ValueError(f"Invalid weight_decay value: {weight_decay}")
defaults = dict(norm_p=norm_p, lr=lr, momentum=momentum, betas=betas, weight_decay=weight_decay, scale=scale, eps=eps, use_nesterov=use_nesterov)
super(MatrixNormSteepestDescent, self).__init__(params, defaults)
def __setstate__(self, state):
super(MatrixNormSteepestDescent, self).__setstate__(state)
@torch.no_grad()
def step(self, closure=None):
"""Performs a single optimization step."""
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
lr = group['lr']
momentum = group['momentum']
beta1, beta2 = group['betas']
weight_decay = group['weight_decay']
scale = group['scale']
eps = group['eps']
norm_p = group['norm_p']
use_nesterov = group['use_nesterov']
for p in group['params']:
if p.grad is None:
continue
grad = p.grad
# Apply decoupled weight decay for all parameters
if weight_decay != 0:
p.add_(p, alpha=-weight_decay * lr)
param_state = self.state[p]
# Check parameter dimension to decide update rule
if p.dim() >= 2:
# --- Matrix Norm Steepest Descent for Tensors with dim >= 2 ---
# Momentum calculation
if momentum != 0:
if 'momentum_buffer' not in param_state:
buf = param_state['momentum_buffer'] = torch.clone(grad).detach()
else:
buf = param_state['momentum_buffer']
buf.mul_(momentum).add_(grad, alpha=1 - momentum)
if use_nesterov:
Gt = torch.lerp(grad, buf, momentum)
else:
Gt = buf
else:
Gt = grad
# original_shape = Gt.shape
# # Matricize the tensor by reshaping: (d1, d_2*d_3*...)
# matrix_Gt = Gt.reshape(original_shape[0], -1)
update_magnitude = 0
update_direction = torch.zeros_like(Gt)
if norm_p == 1:
# Based on the formula for the induced 1-norm update
col_max_abs = torch.max(torch.abs(Gt), dim=-2, keepdim=True)[0]
update_magnitude = torch.sum(col_max_abs)
# Direction is ColumnMax(Gt)
mask = (torch.abs(Gt) >= col_max_abs - 1e-12)
update_direction[mask] = torch.sign(Gt[mask])
elif norm_p == float('inf'):
# Based on the formula for the induced infinity-norm update
row_max_abs = torch.max(torch.abs(Gt), dim=-1, keepdim=True)[0]
update_magnitude = torch.sum(row_max_abs)
# Direction is RowMax(Gt)
mask = (torch.abs(Gt) >= row_max_abs - 1e-12)
update_direction[mask] = torch.sign(Gt[mask])
elif norm_p == 2:
# Implement Muon (similar to the https://arxiv.org/pdf/2502.16982 version since we have the normalization)
try:
update_direction = zeropower_via_newtonschulz5(G=Gt)
except torch.linalg.LinAlgError:
# If NewtonSchulz fails, we skip the update for this parameter to avoid crashing.
update_magnitude = 0
update_direction.fill_(0)
elif norm_p == '2,inf':
# +++ NEW: Row-wise Normalized Gradient Descent +++
# Formula: G_i / ||G_i||_2 for each row i
# 1. Calculate L2 norm for each row (dim=-1)
# Gt shape is likely (Out, In) or (B, Out, In) etc.
# We assume the last dimension is the "input" dimension for summation
# row_norms = torch.norm(Gt, p=2, dim=-1, keepdim=True)
# !!!!!!! here fore embedding layers, we consider using col-wise normalization
# This is also sth interesting: column-wise normalization is 1,2 induced steepest descent
col_norms = torch.norm(Gt, p=2, dim=-1, keepdim=True)
# print(f'The shape of column norm: {col_norms.shape}')
# 2. Normalize gradient rows
# update_direction = Gt / (row_norms + 1e-12)
update_direction = Gt / (col_norms + 1e-12)
else:
raise ValueError(f'Invalid matrix norm choice for p.dim() >= 2. Only 1, 2, and inf are supported.')
# Combine magnitude and direction and reshape back to original tensor shape
# final_update = (update_direction * update_magnitude).reshape(original_shape)
update_direction = update_direction / (torch.linalg.norm(update_direction) + 1e-12)
alpha = -lr * scale * math.sqrt(p.numel())
p.add_(update_direction, alpha=alpha)
else:
# --- Adam Fallback for Params with dim == 1 ---
# Initialize Adam state
if 'step' not in param_state:
param_state['step'] = 0
param_state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
param_state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
exp_avg = param_state['exp_avg']
exp_avg_sq = param_state['exp_avg_sq']
param_state['step'] += 1
step = param_state['step']
# Adam update logic
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
bias_correction1 = 1 - beta1 ** step
bias_correction2 = 1 - beta2 ** step
step_size = lr / bias_correction1
denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
p.addcdiv_(exp_avg, denom, value=-step_size)
return loss |