Warmup Guide — Training Small LLMs

Zero-to-expert primer for Phase 05: everything between "I have a transformer" and "it trained well" — optimization (AdamW, schedules), stability, mixed precision, scaling laws, and the experimental discipline of the nanoGPT lab.

Table of Contents

Chapter 1: The Training Loop, Honestly Stated

The loop is five lines; every line hides a chapter:

for batch in data:                  # Ch. 4: what's a batch, really
    logits = model(x)               # Ch. 5: in what precision
    loss = cross_entropy(logits, y) # Phase 03 Ch. 2: the objective
    loss.backward()                 # autograd (model-accuracy Phase 01)
    clip_grad_norm_(params, 1.0)    # Ch. 6: stability
    optimizer.step(); scheduler.step(); optimizer.zero_grad()   # Ch. 2–3

The phase's real subject is the gap between this loop running and this loop working: the silent failure modes (wrong LR, bad init, precision bugs) that produce a loss curve that goes down — just to a worse place than it should. The defense is the calibration mindset: know what loss value, what curve shape, and what sample quality to expect at each stage, so deviation is informative (Chapter 8).

Chapter 2: Adam and AdamW — Why They Won

SGD's problem for transformers: one global learning rate across parameters whose gradient scales differ by orders of magnitude (embedding rows for rare tokens vs LayerNorm gains). Adam maintains per-parameter statistics — first moment (EMA of gradients, $m$) and second moment (EMA of squared gradients, $v$) — and updates:

$$\theta \mathrel{-}= \eta \cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon}$$

The $\sqrt{\hat{v}}$ denominator is an automatic per-parameter LR: rarely-updated parameters (small $v$) take big steps; noisy ones get damped. Bias correction ($\hat{m} = m/(1-\beta_1^t)$) fixes the zero-initialized EMAs' early-step underestimate — without it the first steps are silently tiny.

AdamW's fix to Adam: in Adam, L2 "weight decay" added to the gradient gets divided by $\sqrt{\hat{v}}$ too — parameters with large gradient history get less regularization, which is backwards. AdamW applies decay decoupled, directly to the weights ($\theta \mathrel{-}= \eta \lambda \theta$), restoring decay's meaning. All modern LLMs: AdamW, $\beta_1{=}0.9$, $\beta_2{=}0.95$ (lower than the 0.999 default — faster-moving curvature estimates suit LLM gradient noise), decay ~0.1 applied to weight matrices but not to biases, LayerNorm parameters, or (usually) embeddings.

The cost nobody mentions until it hurts: two FP32 states per parameter — optimizer memory = 2× model size in FP32, i.e., more than the model itself. This number drives everything in distributed training (Phase 10: ZeRO exists for this) and fine-tuning (Phase 06: LoRA exists substantially for this).

Chapter 3: Learning-Rate Schedules — Warmup and Cosine

The standard LLM schedule is linear warmup → cosine decay:

  • Warmup (a few hundred–few thousand steps, ramping 0 → peak): Adam's $v$ estimate is garbage for the first steps (built from a handful of noisy gradients), and transformer training is most fragile at init — a full-size step at step 1 can spike the loss into a divergence it never recovers from. Warmup lets the optimizer's statistics and the loss landscape's early descent stabilize first.
  • Cosine decay to ~10% of peak: large steps explore early; small steps settle into a minimum late. Empirically robust; the exact shape (cosine vs linear) matters less than (a) peak LR (the single most important hyperparameter — too high diverges or plateaus noisily; too low underfits the budget) and (b) decaying to a low floor by the end of the token budget.
  • The lab's required experiment: an LR range sweep (3e-3 / 6e-4 / 1e-4 on the same config) — the canonical "too hot / right / too cold" triptych of curves is something you should have personally produced once, because you'll be diagnosing it from others' curves forever.

Chapter 4: Batch Size, Gradient Accumulation, and Tokens-Per-Step

  • The meaningful unit is tokens per optimizer step = micro_batch × seq_len × accumulation × (data-parallel ranks). nanoGPT-class runs use ~0.5M tokens/step; report and reason in this unit, not "batch size."
  • Gradient accumulation: run N micro-batch forward/backwards, summing grads, then one optimizer step — mathematically identical to the big batch (within FP non-associativity), trading wall-clock for memory. The classic bug: forgetting to scale the loss by 1/N (or zero_grad placement), silently multiplying the effective LR.
  • Critical batch size intuition (McCandlish et al.): below it, bigger batches are nearly free speedups (gradient noise dominates); above it, diminishing returns — you're averaging already-clean gradients. Small models on small data have small critical batch sizes — another reason the lab's modest config is correct, not just convenient.
  • Batch size and LR interact (bigger batch → cleaner gradient → supports higher LR); change them together or not at all (Chapter 8's one-variable rule has this one sanctioned exception, with the linear-scaling heuristic as the starting point).

Chapter 5: Mixed Precision — BF16 and the Loss-Scale Question

Train in 16-bit where it's safe, 32-bit where it isn't:

  • What stays FP32: master weights (inside AdamW), optimizer states, and the big reductions (loss, norms). What runs 16-bit: matmuls, activations — the bulk of compute and memory.
  • FP16 vs BF16 (the exponent/mantissa trade from the model-accuracy track's Phase 03 Ch. 1): FP16's 5-bit exponent overflows at 65504 and underflows small gradients to zero — hence loss scaling (multiply loss by ~2¹⁵ so gradients shift into representable range, unscale before the step, skip steps on inf/nan). BF16 has FP32's 8-bit exponent — no overflow/underflow drama, no loss scaling, at the cost of ~3 decimal digits of mantissa, which training tolerates. On any hardware that has it (A100+), BF16 is the answer and the lab uses torch.autocast(dtype=bfloat16).
  • The debugging signature to memorize: FP16 run with occasional inf → loss-scaler skips (normal if rare); BF16 run with loss spikes → it's not precision overflow, look at data or LR (this elimination step saves real days).

Chapter 6: Initialization and Stability

Why the lab's init code looks the way it does:

  • Linear layers: normal(0, 0.02); embeddings likewise — small enough that pre-norm residual streams start near-identity.
  • The GPT-2 residual-projection trick: scale the output projections of attention and FFN by $1/\sqrt{2N}$ (N = layer count) — each block adds to the residual stream, and without the scaling the stream's variance grows linearly with depth; the scaling keeps the sum's variance constant. This single line is the difference between 12 layers training smoothly and not.
  • Stability kit, in escalation order: gradient clipping at norm 1.0 (always on; watch the pre-clip norm — its trend is an early-warning instrument), warmup (Ch. 3), then if spikes persist: check data (a pathological document), lower peak LR, raise $\epsilon$, suspect precision last (Ch. 5's elimination).
  • The overfit-one-batch test before any long run: a healthy model+loop drives one small batch to ~zero loss in a few hundred steps. Failure means a wiring bug (shifted targets, mask, LR) — never start a multi-hour run without this two-minute test.

Chapter 7: Scaling Laws — How Big, How Much Data

The empirical regularities that turn "how big a model?" from taste into arithmetic:

  • Kaplan et al. (2020): loss falls as a power law in parameters N, data D, and compute C, over many orders of magnitude — smooth, predictable returns.
  • Chinchilla (2022): for a fixed compute budget $C \approx 6ND$, the optimum is roughly D ≈ 20 N — tokens ≈ 20× parameters. GPT-3 (175B params, 300B tokens) was badly under-trained by this rule; Chinchilla (70B, 1.4T) beat it at the same compute.
  • The inference-cost asterisk (the practitioner's correction): Chinchilla optimizes training compute only. If a model will serve billions of requests, over-training a smaller model far past 20:1 (LLaMA-class models: 100–2000 tokens/param) buys lower inference cost forever at modest training premium — which is the actual logic behind every production "small" model you'll serve in Phase 09.
  • For the lab's scale: a ~10M-param model wants ~200M+ tokens by the 20:1 rule — Shakespeare (~1M) is hopelessly small, which is why the lab also runs OpenWebText- class data and why your Shakespeare model memorizes (watch val loss diverge from train — the overfitting lesson live).

Chapter 8: The Experimental Method

The discipline that separates training-as-science from GPU-warming, the same ledger ethics as the model-accuracy track's capstone (Phase 11 Ch. 3):

  1. One variable per run; config serialized with every checkpoint; seeds fixed (and the run-to-run σ at fixed seed measured once, so you know what differences are real).
  2. Watch validation, not train, loss — and watch samples: fixed prompts generated every N steps; the qualitative arc is a debugging instrument no scalar replaces.
  3. Baseline before improvement: the stock config trained to completion is the reference every change is measured against. "It felt better" without the baseline delta is how teams lose months.
  4. Log the boring numbers: tokens/sec (throughput regressions are bugs), grad norm, LR actually applied. When a run dies at 3am, the logs are the post-mortem.

Lab Walkthrough Guidance

Lab 02 — nanoGPT:

  1. Run the overfit-one-batch test on your Phase 04 model wired into this loop. Only then start real runs.
  2. Shakespeare char-level first (minutes/run): produce the LR triptych (Ch. 3), the batch-size/accumulation equivalence check (same tokens/step two ways → same curve), and the warmup ablation (watch the no-warmup run spike).
  3. Add mixed precision; verify the loss curve matches FP32 within noise and measure the throughput gain — both numbers go in your notes.
  4. Scale to the larger dataset/config; apply Chapter 7's arithmetic to predict where val loss should land relative to the small run before looking. Calibration is the skill; the prediction-vs-actual delta is the lesson.
  5. Keep the ledger (Ch. 8) from run #1 — the lab's deliverable is the table of runs, not the final checkpoint.

Success Criteria

You are ready for Phase 06 when you can, from memory:

  1. Write Adam's update and explain $\sqrt{\hat{v}}$, bias correction, and AdamW's decoupling — plus the 2×-FP32 memory fact and what it later justifies (ZeRO, LoRA).
  2. Justify warmup and cosine decay; sketch the too-hot/right/too-cold triptych.
  3. Compute tokens-per-step for any config and state the accumulation-scaling bug.
  4. Choose BF16 vs FP16 with the exponent argument; explain loss scaling and when it's unnecessary.
  5. Explain the $1/\sqrt{2N}$ residual init and the overfit-one-batch test's purpose.
  6. State Chinchilla's 20:1, its derivation context, and the inference-cost correction that explains over-trained small models.

Interview Q&A

Q: Your loss spikes at step 40K and recovers but lands at a worse plateau. Walk through it. First the data: spikes are most often a pathological batch (corrupt/duplicated document) — find what was sampled at that step (this is why ledgers log data order/ seed). Then optimizer state: a spike corrupts Adam's $v$ (huge squared grads), damping subsequent steps — recovery-to-worse is consistent with poisoned state; mitigations are tighter clipping, lower $\beta_2$, or restart from pre-spike checkpoint with the data fixed. Precision last: BF16 makes overflow unlikely; FP16 scaler logs would show it. The answer's structure — data, optimizer state, precision, in that order — is what's being graded.

Q: Why does everyone use AdamW for transformers when SGD trains ResNets fine? Transformer gradient scales are wildly heterogeneous across parameter types (embedding rows, LayerNorm gains, attention vs FFN matrices) and the loss landscape has sharper curvature anisotropy; per-parameter adaptive steps absorb this where one global LR can't (SGD on transformers needs delicate per-layer LR surgery to compete). Plus decoupled decay gives clean regularization semantics. Cost: 2× FP32 optimizer state — which is a real systems consequence, not a footnote.

Q: With a fixed training budget, would you train a 7B for 1 epoch or a 1.4B for 5× the tokens — and what changes your answer? Chinchilla says match D ≈ 20N for best training-compute loss — compute both options' N:D against that. But the deployment profile dominates: if the model serves at scale, the smaller over-trained model wins on lifetime cost (inference forever beats a small pretraining delta); if it's a research artifact or will be distilled anyway, the bigger model's better loss matters. Saying "Chinchilla, then correct for inference" is the complete answer.

References