Skip to content

llm.schedulers

/home/runner/work/llm-pytorch/llm-pytorch/llm/schedulers.py

Custom learning rate schedules.

LinearWarmupLR 

LinearWarmupLR(
    optimizer: torch.optim.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 
21
22
23
24
25
26
27
28
29
30
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 
32
33
34
35
36
37
38
39
40
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]