wehub/labmlai--annotated_deep_learning_paper_implementations
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:19:01 +08:00
commit 3b90d1192f
2172 changed files with 594509 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
labml_nn/optimizers/__init__.py
+218
View File
@@ -0,0 +1,218 @@
"""
---
title: Optimizers
summary: >
A set of PyTorch implementations/tutorials of popular gradient descent based optimizers.
Currently includes Adam, AMSGrad and RAdam optimizers.
---
# Optimizers
## Optimizer Implementations
* [Adam Optimizer](adam.html)
* [AMSGrad Optimizer](amsgrad.html)
* [Adam Optimizer with warmup](adam_warmup.html)
* [Noam Optimizer](noam.html)
* [Rectified Adam Optimizer](radam.html)
* [AdaBelief Optimizer](ada_belief.html)
* [Sophia-G Optimizer](sophia.html)
This [MNIST example](mnist_experiment.html) uses these optimizers.
## Generic Adaptive Optimizer Base class and Weight Decay
This file defines a common base class for *Adam* and extensions of it.
The base class helps use implement other optimizers with minimal code
because of re-usability.
We also define a special class for L2 weight decay, so that we don't
have to implement it inside each of the optimizers,
and can easily extend to other weight decays like L1 without
changing the optimizers.
Here are some concepts on PyTorch optimizers:
### Parameter groups
PyTorch optimizers group parameters into sets called groups.
Each group can have its own hyper-parameters like learning rates.
In most common cases there will be only one group.
This is when you initialize your optimizer with,
```python
Optimizer(model.parameters())
```
You can define multiple parameter groups when initializing the optimizer:
```python
Optimizer([{'params': model1.parameters()}, {'params': model2.parameters(), 'lr': 2}])
```
Here we pass a list of groups. Each group is a dictionary with its parameters under the key 'params'.
You specify any hyper-parameters as well. If the hyper parameters are not defined they will default
to the optimizer level defaults.
You can access (and even change) these groups, and their hyper-parameters with `optimizer.param_groups`.
Most learning rate schedule implementations I've come across do access this and change 'lr'.
### States
Optimizer maintains states (a dictionary) for each parameter (a tensor), in a dictionary `optimizer.state`.
This is where the optimizer maintains things like exponential averages.
"""
from typing import Dict, Tuple, Any
import torch
from torch import nn
from torch.optim.optimizer import Optimizer
class GenericAdaptiveOptimizer(Optimizer):
"""
## Base class for *Adam* and extensions
"""
def __init__(self, params, defaults: Dict[str, Any], lr: float, betas: Tuple[float, float], eps: float):
"""
### Initialize
* `params` is the collection of parameters or set of parameter groups.
* `defaults` a dictionary of default hyper-parameters
* `lr` is the learning rate, $\alpha$
* `betas` is the tuple $(\beta_1, \beta_2)$
* `eps` is $\epsilon$
"""
# Check the hyper-parameters
if not 0.0 <= lr:
raise ValueError(f"Invalid learning rate: {lr}")
if not 0.0 <= eps:
raise ValueError(f"Invalid epsilon value: {eps}")
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]}")
# Add the hyper-parameters to the defaults
defaults.update(dict(lr=lr, betas=betas, eps=eps))
# Initialize the PyTorch optimizer.
# This will create parameter groups with the default hyper-parameters
super().__init__(params, defaults)
def init_state(self, state: Dict[str, any], group: Dict[str, any], param: nn.Parameter):
"""
### Initialize state for a given parameter tensor
This should be overridden with code to initialize `state` for parameters `param`.
`group` is the parameter group dictionary to which `param` belongs.
"""
pass
def step_param(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor, param: torch.Tensor):
"""
### Take optimizer step on a parameter tensor
This should be overridden and take the optimization step on `param` tensor $\theta$,
where `grad` is the gradient for that parameter, $g_t$,
`state` is the optimizer state dictionary for that parameter, and
`group` is the parameter group dictionary `param` belongs to.
"""
pass
@torch.no_grad()
def step(self, closure=None):
"""
### Optimizer step
We have created a template method that does the common stuff every *Adam* based optimizer needs.
"""
# Calculate loss.
#
# 🤔 I'm not sure when you need this. I guess it's if you define a function that
# calculates the loss, does `loss.backward` and return the loss, instead of calling
# it on your own you could pass it to `optimizer.step`. 🤷‍♂️
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
# Iterate through the parameter groups
for group in self.param_groups:
# Iterate through the parameters in the parameter group
for param in group['params']:
# Skip if the parameter has no gradient
if param.grad is None:
continue
# Get the gradient tensor
grad = param.grad.data
# We don't handle sparse gradients
if grad.is_sparse:
raise RuntimeError('GenericAdaptiveOptimizer does not support sparse gradients,'
' please consider SparseAdam instead')
# Get the state for the parameter
state = self.state[param]
# Initialize the state if state is uninitialized
if len(state) == 0:
self.init_state(state, group, param)
# Take the optimization step on the parameter
self.step_param(state, group, grad, param)
# Return the loss, calculated from closure
return loss
class WeightDecay:
"""
## L2 Weight decay
"""
def __init__(self, weight_decay: float = 0., weight_decouple: bool = True, absolute: bool = False):
"""
### Initialize weight decay
* `weight_decay` is the decay coefficient
* `weight_decouple` is a flag indicating whether to add the weight decay to the gradient or directly
decay from the parameter. If added to the gradient it will go through the normal optimizer update.
* `absolute` this flag indicates whether the weight decay coefficient is absolute. This is applicable
when the decay is performed directly on the parameter. If this is false the actual decay is
`weight_decay`
* `learning_rate`.
"""
# Check hyper-parameters
if not 0.0 <= weight_decay:
raise ValueError(f"Invalid weight_decay value: {weight_decay}")
self.absolute = absolute
self.weight_decouple = weight_decouple
self.weight_decay = weight_decay
def defaults(self):
"""
Return defaults for parameter groups
"""
return dict(weight_decay=self.weight_decay)
def __call__(self, param: torch.nn.Parameter, grad: torch.Tensor, group: Dict[str, any]):
"""
### Perform weight decay and return the gradient
"""
# If we are doing the decay on the parameter directly
if self.weight_decouple:
# If the weight decay coefficient is absolute
if self.absolute:
param.data.mul_(1.0 - group['weight_decay'])
# Otherwise,
else:
param.data.mul_(1.0 - group['lr'] * group['weight_decay'])
# Return the unmodified gradient
return grad
else:
if group['weight_decay'] != 0:
# Add the weight decay to the gradient and return the modified gradient
return grad.add(param.data, alpha=group['weight_decay'])
else:
return grad
labml_nn/optimizers/ada_belief.py
+159
View File
@@ -0,0 +1,159 @@
"""
---
title: AdaBelief optimizer
summary: A simple PyTorch implementation/tutorial of AdaBelief optimizer.
---
# AdaBelief Optimizer
This is based from AdaBelief
[official implementation](https://github.com/juntang-zhuang/Adabelief-Optimizer)
of the paper
[AdaBelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients](https://arxiv.org/abs/2010.07468).
This is implemented in [PyTorch](https://pytorch.org) as an extension to [RAdam](radam.html).
The main difference between Adam optimizer and AdaBelief is that,
how it calculates the adaptive learning rate;
instead of dividing by the exponential moving average of square of the gradients,
AdaBelief divides by the exponential mean of variance.
\begin{align}
m_t &\leftarrow \beta_1 m_{t-1} + (1 - \beta_1) \cdot g_t \\
\textcolor{cyan}{s_t} &\textcolor{cyan}{\leftarrow} \textcolor{cyan}{\beta_2 s_{t-1} + (1 - \beta_2) \cdot (g_t - m_t)^2} \\
\hat{m}_t &\leftarrow \frac{m_t}{1-\beta_1^t} \\
\textcolor{cyan}{\hat{s}_t} &\textcolor{cyan}{\leftarrow} \frac{\textcolor{cyan}{s_t} + \textcolor{red}{\epsilon}}{\textcolor{cyan}{1-\beta_2^t}} \\
\theta_t &\leftarrow \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\textcolor{cyan}{\hat{s}_t}} + \epsilon}
\end{align}
🤔 The paper calculates variance as $(g_t - m_t)^2$,
but I feel it should use the bias corrected momentum
$(g_t - \textcolor{orange}{\hat{m}_t})^2$.
I guess this doesn't affect things much because
bias correction is $\approx 1$ after the initial training steps.
"""
from typing import Dict, Any
import torch
from torch import nn
from labml_nn.optimizers import WeightDecay
from labml_nn.optimizers.radam import RAdam
class AdaBelief(RAdam):
"""
## AdaBelief Optimizer
This class extends from RAdam optimizer defined in [`radam.py`](radam.html).
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16,
weight_decay: WeightDecay = WeightDecay(), amsgrad=False,
degenerate_to_sgd=True,
rectify=True, defaults=None):
"""
### Initialize the optimizer
* `params` is the list of parameters
* `lr` is the learning rate $\alpha$
* `betas` is a tuple of ($\beta_1$, $\beta_2$)
* `eps` is $\hat{\epsilon}$ or $\epsilon$ based on `optimized_update`
* `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html)
* `optimized_update` is a flag whether to optimize the bias correction of the second moment
by doing it after adding $\epsilon$
* `amsgrad` is a flag indicating whether to use AMSGrad or fallback to plain Adam
* `degenerate_to_sgd` whether to use sgd when the rectification term $r_t$ is intractable
* `rectify` is whether to use RAdam update
* `defaults` is a dictionary of default for group values.
This is useful when you want to extend the class `AdaBelief`.
"""
defaults = {} if defaults is None else defaults
super().__init__(params, lr, betas, eps, weight_decay, amsgrad, degenerate_to_sgd, defaults)
self.rectify = rectify
def init_state(self, state: Dict[str, any], group: Dict[str, any], param: nn.Parameter):
"""
### Initialize a parameter state
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `param` is the parameter tensor $\theta_{t-1}$
"""
state['step'] = 0
# Exponential moving average of gradient values
state['exp_avg'] = torch.zeros_like(param, memory_format=torch.preserve_format)
# Exponential moving average of variance
state['exp_avg_var'] = torch.zeros_like(param, memory_format=torch.preserve_format)
# If `amsgrad` flag is `True` for this parameter group, we maintain the maximum of
# exponential moving average of variance
if group['amsgrad']:
# Maintains max of all exp. moving avg. of sq. grad. values
state['max_exp_avg_var'] = torch.zeros_like(param, memory_format=torch.preserve_format)
def get_ms(self, state: Dict[str, Any], group: Dict[str, Any], grad: torch.Tensor):
"""
### Calculate $m_t$ and $s_t$ or $\max(s_1, s_2, ..., s_{t-1}, s_t)$
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$
"""
# Get $\beta_1$ and $\beta_2$
beta1, beta2 = group['betas']
# Get $m_{t-1}$ and $s_{t-1}$
m, s = state['exp_avg'], state['exp_avg_var']
# In-place calculation of $m_t$
# $$m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) \cdot g_t$$
m.mul_(beta1).add_(grad, alpha=1 - beta1)
# Difference between gradient and momentum
grad_residual = grad - m
# In-place calculation of $s_t$
# $$s_t \leftarrow \beta_2 s_{t-1} + (1 - \beta_2) \cdot (g_t - m_t)^2$$
s.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1 - beta2)
# If this parameter group is using `amsgrad`
if group['amsgrad']:
# Get $\max(s_1, s_2, ..., s_{t-1})$.
s_max = state['max_exp_avg_var']
# Calculate $\max(s_1, s_2, ..., s_{t-1}, s_t)$.
torch.maximum(s_max, s, out=s_max)
return m, s_max
else:
# $m_t$ and $s_t$ otherwise
return m, s
def step_param(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor, param: torch.nn.Parameter):
"""
### Take an update step for a given parameter tensor
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$
* `param` is the parameter tensor $\theta_{t-1}$
"""
# Calculate weight decay
grad = self.weight_decay(param, grad, group)
# Get $m_t$ and $v_t$
m, s = self.get_ms(state, group, grad)
# Increment $t$ the number of optimizer steps
state['step'] += 1
if not self.rectify:
# Perform *Adam* update, defined in [`adam.py`](adam.html), with
# $\textcolor{cyan}{s_t} + \textcolor{red}{\epsilon}$ in place of $v_t$.
self.adam_update(state, group, param, m, s + group['eps'])
else:
# Perform *Rectified Adam* update defined in [`radam.py`](radam.html), with
# $\textcolor{cyan}{s_t} + \textcolor{red}{\epsilon}$ in place of $v_t$.
self.r_adam_update(state, group, param, m, s + group['eps'])
labml_nn/optimizers/adam.py
+214
View File
@@ -0,0 +1,214 @@
"""
---
title: Adam Optimizer
summary: A simple PyTorch implementation/tutorial of Adam optimizer
---
# Adam Optimizer
This is a [PyTorch](https://pytorch.org) implementation of popular optimizer *Adam* from paper
[Adam: A Method for Stochastic Optimization](https://arxiv.org/abs/1412.6980).
*Adam* update is,
\begin{align}
m_t &\leftarrow \beta_1 m_{t-1} + (1 - \beta_1) \cdot g_t \\
v_t &\leftarrow \beta_2 v_{t-1} + (1 - \beta_2) \cdot g_t^2 \\
\hat{m}_t &\leftarrow \frac{m_t}{1-\beta_1^t} \\
\hat{v}_t &\leftarrow \frac{v_t}{1-\beta_2^t} \\
\theta_t &\leftarrow \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}
\end{align}
where $\alpha$, $\beta_1$, $\beta_2$ and $\epsilon$ are scalar hyper parameters.
$m_t$ and $v_t$ are first and second order moments.
$\hat{m}_t$ and $\hat{v}_t$ are biased corrected moments.
$\epsilon$ is used as a fix for division by zero error, but also acts as a form of a hyper-parameter
that acts against variance in gradients.
Effective step taken assuming $\epsilon = 0$ is,
$$\Delta t = \alpha \cdot \frac{\hat{m}_t}{\hat{v}_t}$$
This is bounded by,
$$\vert \Delta t \vert \le \alpha \cdot \frac{1 - \beta_1}{\sqrt{1-\beta_2}}$$
when $1-\beta_1 \gt \sqrt{1-\beta_2}$
and
$$\vert \Delta t\vert \le \alpha$$
otherwise.
And in most common scenarios,
$$\vert \Delta t \vert \approx \alpha$$
"""
import math
from typing import Dict, Any, Tuple, Optional
import torch
from labml import tracker
from torch import nn
from labml_nn.optimizers import GenericAdaptiveOptimizer, WeightDecay
class Adam(GenericAdaptiveOptimizer):
"""
## Adam Optimizer
We extend the class `GenericAdaptiveOptimizer` defined in [`__init__.py`](index.html)
to implement the Adam optimizer.
"""
def __init__(self, params,
lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999), eps: float = 1e-16,
weight_decay: WeightDecay = WeightDecay(),
optimized_update: bool = True,
defaults: Optional[Dict[str, Any]] = None):
"""
### Initialize the optimizer
* `params` is the list of parameters
* `lr` is the learning rate $\alpha$
* `betas` is a tuple of ($\beta_1$, $\beta_2$)
* `eps` is $\hat{\epsilon}$ or $\epsilon$ based on `optimized_update`
* `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html)
* `optimized_update` is a flag whether to optimize the bias correction of the second moment
by doing it after adding $\epsilon$
* `defaults` is a dictionary of default for group values.
This is useful when you want to extend the class `Adam`.
"""
defaults = {} if defaults is None else defaults
defaults.update(weight_decay.defaults())
super().__init__(params, defaults, lr, betas, eps)
self.weight_decay = weight_decay
self.optimized_update = optimized_update
def init_state(self, state: Dict[str, any], group: Dict[str, any], param: nn.Parameter):
"""
### Initialize a parameter state
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `param` is the parameter tensor $\theta_{t-1}$
"""
# This is the number of optimizer steps taken on the parameter, $t$
state['step'] = 0
# Exponential moving average of gradients, $m_t$
state['exp_avg'] = torch.zeros_like(param, memory_format=torch.preserve_format)
# Exponential moving average of squared gradient values, $v_t$
state['exp_avg_sq'] = torch.zeros_like(param, memory_format=torch.preserve_format)
def get_mv(self, state: Dict[str, Any], group: Dict[str, Any], grad: torch.Tensor):
"""
### Calculate $m_t$ and and $v_t$
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$
"""
# Get $\beta_1$ and $\beta_2$
beta1, beta2 = group['betas']
# Get $m_{t-1}$ and $v_{t-1}$
m, v = state['exp_avg'], state['exp_avg_sq']
# In-place calculation of $m_t$
# $$m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) \cdot g_t$$
m.mul_(beta1).add_(grad, alpha=1 - beta1)
# In-place calculation of $v_t$
# $$v_t \leftarrow \beta_2 v_{t-1} + (1 - \beta_2) \cdot g_t^2$$
v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
return m, v
def get_lr(self, state: Dict[str, any], group: Dict[str, any]):
"""
### Get learning-rate
This returns the modified learning rate based on the state.
For *Adam* this is just the specified learning rate for the parameter group,
$\alpha$.
"""
return group['lr']
def adam_update(self, state: Dict[str, any], group: Dict[str, any], param: torch.nn.Parameter,
m: torch.Tensor, v: torch.Tensor):
"""
### Do the *Adam* parameter update
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `param` is the parameter tensor $\theta_{t-1}$
* `m` and `v` are the uncorrected first and second moments $m_t$ and $v_t$.
This computes the following
\begin{align}
\theta_t &\leftarrow \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}
\end{align}
Since $\alpha$, $\beta_1$, $\beta_2$ and $\epsilon$ are scalars and others are tensors
we modify this calculation to optimize the computation.
\begin{align}
\theta_t &\leftarrow \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \\
\theta_t &\leftarrow \theta_{t-1} - \alpha \cdot
\frac{m_t / (1-\beta_1^t)}{\sqrt{v_t/(1-\beta_2^t)} + \epsilon} \\
\theta_t &\leftarrow \theta_{t-1} - \alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t} \cdot
\frac{m_t}{\sqrt{v_t} + \hat{\epsilon}} \\
\end{align}
where
$$\hat{\epsilon} = (1-\beta_2^t) \epsilon$$
is what we should specify as the hyper-parameter.
"""
# Get $\beta_1$ and $\beta_2$
beta1, beta2 = group['betas']
# Bias correction term for $\hat{m}_t$, $1 - \beta_1^t$
bias_correction1 = 1 - beta1 ** state['step']
# Bias correction term for $\hat{v}_t$, $1 - \beta_2^t$
bias_correction2 = 1 - beta2 ** state['step']
# Get learning rate
lr = self.get_lr(state, group)
# Whether to optimize the computation
if self.optimized_update:
# $\sqrt{v_t} + \hat{\epsilon}$
denominator = v.sqrt().add_(group['eps'])
# $\alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t}$
step_size = lr * math.sqrt(bias_correction2) / bias_correction1
# $\theta_t \leftarrow \theta_{t-1} - \alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t} \cdot
# \frac{m_t}{\sqrt{v_t} + \hat{\epsilon}}$
param.data.addcdiv_(m, denominator, value=-step_size)
# Computation without optimization
else:
# $\frac{\sqrt{v_t}}{\sqrt{1-\beta_2^t}} + \epsilon$
denominator = (v.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
# $\frac{\alpha}{1-\beta_1^t}$
step_size = lr / bias_correction1
# $\theta_t \leftarrow \theta_{t-1} - \alpha \cdot
# \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$
param.data.addcdiv_(m, denominator, value=-step_size)
def step_param(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor, param: torch.nn.Parameter):
"""
### Take an update step for a given parameter tensor
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$
* `param` is the parameter tensor $\theta_{t-1}$
"""
# Calculate weight decay
grad = self.weight_decay(param, grad, group)
# Get $m_t$ and $v_t$
m, v = self.get_mv(state, group, grad)
# Increment $t$ the number of optimizer steps
state['step'] += 1
# Perform *Adam* update
self.adam_update(state, group, param, m, v)
labml_nn/optimizers/adam_fp16.py
+136
View File
@@ -0,0 +1,136 @@
"""
---
title: Adam Optimizer for Half Precision Training
summary: A simple PyTorch implementation/tutorial of Adam optimizer
---
# Adam Optimizer for Half Precision Training
"""
from typing import Dict, Tuple, Optional, Any
import torch
from torch import nn
from torch.optim import Optimizer
from torch.cuda.amp import grad_scaler
from collections import defaultdict, abc
from labml_nn.optimizers import WeightDecay
from labml_nn.optimizers.adam import Adam
class AdamFP16(Adam):
"""
## Adam Optimizer for Half Precision Training
We extend [Adam Optimizer](adam.html) but use FP32 to store gradients and moments.
"""
def __init__(self, params, lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999), eps: float = 1e-16,
weight_decay: WeightDecay = WeightDecay(), optimized_update: bool = True,
defaults: Optional[Dict[str, Any]] = None):
# Parameter to store 32 bit gradients. This get populated by the `GradScaler` defined below.
self.grad_fp32 = {}
# Call the [Adam Optimizer](adam.html) initializer
super().__init__(params, lr, betas, eps, weight_decay, optimized_update, defaults)
def init_state(self, state: Dict[str, any], group: Dict[str, any], param: nn.Parameter):
"""
### Initialize a parameter state
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `param` is the parameter tensor $\theta_{t-1}$
All the state tensors use FP32.
"""
# This is the number of optimizer steps taken on the parameter, $t$
state['step'] = 0
# Exponential moving average of gradients, $m_t$
state['exp_avg'] = torch.zeros_like(param, memory_format=torch.preserve_format, dtype=torch.float)
# Exponential moving average of squared gradient values, $v_t$
state['exp_avg_sq'] = torch.zeros_like(param, memory_format=torch.preserve_format, dtype=torch.float)
# Maintain a FP32 copy of the parameters
state['fp32_copy'] = param.to(torch.float)
def step_param(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor, param: torch.nn.Parameter):
"""
### Take an update step for a given parameter tensor
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$
* `param` is the parameter tensor $\theta_{t-1}$
"""
# Get the FP32 parameters
param_fp32 = state['fp32_copy']
# Get the FP32 gradients if available
grad_fp32 = self.grad_fp32.get(param, None)
if grad_fp32 is not None:
del self.grad_fp32[param]
grad = grad_fp32
else:
# Otherwise, convert the gradients to FP32
grad = grad.to(torch.float)
# Calculate weight decay
grad = self.weight_decay(param_fp32, grad, group)
# Get $m_t$ and $v_t$
m, v = self.get_mv(state, group, grad)
# Increment $t$ the number of optimizer steps
state['step'] += 1
# Perform *Adam* update
self.adam_update(state, group, param_fp32, m, v)
# Set the parameters
param.data = param_fp32.to(param.dtype)
class GradScalerFP16(grad_scaler.GradScaler):
"""
## Gradient Scaler with half precision gradients
We extend PyTorch gradient scaler to use FP32 gradients.
"""
def _unscale_grads_(self, optimizer: Optimizer, inv_scale: torch.Tensor, found_inf: torch.Tensor,
allow_fp16: bool) -> Dict[torch.device, torch.Tensor]:
per_device_inv_scale = grad_scaler._MultiDeviceReplicator(inv_scale)
per_device_found_inf = grad_scaler._MultiDeviceReplicator(found_inf)
per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list)) # type: ignore[var-annotated]
with torch.no_grad():
# Loop through parameters
for group in optimizer.param_groups:
for param in group["params"]:
# Skip non-trainable parameters
if param.grad is None:
continue
# Not implemented for sparse tensors
if param.grad.is_sparse:
raise NotImplementedError
# If we are using the `AdamFP16` optimizer set `optimizer.grad_fp32[param]` to the FP32 gradients
if isinstance(optimizer, AdamFP16):
grad = param.grad.to(torch.float)
optimizer.grad_fp32[param] = grad
# Otherwise, do not convert the gradients to FP32
else:
grad = param.grad
per_device_and_dtype_grads[grad.device][grad.dtype].append(grad)
# Unscale all the gradients
for device, per_dtype_grads in per_device_and_dtype_grads.items():
for grads in per_dtype_grads.values():
torch._amp_foreach_non_finite_check_and_unscale_(grads,
per_device_found_inf.get(device),
per_device_inv_scale.get(device))
#
return per_device_found_inf._per_device_tensors
labml_nn/optimizers/adam_warmup.py
+61
View File
@@ -0,0 +1,61 @@
"""
---
title: Adam optimizer with warm-up
summary: A simple PyTorch implementation/tutorial of Adam optimizer with warm-up.
---
# Adam Optimizer with Warmup
This extends [AMSGrad optimizer](amsgrad.html) and adds a warmup stage.
"""
from typing import Dict
from labml_nn.optimizers import WeightDecay
from labml_nn.optimizers.amsgrad import AMSGrad
class AdamWarmup(AMSGrad):
"""
## Adam Optimizer with Warmup
This class extends from AMSGrad optimizer defined in [`amsgrad.py`](amsgrad.html).
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16,
weight_decay: WeightDecay = WeightDecay(),
optimized_update: bool = True,
amsgrad=False, warmup=0, defaults=None):
"""
### Initialize the optimizer
* `params` is the list of parameters
* `lr` is the learning rate $\alpha$
* `betas` is a tuple of ($\beta_1$, $\beta_2$)
* `eps` is $\hat{\epsilon}$ or $\epsilon$ based on `optimized_update`
* `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html)
* 'optimized_update' is a flag whether to optimize the bias correction of the second moment
by doing it after adding $\epsilon$
* `amsgrad` is a flag indicating whether to use AMSGrad or fallback to plain Adam
* `warmup` number of warmup steps
* `defaults` is a dictionary of default for group values.
This is useful when you want to extend the class `AdamWarmup`.
"""
defaults = {} if defaults is None else defaults
defaults.update(dict(warmup=warmup))
super().__init__(params, lr, betas, eps, weight_decay, optimized_update, amsgrad, defaults)
def get_lr(self, state: Dict[str, any], group: Dict[str, any]):
"""
### Get learning-rate
$$\alpha \min \bigg(1, \frac{t}{w}\bigg)$$
where $w$ is the number of warmup steps.
"""
# If we are in warmup stage
if group['warmup'] > state['step']:
# A linearly increasing learning rate from $0$ to $\alpha$
return 1e-8 + state['step'] * group['lr'] / group['warmup']
else:
# Constant learning rate $\alpha$
return group['lr']
labml_nn/optimizers/adam_warmup_cosine_decay.py
+97
View File
@@ -0,0 +1,97 @@
"""
---
title: Adam optimizer with warm-up and cosine decay
summary: A PyTorch implementation/tutorial of Adam optimizer with warm-up and cosine decay for GPT.
---
# Adam Optimizer with Warmup and Cosine Decay
This extends [AMSGrad optimizer](adam.html) and adds a warmup stage.
"""
import math
from typing import Dict
from labml_nn.optimizers import WeightDecay
from labml_nn.optimizers.amsgrad import AMSGrad
class AdamWarmupCosineDecay(AMSGrad):
"""
<a id="EmbeddingsWithPositionalEncoding"></a>
## Adam Optimizer with Warmup and Cosine Decay
This class extends from AMSGrad optimizer defined in [`amsgrad.py`](amsgrad.html).
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16,
weight_decay: WeightDecay = WeightDecay(),
optimized_update: bool = True,
amsgrad=False, warmup=0, total_steps=1e10, defaults=None):
"""
### Initialize the optimizer
* `params` is the list of parameters
* `lr` is the learning rate $\alpha$
* `betas` is a tuple of ($\beta_1$, $\beta_2$)
* `eps` is $\hat{\epsilon}$ or $\epsilon$ based on `optimized_update`
* `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html)
* 'optimized_update' is a flag whether to optimize the bias correction of the second moment
by doing it after adding $\epsilon$
* `amsgrad` is a flag indicating whether to use AMSGrad or fallback to plain Adam
* `warmup` number of warmup steps
* `total_steps` total number of steps. Cosine decay reaches 0 at this,
but stays at 10% of `lr` because we take $\alpha * \max(0.1, decay)$
* `defaults` is a dictionary of default for group values.
This is useful when you want to extend the class `AdamWarmup`.
"""
defaults = {} if defaults is None else defaults
defaults.update(dict(warmup=warmup, total_steps=total_steps))
super().__init__(params, lr, betas, eps, weight_decay, optimized_update, amsgrad, defaults)
def get_lr(self, state: Dict[str, any], group: Dict[str, any]):
"""
### Get learning-rate
$$\alpha \min \bigg(1, \frac{t}{w}\bigg)$$
where $w$ is the number of warmup steps.
"""
# If we are in warmup stage
if group['warmup'] > state['step']:
# A linearly increasing learning rate from $0$ to $\alpha$
return 1e-8 + state['step'] * group['lr'] / group['warmup']
else:
# Constant learning rate $\alpha$
progress = (state['step'] - group['warmup']) / max(1, group['total_steps'] - group['warmup'])
return group['lr'] * max(0.1, 0.5 * (1.0 + math.cos(math.pi * progress)))
def _test_lr():
"""
### Plot learning rate for different warmups and model sizes
![Plot of learning rate](noam_lr.png)
"""
import matplotlib.pyplot as plt
import numpy as np
from torch import nn
model = nn.Linear(10, 10)
opt = AdamWarmupCosineDecay(model.parameters(), warmup=5000, lr=1e-4, total_steps=4e6)
steps = 20_000
plt.plot(np.arange(1, steps), [opt.get_lr({'step': i}, opt.defaults) for i in range(1, steps)])
plt.legend(["5000:4e6", "5000:2e6", "5000:1e6"])
plt.title("Learning Rate")
plt.show()
steps = int(6e6)
step_size = 1000
plt.plot(np.arange(1, steps, step_size), [opt.get_lr({'step': i}, opt.defaults) for i in range(1, steps, step_size)])
plt.legend(["5000:4e6", "5000:2e6", "5000:1e6"])
plt.title("Learning Rate")
plt.show()
if __name__ == '__main__':
_test_lr()
labml_nn/optimizers/amsgrad.py
+204
View File
@@ -0,0 +1,204 @@
"""
---
title: AMSGrad Optimizer
summary: A simple PyTorch implementation/tutorial of AMSGrad optimizer.
---
# AMSGrad
This is a [PyTorch](https://pytorch.org) implementation of the paper
[On the Convergence of Adam and Beyond](https://arxiv.org/abs/1904.09237).
We implement this as an extension to our [Adam optimizer implementation](adam.html).
The implementation it self is really small since it's very similar to Adam.
We also have an implementation of the synthetic example described in the paper where Adam fails to converge.
"""
from typing import Dict
import torch
from torch import nn
from labml_nn.optimizers import WeightDecay
from labml_nn.optimizers.adam import Adam
class AMSGrad(Adam):
"""
## AMSGrad Optimizer
This class extends from Adam optimizer defined in [`adam.py`](adam.html).
Adam optimizer is extending the class `GenericAdaptiveOptimizer`
defined in [`__init__.py`](index.html).
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16,
weight_decay: WeightDecay = WeightDecay(),
optimized_update: bool = True,
amsgrad=True, defaults=None):
"""
### Initialize the optimizer
* `params` is the list of parameters
* `lr` is the learning rate $\alpha$
* `betas` is a tuple of ($\beta_1$, $\beta_2$)
* `eps` is $\hat{\epsilon}$ or $\epsilon$ based on `optimized_update`
* `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html)
* 'optimized_update' is a flag whether to optimize the bias correction of the second moment
by doing it after adding $\epsilon$
* `amsgrad` is a flag indicating whether to use AMSGrad or fallback to plain Adam
* `defaults` is a dictionary of default for group values.
This is useful when you want to extend the class `Adam`.
"""
defaults = {} if defaults is None else defaults
defaults.update(dict(amsgrad=amsgrad))
super().__init__(params, lr, betas, eps, weight_decay, optimized_update, defaults)
def init_state(self, state: Dict[str, any], group: Dict[str, any], param: nn.Parameter):
"""
### Initialize a parameter state
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `param` is the parameter tensor $\theta_{t-1}$
"""
# Call `init_state` of Adam optimizer which we are extending
super().init_state(state, group, param)
# If `amsgrad` flag is `True` for this parameter group, we maintain the maximum of
# exponential moving average of squared gradient
if group['amsgrad']:
state['max_exp_avg_sq'] = torch.zeros_like(param, memory_format=torch.preserve_format)
def get_mv(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor):
"""
### Calculate $m_t$ and and $v_t$ or $\max(v_1, v_2, ..., v_{t-1}, v_t)$
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$
"""
# Get $m_t$ and $v_t$ from *Adam*
m, v = super().get_mv(state, group, grad)
# If this parameter group is using `amsgrad`
if group['amsgrad']:
# Get $\max(v_1, v_2, ..., v_{t-1})$.
#
# 🗒 The paper uses the notation $\hat{v}_t$ for this, which we don't use
# that here because it confuses with the Adam's usage of the same notation
# for bias corrected exponential moving average.
v_max = state['max_exp_avg_sq']
# Calculate $\max(v_1, v_2, ..., v_{t-1}, v_t)$.
#
# 🤔 I feel you should be taking / maintaining the max of the bias corrected
# second exponential average of squared gradient.
# But this is how it's
# [implemented in PyTorch also](https://github.com/pytorch/pytorch/blob/19f4c5110e8bcad5e7e75375194262fca0a6293a/torch/optim/functional.py#L90).
# I guess it doesn't really matter since bias correction only increases the value
# and it only makes an actual difference during the early few steps of the training.
torch.maximum(v_max, v, out=v_max)
return m, v_max
else:
# Fall back to *Adam* if the parameter group is not using `amsgrad`
return m, v
def _synthetic_experiment(is_adam: bool):
"""
## Synthetic Experiment
This is the synthetic experiment described in the paper,
that shows a scenario where *Adam* fails.
The paper (and Adam) formulates the problem of optimizing as
minimizing the expected value of a function, $\mathbb{E}[f(\theta)]$
with respect to the parameters $\theta$.
In the stochastic training setting we do not get hold of the function $f$
it self; that is,
when you are optimizing a NN $f$ would be the function on entire
batch of data.
What we actually evaluate is a mini-batch so the actual function is
realization of the stochastic $f$.
This is why we are talking about an expected value.
So let the function realizations be $f_1, f_2, ..., f_T$ for each time step
of training.
We measure the performance of the optimizer as the regret,
$$R(T) = \sum_{t=1}^T \big[ f_t(\theta_t) - f_t(\theta^*) \big]$$
where $\theta_t$ is the parameters at time step $t$, and $\theta^*$ is the
optimal parameters that minimize $\mathbb{E}[f(\theta)]$.
Now lets define the synthetic problem,
\begin{align}
f_t(x) =
\begin{cases}
1010 x, & \text{for } t \mod 101 = 1 \\
-10 x, & \text{otherwise}
\end{cases}
\end{align}
where $-1 \le x \le +1$.
The optimal solution is $x = -1$.
This code will try running *Adam* and *AMSGrad* on this problem.
"""
# Define $x$ parameter
x = nn.Parameter(torch.tensor([.0]))
# Optimal, $x^* = -1$
x_star = nn.Parameter(torch.tensor([-1]), requires_grad=False)
def func(t: int, x_: nn.Parameter):
"""
### $f_t(x)$
"""
if t % 101 == 1:
return (1010 * x_).sum()
else:
return (-10 * x_).sum()
# Initialize the relevant optimizer
if is_adam:
optimizer = Adam([x], lr=1e-2, betas=(0.9, 0.99))
else:
optimizer = AMSGrad([x], lr=1e-2, betas=(0.9, 0.99))
# $R(T)$
total_regret = 0
from labml import monit, tracker, experiment
# Create experiment to record results
with experiment.record(name='synthetic', comment='Adam' if is_adam else 'AMSGrad'):
# Run for $10^7$ steps
for step in monit.loop(10_000_000):
# $f_t(\theta_t) - f_t(\theta^*)$
regret = func(step, x) - func(step, x_star)
# $R(T) = \sum_{t=1}^T \big[ f_t(\theta_t) - f_t(\theta^*) \big]$
total_regret += regret.item()
# Track results every 1,000 steps
if (step + 1) % 1000 == 0:
tracker.save(loss=regret, x=x, regret=total_regret / (step + 1))
# Calculate gradients
regret.backward()
# Optimize
optimizer.step()
# Clear gradients
optimizer.zero_grad()
# Make sure $-1 \le x \le +1$
x.data.clamp_(-1., +1.)
if __name__ == '__main__':
# Run the synthetic experiment is *Adam*.
# You can see that Adam converges at $x = +1$
_synthetic_experiment(True)
# Run the synthetic experiment is *AMSGrad*
# You can see that AMSGrad converges to true optimal $x = -1$
_synthetic_experiment(False)
labml_nn/optimizers/configs.py
+156
View File
@@ -0,0 +1,156 @@
"""
---
title: Configurable optimizer module
summary: This implements a configurable module for optimizers.
---
# Configurable Optimizer
"""
from typing import Tuple
import torch
from labml.configs import BaseConfigs, option, meta_config
from labml_nn.optimizers import WeightDecay
class OptimizerConfigs(BaseConfigs):
"""
<a id="OptimizerConfigs"></a>
## Optimizer Configurations
"""
# Optimizer
optimizer: torch.optim.Adam
# Weight decay
weight_decay_obj: WeightDecay
# Whether weight decay is decoupled;
# i.e. weight decay is not added to gradients
weight_decouple: bool = True
# Weight decay
weight_decay: float = 0.0
# Whether weight decay is absolute or should be multiplied by learning rate
weight_decay_absolute: bool = False
# Whether the adam update is optimized (different epsilon)
optimized_adam_update: bool = True
# Parameters to be optimized
parameters: any
# Learning rate $\alpha$
learning_rate: float = 0.01
# Beta values $(\beta_1, \beta_2)$ for Adam
betas: Tuple[float, float] = (0.9, 0.999)
# Epsilon $\epsilon$ for adam
eps: float = 1e-08
# Momentum for SGD
momentum: float = 0.5
# Whether to use AMSGrad
amsgrad: bool = False
# Number of warmup optimizer steps
warmup: int = 2_000
# Total number of optimizer steps (for cosine decay)
total_steps: int = int(1e10)
# Whether to degenerate to SGD in AdaBelief
degenerate_to_sgd: bool = True
# Whether to use Rectified Adam in AdaBelief
rectify: bool = True
# Model embedding size for Noam optimizer
d_model: int
rho: float
def __init__(self):
super().__init__(_primary='optimizer')
meta_config(OptimizerConfigs.parameters)
@option(OptimizerConfigs.weight_decay_obj, 'L2')
def _weight_decay(c: OptimizerConfigs):
return WeightDecay(c.weight_decay, c.weight_decouple, c.weight_decay_absolute)
@option(OptimizerConfigs.optimizer, 'SGD')
def _sgd_optimizer(c: OptimizerConfigs):
return torch.optim.SGD(c.parameters, c.learning_rate, c.momentum,
weight_decay=c.weight_decay)
@option(OptimizerConfigs.optimizer, 'Adam')
def _adam_optimizer(c: OptimizerConfigs):
if c.amsgrad:
from labml_nn.optimizers.amsgrad import AMSGrad
return AMSGrad(c.parameters,
lr=c.learning_rate, betas=c.betas, eps=c.eps,
optimized_update=c.optimized_adam_update,
weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad)
else:
from labml_nn.optimizers.adam import Adam
return Adam(c.parameters,
lr=c.learning_rate, betas=c.betas, eps=c.eps,
optimized_update=c.optimized_adam_update,
weight_decay=c.weight_decay_obj)
@option(OptimizerConfigs.optimizer, 'AdamW')
def _adam_warmup_optimizer(c: OptimizerConfigs):
from labml_nn.optimizers.adam_warmup import AdamWarmup
return AdamWarmup(c.parameters,
lr=c.learning_rate, betas=c.betas, eps=c.eps,
weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad, warmup=c.warmup)
@option(OptimizerConfigs.optimizer, 'RAdam')
def _radam_optimizer(c: OptimizerConfigs):
from labml_nn.optimizers.radam import RAdam
return RAdam(c.parameters,
lr=c.learning_rate, betas=c.betas, eps=c.eps,
weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad,
degenerated_to_sgd=c.degenerate_to_sgd)
@option(OptimizerConfigs.optimizer, 'AdaBelief')
def _ada_belief_optimizer(c: OptimizerConfigs):
from labml_nn.optimizers.ada_belief import AdaBelief
return AdaBelief(c.parameters,
lr=c.learning_rate, betas=c.betas, eps=c.eps,
weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad,
degenerate_to_sgd=c.degenerate_to_sgd,
rectify=c.rectify)
@option(OptimizerConfigs.optimizer, 'Noam')
def _noam_optimizer(c: OptimizerConfigs):
from labml_nn.optimizers.noam import Noam
return Noam(c.parameters,
lr=c.learning_rate, betas=c.betas, eps=c.eps,
weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad, warmup=c.warmup,
d_model=c.d_model)
@option(OptimizerConfigs.optimizer, 'Sophia')
def _sophia_optimizer(c: OptimizerConfigs):
from labml_nn.optimizers.sophia import Sophia
return Sophia(c.parameters,
lr=c.learning_rate, betas=c.betas, eps=c.eps,
weight_decay=c.weight_decay_obj, rho=c.rho)
@option(OptimizerConfigs.optimizer, 'AdamWarmupCosineDecay')
def _noam_optimizer(c: OptimizerConfigs):
from labml_nn.optimizers.adam_warmup_cosine_decay import AdamWarmupCosineDecay
return AdamWarmupCosineDecay(c.parameters,
lr=c.learning_rate, betas=c.betas, eps=c.eps,
weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad,
warmup=c.warmup, total_steps=c.total_steps)
labml_nn/optimizers/mnist_experiment.py
+132
View File
@@ -0,0 +1,132 @@
"""
---
title: MNIST example to test the optimizers
summary: This is a simple MNIST example with a CNN model to test the optimizers.
---
# MNIST example to test the optimizers
"""
import torch.nn as nn
import torch.utils.data
from labml import experiment, tracker
from labml.configs import option
from labml_nn.helpers.datasets import MNISTConfigs
from labml_nn.helpers.device import DeviceConfigs
from labml_nn.helpers.metrics import Accuracy
from labml_nn.helpers.trainer import TrainValidConfigs, BatchIndex
from labml_nn.optimizers.configs import OptimizerConfigs
class Model(nn.Module):
"""
## The model
"""
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 20, 5, 1)
self.pool1 = nn.MaxPool2d(2)
self.conv2 = nn.Conv2d(20, 50, 5, 1)
self.pool2 = nn.MaxPool2d(2)
self.fc1 = nn.Linear(16 * 50, 500)
self.fc2 = nn.Linear(500, 10)
self.activation = nn.ReLU()
def forward(self, x):
x = self.activation(self.conv1(x))
x = self.pool1(x)
x = self.activation(self.conv2(x))
x = self.pool2(x)
x = self.activation(self.fc1(x.view(-1, 16 * 50)))
return self.fc2(x)
class Configs(MNISTConfigs, TrainValidConfigs):
"""
## Configurable Experiment Definition
"""
optimizer: torch.optim.Adam
model: nn.Module
device: torch.device = DeviceConfigs()
epochs: int = 10
is_save_models = True
model: nn.Module
inner_iterations = 10
accuracy_func = Accuracy()
loss_func = nn.CrossEntropyLoss()
def init(self):
tracker.set_queue("loss.*", 20, True)
tracker.set_scalar("accuracy.*", True)
self.state_modules = [self.accuracy_func]
def step(self, batch: any, batch_idx: BatchIndex):
# Get the batch
data, target = batch[0].to(self.device), batch[1].to(self.device)
# Add global step if we are in training mode
if self.mode.is_train:
tracker.add_global_step(len(data))
# Run the model
output = self.model(data)
# Calculate the loss
loss = self.loss_func(output, target)
# Calculate the accuracy
self.accuracy_func(output, target)
# Log the loss
tracker.add("loss.", loss)
# Optimize if we are in training mode
if self.mode.is_train:
# Calculate the gradients
loss.backward()
# Take optimizer step
self.optimizer.step()
# Log the parameter and gradient L2 norms once per epoch
if batch_idx.is_last:
tracker.add('model', self.model)
tracker.add('optimizer', (self.optimizer, {'model': self.model}))
# Clear the gradients
self.optimizer.zero_grad()
# Save logs
tracker.save()
@option(Configs.model)
def model(c: Configs):
return Model().to(c.device)
@option(Configs.optimizer)
def _optimizer(c: Configs):
"""
Create a configurable optimizer.
We can change the optimizer type and hyper-parameters using configurations.
"""
opt_conf = OptimizerConfigs()
opt_conf.parameters = c.model.parameters()
return opt_conf
def main():
conf = Configs()
conf.inner_iterations = 10
experiment.create(name='mnist_ada_belief')
experiment.configs(conf, {'inner_iterations': 10,
# Specify the optimizer
'optimizer.optimizer': 'Adam',
'optimizer.learning_rate': 1.5e-4})
experiment.add_pytorch_models(dict(model=conf.model))
with experiment.start():
conf.run()
if __name__ == '__main__':
main()
labml_nn/optimizers/noam.py
+88
View File
@@ -0,0 +1,88 @@
"""
---
title: Noam optimizer from Attention is All You Need paper
summary: >
This is a tutorial/implementation of Noam optimizer.
Noam optimizer has a warm-up period and then an exponentially decaying learning rate.
---
# Noam Optimizer
This is the [PyTorch](https://pytorch.org) implementation of optimizer introduced in the paper
[Attention Is All You Need](https://arxiv.org/abs/1706.03762).
"""
from typing import Dict
from labml_nn.optimizers import WeightDecay
from labml_nn.optimizers.amsgrad import AMSGrad
class Noam(AMSGrad):
"""
## Noam Optimizer
This class extends from Adam optimizer defined in [`adam.py`](adam.html).
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16,
weight_decay: WeightDecay = WeightDecay(),
optimized_update: bool = True,
amsgrad=False,
warmup=0, d_model=512, defaults=None):
"""
### Initialize the optimizer
* `params` is the list of parameters
* `lr` is the learning rate $\alpha$
* `betas` is a tuple of ($\beta_1$, $\beta_2$)
* `eps` is $\hat{\epsilon}$ or $\epsilon$ based on `optimized_update`
* `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html)
* 'optimized_update' is a flag whether to optimize the bias correction of the second moment
by doing it after adding $\epsilon$
* `amsgrad` is a flag indicating whether to use AMSGrad or fallback to plain Adam
* `warmup` number of warmup steps
* `d_model` model size; i.e. number of dimensions in the transformer
* `defaults` is a dictionary of default for group values.
This is useful when you want to extend the class `AdamWarmup`.
"""
defaults = {} if defaults is None else defaults
defaults.update(dict(warmup=warmup))
super().__init__(params, lr, betas, eps, weight_decay, optimized_update, amsgrad, defaults)
self.d_model = d_model
def get_lr(self, state: Dict[str, any], group: Dict[str, any]):
"""
### Get learning-rate
$$\alpha \frac{1}{\sqrt{d_{model}}} \min \bigg(\frac{1}{\sqrt{t}}, \frac{t}{w^{3/2}}\bigg)$$
where $w$ is the number of warmup steps.
"""
# $$\min \bigg(\frac{1}{\sqrt{t}}, \frac{t}{w^{3/2}}\bigg)$$
factor = min(state['step'] ** (-0.5), state['step'] * group['warmup'] ** (-1.5))
# $$\alpha \frac{1}{\sqrt{d_{model}}} \min \bigg(\frac{1}{\sqrt{t}}, \frac{t}{w^{3/2}}\bigg)$$
return group['lr'] * self.d_model ** (-0.5) * factor
def _test_noam_lr():
"""
### Plot learning rate for different warmups and model sizes
![Plot of learning rate](noam_lr.png)
"""
import matplotlib.pyplot as plt
import numpy as np
from torch import nn
model = nn.Linear(10, 10)
opts = [Noam(model.parameters(), d_model=512, warmup=4000, lr=1),
Noam(model.parameters(), d_model=512, warmup=8000, lr=1),
Noam(model.parameters(), d_model=2048, warmup=2000, lr=1)]
plt.plot(np.arange(1, 20000), [[opt.get_lr({'step': i}, opt.defaults) for opt in opts] for i in range(1, 20000)])
plt.legend(["512:4000", "512:8000", "2048:2000"])
plt.title("Learning Rate")
plt.show()
if __name__ == '__main__':
_test_noam_lr()
labml_nn/optimizers/performance_test.py
+55
View File
@@ -0,0 +1,55 @@
"""
---
title: Test performance of Adam implementations
summary: This experiment compares performance of Adam implementations.
---
# Performance testing Adam
```text
TorchAdam warmup...[DONE] 222.59ms
TorchAdam...[DONE] 1,356.01ms
MyAdam warmup...[DONE] 119.15ms
MyAdam...[DONE] 1,192.89ms
```
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1ngowaAsADj8VdZfBifu_6L6rtjGoEeoR?usp=sharing)
"""
import torch
import torch.nn as nn
from labml_nn.helpers.device import DeviceInfo
from torch.optim import Adam as TorchAdam
from labml import monit
from labml_nn.optimizers.adam import Adam as MyAdam
from labml_nn.optimizers.mnist_experiment import Model
def test():
device_info = DeviceInfo(use_cuda=True, cuda_device=0)
print(device_info)
inp = torch.randn((64, 1, 28, 28), device=device_info.device)
target = torch.ones(64, dtype=torch.long, device=device_info.device)
loss_func = nn.CrossEntropyLoss()
model = Model().to(device_info.device)
my_adam = MyAdam(model.parameters())
torch_adam = TorchAdam(model.parameters())
loss = loss_func(model(inp), target)
loss.backward()
with monit.section('MyAdam warmup'):
for i in range(100):
my_adam.step()
with monit.section('MyAdam'):
for i in range(1000):
my_adam.step()
with monit.section('TorchAdam warmup'):
for i in range(100):
torch_adam.step()
with monit.section('TorchAdam'):
for i in range(1000):
torch_adam.step()
if __name__ == '__main__':
test()
labml_nn/optimizers/radam.py
+293
View File
@@ -0,0 +1,293 @@
"""
---
title: Rectified Adam (RAdam) optimizer
summary: A simple PyTorch implementation/tutorial of RAdam optimizer.
---
# Rectified Adam (RAdam) optimizer
This implementation is based on
[the official implementation](https://github.com/LiyuanLucasLiu/RAdam)
of the paper
[On the Variance of the Adaptive Learning Rate and Beyond](https://arxiv.org/abs/1908.03265).
We have implemented it in [PyTorch](https://pytorch.org)
as an extension to [our AMSGrad implementation](amsgrad.html)
thus requiring only the modifications to be implemented.
Adam optimizer sometimes converges to a bad local optima during the initial stages of the training;
especially when training transformers.
Researches use warmups to counter this; for the the initial training steps (warm-up stage)
they use a low learning rate.
This paper identifies the problem to be the high variance of adaptive learning rate
during initial stages of training, and counters it using a new rectification term to
reduce variance.
The paper also evaluates two variance reduction mechanisms:
* **Adam-2k**: Only compute the adaptive learning rate ($v_t$ in [Adam](adam.html)) during the first 2k steps,
without changing parameters or calculating momentum ($m_t$).
* **Adam-eps**: Adam with large $\epsilon \approx 10^{-4}$.
## Rectified Adam
Let $\sigma(g_1, ..., g_t)$ and $\psi(g_1, ..., g_t)$ be the functions to calculate
momentum and adaptive learning rate.
For Adam, they are
\begin{align}
\sigma(g_1, ..., g_t) &= \frac{(1 - \beta_1)\sum_{i=1}^t \beta_1^{t-i} g_i}{1 - \beta_1^t} \\
\psi(g_1, ..., g_t) &= \sqrt \frac{1 - \beta_2^t}{(1 - \beta_2)\sum_{i=1}^t \beta_2^{t-i} g_i^2}
\end{align}
### Exponential moving average as simple moving average
The distribution of exponential moving average can be approximated as a simple moving average.
\begin{align}
p\Bigg(\frac{(1-\beta_2) \sum_{i=1}^t \beta_2^{t-i} g_i^2}{1 - \beta_2^t} \Bigg) \approx
p\Bigg(\frac{\sum_{i=1}^{f(t,\beta_2)} g_{t+1-i}^2}{f(t,\beta_2)} \Bigg)
\end{align}
Here we are taking the simple moving average of the last $f(t,\beta_2)$ gradients.
$f(t,\beta_2)$ satisfies the following,
\begin{align}
\frac{(1-\beta_2) \sum_{i=1}^t \beta_2^{t-i} \cdot i}{1 - \beta_2^t} =
\frac{\sum_{i=1}^{f(t,\beta_2)} (t+1-i)}{f(t,\beta_2)}
\end{align}
which gives,
$$f(t,\beta_2) = \frac{2}{1-\beta_2} - 1 - \frac{2 t \beta_2^t}{1 - \beta_2^t}$$
### Scaled inverse chi-squared
From above we have
$$
p\Big( \psi^2(g_1, ..., g_t) \Big) \approx
p\Bigg(\frac{\sum_{i=1}^{f(t,\beta_2)} g_{t+1-i}^2}{f(t,\beta_2)} \Bigg)
$$
where $g_i \sim \mathcal{N}(0, \sigma^2)$.
Note that $sigma$ here is the standard deviation and different from $\sigma(.)$ for momentum.
[Scaled inverse chi-squared](https://en.wikipedia.org/wiki/Scaled_inverse_chi-squared_distribution)
is the distribution of squared inverse of mean of $p$ normal distributions.
$$
p\Bigg(\frac{\sum_{i=1}^{f(t,\beta_2)} g_{t+1-i}^2}{f(t,\beta_2)} \Bigg)
\sim
\text{Scale-inv} \mathcal{X}^2(\rho,\frac{1}{\sigma^2})
$$
where $\rho = f(t,\beta_2)$.
### Rectification
They prove that variance of $\psi(.)$ decreases with $\rho$ when
$\psi^2(.) \sim \text{Scale-inv} \mathcal{X}^2(\rho,\frac{1}{\sigma^2})$.
Therefore the variance is minimized at maximal $\rho$ which is
$\rho_{\infty} = \frac{2}{1-\beta_2} - 1$. Let the minimum variance be $C_{\text{var}}$
In order to ensure that the adaptive learning
rate $\psi(.)$ has consistent variance, we rectify the variance with $r$
\begin{align}
r = \sqrt{\frac{C_{\text{var}}}{Var\big[\psi(.)\big]}}
\end{align}
### Approximating $Var[\psi(.)]$
They estimate $Var[\psi(.)] \approx \frac{Var[\psi^2(.)]}{4 \mathbb{E}[\psi^2(.)}$
based on first order expansion of $\sqrt{\psi^2(.)}$
🤪 I didn't get how it was derived.
From $\text{Scale-inv} \mathcal{X}^2$ distribution we have,
\begin{align}
\mathbb{E}\big[\psi^2(.)\big] &= \frac{\rho / \sigma^2}{\rho-2} \\
Var\big[\psi^2(.)\big] &= \frac{2 \rho / \sigma^4}{(\rho-2)^2 (\rho - 2)}
\end{align}
which gives,
$$
Var[\psi(.)] \approx \frac{\rho}{2(\rho-2)(\rho-4)\sigma^2}
$$
### Rectification term
We have
\begin{align}
r &= \sqrt{\frac{C_{\text{var}}}{Var\big[\psi(.)\big]}} \\
Var[\psi(.)] &\approx \frac{\rho}{2(\rho-2)(\rho-4)\sigma^2}
\end{align}
where $C_{\text{var}}$ is $Var\big[\psi(.)\big]$ for $\rho_\infty$.
Lt $\rho$ and step $t$ be $\rho_t$, and $r_t$ be the rectification term
at step $t$.
\begin{align}
C_{\text{var}} &\approx \frac{\rho_\infty}{2(\rho_\infty-2)(\rho_\infty-4)\sigma^2} \\
Var[\psi(g_1,...,g_t)] &\approx \frac{\rho_t}{2(\rho_t-2)(\rho_t-4)\sigma^2}
\end{align}
This gives,
\begin{align}
r_t &= \sqrt{\frac{(\rho_t-2)(\rho_t-4)\rho_\infty}{(\rho_\infty-2)(\rho_\infty-4)\rho_t}}
\end{align}
"""
import math
from typing import Dict, Optional
import torch
from labml_nn.optimizers import WeightDecay
from labml_nn.optimizers.amsgrad import AMSGrad
class RAdam(AMSGrad):
"""
## Rectified Adam Optimizer
This class extends from AMSAdam optimizer defined in [`amsadam.py`](amsadam.html).
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
weight_decay: WeightDecay = WeightDecay(),
optimized_update: bool = True,
amsgrad=False,
degenerated_to_sgd=True, defaults=None):
"""
### Initialize the optimizer
* `params` is the list of parameters
* `lr` is the learning rate $\alpha$
* `betas` is a tuple of ($\beta_1$, $\beta_2$)
* `eps` is $\hat{\epsilon}$ or $\epsilon$ based on `optimized_update`
* `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html)
* `optimized_update` is a flag whether to optimize the bias correction of the second moment
by doing it after adding $\epsilon$
* `amsgrad` is a flag indicating whether to use AMSGrad or fallback to plain Adam
* `degenerate_to_sgd` whether to use sgd when the rectification term $r_t$ is intractable.
* `defaults` is a dictionary of default for group values.
This is useful when you want to extend the class `RAdam`.
"""
self.degenerated_to_sgd = degenerated_to_sgd
super().__init__(params, lr, betas, eps, weight_decay, optimized_update, amsgrad, defaults)
def step_param(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor, param: torch.nn.Parameter):
"""
### Take an update step for a given parameter tensor
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$
* `param` is the parameter tensor $\theta_{t-1}$
"""
# Calculate weight decay
grad = self.weight_decay(param, grad, group)
# Get $m_t$ and $v_t$; i.e. $\sigma(.)$ and $\psi(.)$ without bias correction
m, v = self.get_mv(state, group, grad)
# Calculate $t$ the number of optimizer steps
state['step'] += 1
# Perform *RAdam* update
self.r_adam_update(state, group, param, m, v)
@staticmethod
def calc_rectification_term(beta2: float, step: int) -> Optional[float]:
"""
### Calculate rectification term $r_t$
"""
# $\beta_2^t$
beta2_t = beta2 ** step
# $$\rho_\infty = \frac{2}{1 - \beta_2} - 1$$
rho_inf = 2 / (1 - beta2) - 1
# $$\rho_t = \frac{2}{1-\beta_2} - 1 - \frac{2 t \beta_2^t}{1-\beta_2^t}$$
rho = rho_inf - 2 * step * beta2_t / (1 - beta2_t)
# $r_t$ is tractable when $\rho_t >= 4$.
# We are being a little more conservative since it's an approximated value
if rho >= 5:
# $$r_t = \sqrt{\frac{(\rho_t-2)(\rho_t-4)\rho_\infty}{(\rho_\infty-2)(\rho_\infty-4)\rho_t}}$$
r2 = (rho - 4) / (rho_inf - 4) * (rho - 2) / rho * rho_inf / (rho_inf - 2)
return math.sqrt(r2)
else:
return None
def r_adam_update(self, state: Dict[str, any], group: Dict[str, any], param: torch.nn.Parameter,
m: torch.Tensor, v: torch.Tensor):
"""
### Do the *RAdam* parameter update
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `param` is the parameter tensor $\theta_{t-1}$
* `m` and `v` are the uncorrected first and second moments $m_t$ and $v_t$;
i.e. $\sigma(.)$ and $\psi(.)$ without bias correction
"""
# Get $\beta_1$ and $\beta_2$
beta1, beta2 = group['betas']
# Bias correction term for $\hat{m}_t$, $1 - \beta_1^t$
bias_correction1 = 1 - beta1 ** state['step']
# Bias correction term for $\hat{v}_t$, $1 - \beta_2^t$
bias_correction2 = 1 - beta2 ** state['step']
r = self.calc_rectification_term(beta2, state['step'])
# Get learning rate
lr = self.get_lr(state, group)
# If $r_t$ is intractable
if r is not None:
# Whether to optimize the computation by combining scalar computations
if self.optimized_update:
# Denominator $\sqrt{v_t} + \hat{\epsilon}$
denominator = v.sqrt().add_(group['eps'])
# Step size $\alpha \sqrt{r_t} * \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t}$
step_size = lr * math.sqrt(bias_correction2) * r / bias_correction1
# Update parameters $\theta_t \leftarrow \theta_{t-1} - \alpha \sqrt{r_t} \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t} \cdot
# \frac{m_t}{\sqrt{v_t} + \hat{\epsilon}}$
param.data.addcdiv_(m, denominator, value=-step_size)
# Computation without optimization
else:
# Denominator $\frac{\sqrt{v_t}}{\sqrt{1-\beta_2^t}} + \epsilon$
denominator = (v.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
# Step size $\frac{\alpha \sqrt{r_t}}{1-\beta_1^t}$
step_size = lr * r / bias_correction1
# Update parameters $\theta_t \leftarrow \theta_{t-1} - \alpha \sqrt{r_t} \cdot
# \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$
param.data.addcdiv_(m, denominator, value=-step_size)
# If $r_t$ is intractable do a SGD with momentum
elif self.degenerated_to_sgd:
# Step size $\frac{\alpha}{1-\beta_1^t}$
step_size = lr / bias_correction1
# Update parameters
# $\theta_t \leftarrow \theta_{t-1} - \alpha \cdot \hat{m}_t$
param.data.add_(m, alpha=-step_size)
def _test_rectification_term():
"""
### Plot $r_t$ against $t$ for various $\beta_2$
![Plot of r_t](radam_r_t.png)
"""
import matplotlib.pyplot as plt
import numpy as np
beta2 = [0.9999, 0.999, 0.99, 0.9, 0.8, 0.6, 0.5]
plt.plot(np.arange(1, 5_000), [[RAdam.calc_rectification_term(b, i) for b in beta2] for i in range(1, 5_000)])
plt.legend(beta2)
plt.title("Optimizer")
plt.show()
if __name__ == '__main__':
_test_rectification_term()
labml_nn/optimizers/readme.md
+10
View File
@@ -0,0 +1,10 @@
# [Optimizers](https://nn.labml.ai/optimizers/index.html)
## Optimizer Implementations
* [Adam Optimizer](https://nn.labml.ai/optimizers/adam.html)
* [AMSGrad Optimizer](https://nn.labml.ai/optimizers/amsgrad.html)
* [Adam Optimizer with warmup](https://nn.labml.ai/optimizers/adam_warmup.html)
* [Noam Optimizer](https://nn.labml.ai/optimizers/noam.html)
* [Rectified Adam Optimizer](https://nn.labml.ai/optimizers/radam.html)
* [AdaBelief Optimizer](https://nn.labml.ai/optimizers/ada_belief.html)
* [Sophia-G Optimizer](https://nn.labml.ai/optimizers/sophia.html)
labml_nn/optimizers/sophia.py
+191
View File
@@ -0,0 +1,191 @@
"""
---
title: Sophia Optimizer
summary: A simple PyTorch implementation/tutorial of Sophia optimizer
---
# Sophia Optimizer
This is a [PyTorch](https://pytorch.org) implementation of *Sophia-G* from paper
[Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://arxiv.org/abs/2305.14342).
Official implementation is available at [Liuhong99/Sophia](https://github.com/Liuhong99/Sophia).
Sophia is more adaptive to heterogeneous curvatures than Adam, more resistant
to non-convexity and rapid change of Hessian than Newton’s method, and also uses a low-cost
pre-conditioner.
Sophia keeps diagonal Hessian estimates with EMA across iterations.
The diagonal Hessian $\hat{h}_t$ is calculated every $k$ steps.
\begin{align}
h_t = \beta_2 h_{t-k} + (1 - \beta_2) \hat{h}_t \ \ \ \ \text{ if } t \text{ mod } k = 1; \text{ else } h_t = h_{t-1}
\end{align}
Sophia uses EMA of gradients $m_t$, only considers positive entries of
the diagonal Hessian and does per-coordinate clipping to the update.
\begin{align}
m_t &\leftarrow \beta_1 m_{t-1} + (1 - \beta_1)g_t \\
\theta_{t + 1} &\leftarrow \theta_t - \eta \cdot \operatorname{clip} \bigg(\frac{m_t}{ \max \{h_t, \epsilon \} }, \rho \bigg)
\end{align}
where $\epsilon$ is a very small value to prevent division by $0$.
### Gauss-Newton-Bartlett (GNB) estimator
\begin{align}
\hat{L}(\theta) &= \frac{1}{B} \sum^{B}_{b=1} \ell_{CE} \big( f(\theta, x_b), \hat{y}_b \big) \\
\hat{h}_t &= B \cdot \nabla_\theta \hat{L} (\theta) \odot \nabla_\theta \hat{L} (\theta)
\end{align}
where $x_b$ are the inputs,
$B$ is the batch size (number of inputs/tokens),
$\ell_{CE}$ is cross entropy loss, and
$\hat{y}_b$ are sampled from the logits $f(\theta, x_b)$.
Note that this hessian estimate is always positive and therefore we
can replace $\max \{h_t, \epsilon \}$ with $h_t + \epsilon$.
Sophia with Gauss-Newton-Bartlett (GNB) estimator is **Sophia-G**
Here is an [experiment](../transformers/basic/with_sophia.html) that uses Sophia-G to train a transformer.
"""
from typing import Dict, Any, Tuple, Optional
import torch
from torch import nn
from labml_nn.optimizers import GenericAdaptiveOptimizer, WeightDecay
class Sophia(GenericAdaptiveOptimizer):
"""
## Sophia-G Optimizer
We extend the class `GenericAdaptiveOptimizer` defined in [`__init__.py`](index.html)
to implement the Sophia optimizer.
"""
def __init__(self, params,
lr: float = 1e-4, betas: Tuple[float, float] = (0.9, 0.95), eps: float = 1e-12,
rho: float = 0.03,
weight_decay: WeightDecay = WeightDecay(),
defaults: Optional[Dict[str, Any]] = None):
"""
### Initialize the optimizer
* `params` is the list of parameters
* `lr` is the maximum learning rate $\eta \rho$
* `betas` is a tuple of ($\beta_1$, $\beta_2$)
* `eps` is $\epsilon$
* `pho` is $\rho$
* `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html)
* `defaults` is a dictionary of default for group values.
This is useful when you want to extend the class `Adam`.
"""
defaults = {} if defaults is None else defaults
defaults.update(weight_decay.defaults())
defaults.update(dict(rho=rho))
super().__init__(params, defaults, lr, betas, eps)
self.weight_decay = weight_decay
def init_state(self, state: Dict[str, any], group: Dict[str, any], param: nn.Parameter):
"""
### Initialize a parameter state
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `param` is the parameter tensor $\theta_{t-1}$
"""
# This is the number of optimizer steps taken on the parameter, $t$
state['step'] = 0
# Exponential moving average of gradients, $m_t$
state['exp_avg'] = torch.zeros_like(param, memory_format=torch.preserve_format)
# Exponential moving average of Hessian diagonal, $h_t$
state['hessian'] = torch.zeros_like(param, memory_format=torch.preserve_format)
def update_hessian(self, n_tokens_training_batch):
"""
### Update the EMA of Hessian diagonal $h_t$
* `n_tokens_training_batch` is the number of tokens/inputs in the batch $B$
\begin{align}
\hat{h}_t &= B \cdot \nabla_\theta \hat{L} (\theta) \odot \nabla_\theta \hat{L} (\theta) \\
h_t &= \beta_2 h_{t-k} + (1 - \beta_2) \hat{h}_t
\end{align}
"""
# Iterate through parameter groups
for group in self.param_groups:
# $\beta_2$
_, beta2 = group['betas']
# Iterate through parameters
for p in group['params']:
# Skip parameters without gradients
if p.grad is None:
continue
# Get optimizer state
state = self.state[p]
# Initialize state if empty
if len(state) == 0:
self.init_state(state, group, p)
# Update EMA Hessian diagonal
#
# \begin{align}
# \hat{h}_t &= B \cdot \nabla_\theta \hat{L} (\theta) \odot \nabla_\theta \hat{L} (\theta) \\
# h_t &= \beta_2 h_{t-k} + (1 - \beta_2) \hat{h}_t
# \end{align}
state['hessian'].mul_(beta2).addcmul_(p.grad, p.grad, value=(1 - beta2) * n_tokens_training_batch)
def step_param(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor, param: torch.nn.Parameter):
"""
### Take an update step for a given parameter tensor
* `state` is the optimizer state of the parameter (tensor)
* `group` stores optimizer attributes of the parameter group
* `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$
* `param` is the parameter tensor $\theta_{t-1}$
We do the following parameter update,
\begin{align}
\theta_{t + 1} &\leftarrow \theta_t - \eta \cdot \operatorname{clip} \bigg(\frac{m_t}{h_t + \epsilon}, \rho \bigg)
\end{align}
"""
# Calculate weight decay
grad = self.weight_decay(param, grad, group)
# Get $\beta_1$ and $\beta_2$
beta1, beta2 = group['betas']
# Get $\rho$
rho = group['rho']
# Get $m_{t-1}$ and $h_{t}$
m, hessian = state['exp_avg'], state['hessian']
# In-place calculation of $m_t$
# $$m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) \cdot g_t$$
m.mul_(beta1).add_(grad, alpha=1 - beta1)
# Increment $t$ the number of optimizer steps
state['step'] += 1
# Get maximum learning rate $\eta \rho$
lr = group['lr']
# $\eta$
eta = lr / rho
# $$\operatorname{clip} \bigg(\frac{m_t}{h_t + \epsilon}, \rho \bigg)$$
ratio = (m / (hessian + group['eps'])).clamp(-rho, rho)
# $$\theta_{t + 1} \leftarrow \theta_t - \eta \cdot \operatorname{clip} \bigg(\frac{m_t}{h_t + \epsilon}, \rho \bigg)$$
param.data.add_(ratio, alpha=-eta)