Skip to content

Schedules, Warmup & Optimizer Tricks

Overview

The optimizer decides where you go; the schedule decides how fast. Most production training uses a boring-but-robust recipe: linear warmup + cosine decay, with a few extras (gradient accumulation, parameter groups, and the occasional optimizer twist) for special cases.

  • Warmup (3-10% of steps): protects against early instability (Adam variance estimates, large initial updates).
  • Cosine decay (or linear decay): strong late-training gains vs constant LR.
  • Parameter groups: different LR per component (embeddings, heads, LoRA, frozen layers).
  • Optimizer choice: AdamW is default; SGD+momentum for some vision; Adamax/RAdam for noisy small-data cases; fused variants for speed.

The single most common LR mistake: no warmup + constant LR → training diverges or plateaus at 50% of potential.

-

The Standard Recipe

import torch, torch.nn as nn, math

model = nn.Linear(8, 4)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)

total_steps = 10_000
warmup_steps = 500 # 5% warmup

def lr_lambda(step):
 if step < warmup_steps:
 return step / warmup_steps # linear warmup
 # cosine decay from warmup end to 0
 progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
 return 0.5 * (1 + math.cos(math.pi * progress))

sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)

for step in range(total_steps):
 loss = model(torch.randn(4, 8)).abs().mean()
 opt.zero_grad(); loss.backward(); opt.step()
 sched.step()

Scheduler Inventory

Scheduler Formula Use
LinearLR / warmup lr * (step/warmup) start of training
CosineAnnealingLR cosine to 0 main schedule
OneCycleLR warmup→high→0 short runs, aggressive
StepLR/MultiStepLR step drops legacy, fine
ReduceLROnPlateau on metric plateau when you have a metric 
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=1e-3,
 total_steps=total_steps,
 pct_start=0.1) # 10% warmup

-

The Easy Correct Way: linear_warmup helper

def warmup_cosine(optimizer, warmup, total):
 def f(step):
 if step < warmup:
 return step / warmup
 p = (step - warmup) / max(1, total - warmup)
 return 0.5 * (1 + math.cos(math.pi * p))
 return torch.optim.lr_scheduler.LambdaLR(optimizer, f)

-

Parameter Groups— Different LRs for Different Parts

model = nn.Sequential(nn.Linear(8, 8), nn.Linear(8, 4))

groups = [
 {"params": model[0].parameters(), "lr": 1e-3}, # backbone
 {"params": model[1].parameters(), "lr": 1e-2}, # head, faster
]
opt = torch.optim.AdamW(groups)

Common uses:

  • High LR for new heads (transfer learning), low for pretrained backbones.
  • LoRA: lr=1e-4 on adapters, 0 (or tiny) on base params.
  • No weight decay on embeddings/bias: {"params": w, "weight_decay": 0.0}.
no_decay = ["bias", "LayerNorm.weight", "embedding"]
decay, no_decay_p = [], []
for name, p in model.named_parameters():
 (no_decay_p if any(k in name for k in no_decay) else decay).append(p)
opt = torch.optim.AdamW([
 {"params": decay, "weight_decay": 0.01},
 {"params": no_decay_p, "weight_decay": 0.0},
])

Optimizer Selection

Optimizer Strength Watch out
AdamW default; robust; decoupled WD memory 2x params
Adam legacy variant wd entangled (avoid)
SGD+momentum vision classic, sharp minima needs schedule care
Adamax large sparse grads less common
Adafactor / 8-bit Adam memory-saving slower convergence, tune
Lion (3rd-party) token efficiency tuning differs
NAdam stable small LR niche

AdamW + warmup/cosine + grad clip = the "boring stack" that beats fancy tricks on most LLM/vision runs.

Common Failure Modes

Symptom Likely fix
Divergence in first 100 steps add/raise warmup; lower peak LR
Plateaus mid-training cosine decay instead of constant
Loss explodes at step 0 with AdamW check WD on embedding/bias; clamp
LR too high after warmup peak_lr = 10-30x effective

Key Takeaways

  • Warmup + cosine decay is the default robust schedule; warmup 3-10% of steps.
  • Parameter groups unlock transfer learning, LoRA, and WD-free embeddings.
  • AdamW (with fused=True on CUDA) is the default; SGD for vision purists.
  • Schedules compose with AMP, clipping, accumulation, and distributed runs.
  • Track LR and loss in logs— schedule bugs are silent until eval.

-