Learning Rate Warmup › Code Implementation

import math


def linear_warmup_cosine_decay(
    step: int,
    warmup_steps: int,
    total_steps: int,
    peak_lr: float,
    min_lr: float = 0.0,
) -> float:
    """
    Compute learning rate at a given step using linear warmup + cosine decay.

    Args:
        step: Current training step (1-indexed)
        warmup_steps: Number of warmup steps
        total_steps: Total number of training steps
        peak_lr: Peak learning rate after warmup
        min_lr: Minimum learning rate at end of training

    Returns:
        Learning rate at this step
    """
    if step <= warmup_steps:
        # Linear warmup phase
        return peak_lr * (step / warmup_steps)
    else:
        # Cosine decay phase
        decay_steps = total_steps - warmup_steps
        completed = step - warmup_steps
        cosine_decay = 0.5 * (1 + math.cos(math.pi * completed / decay_steps))
        return min_lr + (peak_lr - min_lr) * cosine_decay

import torch
import torch.nn as nn
from torch.optim import AdamW


# A minimal model for demonstration
class TinyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(64, 64)

    def forward(self, x):
        return self.linear(x)


model = TinyModel()

# Standard optimizer configuration for transformer training
optimizer = AdamW(
    model.parameters(),
    lr=3e-4,  # This is peak_lr; the scheduler will scale it
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01,
)


# Manually implement the warmup scheduler using LambdaLR
def warmup_cosine_lambda(step):
    warmup = 2000
    total = 100_000
    min_ratio = 0.1  # min_lr / peak_lr
    if step < warmup:
        return step / warmup
    else:
        decay_steps = total - warmup
        completed = step - warmup
        cosine = 0.5 * (1 + math.cos(math.pi * completed / decay_steps))
        return min_ratio + (1.0 - min_ratio) * cosine


scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer, lr_lambda=warmup_cosine_lambda
)

# Hugging Face transformers provides these schedules out of the box.
# The snippet below shows how you would use it in practice.
# We use eval=false here because transformers may not be installed in all environments.
# The API is shown for reference.

# from transformers import get_cosine_schedule_with_warmup
#
# scheduler_hf = get_cosine_schedule_with_warmup(
#     optimizer=optimizer,
#     num_warmup_steps=2000,
#     num_training_steps=100_000
# )
#
# # In your training loop:
# for batch in dataloader:
#     optimizer.zero_grad()
#     loss = model(batch)
#     loss.backward()
#     optimizer.step()
#     scheduler_hf.step()

def scaled_lr_and_warmup(
    base_lr: float,
    base_batch_size: int,
    base_warmup_steps: int,
    new_batch_size: int,
) -> tuple:
    """
    Apply linear scaling rule for batch size changes.

    Returns:
        (new_lr, new_warmup_steps) tuple
    """
    k = new_batch_size / base_batch_size
    new_lr = base_lr * k
    new_warmup = int(base_warmup_steps * k)
    return new_lr, new_warmup


# Baseline configuration
base_lr_scale = 1e-4
base_batch_scale = 256
base_warmup_scale = 1000

# Scaling scenarios
batch_sizes_scale = [256, 512, 1024, 2048, 4096, 8192]