lion-pytorch

When switching from AdamW to Lion, follow these tuning guidelines based on the original paper:

• Learning Rate (lr): A suitable learning rate for Lion is typically 3-10x smaller than that for AdamW. You should change the initial, peak, and end values in your learning rate schedule simultaneously using the same ratio.
• Weight Decay ($\lambda$): Because the effective weight decay is lr * λ, the decoupled weight decay value used for Lion should be 3-10x larger than that used for AdamW to maintain similar strength.
• $\beta_1$ and $\beta_2$: The default values are $\beta_1=0.9$ and $\beta_2=0.99$. If you encounter instability, try using $\beta_1=0.95, \beta_2=0.98$ (and potentially increasing $\epsilon$ to 1e-6).
• Batch Size: It is recommended to use Lion in settings with high batch sizes (64 or above).

You can install the lion-pytorch package using either pip or conda.

$ pip install lion-pytorch

Alternatively, using conda:

$ conda install lion-pytorch

To use a fused kernel for updating parameters via Triton, you must first install the Triton package:

$ pip install triton -U --pre

Then, set use_triton=True when instantiating the Lion optimizer.

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)
)

The Lion implementation includes advanced update mechanisms based on recent research:

1. Cautious Update: Controlled by cautious_factor. When 0 < cautious_factor < 1, the update is scaled based on whether the update direction aligns with the gradient direction. This follows Algorithm 2 in arXiv:2411.16085.
2. Cautious Weight Decay: Controlled by cautious_wd. When enabled, weight decay is applied selectively based on the sign of the update and the parameter (masking), following arXiv:2510.12402.
3. Decoupled Weight Decay: Controlled by decoupled_weight_decay. If set to True, the weight decay is scaled by the initial learning rate to decouple it from the current learning rate.

To use the Lion optimizer, import Lion from lion_pytorch and instantiate it with your model parameters. The usage pattern follows the standard PyTorch optimizer API: step() for updates and zero_grad() to clear gradients.

import torch
from torch import nn
from lion_pytorch import Lion

model = nn.Linear(10, 1)

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

# Standard training loop steps
loss = model(torch.randn(10))
loss.backward()

opt.step()
opt.zero_grad()

# toy model

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()

Use the .step() method to update model parameters based on the computed gradients. Like standard PyTorch optimizers, it supports an optional closure for algorithms that require multiple evaluations of the loss function.

# Standard usage
optimizer.zero_grad()
loss = criterion(output, target)
loss.backward()
optimizer.step()

To use the Lion optimizer, import the Lion class from the lion_pytorch package.

from lion_pytorch import Lion

The Lion optimizer accepts the following arguments in its __init__ method:

• params: An iterable of the parameters to optimize.
• lr (float): The learning rate. Must be greater than 0.
• betas (Tuple[float, float]): A tuple of two beta coefficients for the exponential moving average of gradients. Must be between 0 and 1. Default is (0.9, 0.99).
• weight_decay (float): The weight decay coefficient. Default is 0.0.
• decoupled_weight_decay (bool): If set to True, weight decay is applied in a decoupled manner (similar to AdamW). Default is False.

The Lion class is a PyTorch optimizer implementation. You can initialize it with your model parameters and several hyperparameters to control the update behavior, including learning rate, momentum (betas), weight decay, and specific 'cautious' update modes.

Key arguments:

• params: The parameters to optimize.
• lr: Learning rate (default 1e-4).
• betas: A tuple of (beta1, beta2) for the exponential moving averages (default (0.9, 0.99)).
• weight_decay: Weight decay coefficient (default 0.0).
• cautious_factor: A factor between 0 and 1 used for cautious updates (default 0.0).
• cautious_wd: Boolean indicating whether to use cautious weight decay (default False).
• decoupled_weight_decay: Boolean indicating whether to use decoupled weight decay (default False).

The update_fn provides a Triton-accelerated implementation of the Lion optimizer update step. It is designed to perform the weight decay, momentum update, and sign-based weight update directly on CUDA tensors using a custom Triton kernel.

Requirements:

• triton>=2.2.0 must be installed.
• All input tensors (p, grad, exp_avg) must be on a CUDA device.
• Tensors must not be complex (Triton support for complex tensors is not available).

Parameters:

• p (torch.Tensor): The parameter tensor to be updated.
• grad (torch.Tensor): The gradient tensor.
• exp_avg (torch.Tensor): The momentum running average (exponential moving average) tensor.
• lr (float): Learning rate.
• wd (float): Weight decay coefficient.
• beta1 (float): First momentum coefficient.
• beta2 (float): Second momentum coefficient.