208 lines
7.2 KiB
Python
208 lines
7.2 KiB
Python
""" PyTorch MARS Optimizer
|
|
|
|
Code simplified from https://github.com/AGI-Arena/MARS
|
|
|
|
Paper: MARS: Unleashing the Power of Variance Reduction for Training Large Models - https://arxiv.org/abs/2411.10438
|
|
|
|
@article{yuan2024mars,
|
|
title={MARS: Unleashing the Power of Variance Reduction for Training Large Models},
|
|
author={Yuan, Huizhuo and Liu, Yifeng and Wu, Shuang and Zhou, Xun and Gu, Quanquan},
|
|
journal={arXiv preprint arXiv:2411.10438},
|
|
year={2024}
|
|
}
|
|
"""
|
|
# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
import math
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
from torch.optim.optimizer import Optimizer
|
|
|
|
from ._types import ParamsT
|
|
|
|
|
|
def _mars_single_tensor_step(
|
|
p: torch.Tensor,
|
|
grad: torch.Tensor,
|
|
exp_avg: torch.Tensor,
|
|
exp_avg_sq: torch.Tensor,
|
|
lr: float,
|
|
weight_decay: float,
|
|
beta1: float,
|
|
beta2: float,
|
|
last_grad: torch.Tensor,
|
|
eps: float,
|
|
step: int,
|
|
gamma: float,
|
|
mars_type: str,
|
|
is_grad_2d: bool,
|
|
optimize_1d: bool,
|
|
lr_1d_factor: bool,
|
|
betas_1d: Tuple[float, float],
|
|
caution: bool,
|
|
):
|
|
# optimize_1d ==> use MARS for 1d param, else use AdamW
|
|
if optimize_1d or is_grad_2d:
|
|
one_minus_beta1 = 1. - beta1
|
|
if step == 1:
|
|
# this is a timm addition, making first step more consistent when no grad history, otherwise tests fail
|
|
c_t = grad
|
|
else:
|
|
c_t = (grad - last_grad).mul_(gamma * (beta1 / one_minus_beta1)).add_(grad)
|
|
c_t_norm = torch.norm(c_t)
|
|
if c_t_norm > 1.:
|
|
c_t = c_t / c_t_norm
|
|
exp_avg.mul_(beta1).add_(c_t, alpha=one_minus_beta1)
|
|
|
|
if caution:
|
|
# Apply caution as per 'Cautious Optimizers' - https://arxiv.org/abs/2411.16085
|
|
mask = (exp_avg * grad > 0).to(grad.dtype)
|
|
mask.div_(mask.mean().clamp_(min=1e-3))
|
|
exp_avg = exp_avg * mask
|
|
|
|
if mars_type == "adamw":
|
|
exp_avg_sq.mul_(beta2).addcmul_(c_t, c_t, value=1. - beta2)
|
|
bias_correction1 = 1.0 - beta1 ** step
|
|
bias_correction2 = 1.0 - beta2 ** step
|
|
denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
|
|
update = p * weight_decay + (exp_avg / bias_correction1).div_(denom)
|
|
elif mars_type == "lion":
|
|
update = p * weight_decay + exp_avg.sign()
|
|
else:
|
|
assert False
|
|
p.add_(update, alpha=-lr)
|
|
else:
|
|
beta1_1d, beta2_1d = betas_1d
|
|
exp_avg.mul_(beta1_1d).add_(grad, alpha=1. - beta1_1d)
|
|
exp_avg_sq.mul_(beta2_1d).addcmul_(grad, grad, value=1. - beta2_1d)
|
|
bias_correction1 = 1.0 - beta1_1d ** step
|
|
bias_correction2 = 1.0 - beta2_1d ** step
|
|
denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
|
|
if caution:
|
|
mask = (exp_avg * grad > 0).to(grad.dtype)
|
|
mask.div_(mask.mean().clamp_(min=1e-3))
|
|
exp_avg = exp_avg * mask
|
|
update = p * weight_decay + (exp_avg / bias_correction1).div_(denom)
|
|
p.add_(update, alpha=-(lr * lr_1d_factor))
|
|
return exp_avg, exp_avg_sq
|
|
|
|
|
|
class Mars(Optimizer):
|
|
""" MARS Optimizer
|
|
|
|
Paper: MARS: Unleashing the Power of Variance Reduction for Training Large Models
|
|
https://arxiv.org/abs/2411.10438
|
|
|
|
"""
|
|
def __init__(
|
|
self,
|
|
params: ParamsT,
|
|
lr: float = 3e-3,
|
|
betas: Tuple[float, float] = (0.9, 0.99),
|
|
eps: float = 1e-8,
|
|
weight_decay: float = 0.,
|
|
gamma: float = 0.025,
|
|
mars_type: str = "adamw",
|
|
optimize_1d: bool = False,
|
|
lr_1d_factor: float = 1.0,
|
|
betas_1d: Optional[Tuple[float, float]] = None,
|
|
caution: bool = False
|
|
):
|
|
if not 0.0 <= lr:
|
|
raise ValueError("Invalid learning rate: {}".format(lr))
|
|
if not 0.0 <= eps:
|
|
raise ValueError("Invalid epsilon value: {}".format(eps))
|
|
if not 0.0 <= betas[0] < 1.0:
|
|
raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
|
|
if not 0.0 <= betas[1] < 1.0:
|
|
raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
|
|
assert mars_type in ["adamw", "lion"], "MARS type not supported"
|
|
|
|
defaults = dict(
|
|
lr=lr,
|
|
betas=betas,
|
|
eps=eps,
|
|
weight_decay=weight_decay,
|
|
mars_type=mars_type,
|
|
gamma=gamma,
|
|
optimize_1d=optimize_1d,
|
|
lr_1d_factor=lr_1d_factor,
|
|
betas_1d=betas_1d or betas,
|
|
caution=caution,
|
|
)
|
|
super(Mars, self).__init__(params, defaults)
|
|
|
|
def __setstate__(self, state):
|
|
super(Mars, self).__setstate__(state)
|
|
for group in self.param_groups:
|
|
group.setdefault('caution', False)
|
|
|
|
@torch.no_grad()
|
|
def step(self, closure=None):
|
|
"""Performs a single optimization step.
|
|
|
|
Arguments:
|
|
closure (callable, optional): A closure that reevaluates the model
|
|
and returns the loss.
|
|
"""
|
|
loss = None
|
|
if closure is not None:
|
|
with torch.enable_grad():
|
|
loss = closure()
|
|
|
|
for group in self.param_groups:
|
|
for p in group['params']:
|
|
if p.grad is None:
|
|
continue
|
|
grad = p.grad
|
|
if grad.is_sparse:
|
|
raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
|
|
|
|
state = self.state[p]
|
|
# State initialization
|
|
if len(state) <= 1:
|
|
state['step'] = 0
|
|
# Exponential moving average of gradient values
|
|
state['exp_avg'] = torch.zeros_like(p)
|
|
# Last Gradient
|
|
state['last_grad'] = torch.zeros_like(p)
|
|
# Exponential moving average of squared gradient values
|
|
state['exp_avg_sq'] = torch.zeros_like(p)
|
|
|
|
state['step'] += 1
|
|
step = state['step']
|
|
exp_avg = state['exp_avg']
|
|
exp_avg_sq = state['exp_avg_sq']
|
|
last_grad = state['last_grad']
|
|
lr = group['lr']
|
|
wd = group['weight_decay']
|
|
beta1, beta2 = group['betas']
|
|
is_grad_2d = grad.ndim >= 2
|
|
|
|
# FIXME add multi-tensor (if usage warrants), make more standard
|
|
_mars_single_tensor_step(
|
|
p,
|
|
grad,
|
|
exp_avg,
|
|
exp_avg_sq,
|
|
lr,
|
|
wd,
|
|
beta1,
|
|
beta2,
|
|
last_grad,
|
|
group['eps'],
|
|
step,
|
|
group['gamma'],
|
|
mars_type=group['mars_type'],
|
|
is_grad_2d=is_grad_2d,
|
|
optimize_1d=group['optimize_1d'],
|
|
lr_1d_factor=group['lr_1d_factor'],
|
|
betas_1d=group['betas_1d'],
|
|
caution=group['caution'],
|
|
)
|
|
|
|
state['last_grad'] = grad
|
|
|
|
return loss
|