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.