Warmup — Module 6.2

How This Lesson Fits the Module

Starting training at full learning rate can destabilize large models and transformers—weights and Adam statistics are cold. Warmup linearly ramps LR from near zero to the target over the first few hundred or thousand steps, then hands off to the main scheduler.

Warmup is standard in BERT-style pretraining and pairs naturally with cosine decay and mixed precision.

Prior Lesson Learning Rate Scheduler decays LR over time. Warmup is the opening act—ramp up before decay begins.

Learning Objectives

By the end of this lesson, students should be able to:

Explain why cold-start large LRs cause loss spikes in deep networks.
Implement linear warmup over warmup_steps or warmup_epochs.
Use LambdaLR or OneCycleLR for built-in warmup.
Choose warmup length relative to total training steps.
Step per-batch schedulers inside the training loop (not after epoch).
Combine warmup + cosine decay in one schedule.

Why Warmup Matters

Random initialization plus large gradients in early batches can explode activations or swamp Adam’s bias correction. A small LR lets layers settle; then full LR accelerates learning.

Setting	Typical warmup	Notes
Transformer pretraining	1–10% of total steps	Often linear warmup + cosine
ResNet from scratch	0–5 epochs	Sometimes optional with careful LR
Fine-tuning pretrained CNN	Short or none	Lower base LR already
Large batch training	Longer warmup	Linear scaling rule pairs with warmup

Linear Warmup with LambdaLR

LambdaLR multiplies base LR by a function of epoch. For step-based warmup, switch to a per-batch scheduler or manual LR update.

warmup_epochs = 5
total_epochs = 100

def lr_lambda(epoch):
    if epoch < warmup_epochs:
        return (epoch + 1) / warmup_epochs  # linear 0 → 1
    # cosine decay for remaining epochs (simplified)
    progress = (epoch - warmup_epochs) / (total_epochs - warmup_epochs)
    return 0.5 * (1 + math.cos(math.pi * progress))

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

for epoch in range(total_epochs):
    train_one_epoch(...)
    scheduler.step()

OneCycleLR with Warmup Built In

OneCycleLR ramps LR up during pct_start fraction of training, then decays. Step once per batch inside the inner loop.

steps_per_epoch = len(train_loader)
scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr=1e-3,
    epochs=num_epochs,
    steps_per_epoch=steps_per_epoch,
    pct_start=0.1,  # 10% warmup
)

model.train()
for inputs, targets in train_loader:
    optimizer.zero_grad()
    loss = criterion(model(inputs.to(device)), targets.to(device))
    loss.backward()
    optimizer.step()
    scheduler.step()  # per batch, not per epoch

Critical Mistake — Epoch Scheduler for Batch Warmup

Calling OneCycleLR.step() once per epoch destroys the schedule—LR stays wrong for the entire run. Match step frequency to how the scheduler was designed.

Warmup Length Guidelines

Too short: early instability remains. Too long: wastes compute in a low-LR regime. Start with 5–10% of total steps for transformers; ablate if loss spikes in the first 500 steps.

Engineering Habit — Plot LR Schedule Before Training

Script a dry run: loop scheduler.step() without data and plot LR vs step. Catches off-by-one warmup bugs in minutes.

Knowledge Check

Short Answer: What does warmup do to LR at step 0? Answer: Starts near zero (or very small), ramps to target.
True/False: Warmup is only for transformers. Answer: False—useful for any unstable early training.
Multiple Choice: OneCycleLR.step() is called: (a) per batch, (b) per epoch, (c) once total. Answer: (a).
Short Answer: Typical warmup fraction for BERT-style training? Answer: Often 1–10% of total steps.
Short Answer: Why large batch sizes need longer warmup? Answer: Larger effective step size; more cautious start prevents divergence.
True/False: LambdaLR can encode warmup + decay. Answer: True.
Multiple Choice: pct_start=0.1 in OneCycleLR means: (a) 10% warmup, (b) 10% weight decay, (c) 10 epochs. Answer: (a).
Short Answer: What symptom suggests insufficient warmup? Answer: Loss NaN or spike in first hundreds of steps.
Short Answer: After warmup, what usually follows? Answer: Constant LR or decay (cosine, step, etc.).
Multiple Choice: Fine-tuning with lr=1e-5 often needs: (a) long warmup, (b) little/no warmup, (c) warmup 50% of steps. Answer: (b).

Key Takeaways

Warmup ramps LR gradually to avoid early-training instability.
Use LambdaLR for custom epoch schedules; OneCycleLR for batch-level warmup+decay.
Step frequency must match scheduler design (epoch vs batch).
Next: Gradient Clipping—cap gradient magnitude.

Trainer’s Guide

Hands-on idea: Plot LR curves for 0%, 5%, and 15% warmup on the same architecture; compare step-100 loss.

Discussion prompt: How does warmup interact with Adam’s bias correction in early steps?

What’s Next Stabilize updates further with Gradient Clipping after backward.