### Install lion-pytorch

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/README.md

Install the lion-pytorch library using pip.

```bash
pip install lion-pytorch
```

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

### Basic Lion Optimizer Usage

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/index.md

Demonstrates the basic setup and training loop using the Lion optimizer.

```python
import torch
from torch import nn
from lion_pytorch import Lion

# Create model and optimizer
model = nn.Linear(10, 1)
optimizer = Lion(
    model.parameters(),
    lr=1e-4,          # 3-10x smaller than AdamW
    weight_decay=1e-2
)

# Training loop
for epoch in range(10):
    loss = model(torch.randn(10)).sum()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
```

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

### Install Lion Optimizer

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Install the lion-pytorch package. Use the triton package for GPU acceleration.

```bash
# Basic installation
pip install lion-pytorch
```

```bash
# With Triton for GPU acceleration
pip install lion-pytorch triton>=2.2.0
```

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

### Basic Foreach Lion Example

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/foreach-lion.md

Demonstrates the basic usage of the Foreach Lion optimizer with a sequential PyTorch model. This example shows optimizer initialization and its use within a typical training loop, including zeroing gradients, backpropagation, and stepping the optimizer.

```python
import torch
from torch import nn
from lion_pytorch.foreach import Lion

# Create a larger model to benefit from foreach optimization
model = nn.Sequential(
    nn.Linear(512, 1024),
    nn.ReLU(),
    nn.Linear(1024, 512),
    nn.ReLU(),
    nn.Linear(512, 256)
)

# Initialize foreach variant
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2
)

# Training loop
for epoch in range(10):
    batch = torch.randn(64, 512)
    output = model(batch)
    loss = output.sum()
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

```

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

### Install Lion Optimizer

Source: https://github.com/lucidrains/lion-pytorch/blob/main/README.md

Install the lion-pytorch package using pip or conda.

```bash
pip install lion-pytorch
```

```bash
conda install lion-pytorch
```

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

### Minimal Lion Optimizer Example

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

A basic example demonstrating how to initialize and use the Lion optimizer with a simple PyTorch model.

```python
import torch
from torch import nn
from lion_pytorch import Lion

model = nn.Linear(10, 1)
optimizer = Lion(model.parameters(), lr=1e-4)

loss = model(torch.randn(10)).sum()
optimizer.zero_grad()
loss.backward()
optimizer.step()
```

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

### Install Triton for CUDA Acceleration

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/index.md

Install the Triton library for optional GPU acceleration with Lion.

```bash
pip install triton>=2.2.0
```

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

### Example Usage: Basic Foreach Lion

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/foreach-lion.md

Demonstrates the basic usage of the Foreach Lion optimizer with a sample PyTorch model and training loop.

```APIDOC
### Example Usage: Basic Foreach Lion

```python
import torch
from torch import nn
from lion_pytorch.foreach import Lion

# Create a larger model to benefit from foreach optimization
model = nn.Sequential(
    nn.Linear(512, 1024),
    nn.ReLU(),
    nn.Linear(1024, 512),
    nn.ReLU(),
    nn.Linear(512, 256)
)

# Initialize foreach variant
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2
)

# Training loop
for epoch in range(10):
    batch = torch.randn(64, 512)
    output = model(batch)
    loss = output.sum()
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
```
```

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

### Lion Optimizer Import and Initialization for Different Scenarios

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Provides examples for importing and initializing Lion optimizer variants for stable training, large models, and GPU-optimized training.

```python
# Stable/conservative training
from lion_pytorch.cautious_lion import Lion
optimizer = Lion(model.parameters(), lr=1e-4, cautious_factor=0.5)

# Large models
from lion_pytorch.foreach import Lion
optimizer = Lion(model.parameters(), lr=1e-4)

# GPU-optimized
from lion_pytorch import Lion
optimizer = Lion(model.parameters(), lr=1e-4, use_triton=True)
```

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

### Minimal Lion Optimizer Example

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/README.md

A basic example demonstrating how to use the Lion optimizer with a PyTorch model. This includes model definition, optimizer instantiation, and a single training step.

```python
import torch
from torch import nn
from lion_pytorch import Lion

model = nn.Linear(10, 1)
optimizer = Lion(model.parameters(), lr=1e-4, weight_decay=1e-2)

loss = model(torch.randn(10)).sum()
optimizer.zero_grad()
loss.backward()
optimizer.step()
```

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

### Lion Optimizer Import Examples

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Demonstrates recommended import patterns for the Lion optimizer, including standard imports and aliasing for using multiple variants.

```python
from lion_pytorch import Lion  # Standard variant

# Alternative: Alias for clarity when using multiple variants
from lion_pytorch import Lion as StandardLion
from lion_pytorch.cautious_lion import Lion as CautiousLion
from lion_pytorch.foreach import Lion as ForeachLion
```

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

### Lion Optimizer Standard Usage Example

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/lion.md

Demonstrates how to initialize and use the Lion optimizer with a simple PyTorch model. Includes model creation, optimizer initialization with recommended parameters, and a basic training loop.

```python
import torch
from torch import nn
from lion_pytorch import Lion

# Create a simple model
model = nn.Sequential(
    nn.Linear(64, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

# Initialize optimizer with recommended parameters (3x smaller than AdamW)
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    betas=(0.9, 0.99)
)

# Training loop
for epoch in range(10):
    output = model(torch.randn(32, 64))
    loss = output.sum()
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
```

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

### Triton Kernel Block Setup

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/triton-kernel.md

Calculates the starting index for the current block and generates a mask to handle elements that fall outside the tensor bounds.

```python
pid = tl.program_id(axis=0)  # Block ID
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)  # Element indices
mask = offsets < n_elements  # Boundary mask
```

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

### Example Usage: With Decoupled Weight Decay

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/foreach-lion.md

Shows how to initialize the Foreach Lion optimizer with decoupled weight decay enabled.

```APIDOC
### Example Usage: With Decoupled Weight Decay

```python
from lion_pytorch.foreach import Lion

optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    decoupled_weight_decay=True
)
```
```

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

### Example Usage: Basic Cautious Lion

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/cautious-lion.md

Demonstrates how to initialize and use the Cautious Lion optimizer with its default settings, effectively behaving like the standard Lion optimizer when `cautious_factor` is 0.0.

```APIDOC
### Example Usage: Basic Cautious Lion

```python
import torch
from torch import nn
from lion_pytorch.cautious_lion import Lion

model = nn.Sequential(
    nn.Linear(64, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

# Standard Lion behavior (cautious_factor=0.0)
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2
)

for epoch in range(10):
    output = model(torch.randn(32, 64))
    loss = output.sum()
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
```
```

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

### Lion Optimizer with Triton Acceleration

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/lion.md

Example of initializing the Lion optimizer with Triton acceleration enabled. This requires `pip install triton>=2.2.0` and is beneficial for models running on CUDA-enabled GPUs.

```python
from lion_pytorch import Lion

# Requires: pip install triton>=2.2.0
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    use_triton=True  # Enable CUDA kernel acceleration
)

# Models on GPU will automatically use Triton kernel
output = model(torch.randn(32, 64, device='cuda'))
loss = output.sum()
optimizer.zero_grad()
loss.backward()
optimizer.step()
```

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

### Example Usage: With Cautious Scaling

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/cautious-lion.md

Shows how to enable and utilize the cautious update mechanism by setting a `cautious_factor` greater than 0.0, leading to more conservative parameter updates.

```APIDOC
### Example Usage: With Cautious Scaling

```python
from lion_pytorch.cautious_lion import Lion

# Enable cautious updates with conservative scaling
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    cautious_factor=0.5  # 50% reduction for misaligned updates
)

# Training loop benefits from more conservative parameter updates
# particularly useful for unstable architectures or small batch sizes
output = model(torch.randn(32, 64))
loss = output.sum()
optimizer.zero_grad()
loss.backward()
optimizer.step()
```
```

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

### Example Usage: Decoupled Weight Decay

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/cautious-lion.md

Illustrates the usage of decoupled weight decay in conjunction with the cautious Lion optimizer by setting `decoupled_weight_decay` to `True`.

```APIDOC
### Example Usage: Decoupled Weight Decay

```python
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    cautious_factor=0.5,
    decoupled_weight_decay=True
)
```
```

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

### Integrate Learning Rate Scheduling

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/index.md

Example of integrating a learning rate scheduler with the Lion optimizer.

```python
from torch.optim.lr_scheduler import CosineAnnealingLR

scheduler = CosineAnnealingLR(optimizer, T_max=epochs)

for epoch in range(epochs):
    # training loop
    loss.backward()
    optimizer.step()
    scheduler.step()
```

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

### Training with Triton Kernel

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/triton-kernel.md

Example of using the Lion optimizer with its Triton kernel for parameter updates on a CUDA-enabled GPU.

```python
x = torch.randn(64, 1024, device='cuda')
y = model(x)
loss = y.sum()

optimizer.zero_grad()
loss.backward()
optimizer.step()  # Uses Triton kernel for parameter updates
```

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

### Import Triton and Triton Language

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/triton-kernel.md

Imports the necessary PyTorch and Triton libraries for using the Triton kernel. This snippet is part of the required setup.

```python
import torch
import triton
import triton.language as tl
```

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

### Standard Training Import and Initialization

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Shows how to import the standard Lion optimizer and initialize it for model training.

```python
from lion_pytorch import Lion

optimizer = Lion(model.parameters(), lr=1e-4)
```

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

### Error Handling for Triton Import

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/triton-kernel.md

Includes error handling to check if the Triton library is installed, providing instructions to install it if missing. The module exits if Triton is not found.

```python
try:
    import triton
    import triton.language as tl
except ImportError as e:
    print('triton is not installed, please install by running `pip install triton>=2.2.0`')
    exit()
```

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

### Lion Optimizer with Triton Fused Kernel

Source: https://github.com/lucidrains/lion-pytorch/blob/main/README.md

Utilize the Triton fused kernel for parameter updates by setting `use_triton=True`. Ensure Triton is installed with `pip install triton -U --pre`.

```python
opt = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    use_triton=True # set this to True to use cuda kernel w/ Triton lang (Tillet et al)
)
```

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

### Upgrade PyTorch

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Upgrade your PyTorch installation to version 1.13 or later to resolve `AssertionError` related to foreach functions.

```bash
pip install --upgrade torch
```

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

### Accessing All Lion Variants

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Demonstrates how to import all available Lion optimizer variants for flexible use.

```python
import lion_pytorch
from lion_pytorch.cautious_lion import Lion as CautiousLion
from lion_pytorch.foreach import Lion as ForeachLion

# Use as needed
opt1 = lion_pytorch.Lion(model.parameters())
opt2 = CautiousLion(model.parameters(), cautious_factor=0.5)
opt3 = ForeachLion(model.parameters())
```

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

### Basic Cautious Lion Optimizer Usage

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/cautious-lion.md

Demonstrates the basic initialization and usage of the Cautious Lion optimizer with default parameters, effectively behaving like standard Lion when `cautious_factor` is 0.0. This snippet shows a typical training loop structure.

```python
import torch
from torch import nn
from lion_pytorch.cautious_lion import Lion

model = nn.Sequential(
    nn.Linear(64, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

# Standard Lion behavior (cautious_factor=0.0)
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2
)

for epoch in range(10):
    output = model(torch.randn(32, 64))
    loss = output.sum()
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

```

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

### Lion (Foreach Variant) Constructor

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/foreach-lion.md

Initializes the Foreach Lion optimizer. It takes parameters, learning rate, betas, weight decay, and a flag for decoupled weight decay.

```APIDOC
## Lion (Foreach Variant) Constructor

### Description
Initializes the Foreach Lion optimizer, a performance-optimized variant of Lion that uses PyTorch's `torch._foreach_*` operations for batch tensor updates. This is more efficient for larger models and batches by reducing kernel launch overhead.

### Constructor Signature
```python
class Lion(torch.optim.optimizer.Optimizer):
    def __init__(
        self,
        params,
        lr: float = 1e-4,
        betas: Tuple[float, float] = (0.9, 0.99),
        weight_decay: float = 0.0,
        decoupled_weight_decay: bool = False,
    )
```

### Parameters
#### Constructor Parameters
- **params** (iterable) - Required - Iterable of parameters to optimize or dicts defining parameter groups.
- **lr** (float) - Optional - Learning rate, recommended 3-10x smaller than AdamW. Default: `1e-4`.
- **betas** (Tuple[float, float]) - Optional - Coefficients for gradient exponential moving averages (β₁, β₂). Default: `(0.9, 0.99)`.
- **weight_decay** (float) - Optional - Weight decay coefficient. Default: `0.0`.
- **decoupled_weight_decay** (bool) - Optional - Whether to use decoupled weight decay scaled by initial learning rate. Default: `False`.

### Constructor Validation
- `lr > 0.0` is required.
- All betas must be in the range [0.0, 1.0].
- PyTorch installation must support `torch._foreach_mul_`, `torch._foreach_add_`, `torch._foreach_sign_`, and `torch._foreach_lerp_` operations.
- An error is raised if required foreach operations are unavailable: `'this version of torch does not have the prerequisite foreach operations'`.
```

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

### Alias Multiple Lion Variants

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

When using multiple Lion variants in the same project, alias them to prevent naming conflicts. This example shows how to alias the 'CautiousLion'.

```python
from lion_pytorch.cautious_lion import Lion as CautiousLion
```

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

### Accessing Optimizer State

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Demonstrates how to access and inspect the optimizer's state, such as the current learning rate, weight decay, and momentum buffers.

```python
# View current learning rate
print(optimizer.param_groups[0]['lr'])

# View weight decay
print(optimizer.param_groups[0]['weight_decay'])

# View momentum state for a parameter
param = next(model.parameters())
if param in optimizer.state:
    exp_avg = optimizer.state[param]['exp_avg']
    print(exp_avg.shape)
```

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

### Language Model Training Recipe

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/configuration.md

Recommended configuration for training transformer-based Language Models (LMs) with the Lion optimizer. Emphasizes the importance of a large batch size (256+) for optimal performance.

```python
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.99),
    weight_decay=1e-2,
)
# Large batch size (256+) critical for Lion
train_loader = DataLoader(dataset, batch_size=256)
```

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

### Checkpointing and Resumption with Lion Optimizer

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/index.md

Save and load the state of the model and optimizer for seamless checkpointing and resumption of training.

```python
# Save
checkpoint = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'epoch': epoch
}
torch.save(checkpoint, 'checkpoint.pt')

# Load
checkpoint = torch.load('checkpoint.pt')
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
```

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

### Import Lion Optimizer from Package Root

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/modules.md

Use this import for the standard Lion optimizer implementation. Set `use_triton=True` for high-performance GPU training.

```python
from lion_pytorch import Lion

optimizer = Lion(model.parameters(), lr=1e-4)
```

```python
from lion_pytorch import Lion

optimizer = Lion(model.parameters(), lr=1e-4, use_triton=True)
```

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

### Standard Lion Dependencies

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Lists the dependencies for the standard Lion optimizer implementation.

```text
lion_pytorch.lion_pytorch
├── typing.Tuple
├── typing.Callable
├── torch
├── torch.optim.optimizer.Optimizer
└── lion_pytorch.triton (optional, if use_triton=True)
```

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

### Foreach Lion Dependencies

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Lists the dependencies for the foreach Lion optimizer variant.

```text
lion_pytorch.foreach
├── typing.Tuple
├── typing.Callable
├── torch
├── torch.optim.optimizer.Optimizer
├── torch._foreach_mul_
├── torch._foreach_add_
├── torch._foreach_sign_
└── torch._foreach_lerp_
```

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

### Foreach Lion with Decoupled Weight Decay

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/foreach-lion.md

Shows how to initialize the Foreach Lion optimizer with decoupled weight decay enabled. This is a common configuration for improving model generalization.

```python
from lion_pytorch.foreach import Lion

optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    decoupled_weight_decay=True
)

```

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

### Standard and Foreach Lion Constructor Parameters

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Type signatures for parameters common to the standard and foreach variants of the Lion optimizer.

```python
# Standard and Foreach variants
params: Iterable[torch.nn.Parameter] | Iterable[Dict[str, Any]]
lr: float
betas: Tuple[float, float]
weight_decay: float
decoupled_weight_decay: bool
```

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

### Standard Lion Optimizer Configuration

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Configure the Lion optimizer with standard beta values and weight decay for general use.

```python
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.99),
    weight_decay=1e-2
)
```

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

### Import Lion Optimizer

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Import the Lion optimizer. Recommended for standard usage, or import specific variants like CautiousLion or ForeachLion.

```python
# Standard (recommended)
from lion_pytorch import Lion
```

```python
# Specific variants
from lion_pytorch.cautious_lion import Lion as CautiousLion
from lion_pytorch.foreach import Lion as ForeachLion
```

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

### Basic Training Loop with Lion

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

A standard training loop using the Lion optimizer. Ensure gradients are zeroed, backpropagated, and the optimizer steps.

```python
optimizer = Lion(model.parameters(), lr=1e-4, weight_decay=1e-2)

for epoch in range(epochs):
    for batch in dataloader:
        output = model(batch)
        loss = criterion(output, targets)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
```

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

### Standard Lion Optimizer Initialization

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/modules.md

Initialize the standard Lion optimizer. This is the default choice for most Lion training scenarios. It supports standard PyTorch Optimizer interface.

```python
Lion(params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0, use_triton=False, decoupled_weight_decay=False)
```

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

### Foreach Lion Optimizer Default Configuration

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/configuration.md

Instantiates the Foreach Lion optimizer using default settings for learning rate, betas, weight decay, and decoupled weight decay.

```python
Lion(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.99),
    weight_decay=0.0,
    decoupled_weight_decay=False
)
```

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

### Import Lion Optimizer

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/modules.md

Import the main Lion optimizer class from the lion_pytorch package. This is the standard entry point for most training scenarios.

```python
from lion_pytorch import Lion
```

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

### Cautious Lion Optimizer with Decoupled Weight Decay

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/cautious-lion.md

Demonstrates the initialization of the Cautious Lion optimizer with decoupled weight decay enabled. This option modifies how weight decay is applied during the optimization step.

```python
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    cautious_factor=0.5,
    decoupled_weight_decay=True
)

```

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

### Save and Load Checkpoint

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Save the model, optimizer state, and epoch number to a checkpoint file and load them later to resume training.

```python
# Save checkpoint
checkpoint = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'epoch': epoch,
}
torch.save(checkpoint, 'checkpoint.pth')

# Load checkpoint
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch']
```

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

### Basic Lion Optimizer Usage

Source: https://github.com/lucidrains/lion-pytorch/blob/main/README.md

Instantiate and use the Lion optimizer with a PyTorch model. Ensure to import Lion and pass model parameters, learning rate, and weight decay.

```python
import torch
from torch import nn

model = nn.Linear(10, 1)

# import Lion and instantiate with parameters

from lion_pytorch import Lion

opt = Lion(model.parameters(), lr=1e-4, weight_decay=1e-2)

# forward and backwards

loss = model(torch.randn(10))
loss.backward()

# optimizer step

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

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

### Standard Lion Constructor Parameters

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Additional parameter specific to the standard Lion optimizer.

```python
# Standard variant only
use_triton: bool
```

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

### Standard Lion Optimizer Default Configuration

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/configuration.md

Instantiates the Standard Lion optimizer with default parameters for learning rate, betas, weight decay, Triton usage, and decoupled weight decay.

```python
Lion(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.99),
    weight_decay=0.0,
    use_triton=False,
    decoupled_weight_decay=False
)
```

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

### Move Model and Data to GPU

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Configure your model and input data to run on a single GPU.

```python
# Single GPU
model = model.cuda()
batch = batch.cuda()

optimizer = Lion(model.parameters(), lr=1e-4, use_triton=True)
```

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

### Standard Lion Constructor

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Signature for the standard Lion optimizer. Learning rate is typically 3-10x smaller than AdamW.

```python
Lion(
    params,                          # Model parameters
    lr=1e-4,                         # Learning rate (3-10x smaller than AdamW)
    betas=(0.9, 0.99),              # Momentum coefficients
    weight_decay=0.0,                # Weight decay
    use_triton=False,                # Enable CUDA kernel
    decoupled_weight_decay=False     # Decouple weight decay from LR
)
```

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

### Cautious Lion Optimizer with Cautious Scaling Enabled

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/cautious-lion.md

Shows how to initialize the Cautious Lion optimizer with a non-zero `cautious_factor` to enable conservative parameter updates. This is useful for unstable architectures or small batch sizes. The training loop benefits from more conservative updates.

```python
from lion_pytorch.cautious_lion import Lion

# Enable cautious updates with conservative scaling
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    cautious_factor=0.5  # 50% reduction for misaligned updates
)

# Training loop benefits from more conservative parameter updates
# particularly useful for unstable architectures or small batch sizes
output = model(torch.randn(32, 64))
loss = output.sum()
optimizer.zero_grad()
loss.backward()
optimizer.step()

```

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

### Lion Optimizer Constructor (Standard)

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

The primary export for the Lion optimizer. Recommended for most use cases. Supports optional Triton acceleration.

```python
Lion(
    params,
    lr: float = 1e-4,
    betas: Tuple[float, float] = (0.9, 0.99),
    weight_decay: float = 0.0,
    use_triton: bool = False,
    decoupled_weight_decay: bool = False
)
```

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

### Lion (Cautious Variant) Constructor

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/cautious-lion.md

Initializes the Cautious Lion optimizer. It takes model parameters and various hyperparameters to control the optimization process, including learning rate, betas, weight decay, and the cautious update factor.

```APIDOC
## Class: Lion (Cautious Variant)

```python
class Lion(torch.optim.optimizer.Optimizer):
    def __init__(
        self,
        params,
        lr: float = 1e-4,
        betas: Tuple[float, float] = (0.9, 0.99),
        weight_decay: float = 0.0,
        cautious_factor: float = 0.,
        decoupled_weight_decay: bool = False,
    )
```

### Constructor Parameters

| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| params | iterable | Yes | — | Iterable of parameters to optimize or dicts defining parameter groups |
| lr | float | No | 1e-4 | Learning rate, recommended 3-10x smaller than AdamW |
| betas | Tuple[float, float] | No | (0.9, 0.99) | Coefficients for gradient exponential moving averages (β₁, β₂) |
| weight_decay | float | No | 0.0 | Weight decay coefficient |
| cautious_factor | float | No | 0.0 | Cautious update scaling factor, range [0.0, 1.0]. At 0.0, acts as standard Lion. Higher values apply more conservative updates |
| decoupled_weight_decay | bool | No | False | Whether to use decoupled weight decay scaled by initial learning rate |

### Constructor Validation

- `lr > 0.0` required
- All betas must be in range [0.0, 1.0]
- `cautious_factor` must be in range [0.0, 1.0]
```

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

### Inspect Lion Optimizer State and Configuration

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/configuration.md

Debug optimizer state by printing configuration parameters like learning rate, weight decay, and betas. Access per-parameter state, including momentum buffers.

```python
import json

optimizer = Lion(model.parameters(), lr=1e-4)
optimizer.zero_grad()
loss.backward()
optimizer.step()

# Print optimizer configuration
for group in optimizer.param_groups:
    print(f"LR: {group['lr']}")
    print(f"Weight Decay: {group['weight_decay']}")
    print(f"Betas: {group['betas']}")
    print(f"Num params: {len(group['params'])}")

# Access per-parameter state
for param, state in optimizer.state.items():
    if 'exp_avg' in state:
        print(f"Param shape: {param.shape}, Momentum shape: {state['exp_avg'].shape}")
```

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

### Use Batch-Optimized Lion for Large Models

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/index.md

Instantiate the foreach Lion optimizer variant, optimized for batch processing.

```python
from lion_pytorch.foreach import Lion

optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2
)
```

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

### Cautious Lion Dependencies

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Lists the dependencies for the cautious Lion optimizer variant.

```text
lion_pytorch.cautious_lion
├── typing.Tuple
├── typing.Callable
├── torch
└── torch.optim.optimizer.Optimizer
```

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

### Cautious Lion Optimizer Initialization

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/modules.md

Initialize the cautious variant of the Lion optimizer. Use this when encountering training instability or divergence. It offers conservative updates with adaptive scaling.

```python
Lion(params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0, cautious_factor=0.0, decoupled_weight_decay=False)
```

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

### Configure Lion Optimizer for Unstable Architectures

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/configuration.md

Use this configuration for models that exhibit training difficulties. Set 'cautious_factor' for conservative scaling.

```python
from lion_pytorch.cautious_lion import Lion

optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    cautious_factor=0.5,  # Conservative scaling
)
```

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

### Vision Transformer Training Recipe

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/configuration.md

Recommended configuration for training Vision Transformers (ViT, DeiT, etc.) using the Lion optimizer. It suggests a smaller learning rate and larger weight decay compared to AdamW, along with a cosine decay scheduler.

```python
optimizer = Lion(
    model.parameters(),
    lr=1e-4,              # 3-10x smaller than AdamW
    betas=(0.9, 0.99),
    weight_decay=1e-2,    # Larger than AdamW
)
# Use cosine decay scheduler (cosine annealing)
scheduler = CosineAnnealingLR(optimizer, T_max=epochs)
```

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

### Lion-PyTorch Source Code Organization

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Illustrates the directory structure and main exports for the Lion-PyTorch library.

```text
lion_pytorch/
├── __init__.py                    # Package-level exports
│   └── exports: Lion
├── lion_pytorch.py                # Standard implementation
│   └── exports: Lion, update_fn (internal), exists (internal)
├── cautious_lion.py               # Cautious variant
│   └── exports: Lion, exists (internal)
├── foreach.py                     # Foreach variant
│   └── exports: Lion, exists (internal)
└── triton.py                      # Triton CUDA kernel
    └── exports: update_fn (public), update_fn_kernel (internal)
```

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

### Lion (Standard)

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

The standard implementation of the Lion optimizer. It supports optional Triton acceleration for enhanced performance on compatible hardware.

```APIDOC
## Class Lion (Standard)

### Description
Standard Lion optimizer with optional Triton acceleration.

### Import
`from lion_pytorch import Lion`

### Constructor
```python
Lion(
    params,
    lr: float = 1e-4,
    betas: Tuple[float, float] = (0.9, 0.99),
    weight_decay: float = 0.0,
    use_triton: bool = False,
    decoupled_weight_decay: bool = False
)
```

### Parameters
- **params**: Iterable of `torch.nn.Parameter` or `Dict[str, Any]`
- **lr** (float): Learning rate. Defaults to `1e-4`.
- **betas** (Tuple[float, float]): Coefficients used for computing running averages of gradient and its squared first moment. Defaults to `(0.9, 0.99)`.
- **weight_decay** (float): Weight decay (L2 penalty). Defaults to `0.0`.
- **use_triton** (bool): Whether to use Triton acceleration. Defaults to `False`.
- **decoupled_weight_decay** (bool): Whether to use decoupled weight decay. Defaults to `False`.

### Methods
- **step**(closure: Callable | None = None) -> float | None
  Single optimization step.

### Inherited Methods
- `zero_grad()`
- `state_dict()`
- `load_state_dict()`
- `add_param_group()`

### Inherited Attributes
- `param_groups` (List[Dict]): Parameter groups with LR and other hyperparams.
- `state` (Dict): Per-parameter optimizer state (momentum).
- `defaults` (Dict): Default hyperparameter values.
```

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

### Foreach Lion Optimizer Initialization

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/modules.md

Initialize the foreach variant of the Lion optimizer for performance optimization. This variant uses `torch._foreach_*` operations for batch parameter updates, suitable for large models with PyTorch 1.13+.

```python
Lion(params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0, decoupled_weight_decay=False)
```

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

### Lion Optimizer with Decoupled Weight Decay

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/lion.md

Shows how to configure the Lion optimizer with decoupled weight decay. This ensures the weight decay strength remains consistent even if the learning rate is changed.

```python
# With decoupled weight decay, the effective weight decay strength
# remains consistent regardless of learning rate changes
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    decoupled_weight_decay=True
)
```

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

### Cautious Lion Optimizer Default Configuration

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/configuration.md

Initializes the Cautious Lion optimizer with default values, including a cautious factor of 0.0 for standard Lion behavior.

```python
Lion(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.99),
    weight_decay=0.0,
    cautious_factor=0.0,  # Standard behavior by default
    decoupled_weight_decay=False
)
```

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

### Lion Optimizer Constructor (Cautious Variant)

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

A cautious variant of the Lion optimizer with conservative scaling. Includes an additional `cautious_factor` parameter.

```python
Lion(
    params,
    lr: float = 1e-4,
    betas: Tuple[float, float] = (0.9, 0.99),
    weight_decay: float = 0.0,
    cautious_factor: float = 0.0,
    decoupled_weight_decay: bool = False
)
```

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

### Training Loop with Mixed Precision

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Utilize mixed precision training with `torch.cuda.amp` for potential speedups and reduced memory usage. Enable Triton for GPU acceleration.

```python
from torch.cuda.amp import autocast, GradScaler

optimizer = Lion(model.parameters(), lr=1e-4, weight_decay=1e-2, use_triton=True)
scaler = GradScaler()

for epoch in range(epochs):
    for batch in dataloader:
        with autocast():
            output = model(batch)
            loss = criterion(output, targets)
        
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
```

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

### Triton Kernel Grid Configuration

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/triton-kernel.md

Dynamically computes the grid size for launching the Triton kernel. The grid size is determined by the total number of elements and the auto-tuned block size.

```python
import triton

n_elements = p.numel() # Assuming p, grad, exp_avg have the same number of elements
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
```

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

### Training Loop with Gradient Accumulation

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Implement gradient accumulation to simulate larger batch sizes. Gradients are accumulated over several steps before an optimizer step.

```python
optimizer = Lion(model.parameters(), lr=1e-4, weight_decay=1e-2)
accumulation_steps = 4

for epoch in range(epochs):
    for i, batch in enumerate(dataloader):
        output = model(batch)
        loss = criterion(output, targets) / accumulation_steps
        
        loss.backward()  # Accumulate gradients
        
        if (i + 1) % accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()
```

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

### Lion Optimizer with Parameter Groups

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/index.md

Configure per-layer learning rates and weight decay for the Lion optimizer. Supported overrides include lr, weight_decay, betas, and cautious_factor.

```python
optimizer = Lion([
    {'params': model.early_layers.parameters(), 'lr': 1e-4},
    {'params': model.late_layers.parameters(), 'lr': 5e-5, 'weight_decay': 1e-1},
], lr=1e-4, weight_decay=1e-2)
```

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

### Foreach Lion Constructor

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Signature for the Foreach Lion optimizer variant. This variant is typically used for specific performance optimizations.

```python
from lion_pytorch.foreach import Lion

Lion(
    params,
    lr=1e-4,
    betas=(0.9, 0.99),
    weight_decay=0.0,
    decoupled_weight_decay=False
)
```

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

### Import Cautious Lion Variant

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/modules.md

Import the cautious variant for training unstable or problematic architectures. It includes a `cautious_factor` parameter for fine-tuning.

```python
# Cautious variant
from lion_pytorch.cautious_lion import Lion
optimizer = Lion(model.parameters(), lr=1e-4, cautious_factor=0.5)
```

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

### Cautious Lion Constructor

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Signature for the Cautious Lion optimizer variant. Adjust `cautious_factor` for more conservative updates.

```python
from lion_pytorch.cautious_lion import Lion

Lion(
    params,
    lr=1e-4,
    betas=(0.9, 0.99),
    weight_decay=0.0,
    cautious_factor=0.0,             # 0=standard, 0.5=conservative, 1.0=maximum caution
    decoupled_weight_decay=False
)
```

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

### Enable GPU Acceleration with Triton

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/index.md

Configure the Lion optimizer to use the Triton CUDA kernel for GPU acceleration.

```python
from lion_pytorch import Lion

optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    use_triton=True  # Enable CUDA kernel (requires Triton library)
)
```

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

### Parameter and Momentum Update for Cautious Lion

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/cautious-lion.md

These formulas show how the model parameters and momentum buffers are updated after the cautious scaling is applied. The learning rate is applied to the scaled update.

```mathematica
p ← p - lr * update
m ← β₂ * m + (1 - β₂) * g
```

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

### Debug Triton Usage and Optimizer Configuration

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/quick-reference.md

Check if the Triton kernel is active and inspect the optimizer's learning rate, weight decay, and beta settings.

```python
# Check if using Triton
print("Using Triton:", hasattr(optimizer, 'update_fn') and 
      'triton' in str(optimizer.update_fn))

# Check optimizer configuration
for group in optimizer.param_groups:
    print(f"LR: {group['lr']}")
    print(f"Weight decay: {group['weight_decay']}")
    print(f"Betas: {group['betas']}")
    print(f"Num params: {len(group['params'])}")
```

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

### Lion Optimizer with Parameter Groups

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/configuration.md

Configure the Lion optimizer with different learning rates and weight decay for specific parameter groups. This allows for fine-grained control over optimization for different parts of the model.

```python
optimizer = Lion([
    {'params': model.layer1.parameters(), 'lr': 1e-4},
    {'params': model.layer2.parameters(), 'lr': 1e-5, 'weight_decay': 1e-2},
], lr=1e-4, weight_decay=1e-2)
```

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

### Configure Lion Optimizer for GPU-Optimized Training

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/configuration.md

Achieve maximum performance on modern GPUs by enabling Triton CUDA acceleration. Ensure all tensors are on the GPU.

```python
optimizer = Lion(
    model.parameters(),
    lr=1e-4,
    weight_decay=1e-2,
    use_triton=True,  # CUDA acceleration
)
# Ensure all tensors on GPU
model = model.cuda()
batch = batch.cuda()
```

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

### Cautious Lion Constructor Parameters

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/exports.md

Additional parameter specific to the cautious variant of the Lion optimizer.

```python
# Cautious variant only
cautious_factor: float
```

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

### step Method

Source: https://github.com/lucidrains/lion-pytorch/blob/main/_autodocs/api-reference/cautious-lion.md

Performs a single optimization step. This method applies the cautious update mechanism, adaptively scaling parameter updates based on their alignment with the gradient direction. It can also accept a closure for reevaluating the model.

```APIDOC
### Method: step

```python
@torch.no_grad()
def step(self, closure: Callable | None = None) -> float | None
```

Performs a single optimization step with optional cautious update scaling.

The cautious variant implements Algorithm 2 from "Cautious Optimizers" paper:

1. Weight decay: `p ← p * (1 - lr * wd)`
2. Base update: `update = sign(β₁ * exp_avg + (1 - β₁) * grad)`
3. Cautious scaling (if `cautious_factor < 1.0`):
   - `align_mask = (update * grad) > 0` — identifies aligned elements
   - `scale = where(align_mask, 1.0, cautious_factor)` — full step for aligned, reduced for misaligned
   - `scale /= mean(scale).clamp(min=1e-5)` — normalize to prevent magnitude drift
   - `update *= scale`
4. Parameter update: `p ← p - lr * update`
5. Momentum decay: `exp_avg ← β₂ * exp_avg + (1 - β₂) * grad`

#### Parameters

| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| closure | Callable or None | No | None | A closure that reevaluates the model and returns the loss |

#### Returns

`float | None` — Loss value from closure if provided, otherwise None

#### Behavior Notes

- When `cautious_factor = 0.0`, behaves identically to standard Lion
- When `cautious_factor = 1.0`, applies full cautious scaling to all elements
- Alignment check `(update * grad) > 0` identifies where update direction matches gradient direction
- Scale normalization prevents training instability from accumulated scaling factors
- Uses `@torch.no_grad()` decorator for efficiency
```