Lab 01 — Reverse-Mode Autograd Engine that Trains a Tiny LM

Phase: 03 — Autograd & Training Dynamics Difficulty: ⭐⭐⭐☆☆ (the engine is ~150 lines; the insight is ⭐⭐⭐⭐⭐) Time: 4–5 hours

Every deep-learning framework is, at its core, the thing you build here: a graph of scalar operations, each of which knows its local derivative, plus one reverse pass that applies the chain rule to fill in every gradient at once. This lab is micrograd with teeth — a Value autograd engine, a numerically-stable softmax + cross-entropy with the clean (p − y) gradient, a tiny MLP, SGD and AdamW, global-norm gradient clipping, a warmup+cosine LR schedule, and a deterministic training loop that actually learns XOR with a loss that strictly decreases. When you finish, torch.Tensor.backward() is no longer magic — it's this, vectorized onto a GPU.

What you build

class Value — a scalar node that records data, grad, its parents _prev, and a local _backward closure. Operators + * ** - / neg (and the r-variants) and methods exp / log / tanh / relu, each defining its own local derivative. .backward() builds the reverse topological order, seeds the output grad to 1.0, and runs every closure once.
softmax(values) — max-subtracted, stable for huge logits, inside the graph so gradients flow.
cross_entropy(logits, target_index) — −log(softmax(logits)[target]), built on the stable softmax so log never sees 0; its logit gradient is the famous p − onehot.
Neuron / Layer / MLP(nin, nouts, seed) — a deterministic tiny net with hidden tanh layers and a linear output layer (so the outputs are usable logits), plus .parameters().
sgd_step / adamw_step — the two optimizers, operating directly on Value.grad. AdamW carries 1st/2nd-moment state and applies decoupled weight decay.
clip_grad_norm(params, max_norm) — global-norm clipping; returns the pre-clip norm.
lr_schedule / train(...) — warmup+cosine LR, and a full-batch loop that learns XOR deterministically with a strictly-decreasing loss.

Key concepts

Concept	What to understand
Computation graph	every op creates a node remembering its inputs; the program is the graph
Reverse-mode autodiff	one backward pass yields all gradients — `O(1)` passes, not `O(params)`
The chain rule, locally	each node only needs its local derivative; composition does the rest
Topological order	process a node only after every node that consumes it; grads accumulate (`+=`)
Shared node ⇒ grad sums	a value used twice gets contributions from both paths (the diamond)
Stable softmax	subtract the max before `exp` — identical math, no overflow
`(p − y)` gradient	softmax+CE collapses to "probabilities minus one-hot" — clean and cheap
SGD → Adam → AdamW	momentum, then per-parameter adaptive steps, then decoupled weight decay
Gradient clipping	cap the global norm to survive the occasional giant gradient / loss spike

Files

File	Purpose
`lab.py`	skeleton with `# TODO` markers — your implementation (imports clean before you start)
`solution.py`	complete reference; `python solution.py` runs a worked example
`test_lab.py`	the proof — run it red, make it green
`requirements.txt`	`pytest` only (pure stdlib otherwise; no numpy, no torch)

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v   # against the reference (must be green)
python solution.py                          # the worked example

Success criteria

All tests pass against your implementation (and LAB_MODULE=solution).
You can derive dz/dx and dz/dy for z = x*y + x by hand and match .backward().
test_grad_matches_finite_difference_* passes — the soul test: your analytic gradients match a central finite-difference estimate. If this is green, your chain rule is correct.
You can explain why test_reused_node_accumulates_grad needs += (not =) in every _backward.
test_softmax_stable_for_large_logits passes — you subtracted the max.
test_cross_entropy_softmax_gradient_is_p_minus_onehot passes — you can state why the gradient is p − y without deriving it on the spot.
test_training_loss_strictly_decreases passes — the soul training test: a real net, a real loss, learning a real (non-linearly-separable) task, monotonically.
test_training_is_deterministic passes — same seed → byte-identical loss history (this is why _prev is an ordered tuple, not a set).

How this maps to the real stack

The miniature	The production mechanism	Where to verify it
`Value` + `_backward` closures	`torch.Tensor` with `requires_grad=True` and its `grad_fn` graph	`t = torch.tensor(2., requires_grad=True); (t*t).grad_fn`
`.backward()` topo + chain rule	`loss.backward()` — same reverse-mode pass, vectorized on tensors	PyTorch Autograd mechanics docs
grad accumulation (`+=`)	why you must `optimizer.zero_grad()` every step (grads accumulate)	`optimizer.zero_grad(set_to_none=True)`
`softmax` (max-subtracted)	`torch.softmax` / the fused `F.cross_entropy` (log-sum-exp under the hood)	`torch.nn.functional.cross_entropy`
`cross_entropy` → `(p − y)`	the gradient `F.cross_entropy` computes; logits-in, no manual softmax	the PyTorch CE docs ("input is expected to contain raw logits")
`sgd_step` / `adamw_step`	`torch.optim.SGD` / `torch.optim.AdamW` (decoupled WD)	`torch.optim.AdamW`; the AdamW paper
`clip_grad_norm`	`torch.nn.utils.clip_grad_norm_` — identical global-norm rule, returns pre-clip norm	the clip-grad-norm docs
`lr_schedule`	`get_cosine_schedule_with_warmup` (HF) / `torch.optim.lr_scheduler.CosineAnnealingLR`	HF `transformers` schedulers

Limits of the miniature (say these in the interview): it's scalar autograd — real engines operate on tensors so one node is a whole matmul, not one number, which is the entire reason GPUs help; we keep the whole graph alive every step (no detach/checkpointing, no graph freeing), so memory grows with graph size; there is no mixed precision / loss scaling (fp32 only here); and the train loop is full-batch with no data loader, no minibatch noise, and no validation split.

Extensions (build these on real hardware)

Add Tensor (a 2-D Value analog backed by array) so a "node" is a matmul and you can see why vectorization—not a different algorithm—is what GPUs buy you.
Add momentum SGD and plain Adam (coupled L2) next to AdamW and show the difference decoupling makes on a weight-decayed quadratic.
Add gradient accumulation: backward over k micro-batches before one optimizer step; show the loss curve matches a k×-larger batch.
Add activation checkpointing: free intermediate Values and recompute them in backward; show the time/memory tradeoff.
Port the exact same net to PyTorch and assert your gradients match loss.backward() to 1e-6.

Interview / resume

Talking points: "Walk me through loss.backward() — what is it actually doing?" "Why does one backward pass give every gradient (reverse-mode), and when would forward-mode win?" "Derive the softmax+cross-entropy gradient." "SGD vs Adam vs AdamW — what does decoupled weight decay fix?" "Your loss spiked at step 4000 — what knobs?" (clip, warmup, lower LR, check fp16 overflow).
Resume bullet: Implemented a reverse-mode automatic-differentiation engine from scratch (computation graph, topological backprop, the softmax+cross-entropy (p−y) gradient), with SGD and AdamW optimizers, global-norm gradient clipping, and a warmup+cosine schedule, and used it to train a tiny MLP to a strictly-decreasing, deterministic loss — i.e. rebuilt the core of PyTorch autograd.

Senior AI Engineer — LLMs, VLMs & Agentic Systems