A2 — Backpropagation and AdamW (CPU)

Files: train_gpt2.c
Needs: no GPU
Builds on: A0, A1
Check with: ./grade A2

Overview

You will implement the backward pass of GPT-2 and the AdamW optimizer. By the end of this assignment, ./train_gpt2 will complete a full training loop: forward pass → loss → backward pass → parameter update → repeat, and the loss will actually go down.

Backpropagation is the chain rule applied recursively. Each backward function receives a gradient with respect to its output (dout) and must compute gradients with respect to its inputs (dinp, dweight, dbias) so those can be passed further down the graph.

This is CPU-only and requires no GPU. The CUDA AdamW kernel in llmc/adamw.cuh is part of A3, not A2.

To list your work:

grep -n "TODO (A2)" train_gpt2.c

When these notes and the in-file HINT disagree, the code is authoritative — please report the discrepancy so it can be fixed.

One complication versus a textbook treatment: gradients are accumulated, not assigned. Every backward function does += into its output buffers. This is what makes gradient accumulation across micro-steps possible in A3/A4, and it means the gradient buffers must be explicitly zeroed before each new backward pass.

Where this fits 

flowchart LR
    A0["A0
tokens, batches, PRNG"] --> A1["A1
forward pass (CPU)"]
    A1 --> A2["A2
backward + AdamW (CPU)"]
    A2 --> A3["A3
port to CUDA (1 GPU)"]
    A3 --> A4["A4
scale to N GPUs"]
    A3 --> A5["A5
schedules, clipping, checkpoints"]
    style A2 fill:#2563eb,stroke:#1e40af,color:#fff

You are building the half that makes the network learn: the derivative of every operator in A1, and the AdamW step that turns those derivatives into a weight update. After this assignment ./train_gpt2 is a complete training system.

flowchart TB
    subgraph fwd["forward (A1) — save what backward needs"]
        direction LR
        f1["inp"] --> f2["ln1, ln1_mean, ln1_rstd"] --> f3["qkv"] --> f4["att"] --> f5["fch, fch_gelu"] --> f6["logits, probs"]
    end
    f6 --> loss["loss = -log probs[target]"]
    loss --> b1["dlogits = probs - onehot"]
    subgraph bwd["backward (A2) — consume it in reverse"]
        direction RL
        b2["dlnf"] --> b3["dresidual3"] --> b4["dfch_gelu"] --> b5["datt"] --> b6["dqkv"] --> b7["dencoded"]
    end
    b1 --> b2
    b7 --> upd["encoder_backward → dwte, dwpe"]
    upd --> adamw["AdamW
m, v, bias correction,
decoupled weight decay"]
    adamw --> params["updated weights"]

Why the forward pass saves mean, rstd, att and fch_gelu. Each backward function needs some of what its forward counterpart computed. Recomputing them would be correct but slow, so A1 stashed them in the activation tensors. This is the classic memory-for-compute trade, and it is why activations dominate training memory — A3 revisits it with a recompute flag that trades back the other way.

Why every backward function uses +=. A value used in two places gets gradient from both. wte is the clearest case: it feeds encoder_forward at the bottom of the network and the classifier matmul at the top, so dwte is the sum of two contributions computed at opposite ends of the backward pass. Writing = silently discards one of them. The same convention is what makes gradient accumulation over micro-batches work in A3.

Why AdamW and not SGD. Adam keeps a per-parameter running estimate of the gradient’s mean (m) and uncentred variance (v), and divides by the square root of the latter — so a parameter with consistently small gradients still moves. The “W” is the fix from Loshchilov & Hutter: apply weight decay directly to the parameter rather than adding it to the gradient, so it does not get rescaled by v. That one-line difference is worth a measurable amount of final loss.

Background Reading

Source Read Why it matters here
Karpathy, micrograd engine.py, ~100 lines A scalar autograd engine. If the += convention in Part 2 does not feel obvious, read this first — it is the same idea with the array indexing removed.
Rumelhart, Hinton & Williams 1986, Learning representations by back-propagating errors all 4 pages The original. Short, and worth reading once for the framing of backprop as repeated application of the chain rule over a graph.
Kingma & Ba 2015, Adam §2 and Algorithm 1 The update rule, including the bias-correction terms. Algorithm 1 maps line-for-line onto what you write in Part 4.
Loshchilov & Hutter 2019, Decoupled Weight Decay Regularization §2 and Algorithm 2 Why AdamW is not Adam-plus-L2. Algorithm 2 highlights the one-line difference in green; that line is the whole point.
Bishop, Pattern Recognition and Machine Learning §5.3 Optional. A careful derivation of backpropagation if you want the maths spelled out more slowly than Rumelhart does.

On deriving layernorm_backward yourself. It is the one gradient here that is genuinely fiddly, because mean and rstd both depend on every element of the row — so dinp[i] gets a contribution through inp[i] directly and through both statistics. Working it out on paper for C = 2 first is a good use of twenty minutes. Part 2 gives the two-reduction form the reference uses. —

Gradient Flow Map

The backward pass traverses the computation graph in reverse:

dlosses (= 1.0f / (B*T) per position)
  ↓ crossentropy_softmax_backward
dlogits [B, T, Vp]
  ↓ matmul_backward          (final classifier, weight-shared with wte)
dlnf [B, T, C],  dwte += ...
  ↓ layernorm_backward        (final layernorm)
dresidual3[L-1]
  ↓ (for each layer l = L-1 .. 0)
      matmul_backward          (fcproj: 4C→C)
      gelu_backward
      matmul_backward          (fc: C→4C)
      layernorm_backward       (ln2)
      residual_backward
      matmul_backward          (attproj: C→C)
      attention_backward
      matmul_backward          (qkv: C→3C)
      layernorm_backward       (ln1)
      residual_backward
  ↓
dencoded
  ↓ encoder_backward
dwte +=, dwpe +=

Part 1 — Gradient Infrastructure

gpt2_zero_grad(model)

Set every element of model->grads_memory and model->grads_acts_memory to zero, if they have been allocated. Use memset. This runs at the start of each optimizer step (in a gradient-accumulation setup, not between micro-steps).

Part 2 — Backward Functions

encoder_backward(dwte, dwpe, dout, inp, B, T, C)

dwte[inp[b*T+t], c] += dout[b, t, c]   // += because multiple positions may map to the same token
dwpe[t, c]          += dout[b, t, c]

layernorm_backward(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C)

Let x̂[c] = (inp[c] - mean) * rstd (the normalized value before the affine transform).

dbias[c]   += dout[b,t,c]
dweight[c] += dout[b,t,c] * x̂[c]

For dinp, the cleanest formulation uses two reductions over C. With dnorm[c] = dout[c] * weight[c]:

dnorm_mean      = (1/C) Σ_c dnorm[c]
dnorm_norm_mean = (1/C) Σ_c dnorm[c] * x̂[c]

dinp[c] += rstd * (dnorm[c] - dnorm_mean - x̂[c] * dnorm_norm_mean)

This is algebraically the same as the mean/variance form you may derive by hand, but needs only two passes and is much less error-prone. A common bug is keeping only the first term.

matmul_backward(dinp, dweight, dbias, dout, inp, weight, B, T, C, OC)

dinp[i, c]    += Σ_o dout[i, o] * weight[o, c]
dweight[o, c] += Σ_i dout[i, o] * inp[i, c]
dbias[o]      += Σ_i dout[i, o]        // skip when dbias == NULL

dbias is NULL for the final classifier. Parallelise the two loops separately — the dinp loop parallelises over b,t, the dweight loop over o — so that no two threads accumulate into the same location.

attention_backward(dinp, dpreatt, datt, dout, inp, att, B, T, C, NH)

Work backwards through the four steps of the forward pass. Only t2 <= t1 participates; the masked region contributes nothing.

  1. dV[b,t2,h] += att[b,h,t1,t2] * dout[b,t1,h] and datt[b,h,t1,t2] += Σ_c dout[b,t1,h,c] * V[b,t2,h,c]
  2. Softmax backward: dpreatt[b,h,t1,t2] += att[b,h,t1,t2] * (datt[b,h,t1,t2] - Σ_t att[b,h,t1,t] * datt[b,h,t1,t])
  3. dQ[b,t1,h] += dpreatt[b,h,t1,t2] * K[b,t2,h] / sqrt(HS) and dK[b,t2,h] += dpreatt[b,h,t1,t2] * Q[b,t1,h] / sqrt(HS)

All three write back into the single dinp[B, T, 3C] buffer at the Q/K/V offsets described in A1 Part 5.

gelu_backward(dinp, inp, dout, N)

let x    = inp[i]
    arg  = GELU_SCALING_FACTOR * (x + 0.044715 * x³)
    t    = tanh(arg)
    sech² = 1 / cosh(arg)²          (equivalently 1 - t²)

GELU'(x) = 0.5 * (1 + t) + 0.5 * x * sech² * GELU_SCALING_FACTOR * (1 + 3 * 0.044715 * x²)
dinp[i] += GELU'(x) * dout[i]

Note: gelu_backward uses +=, matching the accumulation convention of every other backward function here (and the fused GPU path in A3, where the residual gradient is already present in dinp).

residual_backward(dinp1, dinp2, dout, N)

dinp1[i] += dout[i]
dinp2[i] += dout[i]

crossentropy_softmax_backward(dlogits, dlosses, probs, targets, B, T, V, Vp)

The fused cross-entropy + softmax backward. The gradient of loss = -log(probs[target]) through the softmax simplifies to:

dlogits[b, t, v] += dlosses[b*T+t] * (probs[b, t, v] - indicator(v == targets[b*T+t]))

Loop v only over [0, V) — the padded columns [V, Vp) must keep a zero gradient — but index with the Vp row stride.

Part 3 — Full Backward Pass

gpt2_backward(model)

  1. Lazily allocate grads_memory / grads_acts_memory on the first call. Error out if mean_loss is -1.0f (forward was run without targets).
  2. Initialise dlosses[b*T+t] = 1.0f / (B*T) — the gradient of the mean loss.
  3. crossentropy_softmax_backward → dlogits
  4. matmul_backward for the classifier: accumulates into dwte (weight-shared with the embedding) and dlnf
  5. layernorm_backward for the final LayerNorm
  6. Loop l = L-1 down to 0, mirroring the forward order in reverse. Naming the gradient buffers by the struct field they live in (grads. for parameters, grads_acts. for activations):
    • residual_backward → splits grads_acts.residual3 into grads_acts.residual2 and grads_acts.fcproj
    • matmul_backward → grads.fcprojw, grads.fcprojb, grads_acts.fch_gelu
    • gelu_backward → grads_acts.fch
    • matmul_backward → grads.fcw, grads.fcb, grads_acts.ln2
    • layernorm_backward → grads.ln2w, grads.ln2b, and back into grads_acts.residual2
    • residual_backward → splits into the layer’s residual input and grads_acts.attproj
    • matmul_backward → grads.attprojw, grads.attprojb, grads_acts.atty
    • attention_backward → grads_acts.qkv (via grads_acts.preatt and grads_acts.att)
    • matmul_backward → grads.qkvw, grads.qkvb, grads_acts.ln1
    • layernorm_backward → grads.ln1w, grads.ln1b, and back into the residual stream
  7. encoder_backward → grads.wte, grads.wpe

For layer l, the “residual stream” gradient is grads_acts.residual3 of layer l-1, or grads_acts.encoded when l == 0. Getting this indexing right is most of the work in gpt2_backward; the per-layer weight pointers follow the same l * <tensor stride> pattern as the forward pass.

Part 4 — AdamW Optimizer

gpt2_update(model, learning_rate, beta1, beta2, eps, weight_decay, t)

Lazily allocate and zero m_memory and v_memory on the first call. Then for each parameter θ with gradient g, where t is the 1-based step count:

m ← β₁·m + (1-β₁)·g
v ← β₂·v + (1-β₂)·g²
m̂ = m / (1 - β₁ᵗ)
v̂ = v / (1 - β₂ᵗ)
θ ← θ - lr · (m̂ / (sqrt(v̂) + ε) + weight_decay · θ)

Decoupled weight decay: apply weight decay directly to the parameter, not through the gradient. This is what distinguishes AdamW from L2-regularized Adam.

The CPU implementation applies a single weight_decay to every parameter, which is what test_gpt2 checks against. Selective decay — 2D tensors (embeddings and weight matrices) decayed, 1D tensors (biases, LayerNorm scales/shifts) not — is introduced with the CUDA optimizer in A3.

Testing 

make test_gpt2
./test_gpt2

test_gpt2 runs 10 forward+backward+update steps against gpt2_124M_debug_state.bin, a fixed batch exported from PyTorch. It is not a finite-difference check — there is no (loss(θ+ε) - loss(θ-ε)) / 2ε anywhere in the harness. It compares your analytic gradients directly against PyTorch’s analytic gradients for the same inputs, which is both faster and stricter.

Shape of the output once A1 and A2 are correct:

OK (LOGITS), max_diff = <small>
LOSS OK: <yours> <expected>
dwte
OK <yours> <expected>      <- first 5 elements of each tensor, printed for eyeballing
...
TENSOR OK, maxdiff = <small>
dwpe
...
TENSOR OK, maxdiff = <small>
... one block per parameter tensor, 16 in all ...
step 0: loss <...> (took <...> ms) OK = 1
...
step 9: loss <...> (took <...> ms) OK = 1
overall okay: 1

overall okay: 1 is the pass condition. Tolerances actually enforced:

Check Tolerance Where
Logits 1e-2 absolute main
Loss (step 0 and each of the 10 steps) 1e-2 absolute main
Every gradient tensor, element-wise 2e-2 absolute check_tensor

The 16 gradient tensors are checked in this order: dwte, dwpe, dln1w, dln1b, dqkvw, dqkvb, dattprojw, dattprojb, dln2w, dln2b, dfcw, dfcb, dfcprojw, dfcprojb, dlnfw, dlnfb. Because the backward pass runs from the output end, the last tensors in that list (dlnfw, dlnfb) are computed first — if only dwte/dwpe are wrong, the bug is in encoder_backward; if everything is wrong, suspect the final classifier or crossentropy_softmax_backward.

The per-step losses at the bottom exercise the optimizer: they must track PyTorch’s within 1e-2 for all 10 steps, which will not happen unless AdamW’s bias correction and step counting are right.

Common Mistakes

  • Forgetting the += convention. Every backward function accumulates. If you use = you will silently drop gradients from other paths — most visibly on dwte, which is written by both the classifier matmul and encoder_backward.
  • Weight sharing. wte is used both in encoder_forward and as the classifier head. Both backward passes must accumulate into dwte.
  • LayerNorm gradient. The formula needs both reductions over C. Keeping only rstd * dnorm[c] gives gradients that look plausible and are wrong.
  • Not zeroing gradients. test_gpt2 calls gpt2_zero_grad between steps; if it is a no-op, step 0 passes and steps 1-9 diverge.
  • Off-by-one on t. AdamW’s bias correction uses a 1-based step; test_gpt2 passes step+1.
  • V vs Vp in the classifier backward. Writing gradient into the padded columns corrupts dwte rows that should stay untouched.

Checking your work 

./grade A2

Runs ./test_gpt2 and maps its 16 gradient tensors onto the table below, so a failure points at the function responsible. Lines marked (inferred) cannot be isolated by the harness and are attributed from the tensors they feed.

Scoring

What ./grade A2 checks, and what each part is worth.

Component Points
gpt2_zero_grad and gradient buffer management 5
encoder_backward matches reference 10
layernorm_backward matches reference 15
matmul_backward matches reference 10
attention_backward matches reference 20
gelu_backward and residual_backward match reference 10
crossentropy_softmax_backward matches reference 10
gpt2_backward wires all backward functions correctly 10
AdamW update — all 10 step losses within 1e-2, overall okay: 1 10
Total 100