Cosine Learning Rate Schedule › Code Implementation
Language AI HandbookAll chapters Cosine Learning Rate Schedule › Code ImplementationOne section · Sections in this chapter9 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We will build the cosine schedule from scratch to make the formula concrete, then show how to use PyTorch's built-in CosineAnnealingLR and CosineAnnealingWarmRestarts schedulers.
From-Scratch Implementation
Let us start by implementing the cosine decay function directly, without any framework.
import math
import numpy as np
def cosine_schedule(eta_max, eta_min, t, T):
"""Compute the cosine-decayed learning rate at step t out of T total steps."""
return eta_min + 0.5 * (eta_max - eta_min) * (1 + math.cos(math.pi * t / T))
def linear_schedule(eta_max, eta_min, t, T):
"""Compute the linearly-decayed learning rate at step t out of T total steps."""
return eta_max - (eta_max - eta_min) * t / T
# Schedule parameters
eta_max = 3e-4
eta_min = 3e-5
T = 1000
steps = np.arange(0, T + 1)
cosine_lrs = np.array([cosine_schedule(eta_max, eta_min, t, T) for t in steps])
linear_lrs = np.array([linear_schedule(eta_max, eta_min, t, T) for t in steps])Schedule parameters: eta_max = 3.0e-04 eta_min = 3.0e-05 T (total steps) = 1000 Learning rates at key steps: Step 0 (start): cosine = 3.00e-04, linear = 3.00e-04 Step 250 (25%): cosine = 2.60e-04, linear = 2.32e-04 Step 500 (50%): cosine = 1.65e-04, linear = 1.65e-04 Step 750 (75%): cosine = 6.95e-05, linear = 9.75e-05 Step 1000 (end): cosine = 3.00e-05, linear = 3.00e-05
Both schedules agree at the start and end, but they diverge significantly in between. At 25% through training, cosine decays more slowly than linear, preserving a higher learning rate for early exploration. At 75% through training, cosine is already below the linear trajectory, giving finer updates for the convergence phase.
Cosine with Warm Restarts
Next, let us implement the warm restart schedule from the SGDR paper.
def cosine_with_restarts(eta_max, eta_min, t, T0, T_mult=1):
"""
Compute the cosine-annealed learning rate with warm restarts at step t.
Args:
eta_max: Peak learning rate at the start of each cycle.
eta_min: Floor learning rate at the end of each cycle.
t: Current step index (global, across all cycles).
T0: Length of the first cycle in steps.
T_mult: Multiplicative factor for cycle lengths (1 = equal, 2 = doubling).
"""
if T_mult == 1:
# All cycles have length T0
t_cycle = t % T0
T_current = T0
else:
# Find which cycle we are in
T_current = T0
t_remaining = t
while t_remaining >= T_current:
t_remaining -= T_current
T_current = int(T_current * T_mult)
t_cycle = t_remaining
return eta_min + 0.5 * (eta_max - eta_min) * (
1 + math.cos(math.pi * t_cycle / T_current)
)
# Compute two restart schedules
T_total = 1000
T0 = 250 # First cycle length
restart_equal = np.array(
[cosine_with_restarts(eta_max, eta_min, t, T0, T_mult=1) for t in steps]
)
restart_double = np.array(
[cosine_with_restarts(eta_max, eta_min, t, T0, T_mult=2) for t in steps]
)Warm restart schedule with T0=250 steps: Equal cycles (T_mult=1): 4 restarts of 250 steps each Restart steps (T_mult=1): [250, 500, 750, 1000] Cycle lengths (T_mult=2): [250, 500] Restart steps (T_mult=2): [250]
The equal-length restart schedule creates four identical cycles of 250 steps, with the learning rate returning to at steps 250, 500, and 750. The doubling schedule creates progressively longer cycles, giving later cycles more time to refine the model.
PyTorch Scheduler Integration
In practice, you will use PyTorch's built-in schedulers rather than reimplementing from scratch. Let us demonstrate how to configure both options.
import torch.nn as nn
import torch.optim as optim
# Create a minimal model and optimizer for demonstration
model = nn.Linear(10, 1)
optimizer = optim.AdamW(model.parameters(), lr=eta_max)
# CosineAnnealingLR: single cycle, no restarts
scheduler_cosine = optim.lr_scheduler.CosineAnnealingLR(
optimizer,
T_max=T, # Number of steps in the half-period (= total decay steps)
eta_min=eta_min, # Floor learning rate
)
# Reset optimizer for second scheduler
optimizer2 = optim.AdamW(model.parameters(), lr=eta_max)
# CosineAnnealingWarmRestarts: multiple cycles with optional doubling
scheduler_restarts = optim.lr_scheduler.CosineAnnealingWarmRestarts(
optimizer2,
T_0=T0, # Steps in the first cycle
T_mult=1, # Equal-length cycles
eta_min=eta_min, # Floor learning rate
)
# Collect learning rates for visualization
pytorch_cosine_lrs = []
pytorch_restart_lrs = []
for step in range(T + 1):
pytorch_cosine_lrs.append(
scheduler_cosine.get_last_lr()[0]
if step > 0
else optimizer.param_groups[0]["lr"]
)
pytorch_restart_lrs.append(
scheduler_restarts.get_last_lr()[0]
if step > 0
else optimizer2.param_groups[0]["lr"]
)
if step < T:
scheduler_cosine.step()
scheduler_restarts.step()
pytorch_cosine_lrs = np.array(pytorch_cosine_lrs)
pytorch_restart_lrs = np.array(pytorch_restart_lrs)PyTorch CosineAnnealingLR: Starting LR: 3.00e-04 LR at step 500: 1.65e-04 Final LR: 3.00e-05 PyTorch CosineAnnealingWarmRestarts: Starting LR: 3.00e-04 LR at step 250 (first restart): 3.00e-04 LR at step 500 (second restart): 3.00e-04 Final LR: 3.00e-04
The PyTorch CosineAnnealingLR takes T_max as the half-period of the cosine, not the full decay duration. When you want a single decay from to over steps, pass T_max=T. When you want restarts, use CosineAnnealingWarmRestarts with T_0 matching your desired cycle length.
One important note about CosineAnnealingLR: PyTorch's implementation uses T_max as a half-period, meaning that after steps the learning rate reaches , but if you continue calling scheduler.step() beyond that point, the learning rate will start increasing again toward . In practice, you want to stop calling the scheduler after the final step, or set the training loop to end exactly at steps.
Combining Warmup with Cosine Decay
The most common configuration in large-scale training combines a linear warmup phase with a subsequent cosine decay. Let us implement this combined schedule.
def warmup_cosine_schedule(eta_max, eta_min, t, T_warmup, T_total):
"""
Compute the learning rate at step t for a warmup-then-cosine schedule.
Args:
eta_max: Peak learning rate reached at end of warmup.
eta_min: Floor learning rate at end of training.
t: Current step index.
T_warmup: Number of warmup steps.
T_total: Total training steps.
"""
if t < T_warmup:
# Linear warmup from 0 to eta_max
return eta_max * t / T_warmup
else:
# Cosine decay from eta_max to eta_min
t_decay = t - T_warmup
T_decay = T_total - T_warmup
return eta_min + 0.5 * (eta_max - eta_min) * (
1 + math.cos(math.pi * t_decay / T_decay)
)
# Configuration matching typical LLM training (proportions)
T_warmup_steps = 100 # 10% of training for warmup
T_total_steps = 1000
warmup_cosine_lrs = np.array(
[
warmup_cosine_schedule(
eta_max, eta_min, t, T_warmup_steps, T_total_steps
)
for t in range(T_total_steps + 1)
]
)Warmup + Cosine schedule (100 warmup steps, 1000 total): Step 0 (cold start): 0.00e+00 Step 50 (mid warmup): 1.50e-04 Step 100 (peak, warmup end): 3.00e-04 Step 550 (midpoint decay): 1.65e-04 Step 1000 (end): 3.00e-05
This combined schedule is what you will see in virtually every modern LLM training recipe. The warmup period prevents the unstable early-training loss spikes discussed in the Learning Rate Warmup chapter, and the cosine decay efficiently brings the learning rate to its minimum over the bulk of training.
Key Parameters
The key parameters for cosine learning rate schedules are:
- eta_max (peak learning rate): The starting learning rate after warmup. Typical values range from
1e-4to3e-4for large language models using AdamW. - eta_min (minimum learning rate): The floor learning rate at the end of training. A common choice is
eta_max / 10or a fixed small value like1e-6. - T_max / T_total (decay steps): Total steps for the cosine decay. Usually set to the full training duration minus warmup steps, or slightly longer to avoid the zero-LR degenerate case.
- T_warmup (warmup steps): Steps for linear warmup. Commonly 1-5% of total training steps.
- T_0 (cycle length for restarts): Number of steps in the first restart cycle. Only relevant when using warm restarts.
- T_mult (cycle multiplier): Multiplicative factor for restart cycle lengths. Common values are 1 (equal cycles) and 2 (doubling cycles).
Visualizing the Effect of Minimum Learning Rate
One underappreciated aspect of the cosine schedule is how dramatically the minimum learning rate changes the character of the decay. Let us compare several values to see the differences clearly.
The visual differences are most pronounced in the final third of training. With , the learning rate decays all the way to zero, which can cause the optimizer to stall. With (30% of peak), the decay is much more modest, and the model retains a useful update signal even at the end. The LLaMA choice of 10% sits in a practical middle ground: the floor is small enough that the parameter updates approach convergence, but large enough to prevent complete optimizer shutdown.
Using torch.optim.lr_scheduler.LambdaLR for Custom Schedules
If you need a warmup-plus-cosine schedule in PyTorch and want full control over the exact formula, LambdaLR lets you define the multiplier function directly. This approach is also compatible with Hugging Face's transformers library, which uses LambdaLR internally.
import torch.nn as nn
import torch.optim as optim
def get_cosine_schedule_with_warmup(
optimizer, num_warmup_steps, num_training_steps, eta_min_fraction=0.1
):
"""
Create a cosine schedule with warmup using LambdaLR.
Args:
optimizer: The optimizer to schedule.
num_warmup_steps: Steps for linear warmup.
num_training_steps: Total training steps.
eta_min_fraction: Minimum LR as fraction of peak LR (default: 10%).
Returns:
LambdaLR scheduler with cosine-plus-warmup multiplier.
"""
def lr_lambda(current_step):
if current_step < num_warmup_steps:
return float(current_step) / float(max(1, num_warmup_steps))
progress = float(current_step - num_warmup_steps) / float(
max(1, num_training_steps - num_warmup_steps)
)
cosine_factor = 0.5 * (1.0 + math.cos(math.pi * progress))
# Scale to [eta_min_fraction, 1.0]
return eta_min_fraction + (1.0 - eta_min_fraction) * cosine_factor
return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
# Test the custom schedule
model_demo = nn.Linear(10, 1)
opt_demo = optim.AdamW(model_demo.parameters(), lr=3e-4)
sched_demo = get_cosine_schedule_with_warmup(
opt_demo, num_warmup_steps=100, num_training_steps=1000
)
lambda_lrs = []
for step in range(1001):
lambda_lrs.append(opt_demo.param_groups[0]["lr"])
if step < 1000:
sched_demo.step()
lambda_lrs = np.array(lambda_lrs)Custom LambdaLR warmup-cosine schedule: Step 0 (start): 0.00e+00 Step 50 (mid warmup): 1.50e-04 Step 100 (peak): 3.00e-04 Step 550 (midpoint decay): 1.65e-04 Step 1000 (end): 3.00e-05 Ratio eta_end / eta_peak: 0.10
The LambdaLR approach gives you a single, auditable function for the entire schedule. The multiplier function returns values between eta_min_fraction and 1.0, and these multiply against the optimizer's base learning rate. If the base rate is 3e-4 and eta_min_fraction=0.1, the floor will be 3e-5, consistent with the LLaMA configuration.