Hitchhiker's Guide — Autograd & Training Dynamics

The compressed practitioner tour. If WARMUP.md is the professor, this is the senior who leans over and says "here's what you actually need to remember about loss.backward() and the loop."

The 30-second mental model

A model's program is a computation graph; every op is a node that knows its local derivative. loss.backward() walks that graph in reverse topological order applying the chain rule, so one pass fills every gradient (reverse-mode wins because nets are many-params-in, one-loss-out). Gradients accumulate (+=) — that's why you zero_grad() every step. Then an optimizer turns grads into a step: SGD → momentum → Adam → AdamW (decoupled weight decay). Wrap it with **warmup

• cosine** LR, clip the global grad norm to survive spikes, and pick a precision (bf16 ≫ fp16 for training because of range). softmax+CE has the cleanest gradient in ML: p − y.

The numbers to tattoo on your arm

Back-of-envelope one-liners

# "Does my autograd work?" → gradcheck against central finite differences
analytic == (L(x+h) - L(x-h)) / (2h)   within ~1e-4   → correct

# "Why is one backward pass enough?" → reverse-mode, scalar loss
N params, 1 loss  → reverse-mode: O(1) passes for all N grads (forward-mode: O(N))

# "Loss spiked at step 4000" → the ladder
LR too high post-warmup? grad-norm spiking (clip!)? fp16 overflow→NaN (bf16/loss-scale)? bad batch?

# "Batch won't fit" → gradient accumulation
want batch 256, fit 32 → 8 micro-batches, accumulate grads, ONE step  (time cost, not memory)

# softmax+CE gradient (no derivation needed at the whiteboard)
d(CE)/d(logit_i) = softmax(logit)_i - onehot_i

The framework one-liners (where these live in real tools)

import torch, torch.nn.functional as F
x = torch.tensor(2.0, requires_grad=True)
y = (x*x + x); y.backward(); x.grad          # = 5.0  (your Value engine, vectorized)

loss = F.cross_entropy(logits, target)        # logits in! softmax+log-sum-exp fused → (p-y) grad
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)  # decoupled WD
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=total)    # or HF warmup+cosine

opt.zero_grad(set_to_none=True)               # grads accumulate — clear them every step
loss.backward()                               # reverse-mode autodiff
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)        # returns pre-clip norm
opt.step(); sched.step()

with torch.autocast("cuda", dtype=torch.bfloat16):  # mixed precision; bf16 ⇒ no loss scaling
    loss = model(batch)
torch.autograd.gradcheck(fn, inputs)          # the finite-difference referee

War stories

• The forgotten zero_grad. Loss looked like it was learning then chaotically diverged. Cause: no optimizer.zero_grad() — gradients accumulated across steps, so each update was the sum of all prior gradients. One line fixed a week of confusion. Grads accumulate by design.
• The 2 a.m. NaN. fp16 run, loss → NaN around step 1500. Small gradients had underflowed to zero and a big activation overflowed to inf. Fix: enable dynamic loss scaling (or just switch to bf16, whose fp32-range exponent dodged the whole problem).
• The AdamW-vs-Adam regression. Team "added weight decay" by setting Adam's weight_decay and saw worse generalization than expected. The L2-in-the-gradient was being rescaled by Adam's adaptive denominator. Switching to AdamW (decoupled) recovered the expected regularization.
• The loss spike that clipping ate. Long training, periodic loss spikes from rare hard batches. Logging the pre-clip gradient norm showed the spikes coming; turning on global-norm clipping at 1.0 flattened them without changing the final loss — direction preserved, length capped.
• The double-softmax. Someone fed softmax(logits) into F.cross_entropy (which softmaxes again). Training "worked" but underperformed for weeks. CE takes logits, not probabilities.

Vocabulary (rapid-fire)

• Computation graph — nodes = ops, edges = data dependencies; differentiated by walking backward.
• Reverse-mode autodiff / backprop — one backward pass → all gradients (the right mode for nets).
• Local derivative — an op's own ∂out/∂in; the chain rule composes them.
• Topological order — process a node only after every consumer; lets grads finish before use.
• Grad accumulation — += into .grad: required within a pass (shared nodes), exploited across micro-batches, and the reason zero_grad exists.
• (p − y) — the softmax+cross-entropy gradient w.r.t. logits.
• Moments (Adam) — 1st = mean of grads (momentum), 2nd = mean of grad² (per-param scale).
• Decoupled weight decay (AdamW) — shrink the weight directly, not via the gradient.
• Bias correction — 1/(1−β^t) fixup for zero-initialized moments.
• Warmup / cosine — ramp the LR up, then anneal it down smoothly.
• Global-norm clipping — cap ‖grad‖₂ to a threshold; keeps direction, bounds step length.
• fp16 / bf16 / loss scaling / AMP — 16-bit training; bf16 has the range, fp16 needs loss scaling.
• gradcheck — verify autograd against central finite differences.

Beginner mistakes

• Forgetting optimizer.zero_grad() (grads accumulate across steps).
• Writing softmax without max-subtraction (overflow on large logits).
• Feeding probabilities to cross_entropy instead of logits (double softmax).
• Overwriting (=) instead of accumulating (+=) gradients in a custom backward (shared nodes break).
• Treating AdamW as "Adam + L2" — the decoupling is the point.
• Skipping warmup on a transformer with Adam (early divergence).
• Using fp16 without loss scaling and chasing the resulting NaNs (reach for bf16).
• Not logging the gradient norm — then having no idea why the loss spiked.

The one thing to take away

Before you trust any training run, be able to (1) explain loss.backward() as "reverse-topo chain rule, one pass, all grads," (2) verify a gradient against finite differences, and (3) name your optimizer/schedule/clip/precision and why. If you can, you can debug any loss curve. If you can't, you're praying to a black box.