llm.schedulers

llm/schedulers.py

Custom learning rate schedules.

LinearWarmupLR

LinearWarmupLR(
    optimizer: Optimizer,
    total_steps: int,
    warmup_steps: int = 0,
    last_epoch: int = -1,
)

Bases: _LRScheduler

Linear warmup and decay LR scheduler.

Source: ColossalAI

Parameters:

• optimizer (Optimizer) –

  Optimizer to adjust learning rate of.

• total_steps (int) –

  Total training steps.

• warmup_steps (int, default: 0 ) –

  Steps to linearly warmup the learning rate.

• last_epoch (int, default: -1 ) –

  Optional last epoch.

Source code in llm/schedulers.py

def __init__(
    self,
    optimizer: torch.optim.Optimizer,
    total_steps: int,
    warmup_steps: int = 0,
    last_epoch: int = -1,
) -> None:
    self.total_steps = total_steps
    self.warmup_steps = warmup_steps
    super().__init__(optimizer, last_epoch)

get_lr

get_lr() -> list[float]

Compute the current learning rate.

Source code in llm/schedulers.py

def get_lr(self) -> list[float]:
    """Compute the current learning rate."""
    if self.last_epoch < self.warmup_steps:
        factor = (self.last_epoch + 1) / (self.warmup_steps + 1)
        return [base_lr * factor for base_lr in self.base_lrs]
    else:
        factor = self.total_steps - self.last_epoch
        factor = factor / (self.total_steps - self.warmup_steps)
        return [base_lr * factor for base_lr in self.base_lrs]