### Install optimi

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/index.md

Install the optimi library using pip.

```bash
pip install torch-optimi
```

--------------------------------

### Initialize Optimizer and Model for Gradient Release

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizer_accumulation.md

Initialize any optimi optimizer with `gradient_release=True` and call `prepare_for_gradient_release` on both the model and optimizer. This setup is required for optimizer accumulation.

```python
import torch
from torch import nn
from optimi import AdamW

# create or cast model in low precision (bfloat16)
model = nn.Linear(20, 1, dtype=torch.bfloat16)

# initialize any optimi optimizer with `gradient_release=True`
# and call `prepare_for_gradient_release` on model and optimizer
opt = AdamW(model.parameters(), lr=1e-3, gradient_release=True)
prepare_for_gradient_release(model, opt)

# update model parameters every four steps after accumulating
# gradients directly into the optimizer states
accumulation_steps = 4

# setup a learning rate scheduler for gradient accumulation
scheduler = CosineAnnealingLR(opt, ...)
```

--------------------------------

### Initialize Model in Low Precision

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/kahan_summation.md

Prepare your model by casting its layers to a low-precision dtype like bfloat16. This setup is necessary before applying optimi optimizers with Kahan summation.

```python
import torch
from torch import nn
from optimi import AdamW

# create or cast some model layers in low precision (bfloat16)
model = nn.Linear(20, 1, dtype=torch.bfloat16)

```

--------------------------------

### Gradient Accumulation Example

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/index.md

This snippet shows how to configure and use Optimi for gradient accumulation. Gradients are accumulated into optimizer states, and the optimizer step is performed automatically after `loss.backward()` when `optimizer_accumulation` is set appropriately. The learning rate scheduler is stepped only when parameters are updated.

```python
accumulation_steps = 4
scheduler = CosineAnnealingLR(opt, ...)

for idx, batch in enumerate(dataloader):
    opt.optimizer_accumulation = (idx+1) % accumulation_steps != 0
    loss = model(batch)
    loss.backward()
    if not opt.optimizer_accumulation:
        scheduler.step()

remove_gradient_release(model)
```

--------------------------------

### Initialize AdamW Optimizer with Triton

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/triton.md

Demonstrates how to initialize the AdamW optimizer. Models on supported GPUs default to using the Triton backend. You can explicitly enable it by setting `triton=True`.

```python
import torch
from torch import nn
from optimi import AdamW

# create model
model = nn.Linear(20, 1, device="cuda")

# models on a supported GPU will default to `triton=True`
opt = AdamW(model.parameters(), lr=1e-3)

# or initialize any optimi optimizer with `triton=True`
opt = AdamW(model.parameters(), lr=1e-3, triton=True)

# forward and backward
loss = model(torch.randn(20))
loss.backward()

# optimizer step is the Triton implementation
opt.step()
opt.zero_grad()
```

--------------------------------

### Standard Optimizer Step with Optimi

Source: https://github.com/warner-benjamin/optimi/blob/main/README.md

Use Optimi optimizers for standard training loops. The optimizer step and zero_grad are handled automatically after `loss.backward()`. Learning rate schedulers are stepped as usual.

```python
scheduler = CosineAnnealingLR(opt, ...)

loss = model(torch.randn(20, dtype=torch.bfloat16))
loss.backward()

scheduler.step()

remove_gradient_release(model)
```

--------------------------------

### SGD with Kahan Summation Algorithm

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/kahan_summation.md

This pseudocode illustrates the SGD algorithm with Kahan summation, showing the compensation buffer 'k' for mitigating precision errors.

```latex
\begin{aligned}
    &\rule{90mm}{0.4pt}\\
    &\hspace{2mm} \textcolor{#009ddb}{\textbf{SGD}} \: \textcolor{#9a3fe4}{\text{with Kahan summation}}\\
    &\hspace{5mm} \text{inputs} : \bm{\theta}_0 \: \text{(params)}; \: f(\bm{\theta}) \text{(objective)};\\
    &\hspace{17.25mm} \gamma_t \:\text{(learning rate at } t \text{)}; \: \lambda \: \text{(weight decay)}\\
    &\hspace{5mm} \text{initialize} : \textcolor{#9a3fe4}{\bm{k}_{0} \leftarrow \bm{0}}\\
\[-0.5em]
    &\rule{90mm}{0.4pt}\\
    &\hspace{5mm} \textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}	:
        &	\hspace{10mm} \bm{g}_t \leftarrow \nabla_{\theta} f_t(\bm{\theta}_{t-1}) - \lambda\bm{\theta}_{t-1}\\
        &	\hspace{10mm} \textcolor{#009ddb}{\bm{\theta}_t \leftarrow \bm{\theta}_{t-1} - \gamma_t\bm{g}_t}\\
        &	\hspace{10mm} \textcolor{#9a3fe4}{\bm{u}_t \leftarrow \bm{k}_{t-1} - \gamma_t\bm{g}_t}\\
        &	\hspace{10mm} \textcolor{#9a3fe4}{\bm{\theta}_t \leftarrow \bm{\theta}_{t-1} + \bm{u}_t}\\
        &	\hspace{10mm} \textcolor{#9a3fe4}{\bm{k}_t \leftarrow \bm{u}_t + (\bm{\theta}_{t-1} - \bm{\theta}_t)}\\
\[-0.5em]
    &\rule{90mm}{0.4pt}\\
\end{aligned}

```

--------------------------------

### Initialize and Use ForEach AdamW Optimizer

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/foreach.md

Demonstrates how to initialize an AdamW optimizer with `foreach=True` for potential performance gains. Ensure PyTorch 2.1+ and a CUDA device are used. Set `foreach=False` if gradients are needed between `step()` and `zero_grad()`.

```python
import torch
from torch import nn
from optimi import AdamW

# create model
model = nn.Linear(20, 1, device="cuda")

# initialize any optmi optimizer with `foreach=True`
opt = AdamW(model.parameters(), lr=1e-3, foreach=True)

# forward and backward
loss = model(torch.randn(20))
loss.backward()

# optimizer step is the foreach implementation
opt.step()
opt.zero_grad()
```

--------------------------------

### Usage with Gradient Release

Source: https://github.com/warner-benjamin/optimi/blob/main/README.md

Initialize an optimi optimizer with gradient release enabled and prepare the model and optimizer for it.

```python
# initialize any optimi optimizer with `gradient_release=True`
# and call `prepare_for_gradient_release` on model and optimizer
opt = AdamW(model.parameters(), lr=1e-3, gradient_release=True)
prepare_for_gradient_release(model, opt)
```

--------------------------------

### Initialize and Use Gradient Release

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/gradient_release.md

Initialize an optimi optimizer with `gradient_release=True` and call `prepare_for_gradient_release` on the model and optimizer. The backward pass will then automatically perform the optimizer step and clear gradients. Learning rate schedulers can be stepped as usual. Gradient release hooks can be removed when training is complete.

```python
import torch
from torch import nn
from optimi import AdamW

# create or cast model in low precision (bfloat16)
model = nn.Linear(20, 1, dtype=torch.bfloat16)

# initialize any optimi optimizer with `gradient_release=True`
# and call `prepare_for_gradient_release` on model and optimizer
opt = AdamW(model.parameters(), lr=1e-3, gradient_release=True)
prepare_for_gradient_release(model, opt)

# setup a learning rate scheduler like normal
scheduler = CosineAnnealingLR(opt, ...)

# calling backward on the model will peform the optimzier step
loss = model(torch.randn(20, dtype=torch.bfloat16))
loss.backward()

# optimizer step and zero_grad are no longer needed, and will
# harmlessly no-op if called by an existing training framework
# opt.step()
# opt.zero_grad()

# step the learning rate scheduler like normal
scheduler.step()

# optionally remove gradient release hooks when done training
remove_gradient_release(model)
```

--------------------------------

### Usage with Gradient Release

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/index.md

Initialize an optimi optimizer with gradient release enabled and prepare the model and optimizer. The backward pass automatically performs the optimizer step.

```python
# initialize any optimi optimizer with `gradient_release=True`
# and call `prepare_for_gradient_release` on model and optimizer
opt = AdamW(model.parameters(), lr=1e-3, gradient_release=True)
prepare_for_gradient_release(model, opt)

# setup a learning rate scheduler like normal
scheduler = CosineAnnealingLR(opt, ...)

# calling backward on the model will peform the optimzier step
loss = model(torch.randn(20, dtype=torch.bfloat16))
loss.backward()

# optimizer step and zero_grad are no longer needed, and will
# harmlessly no-op if called by an existing training framework
# opt.step()
# opt.zero_grad()

# step the learning rate scheduler like normal
scheduler.step()

# optionally remove gradient release hooks when done training
remove_gradient_release(model)
```

--------------------------------

### Optimizer with Gradient Accumulation

Source: https://github.com/warner-benjamin/optimi/blob/main/README.md

Configure Optimi optimizers with `gradient_release=True` and prepare the model and optimizer using `prepare_for_gradient_release`. Gradients are accumulated into optimizer states, and parameter updates occur every `accumulation_steps`.

```python
opt = AdamW(model.parameters(), lr=1e-3, gradient_release=True)
prepare_for_gradient_release(model, opt)

accumulation_steps = 4

scheduler = CosineAnnealingLR(opt, ...)

for idx, batch in enumerate(dataloader):
    opt.optimizer_accumulation = (idx+1) % accumulation_steps != 0

    loss = model(batch)
    loss.backward()

    if not opt.optimizer_accumulation:
        scheduler.step()

remove_gradient_release(model)
```

--------------------------------

### Usage with Optimizer Accumulation

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/index.md

Initialize an optimi optimizer with gradient release enabled and prepare the model and optimizer for optimizer accumulation. Model parameters are updated every four steps.

```python
# initialize any optimi optimizer with `gradient_release=True`
# and call `prepare_for_gradient_release` on model and optimizer
opt = AdamW(model.parameters(), lr=1e-3, gradient_release=True)
prepare_for_gradient_release(model, opt)

# update model parameters every four steps after accumulating

```

--------------------------------

### Usage with Kahan Summation and Decoupled Weight Decay

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/index.md

Initialize an optimi optimizer with Kahan summation and fully decoupled weight decay. Kahan summation is automatically enabled for BFloat16 models and inputs.

```python
import torch
from torch import nn
from optimi import AdamW

# create or cast model in low precision (bfloat16)
model = nn.Linear(20, 1, dtype=torch.bfloat16)

# initialize any optimi optimizer with parameters & fully decoupled weight decay
# Kahan summation is automatically enabled since model & inputs are bfloat16
opt = AdamW(model.parameters(), lr=1e-3, weight_decay=1e-5, decouple_lr=True)

# forward and backward, casting input to bfloat16 if needed
loss = model(torch.randn(20, dtype=torch.bfloat16))
loss.backward()

# optimizer step
opt.step()
opt.zero_grad()
```

--------------------------------

### Adam Optimizer Initialization

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizers/adam.md

Initializes the Adam optimizer. This is the primary way to use the Adam optimizer in optimi. It supports various hyperparameters for fine-tuning the optimization process.

```python
optimi.adam.Adam
```

--------------------------------

### Usage with PyTorch-style Weight Decay

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/index.md

Initialize an optimi optimizer with PyTorch-style weight decay for float32 or mixed precision training.

```python
# create model
model = nn.Linear(20, 1)

# initialize any optimi optimizer with parameters
opt = AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
```

--------------------------------

### AdamW Optimizer

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizers/adamw.md

Instantiates the AdamW optimizer from the optimi library. Default hyperparameters are provided, but can be overridden.

```python
from optimi import AdamW

# Default hyperparameters
optimizer = AdamW()

# Custom hyperparameters
optimizer = AdamW(lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01, decouple_lr=False)
```

--------------------------------

### Ranger Optimizer Usage

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizers/ranger.md

Instantiate the Ranger optimizer from the optimi library. It supports decoupled weight decay options.

```python
from optimi.ranger import Ranger

# Example usage with default parameters
optimizer = Ranger(params, lr=1e-3, weight_decay=1e-2)

# Example with decoupled weight decay
optimizer_decouple_wd = Ranger(params, lr=1e-3, weight_decay=1e-2, decouple_wd=True)

# Example with fully decoupled weight decay
optimizer_decouple_lr = Ranger(params, lr=1e-3, weight_decay=1e-4, decouple_lr=True)
```

--------------------------------

### Initialize AdamW with Fully Decoupled Weight Decay

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/fully_decoupled_weight_decay.md

Initialize an optimi AdamW optimizer with `decouple_lr=True` to enable fully decoupled weight decay. Note that the `weight_decay` value is typically lower than the default when using this feature.

```python
import torch
from torch import nn
from optimi import AdamW

# create model
model = nn.Linear(20, 1, dtype=torch.bfloat16)

# initialize any optimi optimizer useing `decouple_lr=True` to enable fully
# decoupled weight decay. note `weight_decay` is lower then the default of 1e-2
opt = AdamW(model.parameters(), lr=1e-3, weight_decay=1e-5, decouple_lr=True)

# model is optimized using fully decoupled weight decay
loss = model(torch.randn(20, dtype=torch.bfloat16))
loss.backward()

opt.step()
opt.zero_grad()
```

--------------------------------

### StableAdamW Algorithm

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizers/stableadamw.md

This is the mathematical formulation for the StableAdamW algorithm, detailing the steps for parameter updates, including RMS calculation and learning rate scaling.

```latex
\begin{aligned}
    &\rule{100mm}{0.4pt}\\
    &\hspace{2mm} \textbf{\textcolor{#9a3fe4}{Stable}AdamW} \\
    &\hspace{5mm} \text{inputs} : \bm{\theta}_0 \: \text{(params)}; \: f(\bm{\theta}) \text{(objective)}; \: \gamma_t \:\text{(learning rate at } t 	ext{)}; \\
    &\hspace{17.25mm} \beta_1, \beta_2 \: \text{(betas)}; \: \lambda \: \text{(weight decay)}; \: \epsilon \: \text{(epsilon)}\\
    &\text{initialize} : \bm{m}_{0} \leftarrow \bm{0}; \: \bm{v}_{0} \leftarrow \bm{0}\\[-0.5em]
    &\rule{100mm}{0.4pt}\\
    &\hspace{5mm} \textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}	ext{:}\\
        &\hspace{10mm} \bm{g}_t \leftarrow \nabla_{\theta} f_t(\bm{\theta}_{t-1})\\\n        &\hspace{10mm} \bm{m}_t \leftarrow \beta_1 \bm{m}_{t-1} + (1 - \beta_1) \bm{g}_t\\
        &\hspace{10mm} \bm{v}_t \leftarrow \beta_2 \bm{v}_{t-1} + (1 - \beta_2) \bm{g}^2_t\\\n        &\hspace{10mm} \hat{\bm{m}}_t \leftarrow \bm{m}_t/(1 - \beta_1^t)\\
        &\hspace{10mm} \hat{\bm{v}}_t \leftarrow \bm{v}_t/(1 - \beta_2^t)\\
        &\hspace{10mm} \textcolor{#9a3fe4}{\textbf{RMS}_t \leftarrow  \sqrt{\mathbb{E[\bm{g}^2_t/\text{max}(\bm{v}_t, \epsilon^2)]}}}\\
        &\hspace{10mm} \textcolor{#9a3fe4}{\bm{\eta}_t \leftarrow  \gamma_t/\text{max}(1,\textbf{RMS}_t)}\\
        &\hspace{10mm} \bm{\theta}_t \leftarrow \bm{\theta}_{t-1} - \textcolor{#9a3fe4}{\bm{\eta}_t} \bigl( \hat{\bm{m}}_t / (\sqrt{\hat{\bm{v}}_t} + \epsilon) + \lambda\bm{\theta}_{t-1} \bigr)\\[-0.5em]
    &\rule{100mm}{0.4pt}\\
\end{aligned}

```

--------------------------------

### Initialize AdamW with Kahan Summation

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/kahan_summation.md

Initialize an AdamW optimizer with Kahan summation enabled. This is useful when model layers use low-precision types like bfloat16.

```python
opt = AdamW(model.parameters(), lr=1e-3)
```

--------------------------------

### Optimizer Accumulation Training Loop

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizer_accumulation.md

Demonstrates how to implement optimizer accumulation within a PyTorch training loop. Set `optimizer_accumulation` to control gradient accumulation or parameter updates. The optimizer step and zero_grad are handled automatically.

```python
for idx, batch in enumerate(dataloader):
    # `optimizer_accumulation=True` accumulates gradients into
    # optimizer states. set `optimizer_accumulation=False` to
    # update parameters by performing a full gradient release step
    opt.optimizer_accumulation = (idx+1) % accumulation_steps != 0

    # calling backward on the model will peform the optimizer step
    # either accumulating gradients or updating model parameters
    loss = model(batch)
    loss.backward()

    # optimizer step and zero_grad are no longer needed, and will
    # harmlessly no-op if called by an existing training framework
    # opt.step()
    # opt.zero_grad()

    # step the learning rate scheduler after accumulating gradients
    if not opt.optimizer_accumulation:
        scheduler.step()

# optionally remove gradient release hooks when done training
remove_gradient_release(model)
```

--------------------------------

### Create Parameter Groups with Weight Decay

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/utils.md

Separates model parameters into groups with and without weight decay. Useful for fine-tuning or when specific layers should not have weight decay applied.

```python
params = param_groups_weight_decay(model, weight_decay=1e-5)
optimizer = StableAdamW(params, decouple_lr=True)
```

--------------------------------

### Convert Model to Low Precision and Move to Device

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/utils.md

Casts model parameters to a specified low-precision data type and moves the model to a target device (e.g., 'cuda') in a single operation. Simplifies device and precision management.

```python
to_low_precision(model, dtype=torch.bfloat16, device="cuda")
```

--------------------------------

### Adan Optimizer Implementation

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizers/adan.md

This snippet shows the Adan optimizer class. It supports standard and fully decoupled weight decay. Note that training with Adan in bfloat16 can be noisier than in float32 or mixed precision.

```python
class Adan(Optimizer):
    """Adan: ADAptive Nesterov Momentum

    Args:
        params (iterable): iterable of parameters to optimize or dicts defining parameter groups
        lr (float, optional): learning rate (default: 1e-3)
        betas (Tuple[float, float, float], optional): coefficients used for computing
            running averages of gradient and its square (default: (0.9, 0.92, 0.99))
        eps (float, optional): term added to the denominator to improve
            numerical stability (default: 1e-8)
        weight_decay (float, optional): weight decay (default: 0.02)
        decouple_lr (bool, optional): Whether to use fully decoupled weight decay. (default: False)
    """

    def __init__(self, params, lr=1e-3, betas=(0.9, 0.92, 0.99), eps=1e-8,
                 weight_decay=0.02, decouple_lr=False):
        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate value: {}".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]))
        if not 0.0 <= betas[2] < 1.0:
            raise ValueError("Invalid beta parameter at index 2: {}".format(betas[2]))
        if weight_decay < 0.0:
            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))

        defaults = dict(lr=lr, betas=betas, eps=eps,
                        weight_decay=weight_decay, decouple_lr=decouple_lr)
        super(Adan, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(Adan, self).__setstate__(state)
        for group in self.param_groups:
            group.setdefault('exp_avg', collections.defaultdict(torch.Tensor))

    @torch.no_grad()
    def step(self, closure=None):
        """Performs a single optimization step.

        Args:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss. (default: None)
        """
        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

                # State initialization
                if p.shape not in self.state[p].keys():
                    self.state[p]['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    self.state[p]['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    self.state[p]['exp_avg_diff'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    self.state[p]['step_num'] = torch.zeros(1, dtype=torch.long)

                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError('Adan does not support sparse gradients, please consider AdamW or SparseAdam instead')

                state = self.state[p]
                step_num = state['step_num'].item() + 1
                state['step_num'].fill_(step_num)

                # Hyperparameter
                beta1, beta2, beta3 = group['betas']
                lr = group['lr']
                eps = group['eps']
                weight_decay = group['weight_decay']
                decouple_lr = group['decouple_lr']

                # Update biased first moment estimate
                exp_avg = state['exp_avg']
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)

                # Update biased second moment estimate
                exp_avg_sq = state['exp_avg_sq']
                if step_num > 1:
                    exp_avg_diff = state['exp_avg_diff']
                    exp_avg_diff.mul_(beta2).add_(grad - p.grad_prev, alpha=1 - beta2)
                else:
                    exp_avg_diff = torch.zeros_like(p, memory_format=torch.preserve_format)

                # Update Nesterov momentum
                exp_avg_diff_nesterov = exp_avg_diff.clone()
                exp_avg_diff_nesterov.mul_(beta2)
                exp_avg_diff_nesterov.add_(exp_avg, alpha=1 - beta1)

                # Update third moment estimate
                exp_avg_diff_third = state['exp_avg_diff']
                exp_avg_diff_third.mul_(beta3).add_(exp_avg_diff_nesterov, alpha=1 - beta3)

                # Bias correction
                bias_correction1 = 1 - beta1 ** step_num
                bias_correction2 = 1 - beta2 ** step_num
                bias_correction3 = 1 - beta3 ** step_num

                # Update parameters
                step_size = lr / bias_correction2
                bias_exp_avg_diff_third = exp_avg_diff_third / bias_correction3
                step_size /= (torch.sqrt(bias_exp_avg_diff_third) + eps)
                step_size *= (exp_avg / bias_correction1 + beta2 * exp_avg_diff / bias_correction2)

                if weight_decay != 0:
                    if decouple_lr:
                        p.mul_(1.0 - lr * weight_decay)
                    else:
                        p.mul_(1.0 - group['lr'] * weight_decay)

                p.data.add_(step_size, alpha=-1)
                p.grad_prev = grad.clone() # Store grad for next step

        return loss
```

--------------------------------

### optimi.lion.Lion

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizers/lion.md

The Lion optimizer class from the optimi library. It implements the Evolved Sign Momentum algorithm, offering memory savings and potentially faster convergence.

```APIDOC
## Class: optimi.lion.Lion

### Description
Implements the Lion optimizer, which uses a sign-based momentum update for efficient training. It tracks gradient moving averages to reduce memory usage compared to optimizers like AdamW.

### Algorithm
Lion: Evolved Sign Momentum.

$$ \
\begin{aligned} \
    &\\rule{100mm}{0.4pt}\\
    &\hspace{2mm} \textbf{Lion} \
    &\hspace{5mm} \text{inputs} : \bm{\theta}_0 \: \text{(params)}; \: f(\bm{\theta}) \text{(objective)}; \: \gamma_t \:\text{(learning rate at } t \text{)}; \\
    &\hspace{17.25mm} \beta_1, \beta_2 \: \text{(betas)}; \: \lambda \: \text{(weight decay)}\
    &\hspace{5mm} \text{initialize} : \bm{m}_{0} \leftarrow \bm{0}\\[-0.5em] \
    &\rule{100mm}{0.4pt}\\\n    &\hspace{5mm} \textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}\
        &\hspace{10mm} \bm{g}_t \leftarrow \nabla_{\theta} f_t(\bm{\theta}_{t-1})\\\n        &\hspace{10mm} \bm{u} \leftarrow \beta_1 \bm{m}_{t-1} + (1 - \beta_1) \bm{g}_t\\ \
        &\hspace{10mm} \bm{m}_t \leftarrow \beta_2 \bm{m}_{t-1} + (1 - \beta_2) \bm{g}_t\\\n        &\hspace{10mm} \bm{\theta}_t \leftarrow \bm{\theta}_{t-1} - \gamma_t \bigl(\text{sign}(\bm{u}) + \lambda\bm{\theta}_{t-1} \bigr)\\[-0.5em] \
    &\rule{100mm}{0.4pt}\\\n\end{aligned}
$$

Note: This implementation also supports fully decoupled weight decay, which is not explicitly shown in the algorithm above.
```

--------------------------------

### StableAdamW Optimizer Class

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizers/stableadamw.md

This snippet shows the StableAdamW optimizer class from the optimi library. It's a drop-in replacement for AdamW, offering improved training stability.

```python
from optimi.stableadamw import StableAdamW
```

--------------------------------

### AdamW Optimizer Class

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizers/adamw.md

The AdamW optimizer class from the optimi library. It implements Adam with decoupled weight decay, offering improved training dynamics compared to standard Adam. It also supports fully decoupled weight decay as an option.

```APIDOC
## Class optimi.adamw.AdamW

### Description
Implements the AdamW optimization algorithm, which decouples weight decay from the gradient updates. This approach often leads to better generalization and convergence.

### Parameters
- **params** (iterable) - Iterable of parameters to optimize or dicts defining parameter groups.
- **lr** (float, optional) - Learning rate. Defaults to 1e-3.
- **betas** (Tuple[float, float], optional) - Coefficients used for computing running averages of gradient and its square. Defaults to (0.9, 0.99).
- **eps** (float, optional) - Term added to the denominator to improve numerical stability. Defaults to 1e-6.
- **weight_decay** (float, optional) - Weight decay (L2 penalty). Defaults to 0.01.
- **decouple_lr** (bool, optional) - If True, uses fully decoupled weight decay. Defaults to False.

### Example
```python
from optimi.adamw import AdamW

# Assuming 'model' is your PyTorch model
optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.01, decouple_lr=True)
```
```

--------------------------------

### SGD Optimizer Class

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizers/sgd.md

This snippet shows the SGD optimizer class from the optimi library. It supports various hyperparameters for customization.

```python
from optimi.sgd import SGD

# Example usage (assuming model and data are defined)
# optimizer = SGD(params=model.parameters(), lr=0.01, dampening=True, torch_init=True, decouple_wd=True, decouple_lr=True)
```

--------------------------------

### RAdam Optimizer Class

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/optimizers/radam.md

This snippet shows the RAdam optimizer class from the optimi library. It is designed to improve Adam's convergence by stabilizing the adaptive learning rate. It supports decoupled weight decay and fully decoupled weight decay.

```python
from optimi.radam import RAdam
```

--------------------------------

### Disable Kahan Summation

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/kahan_summation.md

To disable Kahan summation, pass `kahan_summation=False` during optimizer initialization.

```python
# To disable Kahan Summation pass kahan_summation=False on optimizer initialization.
```

--------------------------------

### Keep Specific Modules in Float32 during Low Precision Conversion

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/utils.md

Allows specific modules (e.g., LayerNorm) to remain in float32 while the rest of the model is cast to a lower precision. This can help maintain numerical stability for certain layers.

```python
to_low_precision(model, dtype=torch.bfloat16, fp32_modules=(nn.LayerNorm,))
```

--------------------------------

### Convert Model to Low Precision

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/utils.md

Casts model parameters to a specified low-precision data type (e.g., bfloat16). Useful for reducing memory footprint and potentially increasing speed on compatible hardware.

```python
from torch import nn
from optimi.utils import to_low_precision

model = nn.Sequential(
    nn.Embedding(1000, 128),
    nn.Linear(128, 128),
)

to_low_precision(model, dtype=torch.bfloat16)
```

--------------------------------

### Exclude Additional Layers from Weight Decay

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/utils.md

Allows specifying additional layer names or substrings to exclude from weight decay, beyond bias and normalization layers. Useful for layers like token embeddings.

```python
class MiniLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok_embeddings = nn.Embedding(1000, 20)
        self.pos_embeddings = nn.Embedding(100, 20)
        self.norm = nn.LayerNorm(20)
        self.layer1 = nn.Linear(20, 30)
        self.layer2 = nn.Linear(30, 1000)

model = MiniLM()

# Exclude token embeddings from weight decay in addition to bias and normalization layers
params = param_groups_weight_decay(
    model,
    weight_decay=1e-5,
    additional_layers=["tok_embeddings"]
)
```

--------------------------------

### Disable Rotary Buffer Preservation during Low Precision Conversion

Source: https://github.com/warner-benjamin/optimi/blob/main/docs/utils.md

Disables the preservation of RoPE/rotary buffers, allowing them to also be cast to the specified low-precision data type. Use when buffer preservation is not required.

```python
to_low_precision(model, dtype=torch.bfloat16, fp32_buffers=None)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.