A1 — GPT-2 Forward Pass (CPU)

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

Overview

You will implement the complete forward pass of GPT-2 in C. By the end of this assignment, ./train_gpt2 will load a real GPT-2 checkpoint, run a forward pass, and compute a cross-entropy loss that matches PyTorch to floating-point precision.

The forward pass has this structure:

tokens → Embedding (wte + wpe)
       → [LayerNorm → QKV → Attention → Projection → Residual] × L
       → LayerNorm
       → Linear (logits)
       → Softmax + CrossEntropy → loss

Every function you implement is a mathematical operation on CPU arrays. This is CPU-only and requires no GPU — the CUDA port is A3.

To list your work:

grep -n "TODO (A1)" 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.

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 A1 fill:#2563eb,stroke:#1e40af,color:#fff

You are building the forward half of a transformer: the arithmetic that turns a sequence of token ids into a probability distribution over the next token, and a scalar loss saying how wrong it was. A2 differentiates exactly these functions, and A3 ports exactly these functions to CUDA — so the structure you build here is the structure you live with for the rest of the exercises.

One transformer block, which repeats L = 12 times in GPT-2 124M:

flowchart TB
    x["residual stream
[B,T,C]"] --> ln1["layernorm_forward → ln1"]
    ln1 --> qkv["matmul_forward
C → 3C"]
    qkv --> attn["attention_forward
causal, NH heads"]
    attn --> proj["matmul_forward
C → C"]
    proj --> add1(("+"))
    x --> add1
    add1 --> r2["residual2 [B,T,C]"]
    r2 --> ln2["layernorm_forward → ln2"]
    ln2 --> fc["matmul_forward
C → 4C"]
    fc --> gelu["gelu_forward"]
    gelu --> fcproj["matmul_forward
4C → C"]
    fcproj --> add2(("+"))
    r2 --> add2
    add2 --> out["residual3 [B,T,C]"]

The two + nodes are the residual connections, and they are the reason the network trains at all: they give the gradient a path from the loss back to layer 0 that does not pass through any matmul. Notice that each block reads the residual stream, computes a correction, and adds it back — the stream itself is never overwritten. In A2 that shape is what makes residual_backward a simple fan-out of the same gradient to two places.

Attention, the one part that mixes information across positions:

flowchart LR
    inp["inp[B,T,3C]"] --> Q["Q[b,t,h]"]
    inp --> K["K[b,t,h]"]
    inp --> V["V[b,t,h]"]
    Q --> dot["preatt = Q·K / sqrt(HS)
[B,NH,T,T]"]
    K --> dot
    dot --> mask["causal mask
t2 > t1 excluded"]
    mask --> sm["softmax over t2
→ att"]
    sm --> wsum["out = Σ att · V
[B,T,C]"]
    V --> wsum

Everything else in the model is position-wise: layernorm, gelu, the matmuls and the residuals all treat each (b,t) independently. Attention is the only operator where position t1 reads position t2, and the causal mask is what stops it reading the future — without it the model could see the answer it is being asked to predict, and the loss would collapse to near zero.

Background Reading

Source	Read	Why it matters here
The structs at the top of `train_gpt2.c`	`GPT2Config`, `ParameterTensors`, `ActivationTensors`	The shape of every tensor you will touch, with sizes in the comments. Read this before writing a line.
Vaswani et al. 2017, Attention Is All You Need	§3.1–3.3	The block structure and multi-head attention. Note GPT-2 differs from the paper: it is decoder-only, it puts LayerNorm before each sublayer rather than after, and it is causal.
Radford et al. 2019, Language Models are Unsupervised Multitask Learners	§2.3 and Table 2	The specific GPT-2 shapes you are implementing: 12 layers, 12 heads, 768 channels, 1024 context, 50257 vocab.
Ba et al. 2016, Layer Normalization	§3	Why normalise over channels rather than over the batch — and therefore why `mean` and `rstd` are per `(b,t)` and not per channel.
Hendrycks & Gimpel 2016, Gaussian Error Linear Units	§2, eq. 3	GELU and its `tanh` approximation. GPT-2 uses the approximation, which is the formula in Part 6 — the exact `erf` form gives slightly different numbers and will fail the tests.
Xiong et al. 2020, On Layer Normalization in the Transformer Architecture	§3 and Fig. 1	Optional but clarifying. Explains pre-LN vs post-LN and why pre-LN (what GPT-2 and this code use) trains without a warmup-sensitive gradient explosion at initialisation.

A note on `V` vs `Vp`

GPT-2’s real vocabulary is V = 50257. The checkpoint stores a padded vocabulary Vp = 50304 so the classifier matmul hits nicer memory alignment. This distinction is the single most common source of A1 bugs:

Parameter and activation tensors are allocated with Vp. wte is [Vp, C], logits and probs are [B, T, Vp].
Math only ever runs over the first V entries. The softmax normaliser, the max, and the cross-entropy lookup all stop at V.
The padded slots [V, Vp) must be explicitly zeroed in softmax_forward, so they contribute nothing downstream.

Vp is read from the checkpoint header — you do not compute it.

Tensor Layout Reference

All tensors are stored in row-major (C) order. B = batch size, T = sequence length, C = channels, NH = number of heads, HS = head size = C/NH, L = number of layers, maxT = config.max_seq_len.

Parameters (ParameterTensors, 16 tensors):

Tensor	Shape
`wte`	`[Vp, C]`
`wpe`	`[maxT, C]`
`ln1w`, `ln1b`	`[L, C]`
`qkvw`	`[L, 3C, C]`
`qkvb`	`[L, 3C]`
`attprojw`	`[L, C, C]`
`attprojb`	`[L, C]`
`ln2w`, `ln2b`	`[L, C]`
`fcw`	`[L, 4C, C]`
`fcb`	`[L, 4C]`
`fcprojw`	`[L, C, 4C]`
`fcprojb`	`[L, C]`
`lnfw`, `lnfb`	`[C]`

Activations (ActivationTensors, 23 tensors — note the names, they differ from the CUDA version you will meet in A3):

Tensor	Shape	Notes
`encoded`	`[B, T, C]`	Output of embedding layer
`ln1`	`[L, B, T, C]`	LayerNorm 1 output
`ln1_mean`, `ln1_rstd`	`[L, B, T]`	Saved for the backward pass
`qkv`	`[L, B, T, 3C]`	Q/K/V after projection
`atty`	`[L, B, T, C]`	Attention output
`preatt`	`[L, B, NH, T, T]`	Attention scores before softmax
`att`	`[L, B, NH, T, T]`	Attention weights after softmax
`attproj`	`[L, B, T, C]`	Attention output projection
`residual2`	`[L, B, T, C]`	After attention residual
`ln2`	`[L, B, T, C]`	LayerNorm 2 output
`ln2_mean`, `ln2_rstd`	`[L, B, T]`	Saved for the backward pass
`fch`	`[L, B, T, 4C]`	After MLP fc
`fch_gelu`	`[L, B, T, 4C]`	After GeLU
`fcproj`	`[L, B, T, C]`	MLP output projection
`residual3`	`[L, B, T, C]`	After MLP residual
`lnf`	`[B, T, C]`	Final LayerNorm output
`lnf_mean`, `lnf_rstd`	`[B, T]`	Saved for the backward pass
`logits`	`[B, T, Vp]`	Classifier output
`probs`	`[B, T, Vp]`	Softmax of logits
`losses`	`[B, T]`	Per-position cross-entropy

preatt and att are separate tensors: the causal mask and the scaled scores land in preatt, the normalised weights in att. A2 needs both.

Part 1 — Checkpoint Loading

gpt2_build_from_checkpoint(model, checkpoint_path)

fread a 256-int header.
Validate header[0] == 20240326 (magic). If it does not match, print an error and exit.
Validate header[1] == 3. Version 3 is “fp32, padded vocab”. Older checkpoints used version 1 and are not supported — the correct response is a fatal error pointing the user at python train_gpt2.py.

Read the config:

max_seq_len       = header[2]
vocab_size        = header[3]   // V  = 50257
num_layers        = header[4]
num_heads         = header[5]
channels          = header[6]
padded_vocab_size = header[7]   // Vp = 50304

Call fill_in_parameter_sizes, sum the sizes into num_parameters, then malloc_and_point_parameters.
fread the weight data directly into model->params_memory.
Initialise the remaining bookkeeping fields (acts_memory, grads_memory, … to NULL, batch_size/seq_len to 0, mean_loss to -1.0f).

Part 2 — Embedding Layer

encoder_forward(out, inp, wte, wpe, B, T, C)

For each (b, t) position, add the token embedding and the position embedding:

out[b, t, c] = wte[inp[b*T+t], c] + wpe[t, c]

Part 3 — Layer Normalization

layernorm_forward(out, mean, rstd, inp, weight, bias, B, T, C)

For each (b, t) position, compute the mean and reciprocal standard deviation over the C channels, then normalize and affine-transform:

μ = (1/C) Σ_c inp[b,t,c]
σ² = (1/C) Σ_c (inp[b,t,c] - μ)²
rstd = 1 / sqrt(σ² + ε)       where ε = 1e-5
out[b,t,c] = (inp[b,t,c] - μ) * rstd * weight[c] + bias[c]

Store mean[b*T+t] and rstd[b*T+t] — they are needed in the backward pass (A2), and the gradient check will fail in confusing ways if you skip them.

Part 4 — Matrix Multiplication

matmul_forward_naive(out, inp, weight, bias, B, T, C, OC)
matmul_forward(out, inp, weight, bias, B, T, C, OC)

Computes out = inp @ weight.T + bias. inp has shape [B*T, C], weight has shape [OC, C] (stored transposed relative to the standard convention), out has shape [B*T, OC]:

out[i, o] = Σ_c inp[i, c] * weight[o, c] + bias[o]

bias may be NULL, in which case it contributes nothing.

Implement matmul_forward_naive first as a straightforward triple loop. Then implement matmul_forward — it has the same semantics but should be loop-tiled and OpenMP-parallel, because this is where essentially all of the runtime goes. A correct matmul_forward that just calls matmul_forward_naive will pass the numerical tests but is slow enough to make the rest of this exercise unpleasant.

This is used for: the QKV projection (C → 3C), the attention output projection (C → C), the MLP layers (C → 4C and 4C → C), and the final classifier head (C → Vp).

Part 5 — Attention

attention_forward(out, preatt, att, inp, B, T, C, NH)

inp is the concatenated Q, K, V with shape [B, T, 3C]. For position t and head h (head size HS = C/NH), the three vectors live at:

Q = inp[b, t, h*HS       .. h*HS + HS)
K = inp[b, t, C + h*HS   .. C + h*HS + HS)
V = inp[b, t, 2C + h*HS  .. 2C + h*HS + HS)

Then for each head:

Scaled dot-product scores for t2 <= t1: preatt[b,h,t1,t2] = (Q[b,t1,h] · K[b,t2,h]) / sqrt(HS)
Causal masking: positions t2 > t1 are never computed. Leave preatt untouched there and write att[b,h,t1,t2] = 0. (Do not write -1e10 into preatt and rely on exp underflow — A2’s backward pass reads these slots.)
Softmax over t2 ∈ [0, t1] → att[b,h,t1,:]
Weighted sum of values: out[b,t1,h*HS + c] = Σ_{t2<=t1} att[b,h,t1,t2] * V[b,t2,h*HS + c]

Use the numerically stable softmax: subtract the row max before exponentiating.

Part 6 — Activation Functions, Residuals, Softmax

gelu_forward(out, inp, N)

GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))

The scaling factor sqrt(2/π) ≈ 0.7978845608 is defined as GELU_SCALING_FACTOR.

residual_forward(out, inp1, inp2, N)

out[i] = inp1[i] + inp2[i]

softmax_forward(probs, logits, B, T, V, Vp)

For each (b, t) row of length Vp: take the max over the first V entries, exponentiate logit - max over the first V, divide by their sum, and write 0.0f into probs[b, t, V .. Vp).

Part 7 — Loss

crossentropy_forward(losses, probs, targets, B, T, Vp)

losses[b*T+t] = -log(probs[b, t, targets[b*T+t]])

The row stride is Vp, even though targets are always < V.

Part 8 — Full Forward Pass

gpt2_forward(model, inputs, targets, B, T)

Wire together all of the above. On the first call this also has to allocate the activation buffers and record B/T; on later calls it must check the batch shape still fits.

encoder_forward → encoded
For each layer l = 0..L-1 (the residual stream input is encoded for l == 0, else residual3[l-1]):
- layernorm_forward → ln1, ln1_mean, ln1_rstd
- matmul_forward → qkv (C → 3C)
- attention_forward → atty, preatt, att
- matmul_forward → attproj (C → C)
- residual_forward → residual2
- layernorm_forward → ln2, ln2_mean, ln2_rstd
- matmul_forward → fch (C → 4C)
- gelu_forward → fch_gelu
- matmul_forward → fcproj (4C → C)
- residual_forward → residual3
layernorm_forward → lnf, lnf_mean, lnf_rstd
matmul_forward → logits (bias is NULL; the classifier reuses wte)
softmax_forward → probs
If targets != NULL: crossentropy_forward → losses, then set model->mean_loss to their average. Otherwise set model->mean_loss = -1.0f.

Testing

# fetch the checkpoint and the PyTorch reference state
chmod u+x ./dev/download_starter_pack.sh && ./dev/download_starter_pack.sh
# ...or regenerate them yourself:
#   pip install -r requirements.txt && python train_gpt2.py

make test_gpt2
./test_gpt2

test_gpt2 loads gpt2_124M.bin plus gpt2_124M_debug_state.bin (a fixed batch with PyTorch’s logits, loss, and gradients) and runs 10 steps. A1 only affects the first three checks — the gradient checks below them are A2’s job and will fail until then.

The shape of the output, once A1 is correct (exact values depend on your checkpoint):

[GPT-2]
max_seq_len: 1024
vocab_size: 50257
padded_vocab_size: 50304
num_layers: 12
num_heads: 12
channels: 768
num_parameters: 124475904
<10 lines of "expected_logit, your_logit" pairs>
OK (LOGITS), max_diff = <small>
LOSS OK: <yours> <expected>

Note there is no per-layer [layernorm_forward] OK line — test_gpt2 checks the model end to end, not layer by layer. If OK (LOGITS) fails, the printed first-mismatch index is your only localisation hint; bisect by commenting out layers in gpt2_forward.

Tolerances actually enforced by test_gpt2.c:

Check	Tolerance
Logits	`1e-2` absolute, per element, over the first `V` columns
Loss	`1e-2` absolute

You can also run the trainer end to end:

OMP_NUM_THREADS=8 ./train_gpt2

With a correct forward pass, the numbers at the top of the run should land near the reference output in the repo README: val loss 5.252026, step 0: train loss 5.356189. The loss will not go down over subsequent steps until A2 is done — gpt2_backward and gpt2_update are A2 stubs.

Checking your work

./grade A1

Runs ./test_gpt2 and scores it against the table below. Note that the harness is end to end, so a single wrong layer fails the whole forward-pass line; the grader says so rather than pretending to localise the bug.

Scoring

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

Component	Points
`gpt2_build_from_checkpoint` (magic, version 3, `header[2..7]` incl. `padded_vocab_size`)	10
`encoder_forward` matches reference	10
`layernorm_forward` matches reference (mean/rstd stored correctly)	15
`matmul_forward` / `matmul_forward_naive` match reference	15
`attention_forward` matches reference (causal masking, stable softmax)	20
`gelu_forward`, `residual_forward`, `softmax_forward` match reference (incl. `V`/`Vp` padding)	10
`crossentropy_forward` matches reference	5
`gpt2_forward` — logits and loss within `1e-2` of PyTorch on the real checkpoint	10
Code compiles cleanly; no undefined behaviour (`-fsanitize=address,undefined`)	5
Total	100