Phase 03 — Autograd & Training Dynamics

The phase where the model stops being a fixed function and starts to learn. Phase 02 built the forward pass — attention, RoPE, a GPT that turns tokens into logits. This phase builds the machinery that makes those weights correct: reverse-mode automatic differentiation (the thing behind loss.backward()), the optimizers that turn gradients into updates, and the training-loop dynamics — learning-rate schedules, gradient clipping, numerical precision — that decide whether a run converges, plateaus, or blows up at 2 a.m. on step 4000.

Why this phase exists

Ask a candidate "what does loss.backward() do?" and you separate the people who use PyTorch from the people who understand it. The honest answer is small and beautiful: the framework records every operation as a node in a computation graph, each node knows its own local derivative, and one reverse pass over the graph applies the chain rule to produce every gradient at once. That is the entire idea of backpropagation, and it is the engine under every model you will ever train. Everything else in training — Adam, weight decay, warmup, clipping, mixed precision — is a refinement on top of "compute the gradient, take a step."

You cannot reason about a training run you do not understand. When the loss spikes, when it won't go down, when fp16 produces NaN, when Adam diverges but SGD doesn't — the engineer who built the autograd tape and the optimizer by hand knows where to look. The five things to own here:

1. Reverse-mode autodiff — why one backward pass yields all gradients, and how the topological order + grad accumulation make it work (and why forward-mode is the opposite tradeoff).
2. The softmax + cross-entropy gradient — the loss that trains every classifier and LM, and why it collapses to the clean p − y.
3. Optimizers: SGD → momentum → Adam → AdamW — the moment estimates, bias correction, and why decoupled weight decay (the "W") was a real fix, not a rebrand.
4. Training dynamics — warmup + cosine schedules, gradient clipping, loss spikes: the knobs that decide convergence.
5. Numerical precision — fp32 / fp16 / bf16, mixed precision, loss scaling, gradient accumulation: the difference between a run that fits and one that OOMs or NaNs.

Concept map

                         ┌──────────────────────────────────────────┐
                         │  Forward pass (Phase 02): tokens → logits │
                         └──────────────────────────────────────────┘
                                            │  loss = cross_entropy(logits, target)
                                            ▼
                    ┌───────────────────────────────────────────────┐
                    │     COMPUTATION GRAPH  (every op = a node)      │
                    │   each node knows its LOCAL derivative          │
                    └───────────────────────────────────────────────┘
                                            │  .backward()
                          reverse topological order + chain rule
                                            ▼
                    ┌───────────────────────────────────────────────┐
                    │   GRADIENTS for every parameter (one pass)      │
                    └───────────────────────────────────────────────┘
                          │              │                │
                  clip_grad_norm   lr_schedule       OPTIMIZER
                  (tame spikes)    (warmup+cosine)   SGD→Adam→AdamW
                          └──────────────┴────────────────┘
                                            ▼
                              θ ← θ − lr · update(grad)   ← one step
                                            │
                            repeat → loss strictly decreases
                                            │
                 (precision: fp32 / fp16+loss-scaling / bf16 ; grad accumulation)

The lab

The lab is a runnable, test-verified miniature — see the lab standard. Run it red (pytest test_lab.py), make it green, then read solution.py.

Anatomy of one training step

Every training loop in existence — from this lab to a 70B run on a thousand GPUs — is the same six moves. Burn this into muscle memory; the rest of the phase is just how each move works:

1. zero_grad        clear last step's gradients (.grad accumulates by design)
2. forward          inputs → logits → loss        (builds the computation graph)
3. backward         loss.backward()               (reverse-mode: one pass, all grads)
4. clip             clip_grad_norm(params, τ)      (cap step length, keep direction)
5. step             optimizer.step()               (SGD / AdamW turns grad → update)
6. schedule         lr_scheduler.step()            (warmup + cosine: how big the step)

Miss step 1 and gradients pile up (the classic silent bug). Get step 3's order wrong and you read half-finished gradients. Skip step 4 on a bad batch and the loss spikes. The lab implements all six.

Integrated scenario ideas

• Explain loss.backward() on a whiteboard: draw the graph for loss = CE(W·x, y), label each node's local derivative, and walk the reverse pass that fills W.grad. This is the autograd interview question.
• Debug a loss spike: a run is fine until step 4000, then loss explodes. Enumerate causes (LR too high after warmup, a bad batch, fp16 overflow, no grad clipping) and the fix order.
• Defend AdamW over Adam: show on a weight-decayed objective why coupling L2 into the gradient (Adam) interacts badly with the adaptive denominator, and why decoupling (AdamW) fixes it.
• Size a precision decision: when do you reach for bf16 vs fp16+loss-scaling, and what does gradient accumulation buy you when the batch won't fit?

Deliverables checklist

• lab.py passes pytest test_lab.py -v (and LAB_MODULE=solution does too).
• You can derive the gradient of z = x*y + x by hand and trace it through .backward().
• You can explain why reverse-mode gives all gradients in one pass and when forward-mode wins.
• You can derive the softmax + cross-entropy gradient and state why it equals p − y.
• You can explain SGD vs Adam vs AdamW, the moment estimates, bias correction, and decoupled weight decay — without notes.
• You can explain warmup, cosine decay, gradient clipping, and the precision/loss-scaling story.
• You can tie every piece of the Value engine to its torch equivalent.

Key takeaways

• One backward pass, all gradients. Reverse-mode autodiff is the whole trick: record the graph, give each op a local derivative, walk it backward applying the chain rule. Gradients accumulate, which is exactly why you zero_grad() every step.
• softmax + cross-entropy is p − y. The most common loss in the field has the cleanest gradient in the field — probabilities minus the one-hot target. Know the derivation cold.
• AdamW ≠ Adam + L2. Decoupled weight decay is a genuine correction; it's the default optimizer for transformers for a reason. Momentum, adaptive per-parameter steps, bias correction — own each.
• Training dynamics are engineering, not luck. Warmup avoids early divergence, cosine anneals to a minimum, clipping survives the giant gradient, and precision/loss-scaling/accumulation decide whether the run fits the hardware at all.
• This is PyTorch. torch.Tensor.requires_grad + loss.backward() is the Value engine you built, vectorized onto tensors and a GPU. The algorithm is identical; only the data type changed.

Next: Phase 04 — Vision-Language Models & Multimodal Fusion.