Bells & Whistles: Warmup, Cosine Decay & Gradient Clipping

The ch5 training loop works — but production LLM training adds three stabilizers on top: ramp the learning rate up gently, wind it down smoothly, and cap runaway gradients. This appendix adds all three to the train_model_simple skeleton.

D.1 Learning rate warmup

Warmup gradually increases the learning rate from a very low initial_lr to the user-specified maximum peak_lr, so training starts with small, non-destabilizing weight updates. Typical warmup length: 0.1%–20% of total steps.

</> Linear warmup — the schedule inside the loop

n_epochs = 15
initial_lr = 0.0001
peak_lr = 0.01

total_steps = len(train_loader) * n_epochs
warmup_steps = int(0.2 * total_steps)   # 20% warmup

lr_increment = (peak_lr - initial_lr) / warmup_steps  #A

global_step = -1
track_lrs = []
optimizer = torch.optim.AdamW(model.parameters(), weight_decay=0.1)

for epoch in range(n_epochs):
  for input_batch, target_batch in train_loader:
      optimizer.zero_grad()
      global_step += 1

      if global_step < warmup_steps:
          lr = initial_lr + global_step * lr_increment  #B
      else:
          lr = peak_lr

      for param_group in optimizer.param_groups:  #C
          param_group["lr"] = lr
      track_lrs.append(optimizer.param_groups[0]["lr"])

      # Calculate loss and update weights
      # ...

#A Divide the climb (peak − initial) evenly across the warmup steps.
#B A straight line from initial_lr to peak_lr, one increment per step; after warmup the rate holds at the peak.
#C The optimizer's lr lives in its param_groups — writing it there each step is how ANY schedule is applied in PyTorch.

D.2 Cosine decay

After warmup, cosine decay lowers the rate from the peak toward min_lr along a half-cosine — slowing learning as the model improves, reducing the risk of overshooting minima late in training:

$lr = lr_{min} + (lr_{peak} - lr_{min}) \cdot \tfrac{1}{2}\left(1 + \cos(\pi \cdot \text{progress})\right)$

symbol	meaning	in the demo
$\text{progress}$	fraction of the post-warmup training completed (0 → 1)	(step − warmup) / (total − warmup)
$lr_{peak}$	the maximum rate warmup climbed to	0.01
$lr_{min}$	the floor the rate decays toward	0.1 · initial_lr = 1e-5
$\tfrac{1}{2}(1{+}\cos \pi p)$	a smooth 1 → 0 ramp (the half-cosine)	—

Both schedules, plotted with the demo’s exact constants — hover to read the rate at any point of training:

warmup → constant peak (D.1)warmup → cosine decay (D.2)

hover the plot to read exact values

x = fraction of training completed. Both schedules share the 20% linear warmup from 1e-4 to 0.01; D.1 then holds the peak, while D.2's half-cosine glides down to min_lr = 1e-5. Cosine is preferred for its smoothness; linear decay also appears in practice (e.g. OLMo).

</> Warmup + cosine annealing

import math

min_lr = 0.1 * initial_lr

for epoch in range(n_epochs):
  for input_batch, target_batch in train_loader:
      optimizer.zero_grad()
      global_step += 1

      if global_step < warmup_steps:
          lr = initial_lr + global_step * lr_increment  # linear warmup
      else:
          progress = ((global_step - warmup_steps) /
                      (total_training_steps - warmup_steps))  #A
          lr = min_lr + (peak_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * progress))

      for param_group in optimizer.param_groups:
          param_group["lr"] = lr
      track_lrs.append(lr)
      # ...

#A progress runs 0 → 1 over the post-warmup steps; cos(π·progress) runs 1 → −1, so the bracket runs 1 → 0 — the smooth glide of the plot above.

D.3 Gradient clipping

Gradient clipping caps the size of updates: if the gradients’ L2 norm (Euclidean length) exceeds max_norm, the whole gradient is scaled down to that norm:

$\|\mathbf{v}\|_2 = \sqrt{v_1^2 + v_2^2 + \dots + v_n^2}$

Work the book’s example — a 2×2 gradient matrix clipped to max_norm 1:

The raw gradients

Suppose backward() produced the gradient matrix G = [[1, 2], [2, 4]] for some parameter. Individually modest numbers — but their combined length is what matters for the update size.

1 / 3

Norm clipping, by hand: measure the gradient's length, compare to the cap, rescale uniformly. Direction is preserved; only the magnitude shrinks.

</> Clipping in practice — before and after

loss = calc_loss_batch(input_batch, target_batch, model, device)
loss.backward()  #A

def find_highest_gradient(model):
  max_grad = None
  for param in model.parameters():
      if param.grad is not None:
          grad_values = param.grad.data.flatten()
          max_grad_param = grad_values.max()
          if max_grad is None or max_grad_param > max_grad:
              max_grad = max_grad_param
  return max_grad

print(find_highest_gradient(model))
# tensor(0.0411)

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  #B
print(find_highest_gradient(model))
# tensor(0.0185)

#A backward() fills every parameter's .grad — clipping must happen AFTER this and BEFORE optimizer.step() consumes the gradients.
#B One call rescales the global gradient in place; the largest single gradient drops from 0.0411 to 0.0185.

D.4 The modified training function

All three refinements, folded into ch5’s loop:

</> train_model — train_model_simple + warmup + cosine + clipping

from previous_chapters import evaluate_model, generate_and_print_sample

BOOK_VERSION = True

def train_model(model, train_loader, val_loader, optimizer, device,
              n_epochs, eval_freq, eval_iter, start_context, tokenizer,
              warmup_steps, initial_lr=3e-05, min_lr=1e-6):

  train_losses, val_losses, track_tokens_seen, track_lrs = [], [], [], []
  tokens_seen, global_step = 0, -1

  peak_lr = optimizer.param_groups[0]["lr"]  #A
  total_training_steps = len(train_loader) * n_epochs
  lr_increment = (peak_lr - initial_lr) / warmup_steps

  for epoch in range(n_epochs):
      model.train()
      for input_batch, target_batch in train_loader:
          optimizer.zero_grad()
          global_step += 1

          # Adjust the learning rate (warmup or cosine annealing)
          if global_step < warmup_steps:
              lr = initial_lr + global_step * lr_increment
          else:
              progress = ((global_step - warmup_steps) /
                          (total_training_steps - warmup_steps))
              lr = min_lr + (peak_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * progress))

          for param_group in optimizer.param_groups:
              param_group["lr"] = lr
          track_lrs.append(lr)

          loss = calc_loss_batch(input_batch, target_batch, model, device)
          loss.backward()

          # Gradient clipping after the warmup phase
          if BOOK_VERSION:
              if global_step > warmup_steps:  #B
                  torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
          else:
              if global_step >= warmup_steps:
                  torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

          optimizer.step()
          tokens_seen += input_batch.numel()

          if global_step % eval_freq == 0:
              train_loss, val_loss = evaluate_model(
                  model, train_loader, val_loader, device, eval_iter)
              train_losses.append(train_loss)
              val_losses.append(val_loss)
              track_tokens_seen.append(tokens_seen)
              print(f"Ep {epoch+1} (Iter {global_step:06d}): "
                    f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")

      generate_and_print_sample(model, tokenizer, device, start_context)

  return train_losses, val_losses, track_tokens_seen, track_lrs

# Usage — note the corrected settings:
peak_lr = 0.001  #C
optimizer = torch.optim.AdamW(model.parameters(), lr=peak_lr, weight_decay=0.1)

#A A tidy convention: peak_lr is READ from the optimizer (whatever lr it was constructed with) rather than passed separately.
#B Erratum, flagged in the notes: `>` skips clipping on the ONE step where global_step equals warmup_steps; the repo's corrected version uses `>=`. The BOOK_VERSION flag lets you reproduce either.
#C Second erratum: the book originally set 5e-4 by mistake AND omitted the lr= assignment when constructing AdamW — corrected here to peak_lr = 0.001, passed explicitly.

Key idea — schedule the rate, cap the step

Warmup protects the fragile start, cosine decay protects the delicate finish, and clipping insures every step in between against gradient spikes — three small additions that turn the teaching loop of ch5 into something shaped like production LLM training. One appendix remains: fine-tuning without touching most of the weights — LoRA.

📝 Check yourself: training-loop refinements

0 / 5

1.What problem does learning rate warmup address?
2.After warmup, cosine decay takes the LR from peak toward min_lr along a half-cosine. Why decay at all?
3.Predict: gradients G = [[1, 2], [2, 4]] are clipped with max_norm = 1. What happens?
4.Where exactly in the training step does gradient clipping happen?
5.The notes flag that the book's `if global_step > warmup_steps:` differs from the repo's `>=`. What's the practical consequence?