# utils.schedulers – Axolotl

# utils.schedulers

`utils.schedulers`

Module for custom LRScheduler class

## Classes

Name | Description
--- | ---
InterpolatingLogScheduler | A scheduler that interpolates learning rates in a logarithmic fashion
JaggedLRRestartScheduler | Wraps another scheduler to apply per-lora-restart learning rate warmups.
RexLR | Reflected Exponential (REX) learning rate scheduler.

### InterpolatingLogScheduler

    utils.schedulers.InterpolatingLogScheduler(
        optimizer,
        num_steps,
        min_lr,
        max_lr,
        last_epoch=-1,
    )

A scheduler that interpolates learning rates in a logarithmic fashion

### JaggedLRRestartScheduler

    utils.schedulers.JaggedLRRestartScheduler(
        optimizer,
        inner_schedule,
        jagged_restart_steps,
        jagged_restart_warmup_steps,
        jagged_restart_anneal_steps=1,
        min_lr_scale=0.001,
    )

Wraps another scheduler to apply per-lora-restart learning rate warmups.

### RexLR

    utils.schedulers.RexLR(
        optimizer,
        max_lr,
        min_lr,
        total_steps=0,
        num_warmup_steps=0,
        last_step=0,
    )

Reflected Exponential (REX) learning rate scheduler.

- Original implementation: https://github.com/IvanVassi/REX_LR
- Original license: Apache 2.0
- Based on: https://arxiv.org/abs/2107.04197

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
optimizer | torch.optim.Optimizer | The optimizer to schedule the learning rate for. | required
max_lr | float | The maximum learning rate. | required
min_lr | float | The minimum learning rate. | required
total_steps | int | The total number of training steps. | `0`
num_warmup_steps | int | The number of warmup steps. | `0`
last_step | int | The index of last step. | `0`

## Functions

Name | Description
--- | ---
get_cosine_schedule_with_min_lr |
get_cosine_schedule_with_quadratic_warmup | Create a schedule with a learning rate that decreases following the values of the cosine function between the
get_cosine_schedule_with_warmup_decay_constant | Implementation of Continual Pre-Training of Large Language Models: How to (re)warm your model? (https://arxiv.org/pdf/2308.04014.pdf)

### get_cosine_schedule_with_min_lr

    utils.schedulers.get_cosine_schedule_with_min_lr(
        optimizer,
        num_warmup_steps,
        num_training_steps,
        min_lr_ratio=0.0,
    )

#### Create a learning rate schedule which has

- linear warmup from 0 -> `max_lr` over `num_warmup_steps`
- cosine learning rate annealing from `max_lr` -> `min_lr` over `num_training_steps`

### get_cosine_schedule_with_quadratic_warmup

    utils.schedulers.get_cosine_schedule_with_quadratic_warmup(
        optimizer,
        num_warmup_steps,
        num_training_steps,
        num_cycles=0.5,
        last_epoch=-1,
    )

Create a schedule with a learning rate that decreases following the values of the cosine function between the initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the initial lr set in the optimizer.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
optimizer | [`~torch.optim.Optimizer`] | The optimizer for which to schedule the learning rate. | required
num_warmup_steps | `int` | The number of steps for the warmup phase. | required
num_training_steps | `int` | The total number of training steps. | required
num_cycles | `float`, optional, defaults to 0.5 | The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 following a half-cosine). | `0.5`
last_epoch | `int`, optional, defaults to -1 | The index of the last epoch when resuming training. | `-1`

#### Return

`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.

### get_cosine_schedule_with_warmup_decay_constant

    utils.schedulers.get_cosine_schedule_with_warmup_decay_constant(
        optimizer,
        num_warmup_steps,
        num_training_steps,
        constant_lr_ratio,
        min_lr_ratio,
        num_cycles=0.5,
        last_epoch=-1,
    )

Implementation of Continual Pre-Training of Large Language Models: How to (re)warm your model? (https://arxiv.org/pdf/2308.04014.pdf) Create a schedule with a learning rate that decreases following the values of the cosine function between the initial lr set in the optimizer to min_lr_ratio until num_training_steps * constant_lr_ratio, after constant_rate returns constant value of min_rate , after a warmup period during which it increases linearly between 0 and the initial lr set in the optimizer.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
optimizer | [`~torch.optim.Optimizer`] | The optimizer for which to schedule the learning rate. | required
num_warmup_steps | `int` | The number of steps for the warmup phase. | required
num_training_steps | `int` | The total number of training steps. | required
constant_lr_ratio | float | (`float`): The ratio of num_training_steps to decrease by cosine function. | required

min_lr_ratio

float

(`float): The ratio of maximum learning rate for cosine function to decay to minimum learning rate. | _required_ | | num_cycles |`float`, *optional*, defaults to 0.5 | The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 following a half-cosine). |`0.5`| | last_epoch |`int`, *optional*, defaults to -1 | The index of the last epoch when resuming training. |`-1`

#### Return

`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.