35. Training a Transformer Language Model: Backpropagation Through Causal Attention, Layer Norms, a Feed-Forward Network and Embeddings

What you will understand: how to train, in Mountain Goat, a small but complete causal transformer language model: token embeddings plus positions, a layer-normalized block of causal self-attention and a feed-forward network with residual connections, a final layer norm and a projection to the vocabulary, trained to predict the next token. Chapter 30 built the forward pass of such a model with random weights; Chapters 33 and 34 derived the backward pass piece by piece for a classifier. This chapter joins the pieces and adds the four that were missing: the causal mask inside the attention backward pass, many sequences stacked into one matrix, the per-position (masked) cross-entropy, and the embedding gradient. The model has 648 weights in 16 matrices; the hand-derived gradient of all 16 matrices is checked against an independent Python calculation, and against finite differences. After 200 steps the model predicts every token it is supposed to be able to predict on all 64 possible sequences, and the page then shows why you should not take that result as more than it is (training turns out to be chaotic).

What you need to know first: Chapter 33 (the attention backward pass), Chapter 34 (layer norm, ReLU and residual backward passes), Chapter 31 (cross-entropy) and Chapter 30 (the transformer's forward pass). No new compiler operation is needed: the programs run on Chapter 32's compiler.

Compile and run

cd docs/part32/code && ./build.sh                 # once: the newest compiler (this chapter adds nothing to it)
cd ../../part35/code
python3 make_examples.py                          # (re)writes the examples from lm_lib.py
../../part32/code/mgc run examples/03_train_10_steps.mg    # 10 training steps (under a minute)
../../part32/code/mgc run examples/04_train_200_steps.mg   # the full 200 steps: a 12,455-line program, about 8 minutes
./run_examples.sh > examples_out.txt              # every example and every error example on this page (about 8 minutes)
./check_lm.py                                     # the quick check: gradients, finite differences, 10 training steps (about 25 seconds)
./check_lm.py --full > check_lm_out.txt           # the page's 200 steps against the reference, and the experiment (about 8 minutes)
./sensitivity.py > sensitivity_out.txt            # Python only: how sensitive is training to a 1e-12 change? (about 15 seconds)
./lr_sweep.py > lr_sweep_out.txt                  # Python only: learning rates and starting weights (a few minutes)
./model_mutation.py > model_mutation_out.txt      # runs 18 WRONG programs through the quick check (about ten minutes)
cd ../../part15/code && ./run_lit.sh              # the whole test suite

Every listing and output on this page comes from these commands.

Why this page shows FAIL lines and error messages, and why that is good

The three error examples are programs the compiler must reject (each ends exit status 1). The mutation output is a negative control: deliberately wrong programs run against the checker, where "caught" is the wanted result and "NOT CAUGHT" would be a gap (one is reported, and explained). The -- lines in the full check are not failures: they mark values the page deliberately does not compare digit for digit, for a reason the page measures.

The task: copy the token from two positions back

A sequence has seven tokens from a vocabulary of four (0, 1, 2, 3). The first three are chosen freely and then the sequence repeats them: a b c a b c a. For example 0 1 2 0 1 2 0 and 3 3 1 3 3 1 3. There are 4³ = 64 possible sequences. A language model is trained to predict, at each position, the next token. In these sequences the next token is predictable only at positions 2 to 5: there, token i+1 equals token i−2 (the model has seen a b c and must say a, then b, then c, then a). At positions 0 and 1 the next token is a free choice, and position 6 has no next token, so those three positions are not scored. That leaves 4 scored predictions per sequence, and 256 scored predictions over all 64 sequences. Guessing gives 25%; so does always answering the most common token.

16 random sequences train the model (their first triples are 222 101 130 221 322 312 200 302 123 232 211 220 133 131 001 310), 16 others are held out (121 212 331 223 311 010 112 002 032 313 230 332 033 321 213 012), and the remaining 32 are seen only by the exhaustive check at the end. A model that memorizes its 16 sequences scores 64 of 64 on them and learns nothing about the rest; a model that has found the rule gets 256 of 256.

The model

The tokens are one-hot rows. For a sequence of 7 tokens, the matrices below have 7 rows (N = 7); the width is D = 8, the feed-forward width H = 16, the vocabulary V = 4:

step	what it computes	shape
x0 = x @ e + ps	embed each token (e is learned, 4×8), add the fixed sinusoidal position vectors ps	N×8
u1 = ln_out(x0, g1, n1)	layer norm 1	N×8
q, k, v = u1 @ wq, u1 @ wk, u1 @ wv	queries, keys, values (one head, width 8)	N×8
a = softmax_rows(q @ kᵀ / √8 + mask)	causal attention weights	N×N
o = a @ v; x1 = x0 + o @ wo	mix values; project; residual	N×8
u2 = ln_out(x1, g2, n2)	layer norm 2	N×8
pre = u2 @ w1 + c1; r = relu(pre)	feed-forward, first layer	N×16
x2 = x1 + (r @ w2 + c2)	feed-forward, second layer; residual	N×8
uf = ln_out(x2, gf, nf); lg = uf @ wu	final layer norm; scores for the 4 next-token candidates	N×4

The loss is the average over the scored positions of −log softmax(lg)[correct next token]. Position vectors are fixed (sines and cosines of the position, never trained), so only the 16 matrices e, g1, n1, wq, wk, wv, wo, g2, n2, w1, c1, w2, c2, gf, nf, wu are learned (32 + 8 + 8 + 4·64 + 8 + 8 + 128 + 16 + 128 + 8 + 8 + 8 + 32 = 648 numbers).

Many sequences, one matrix, and the causal mask

Mountain Goat has matrices and no batch dimension, so 16 sequences are stacked: 16 × 7 = 112 rows, and attention is computed for all 112 rows at once, a 112×112 score matrix. Row r must be allowed to look only at columns in its own sequence (no leaking between sequences) and only at positions not after itself (causal: no looking at the answer). The mask adds 0 to allowed entries and −10⁹ to the rest, and softmax turns those into exactly 0. Here is the idea at a size you can read, two sequences of three tokens (6 rows), with every score equal to 1 so that all structure comes from the mask:

# Two sequences of three tokens are stacked into one 6-row matrix, so attention is computed over all 6 rows at once. The mask must let row r look at
# column c only if c is in the SAME sequence and c <= r. Entries that must not be seen get -1000000000, which softmax turns into exactly 0.
# (Rows 0-2 are sequence one, rows 3-5 sequence two.) The scores are made up: each is 1, so without the mask every row would be uniform.
def softmax_rows(m: tensor[6x6]) = exp(m - row_max(m)) / row_sum(exp(m - row_max(m)))
let mask = [[0, -1000000000, -1000000000, -1000000000, -1000000000, -1000000000], [0, 0, -1000000000, -1000000000, -1000000000, -1000000000], [0, 0, 0, -1000000000, -1000000000, -1000000000], [-1000000000, -1000000000, -1000000000, 0, -1000000000, -1000000000], [-1000000000, -1000000000, -1000000000, 0, 0, -1000000000], [-1000000000, -1000000000, -1000000000, 0, 0, 0]]
let scores = [[1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]]
print softmax_rows(scores + mask)
# Each row sums to 1; row 0 is all on itself, row 2 is 1/3 each on rows 0-2, row 5 is 1/3 each on rows 3-5, and nothing leaks between the two sequences.
print row_sum(softmax_rows(scores + mask))

################ mgc run 01_stacked_causal_mask.mg
  | # Two sequences of three tokens are stacked into one 6-row matrix, so attention is computed over all 6 rows at once. The mask must let row r look at
  | # column c only if c is in the SAME sequence and c <= r. Entries that must not be seen get -1000000000, which softmax turns into exactly 0.
  | # (Rows 0-2 are sequence one, rows 3-5 sequence two.) The scores are made up: each is 1, so without the mask every row would be uniform.
  | def softmax_rows(m: tensor[6x6]) = exp(m - row_max(m)) / row_sum(exp(m - row_max(m)))
  | let mask = [[0, -1000000000, -1000000000, -1000000000, -1000000000, -1000000000], [0, 0, -1000000000, -1000000000, -1000000000, -1000000000], [0, 0, 0, -1000000000, -1000000000, -1000000000], [-1000000000, -1000000000, -100000
  | let scores = [[1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]]
  | print softmax_rows(scores + mask)
  | # Each row sums to 1; row 0 is all on itself, row 2 is 1/3 each on rows 0-2, row 5 is 1/3 each on rows 3-5, and nothing leaks between the two sequences.
  | print row_sum(softmax_rows(scores + mask))
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [6, 6] strides = [6, 1] data = 
[[1,   0,   0,   0,   0,   0], 
 [0.5,   0.5,   0,   0,   0,   0], 
 [0.333333,   0.333333,   0.333333,   0,   0,   0], 
 [0,   0,   0,   1,   0,   0], 
 [0,   0,   0,   0.5,   0.5,   0], 
 [0,   0,   0,   0.333333,   0.333333,   0.333333]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [6, 1] strides = [1, 1] data = 
[[1], 
 [1], 
 [1], 
 [1], 
 [1], 
 [1]]
  ----- exit status 0

Row 0 attends to itself only; row 2 splits its attention equally over rows 0 to 2; row 5 over rows 3 to 5; nothing crosses from rows 0–2 to rows 3–5. The real mask is the same pattern with blocks of 7 inside 112×112 (a literal in the programs).

The loss counts only some rows

Of the 112 stacked rows only 4 in 7 are predictions. The program carries a column mr (1 on scored rows, 0 elsewhere) and a one-hot matrix y that is all zeros on unscored rows, so that unscored rows add nothing to the loss, and mr also zeros their gradient:

# Only some rows of the stacked matrix are predictions. Here a 4-row logits matrix over 3 words, of which rows 1 and 3 are targets (mr is 1 on them, 0 elsewhere).
# y holds the correct next word one-hot on the target rows and zeros elsewhere. The loss averages over the 2 target rows (hence the 0.5), and the gradient with
# respect to the logits is (softmax(z) * mr - y) * 0.5: the rows that are not targets get exactly zero gradient.
def log_softmax(z: tensor[4x3]) = z - row_max(z) - log(row_sum(exp(z - row_max(z))))
let z = [[1, 2, 3], [0, 0, 0], [5, 1, 1], [1, 3, 2]]
let mr = [[0], [1], [0], [1]]
let y = [[0, 0, 0], [0, 1, 0], [0, 0, 0], [0, 0, 1]]
print col_sum(row_sum(y * log_softmax(z))) * -0.5
print (exp(log_softmax(z)) * mr - y) * 0.5
# Second row: three equal logits give each word 1/3; the correct word is the middle one, so that row's gradient is (1/3, 1/3 - 1, 1/3) / 2 = (0.1667, -0.3333, 0.1667).

################ mgc run 02_masked_cross_entropy.mg
  | # Only some rows of the stacked matrix are predictions. Here a 4-row logits matrix over 3 words, of which rows 1 and 3 are targets (mr is 1 on them, 0 elsewhere).
  | # y holds the correct next word one-hot on the target rows and zeros elsewhere. The loss averages over the 2 target rows (hence the 0.5), and the gradient with
  | # respect to the logits is (softmax(z) * mr - y) * 0.5: the rows that are not targets get exactly zero gradient.
  | def log_softmax(z: tensor[4x3]) = z - row_max(z) - log(row_sum(exp(z - row_max(z))))
  | let z = [[1, 2, 3], [0, 0, 0], [5, 1, 1], [1, 3, 2]]
  | let mr = [[0], [1], [0], [1]]
  | let y = [[0, 0, 0], [0, 1, 0], [0, 0, 0], [0, 0, 1]]
  | print col_sum(row_sum(y * log_softmax(z))) * -0.5
  | print (exp(log_softmax(z)) * mr - y) * 0.5
  | # Second row: three equal logits give each word 1/3; the correct word is the middle one, so that row's gradient is (1/3, 1/3 - 1, 1/3) / 2 = (0.1667, -0.3333, 0.1667).
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.25311]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 3] strides = [3, 1] data = 
[[0,   0,   0], 
 [0.166667,   -0.333333,   0.166667], 
 [0,   0,   0], 
 [0.0450153,   0.33262,   -0.377636]]
  ----- exit status 0

The loss 1.25311 is the mean of ln 3 = 1.0986 (three equal logits) and 1.4076 (the correct word has probability 0.245). The rows 0 and 2 of the gradient are exactly zero; row 1 is (1/3, 1/3 − 1, 1/3)/2. In the real model the factor is 1/64, one over the number of scored predictions in the batch.

The backward pass: what is new

Everything below dl (the gradient of the loss with respect to the scores) is a rule met before, applied in reverse order of the forward pass. The new rows:

• The scores' gradient, masked: dl = (exp(log_softmax(lg)) * mr − y) / 64: softmax minus one-hot (Chapter 31), zeroed on unscored rows.
• Through the final layer norm and the two residuals: ln_back (Chapter 34) and dx1 = dx2 + ln_back(du2, …): where a residual joins, the gradient paths add.
• Through the ReLU: multiply by ge(pre, [[0]]) (Chapter 34).
• Through softmax attention: da = do @ vᵀ, then ds = a * (da − row_sum(a * da)) (Chapter 33). The mask needs no special treatment: it is a constant added before the softmax, so it has no gradient, and the masked entries of a are exactly 0, so ds is exactly 0 there: no gradient flows to a position that was not seen. This is the one place where one might expect an extra term, and there is none.
• Three paths into u1: u1 feeds q, k and v, so du1 = dq @ wqᵀ + dk @ wkᵀ + dv @ wvᵀ, three terms that add. (dq = ds @ k / √8, dk = dsᵀ @ q / √8, dv = aᵀ @ do.)
• The embedding gradient: x is one-hot, so x0 = x @ e + ps is a lookup; its weight gradient is transpose(x) @ dx0: each row of e receives the sum of dx0 over the positions where that token occurs (a scatter-add written as a matrix product). The positions ps are constants, so they receive nothing.

The definitions

One training step as Mountain Goat text (the lets of the forward pass, the backward pass, the update); in the generated training programs each step's names carry a suffix and every step is unrolled, which is why the 200-step program has 12,455 lines:

# A one-block, one-head, causal transformer language model and its hand-derived backward pass. B = 16 sequences of 7 tokens are stacked into one
# matrix of 112 rows; attention is over all 112 rows with a mask that allows a position to see only earlier positions of its OWN sequence.
# (This listing is one training step with every step's suffix removed: the parameters are named after their matrices, e.g. e is the embedding. In the training
# programs each step's lets carry a suffix such as _s7, and the parameters are numbered per step, e0 .. e200.)

def softmax_rows(m: tensor[112x112]) = exp(m - row_max(m)) / row_sum(exp(m - row_max(m)))
def ln_sigma(h: tensor[112x8]) = sqrt(row_mean((h - row_mean(h)) * (h - row_mean(h))) + 1e-05)
def ln_hat(h: tensor[112x8]) = (h - row_mean(h)) / ln_sigma(h)
def ln_out(h: tensor[112x8], g: tensor[1x8], t: tensor[1x8]) = ln_hat(h) * g + t
def ln_back(d: tensor[112x8], h: tensor[112x8], g: tensor[1x8]) = ((d * g) - row_mean(d * g) - ln_hat(h) * row_mean((d * g) * ln_hat(h))) / ln_sigma(h)
def log_softmax(z: tensor[112x4]) = z - row_max(z) - log(row_sum(exp(z - row_max(z))))

# ---- the forward pass, one let per step
let x0_ = x @ e + ps
let u1_ = ln_out(x0_, g1, n1)
let q_ = u1_ @ wq
let k_ = u1_ @ wk
let v_ = u1_ @ wv
let a_ = softmax_rows(q_ @ transpose(k_) * 0.35355339059327373 + mk)
let o_ = a_ @ v_
let x1_ = x0_ + o_ @ wo
let u2_ = ln_out(x1_, g2, n2)
let pre_ = u2_ @ w1 + c1
let r_ = relu(pre_)
let x2_ = x1_ + (r_ @ w2 + c2)
let uf_ = ln_out(x2_, gf, nf)
let lg_ = uf_ @ wu
# ---- the loss, and the number of correct predictions
# loss    = col_sum(row_sum(y * log_softmax(lg_))) * -0.015625
# correct = col_sum(row_sum(ge(lg_, row_max(lg_)) * y))

# ---- the backward pass, one let per step, working backwards from the loss
let dl_ = (exp(log_softmax(lg_)) * mr - y) * 0.015625
let Gwu_ = transpose(uf_) @ dl_
let duf_ = dl_ @ transpose(wu)
let Gnf_ = col_sum(duf_)
let Ggf_ = col_sum(duf_ * ln_hat(x2_))
let dx2_ = ln_back(duf_, x2_, gf)
let Gc2_ = col_sum(dx2_)
let Gw2_ = transpose(r_) @ dx2_
let dpre_ = (dx2_ @ transpose(w2)) * ge(pre_, [[0]])
let Gc1_ = col_sum(dpre_)
let Gw1_ = transpose(u2_) @ dpre_
let du2_ = dpre_ @ transpose(w1)
let Gn2_ = col_sum(du2_)
let Gg2_ = col_sum(du2_ * ln_hat(x1_))
let dx1_ = dx2_ + ln_back(du2_, x1_, g2)
let Gwo_ = transpose(o_) @ dx1_
let do_ = dx1_ @ transpose(wo)
let dv_ = transpose(a_) @ do_
let da_ = do_ @ transpose(v_)
let ds_ = a_ * (da_ - row_sum(a_ * da_))
let dq_ = ds_ @ k_ * 0.35355339059327373
let dk_ = transpose(ds_) @ q_ * 0.35355339059327373
let Gwq_ = transpose(u1_) @ dq_
let Gwk_ = transpose(u1_) @ dk_
let Gwv_ = transpose(u1_) @ dv_
let du1_ = dq_ @ transpose(wq) + dk_ @ transpose(wk) + dv_ @ transpose(wv)
let Gn1_ = col_sum(du1_)
let Gg1_ = col_sum(du1_ * ln_hat(x0_))
let dx0_ = dx1_ + ln_back(du1_, x0_, g1)
let Ge_ = transpose(x) @ dx0_

# ---- the update, one per matrix (lr = 1.0)
let e_new = e - Ge_ * 1.0
let g1_new = g1 - Gg1_ * 1.0
let n1_new = n1 - Gn1_ * 1.0
let wq_new = wq - Gwq_ * 1.0
let wk_new = wk - Gwk_ * 1.0
let wv_new = wv - Gwv_ * 1.0
let wo_new = wo - Gwo_ * 1.0
let g2_new = g2 - Gg2_ * 1.0
let n2_new = n2 - Gn2_ * 1.0
let w1_new = w1 - Gw1_ * 1.0
let c1_new = c1 - Gc1_ * 1.0
let w2_new = w2 - Gw2_ * 1.0
let c2_new = c2 - Gc2_ * 1.0
let gf_new = gf - Ggf_ * 1.0
let nf_new = nf - Gnf_ * 1.0
let wu_new = wu - Gwu_ * 1.0

Why one let per step and not nested functions. The first version of this chapter wrote each gradient as a function of the parameters (as in Chapter 34). For this model mgc was killed (exit status −9; running out of memory is the likely cause, not confirmed beyond that): a nested function is expanded in place at every use, so a quantity used several times is recomputed several times, the generated code never frees memory, and the cost multiplied at every level. Binding every intermediate result with let computes it once. This is a property of this compiler (no common-subexpression elimination, no freeing of temporaries), not of the method.

Training

The program for the page, examples/04_train_200_steps.mg, runs 200 steps of gradient descent with learning rate 1.0 from seeded starting weights, printing the loss and the number of correct predictions (of 64) on the 16 training sequences at checkpoints 0, 1, 10, 25, 50, 100, 150 and 200; then the held-out loss and count; then the number correct on each of four chunks of 16 of all 64 possible sequences; then the attention weights of the first 7 rows (the first training sequence). First, the same program stopped at 10 steps (03_train_10_steps.mg, which the test suite runs):

################ mgc run 03_train_10_steps.mg
  | (785 lines: the definitions are in lm_lib.py; the file holds the data, the unrolled training steps and the prints, written by make_examples.py)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.42919]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[20]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.28756]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[29]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.965858]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[39]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.58847]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[20]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[11]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[31]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[39]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[20]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 
[[1,   0,   0,   0,   0,   0,   0], 
 [0.416434,   0.583566,   0,   0,   0,   0,   0], 
 [0.375767,   0.356487,   0.267746,   0,   0,   0,   0], 
 [0.445018,   0.32708,   0.141979,   0.0859235,   0,   0,   0], 
 [0.436527,   0.296865,   0.121964,   0.0729565,   0.0716872,   0,   0], 
 [0.334917,   0.24728,   0.132274,   0.0910244,   0.0858113,   0.108694,   0], 
 [0.252324,   0.205006,   0.132171,   0.0983111,   0.08894,   0.100532,   0.122715]]
  ----- exit status 0

(The first six numbers are loss and number correct at steps 0, 1 and 10; the next two are the held-out loss and count; then the four chunk counts; then the attention matrix, which at this early stage is spread broadly: row 6 puts 0.25, 0.21, 0.13, … on the seven positions.) The loss falls from 1.429 to 0.966 in ten steps; the number correct rises from 20 of 64 (chance is 16) to 39. And the full run:

################ mgc run 04_train_200_steps.mg
  | (12455 lines: the definitions are in lm_lib.py; the file holds the data, the unrolled training steps and the prints, written by make_examples.py)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.42919]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[20]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.28756]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[29]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.965858]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[39]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.727803]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[45]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.525403]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[53]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.0356118]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[64]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.00192844]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[64]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.000917701]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[64]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.00759067]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[64]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[64]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[64]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[64]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[64]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 
[[1,   0,   0,   0,   0,   0,   0], 
 [0.990166,   0.00983402,   0,   0,   0,   0,   0], 
 [0.977195,   0.0224568,   0.000348693,   0,   0,   0,   0], 
 [0.429867,   0.518344,   0.0497861,   0.00200291,   0,   0,   0], 
 [0.00284028,   0.10505,   0.803391,   0.0871705,   0.00154816,   0,   0], 
 [3.99338e-05,   0.00111138,   0.149104,   0.783051,   0.0653207,   0.00137268,   0], 
 [1.24684e-05,   1.48271e-05,   0.00118552,   0.187334,   0.785065,   0.0261648,   0.000223552]]
  ----- exit status 0

after step	0	1	10	25	50	100	150	200
training loss	1.429	1.288	0.966	0.728	0.525	0.0356	0.0019	0.0009
correct of 64	20	29	39	45	53	64	64	64

• It fits the training data: loss 0.0009, all 64 scored predictions right.
• It generalizes, on this run: the 16 held-out sequences give loss 0.0076 and 64 of 64, and every one of the four chunks of 16 of all 64 possible sequences gives 64 of 64: 256 of 256.
• The attention has learned where to look. The last matrix (the first training sequence, 2 2 2 2 2 2 2) shows each row attending to the position two before itself: row 2 puts 0.977 on position 0, row 3 puts 0.518 on position 1 (and 0.430 on position 0), row 4 puts 0.803 on position 2, row 5 puts 0.783 on position 3, row 6 puts 0.785 on position 4. Position i is predicting token i+1, which equals token i−2, and that is where it looks. This is a positional pattern; because this sequence has the same token everywhere, it cannot show whether the attention also depends on the token content, and no other sequence's attention was printed, so that is not established.
• The loss did not fall smoothly. It is 0.525 at step 50, rises to about 0.58 at step 75 (from the Python run, below), and has collapsed to 0.036 by step 100. It is the shape of a model finding the trick, and it has a consequence, next.

Training is chaotic: a measured caution

check_lm.py compares the Mountain Goat training with the independent Python trainer. The comparison is digit for digit up to step 50 (losses and counts equal to a relative 10⁻⁵, the six digits mgc prints) and then stops being possible. At step 100 the program's loss is 0.0356 and the Python reference's is 0.0346; the first version of the checker, which demanded agreement at every checkpoint, failed there. Was something wrong? The Python-only sensitivity.py answers by training the same Python code twice, once as is and once with a single starting weight changed by 10⁻¹²:

 step   train loss  nudged by 1e-12   held-out loss     nudged
    0     1.429187         1.429187        1.478939   1.478939
   25     0.727803         0.727803        1.666671   1.666671
   50     0.525403         0.525403        1.564036   1.564036
   60     0.631461         0.631461        1.574613   1.574613
   70     0.208339         0.208347        1.179501   1.179517
   75     0.581304         0.580148        1.207885   1.206682
   80     0.602086         0.599439        1.239171   1.236996
   90     0.376693         0.322604        0.808559   0.995614
  100     0.034635         0.211520        0.251165   0.896438
  125     0.003914         0.041846        0.042737   0.279762
  150     0.001926         0.003073        0.020251   0.167249
  200     0.000912         0.001132        0.008777   0.154761

A change of one part in a trillion in one weight is invisible for the first 60 steps and then grows until the two runs differ completely: at step 100 the loss is 0.035 in one run and 0.21 in the other; at step 200 the held-out loss is 0.0088 in one and 0.155 in the other. Two correct implementations differ at about that level (different order of floating-point operations), so exact agreement cannot be expected after the transition. What this does and does not mean:

• The backward pass is right (all 16 gradient matrices equal the reference at three starting points, to 10⁻⁵, and agree with finite differences), and the training agrees digit for digit while the run is still stable. Divergence after the transition is a property of the training, not of the gradient.
• Where the check cannot compare, it checks claims instead: from step 150 on, the program's loss must be below 0.01 with 64 of 64 right; and the final result is judged by the experiment's claims. The -- lines in the output below are the values left uncompared.
• "256 of 256" is a statement about this particular run, not about the model class. The nudged Python run has the same training loss to within 0.0002 and a held-out loss 18 times larger.

How reliable is this? A Python-only sweep

lr_sweep.py (plain Python, not mgc) trains for 200 steps with four learning rates and four starting-weight seeds, and reports the final training loss and the number correct of all 256 scored predictions:

   lr   per starting-weight seed 1..4: (final training loss, correct of 256)
 0.25   (0.00843, 228)  (0.0056, 239)  (0.00273, 225)  (0.011, 235)
  0.5   (0.00144, 249)  (0.00218, 234)  (0.00125, 231)  (0.0494, 189)
  1.0   (0.000912, 256)  (0.00334, 239)  (0.00118, 249)  (0.00419, 239)
  2.0   (1.15, 104)  (0.723, 134)  (0.836, 124)  (0.862, 150)

• Learning rate 1.0 (the page's) always drives the training loss below 0.005, but the number correct over all 64 sequences is 256, 239, 249, 239: seed 1, the page's, is the best of the four. A model can fit all 16 training sequences and still be wrong on 5 to 7% of the scored predictions the data never showed.
• 0.25 and 0.5 are slower and do no better on unseen sequences (225–239 at 0.25; 189–249 at 0.5, where one seed's training loss is still 0.049). The sweep did not run longer.
• 2.0 does not learn at all (loss above 0.7 after 200 steps; no overflow, it just does not settle).
• Seeds differ, and (as the sensitivity table shows) so can two correct implementations on one seed. With 16 training sequences out of 64, which sequences are held out is a further source of variation that was not varied.

Mistakes the compiler reports

################ mgc run errors/e1_mask_of_the_wrong_size.mg
  | # The mask must be as large as the score matrix: 6 x 6 here. A 4 x 4 mask does not fit, and the compiler says so.
  | def softmax_rows(m: tensor[6x6]) = exp(m - row_max(m)) / row_sum(exp(m - row_max(m)))
  | let scores = [[1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]]
  | let mask = [[0, -1000000000, -1000000000, -1000000000], [0, 0, -1000000000, -1000000000], [0, 0, 0, -1000000000], [0, 0, 0, 0]]
  | print softmax_rows(scores + mask)
  ----- output
errors/e1_mask_of_the_wrong_size.mg:5: error: cannot add shapes 6x6 and 4x4: dimensions 6 and 4 differ and neither is 1
  ----- exit status 1

################ mgc run errors/e2_missing_argument.mg
  | # ln_back takes three arguments (the upstream gradient, the layer's input and the scale). Two is an arity error, reported before any code exists.
  | def ln_sigma(h: tensor[2x4]) = sqrt(row_mean((h - row_mean(h)) * (h - row_mean(h))) + 0.00001)
  | def ln_hat(h: tensor[2x4]) = (h - row_mean(h)) / ln_sigma(h)
  | def ln_back(d: tensor[2x4], h: tensor[2x4], g: tensor[1x4]) = ((d * g) - row_mean(d * g) - ln_hat(h) * row_mean((d * g) * ln_hat(h))) / ln_sigma(h)
  | print ln_back([[1, 2, 3, 4], [4, 3, 2, 1]], [[1, 2, 4, 8], [8, 4, 2, 1]])
  ----- output
errors/e2_missing_argument.mg:5: error: ln_back takes 3 arguments, got 2
  ----- exit status 1

################ mgc run errors/e3_stacked_rows_do_not_match.mg
  | # The tokens matrix has one row per token of every stacked sequence (6 here); the mask-row vector must have the same 6 rows. A 4-row one cannot multiply it.
  | let logits = [[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]
  | let mr = [[0], [1], [0], [1]]
  | print logits * mr
  ----- output
errors/e3_stacked_rows_do_not_match.mg:4: error: cannot multiply shapes 6x3 and 4x1: dimensions 6 and 4 differ and neither is 1
  ----- exit status 1

A mask of the wrong size for its scores, a call with a missing argument, and a row-selection column that does not match the stacked rows. As always, the shapes are types.

Tests

Three new lit files in test/lm35/ (the suite is now 141 tests; no verifiers or ir test because no compiler operation was added):

Test	What it checks
examples	the two small examples (worked out above) and the loss and correct count of the 10-step training run
errors	the three mistakes give the exact messages
against-reference	check_lm.py in its quick mode, below (about 25 seconds). The 200-step run is not in CI because it takes about 8 minutes; check_lm.py --full runs it

check_lm.py compares with an independent Python implementation (plain lists, one sequence and one position at a time; every gradient derived element by element) and checks, in order:

1. The gradients, for three starting-weight sets: the loss and all 16 gradient matrices equal the reference's, and for the first set two weights of every matrix (the first, and the one with the largest gradient: 32 weights) agree with a finite-difference estimate (step 0.001). Checking every one of the 648 would take 1,296 more forward passes in one program and has not been done.
2. Training: loss and correct count at the checkpoints up to step 50 (and, in the quick mode, up to step 10), the held-out loss and count, the four chunk counts and the attention pattern equal the reference's; every attention row sums to 1 and puts exactly 0 after the diagonal.
3. The experiment (--full): the 200-step run fits the training sequences (loss below 0.002), gets all 16 held-out and all 64 possible sequences right, with the later checkpoints held to their claims rather than to digits.

-- the hand-derived backward pass against the independent per-position gradient, and against finite differences
  ok   starting weights #1: the loss and all 16 gradient matrices equal the reference
  ok   starting weights #1: 32 weights (two of every matrix) agree with finite differences (largest difference 1.6e-04)
  ok   starting weights #2: the loss and all 16 gradient matrices equal the reference
  ok   starting weights #3: the loss and all 16 gradient matrices equal the reference
-- training 200 steps against the independent Python trainer
  ok   step 0: loss 1.429187 and 20 of 64 correct equal the reference
  ok   step 1: loss 1.287565 and 29 of 64 correct equal the reference
  ok   step 10: loss 0.965858 and 39 of 64 correct equal the reference
  ok   step 25: loss 0.727803 and 45 of 64 correct equal the reference
  ok   step 50: loss 0.525403 and 53 of 64 correct equal the reference
  --   step 100: loss 0.035612, 64 of 64 correct (the reference gives 0.034635, 64: inside the chaotic phase, not compared)
  ok   step 150: loss 0.001928 (below 0.01) and 64 of 64 correct (the reference's run, 0.001926, differs by rounding-level amplification; not compared digit for digit)
  ok   step 200: loss 0.000918 (below 0.01) and 64 of 64 correct (the reference's run, 0.000912, differs by rounding-level amplification; not compared digit for digit)
  --   the 16 held-out sequences: 64 of 64 correct, loss 0.007591 (the reference run gives 0.008777; compared by the experiment's claims below)
  ok   the number correct in each of the four chunks of 16 of all 64 sequences, [64, 64, 64, 64], equals the reference's [64, 64, 64, 64] (or the run is past the tight region and the experiment's claims below decide)
  ok   every attention row sums to 1 and puts exactly 0 on every later position (the mask works)
-- the experiment
  ok   it fits its 16 training sequences: loss 0.000918, 64 of 64 predictions right
  ok   it gets all 16 held-out sequences right (64 of 64 predictions)
  ok   it gets all 64 possible sequences right: 256 of 256 predictions
  ok   for the first training sequence [2, 2, 2, 2, 2, 2, 2], position 5 (predicting token 2) attends mostly to position 3 (weight 0.783); the page reports the whole pattern
all checks pass

A repaired check. The finite-difference check first failed (largest difference 2.8·10⁻³ against a threshold of 10⁻³). The gradient was not wrong: the loss is about 1.4 and mgc prints six significant digits, so each printed loss carries an error of about 5·10⁻⁶, which divided by the step 0.002 becomes about 2.5·10⁻³. The program now forms the difference of the two losses inside Mountain Goat and prints that, so that the six digits are spent on the small difference, and the largest difference fell to 1.6·10⁻⁴.

Are the checks good enough? Run eighteen wrong programs.

model_mutation.py writes eighteen versions of the program, each with one mistake, and runs the quick checker on each:

NOTE: this script deliberately runs wrong backward passes. 'caught' lines are EXPECTED: they show the checker can detect the mistake.
baseline (nothing broken): checker passes (expected)
layer-norm backward misses the mean-subtraction term: caught by 9 failing checks (the reference gradient, finite differences, the reference training run, other)
layer-norm backward misses the projection term: caught by 9 failing checks (the reference gradient, finite differences, the reference training run, other)
layer-norm backward forgets to divide by sigma: caught by 9 failing checks (the reference gradient, finite differences, the reference training run, other)
softmax backward misses the row-sum correction (ds = a * da): caught by 4 failing checks (the reference gradient, finite differences)
the 1/sqrt(8) score scale is missing from dq: caught by 4 failing checks (the reference gradient, finite differences)
dk forgets to transpose ds: caught by 4 failing checks (the reference gradient, finite differences)
the gradient of the cross-entropy is not masked to the target rows: caught by 9 failing checks (the reference gradient, finite differences, the reference training run, other)
the loss's 1/64 is 1/32 in the backward pass only: caught by 9 failing checks (the reference gradient, finite differences, the reference training run, other)
the relu mask is omitted (the gradient passes everywhere): caught by 4 failing checks (the reference gradient, finite differences)
the relu mask is the wrong way round: caught by 4 failing checks (the reference gradient, finite differences)
the residual path is dropped from dx1: caught by 4 failing checks (the reference gradient, finite differences)
the residual path is dropped from dx0 (the embedding gradient): caught by 4 failing checks (the reference gradient, finite differences)
the value path is dropped from du1 (attention's three inputs become two): caught by 4 failing checks (the reference gradient, finite differences)
the forward pass lets every position see ALL 112 rows (no mask; the backward pass is unchanged): caught by 10 failing checks (the reference gradient, the reference training run, other, the mask check)
the position encodings are left out of the forward pass only: caught by 10 failing checks (the reference gradient, finite differences, the reference training run, other)
the layer-norm epsilon is 0.1 in the forward AND backward pass (consistent with itself, not with the reference): caught by 9 failing checks (the reference gradient, the reference training run, other)
naive softmax (no row-max subtraction); algebraically identical, and the masked entries are exp(-1e9) = 0: NOT CAUGHT
the training step has the wrong learning rate (0.5 instead of 1.0): caught by 5 failing checks (the reference training run, other)

Seventeen of eighteen are caught. What the pattern shows:

• Thirteen are mistakes in the backward pass alone (the three layer-norm errors, softmax's missing row-sum correction, the dropped 1/√8 in dq, dk without the transpose, the loss gradient not masked to the scored rows, 1/32 instead of 1/64, the ReLU mask omitted or reversed, the residual dropped from dx1 or from dx0, the value path dropped from du1). Each is caught by the reference gradient and by the finite differences, which see only the loss.
• The forward-only changes are caught by the reference, not by finite differences. The forward pass with no mask (every position sees all 112 rows) is also caught by the mask check and the training run; the layer-norm epsilon of 0.1 in forward and backward is consistent with itself, so finite differences cannot see it, and only the comparison with the reference does (as in Chapter 34). The wrong learning rate changes only training, and only the comparison of the training run notices.
• A caveat about one line. The mutant that drops the position vectors from the forward pass is reported as caught by finite differences too; that is partly an artifact of the mutation script, which changed the program's main forward pass but not the perturbed copies used for the embedding's finite differences. It is caught regardless, by the reference.
• One is NOT CAUGHT: the softmax without the row-maximum subtraction. It is algebraically the same function, and with scores of this size (and the masked entries exp(−10⁹) = 0 exactly) it gives the same numbers. The reason for subtracting the maximum is to avoid overflow for large scores, and no check here has large scores. This is a gap in the checks, not in the program: a test with scores of a few hundred would catch it, and none was written.

There is no compiler-mutation run for this chapter: no compiler code changed.

The full suite:

-- Testing: 141 tests, 4 workers --
PASS: IR Forge / Mountain Goat :: attention33/examples.mlir (1 of 141)
PASS: IR Forge / Mountain Goat :: attention33/against-reference.mlir (2 of 141)
PASS: IR Forge / Mountain Goat :: generalize32/against-reference.mlir (3 of 141)
PASS: IR Forge / Mountain Goat :: transformer30/against-reference.mlir (4 of 141)
PASS: IR Forge / Mountain Goat :: lm35/against-reference.mlir (5 of 141)
PASS: IR Forge / Mountain Goat :: contraction/tensor-contraction.mlir (6 of 141)
PASS: IR Forge / Mountain Goat :: contraction28/against-definition.mlir (7 of 141)
PASS: IR Forge / Mountain Goat :: attention29/against-reference.mlir (8 of 141)
PASS: IR Forge / Mountain Goat :: learning/gradient-descent.mlir (9 of 141)
PASS: IR Forge / Mountain Goat :: lm35/examples.mlir (10 of 141)
PASS: IR Forge / Mountain Goat :: bigram31/against-reference.mlir (11 of 141)
PASS: IR Forge / Mountain Goat :: generalize32/examples.mlir (12 of 141)
PASS: IR Forge / Mountain Goat :: tour/gallery.mlir (13 of 141)
PASS: IR Forge / Mountain Goat :: attention29/examples.mlir (14 of 141)
PASS: IR Forge / Mountain Goat :: contraction28/examples.mlir (15 of 141)
PASS: IR Forge / Mountain Goat :: broadcast/values.mlir (16 of 141)
PASS: IR Forge / Mountain Goat :: tour/more-examples.mlir (17 of 141)
PASS: IR Forge / Mountain Goat :: bigram31/examples.mlir (18 of 141)
PASS: IR Forge / Mountain Goat :: tour/tour-examples.mlir (19 of 141)
PASS: IR Forge / Mountain Goat :: transformer30/examples.mlir (20 of 141)
PASS: IR Forge / Mountain Goat :: loop-order/bitwise-identical.mlir (21 of 141)
PASS: IR Forge / Mountain Goat :: loop-order/results-identical.mlir (22 of 141)
PASS: IR Forge / Mountain Goat :: perf/options-same-results.mlir (23 of 141)
PASS: IR Forge / Mountain Goat :: broadcast/cpp-interop-reductions.mlir (24 of 141)
PASS: IR Forge / Mountain Goat :: contraction/mismatch-trap.mlir (25 of 141)
PASS: IR Forge / Mountain Goat :: loop-order/cpp-ikj.mlir (26 of 141)
PASS: IR Forge / Mountain Goat :: ops/values.mlir (27 of 141)
PASS: IR Forge / Mountain Goat :: frontend/cpp-interop.mlir (28 of 141)
PASS: IR Forge / Mountain Goat :: ops/elementwise-mismatch-aborts.mlir (29 of 141)
PASS: IR Forge / Mountain Goat :: tour/scientific-literals.mlir (30 of 141)
PASS: IR Forge / Mountain Goat :: ops/cpp-interop-ops.mlir (31 of 141)
PASS: IR Forge / Mountain Goat :: broadcast/networks.mlir (32 of 141)
PASS: IR Forge / Mountain Goat :: execution/gpu-host-with-stub-runtime.mlir (33 of 141)
PASS: IR Forge / Mountain Goat :: ops/precedence-and-identities.mlir (34 of 141)
PASS: IR Forge / Mountain Goat :: frontend/more-examples.mlir (35 of 141)
PASS: IR Forge / Mountain Goat :: ops/negate-and-dynamic-divide.mlir (36 of 141)
PASS: IR Forge / Mountain Goat :: broadcast/dynamic-reductions.mlir (37 of 141)
PASS: IR Forge / Mountain Goat :: ops/composed-programs.mlir (38 of 141)
PASS: IR Forge / Mountain Goat :: ops/dynamic-ops.mlir (39 of 141)
PASS: IR Forge / Mountain Goat :: frontend/runtime-mismatch-aborts.mlir (40 of 141)
PASS: IR Forge / Mountain Goat :: broadcast/identities.mlir (41 of 141)
PASS: IR Forge / Mountain Goat :: loop-order/flag-reaches-pass.mlir (42 of 141)
PASS: IR Forge / Mountain Goat :: chapters/ch09-execution-engine.mlir (43 of 141)
PASS: IR Forge / Mountain Goat :: execution/mismatch-aborts.mlir (44 of 141)
PASS: IR Forge / Mountain Goat :: assert-stderr/message-on-stderr.mlir (45 of 141)
PASS: IR Forge / Mountain Goat :: frontend/dynamic-runs.mlir (46 of 141)
PASS: IR Forge / Mountain Goat :: perf/flags-reach-clang.mlir (47 of 141)
PASS: IR Forge / Mountain Goat :: tour/tour-errors.mlir (48 of 141)
PASS: IR Forge / Mountain Goat :: execution/mismatch-columns-aborts.mlir (49 of 141)
PASS: IR Forge / Mountain Goat :: contraction28/errors.mlir (50 of 141)
PASS: IR Forge / Mountain Goat :: execution/dynamic-shapes-run.mlir (51 of 141)
PASS: IR Forge / Mountain Goat :: ops/ptx-new-ops.mlir (52 of 141)
PASS: IR Forge / Mountain Goat :: perf/tiling-structure.mlir (53 of 141)
PASS: IR Forge / Mountain Goat :: ops/matmul-mismatch-aborts.mlir (54 of 141)
PASS: IR Forge / Mountain Goat :: frontend/hello-and-static.mlir (55 of 141)
PASS: IR Forge / Mountain Goat :: execution/static-add-runs.mlir (56 of 141)
PASS: IR Forge / Mountain Goat :: ops/division-ieee.mlir (57 of 141)
PASS: IR Forge / Mountain Goat :: bigram31/errors.mlir (58 of 141)
PASS: IR Forge / Mountain Goat :: dynamic-loops/runtime-check-survives-transforms.mlir (59 of 141)
PASS: IR Forge / Mountain Goat :: dynamic-loops/tiling-runs.mlir (60 of 141)
PASS: IR Forge / Mountain Goat :: transformer30/errors.mlir (61 of 141)
PASS: IR Forge / Mountain Goat :: broadcast/ptx-layer.mlir (62 of 141)
PASS: IR Forge / Mountain Goat :: assert-stderr/valid-shapes-still-pass.mlir (63 of 141)
PASS: IR Forge / Mountain Goat :: attention33/errors.mlir (64 of 141)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-after-tiling.mlir (65 of 141)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-runs.mlir (66 of 141)
PASS: IR Forge / Mountain Goat :: generalize32/errors.mlir (67 of 141)
PASS: IR Forge / Mountain Goat :: assert-stderr/nested-assert-runs.mlir (68 of 141)
PASS: IR Forge / Mountain Goat :: execution/mixed-static-dynamic.mlir (69 of 141)
PASS: IR Forge / Mountain Goat :: ffn34/errors.mlir (70 of 141)
PASS: IR Forge / Mountain Goat :: gpu/host-runtime-calls.mlir (71 of 141)
PASS: IR Forge / Mountain Goat :: assert-stderr/message-bytes-exact.mlir (72 of 141)
PASS: IR Forge / Mountain Goat :: contraction28/ir.mlir (73 of 141)
PASS: IR Forge / Mountain Goat :: attention29/errors.mlir (74 of 141)
PASS: IR Forge / Mountain Goat :: gpu/ptx-default.mlir (75 of 141)
PASS: IR Forge / Mountain Goat :: lm35/errors.mlir (76 of 141)
PASS: IR Forge / Mountain Goat :: loop-order/structure.mlir (77 of 141)
PASS: IR Forge / Mountain Goat :: gpu/ptx-bare-pointers.mlir (78 of 141)
PASS: IR Forge / Mountain Goat :: gpu/pass-order-crash.mlir (79 of 141)
PASS: IR Forge / Mountain Goat :: execution/strided-view.mlir (80 of 141)
PASS: IR Forge / Mountain Goat :: gpu/host-runtime-calls-bare.mlir (81 of 141)
PASS: IR Forge / Mountain Goat :: frontend/ptx-from-source.mlir (82 of 141)
PASS: IR Forge / Mountain Goat :: gpu/ptx-index32.mlir (83 of 141)
PASS: IR Forge / Mountain Goat :: chapters/ch01-pipeline.mlir (84 of 141)
PASS: IR Forge / Mountain Goat :: ops/front-end-errors.mlir (85 of 141)
PASS: IR Forge / Mountain Goat :: frontend/front-end-errors.mlir (86 of 141)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-needs-constant-trip-counts.mlir (87 of 141)
PASS: IR Forge / Mountain Goat :: broadcast/front-end-errors.mlir (88 of 141)
PASS: IR Forge / Mountain Goat :: broadcast/roundtrip-and-lowering.mlir (89 of 141)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-full-needs-constant-trip-count.mlir (90 of 141)
PASS: IR Forge / Mountain Goat :: transformer30/ir.mlir (91 of 141)
PASS: IR Forge / Mountain Goat :: chapters/ch02-example-programs.mlir (92 of 141)
PASS: IR Forge / Mountain Goat :: assert-stderr/guard-nested-assert.mlir (93 of 141)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice-dynamic-mixed.mlir (94 of 141)
PASS: IR Forge / Mountain Goat :: gpu/nvvm-lowering.mlir (95 of 141)
PASS: IR Forge / Mountain Goat :: ops/matmul-lowering-structure.mlir (96 of 141)
PASS: IR Forge / Mountain Goat :: bigram31/ir.mlir (97 of 141)
PASS: IR Forge / Mountain Goat :: ops/matmul-inner-mismatch.mlir (98 of 141)
PASS: IR Forge / Mountain Goat :: broadcast/verifier-broadcast.mlir (99 of 141)
PASS: IR Forge / Mountain Goat :: verifier/transpose-bad-shape.mlir (100 of 141)
PASS: IR Forge / Mountain Goat :: frontend/emit-mlir.mlir (101 of 141)
PASS: IR Forge / Mountain Goat :: generalize32/ir.mlir (102 of 141)
PASS: IR Forge / Mountain Goat :: verifier/dynamic-accepted.mlir (103 of 141)
PASS: IR Forge / Mountain Goat :: gpu/kernel-outlining.mlir (104 of 141)
PASS: IR Forge / Mountain Goat :: verifier/add-static-mismatch.mlir (105 of 141)
PASS: IR Forge / Mountain Goat :: assert-stderr/run-twice-no-collision.mlir (106 of 141)
PASS: IR Forge / Mountain Goat :: assert-stderr/static-unchanged.mlir (107 of 141)
PASS: IR Forge / Mountain Goat :: dynamic-loops/tiling-structure.mlir (108 of 141)
PASS: IR Forge / Mountain Goat :: ops/matmul-bad-result-shape.mlir (109 of 141)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-declines.mlir (110 of 141)
PASS: IR Forge / Mountain Goat :: gpu/cuda-comparison.mlir (111 of 141)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-fires-when-static.mlir (112 of 141)
PASS: IR Forge / Mountain Goat :: lowering/bufferize-static-add.mlir (113 of 141)
PASS: IR Forge / Mountain Goat :: attention29/ir.mlir (114 of 141)
PASS: IR Forge / Mountain Goat :: verifier/unranked-rejected.mlir (115 of 141)
PASS: IR Forge / Mountain Goat :: verifier/transpose-rank3.mlir (116 of 141)
PASS: IR Forge / Mountain Goat :: loops/tiling.mlir (117 of 141)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-epilogue.mlir (118 of 141)
PASS: IR Forge / Mountain Goat :: broadcast/verifier-reduce.mlir (119 of 141)
PASS: IR Forge / Mountain Goat :: loops/fusion.mlir (120 of 141)
PASS: IR Forge / Mountain Goat :: lowering/dynamic-transpose-no-check.mlir (121 of 141)
PASS: IR Forge / Mountain Goat :: bigram31/verifiers.mlir (122 of 141)
PASS: IR Forge / Mountain Goat :: loops/unroll-full.mlir (123 of 141)
PASS: IR Forge / Mountain Goat :: lowering/bufferize-dynamic-add.mlir (124 of 141)
PASS: IR Forge / Mountain Goat :: ops/dynamic-accepted.mlir (125 of 141)
PASS: IR Forge / Mountain Goat :: assert-stderr/rewrite-structure.mlir (126 of 141)
PASS: IR Forge / Mountain Goat :: ops/elementwise-static-mismatch.mlir (127 of 141)
PASS: IR Forge / Mountain Goat :: canonicalize/fold-constant-add.mlir (128 of 141)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice.mlir (129 of 141)
PASS: IR Forge / Mountain Goat :: ops/scalar-bad-op.mlir (130 of 141)
PASS: IR Forge / Mountain Goat :: lowering/static-add-affine.mlir (131 of 141)
PASS: IR Forge / Mountain Goat :: lowering/mixed-add-one-check.mlir (132 of 141)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice-dynamic-exact.mlir (133 of 141)
PASS: IR Forge / Mountain Goat :: loop-order/bad-option.mlir (134 of 141)
PASS: IR Forge / Mountain Goat :: transformer30/verifiers.mlir (135 of 141)
PASS: IR Forge / Mountain Goat :: lowering/dynamic-add-affine.mlir (136 of 141)
PASS: IR Forge / Mountain Goat :: contraction28/verifiers.mlir (137 of 141)
PASS: IR Forge / Mountain Goat :: attention29/verifiers.mlir (138 of 141)
PASS: IR Forge / Mountain Goat :: generalize32/verifiers.mlir (139 of 141)
PASS: IR Forge / Mountain Goat :: ffn34/examples.mlir (140 of 141)
PASS: IR Forge / Mountain Goat :: ffn34/against-reference.mlir (141 of 141)

Testing Time: 142.50s

Total Discovered Tests: 141
  Passed: 141 (100.00%)

Limits and what is not established

• One toy language. Seven tokens, four symbols, a rule that a single attention head can implement by looking two positions back. The result does not show that the model could learn anything harder, and nothing was tried on text.
• Not Chapter 30's transformer. This model has one head, one block and a single D = 8 width. Chapter 30's has two heads and a different layout; its backward pass is still not derived and its weights are still random. Two heads and stacked blocks need a split-and-merge of heads in the backward pass and a second set of residual paths, not done here.
• The result is one run in a chaotic regime. The measured sensitivity (a 10⁻¹² nudge changes the step-100 loss by a factor of six) means this exact run cannot be reproduced digit for digit by a differently ordered computation after about step 70. In the sweep's 12 runs at rates 0.25 to 1.0 the training loss ends below 0.05, and the number correct of all 256 ranges from 189 to 256; the 256 of 256 is the best case. Whether the spread would narrow with more steps was not tested.
• The mask-and-content question is open. The attention pattern printed is for one all-identical sequence. That the model looks two positions back is shown there; that it copies the content it finds there is only indirectly supported by 256 of 256 correct predictions on all 64 sequences, and the attention for other sequences was not looked at.
• Finite differences cover two weights per matrix, 32 of 648, for one starting-weight set; the gradient-versus-reference comparison covers all of them for three. The agreement is limited to about 10⁻⁴ by the step and the printed digits.
• The softmax-without-max gap (above) means overflow robustness is untested.
• The training program is large and slow. 12,455 lines and about 8 minutes (one measurement, this container) because the compiler has no loops over steps for this program shape, no common-subexpression elimination, and never frees temporaries. That is a limitation of this compiler, not of the method; no attempt was made to speed it up.
• Generation is not shown. The model is trained to predict the next token; this chapter does not sample from it or measure tie-breaking, which Chapter 32 did for a bigram model.
• Platform: x86-64 Linux, LLVM 18.1.3.

Reproducing

cd docs/part32/code && ./build.sh
cd ../../part35/code
python3 make_examples.py && ./run_examples.sh > examples_out.txt
./check_lm.py --full > check_lm_out.txt
./sensitivity.py > sensitivity_out.txt
./lr_sweep.py > lr_sweep_out.txt
./model_mutation.py > model_mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 141 tests

Chapter summary

• A one-block, one-head, causal transformer language model (embeddings, layer norms, attention, feed-forward, residuals, vocabulary projection; 648 weights) is trained in Mountain Goat by hand-derived backpropagation. The task is to repeat a triple; after 200 steps the loss is 0.0009 and every scored prediction on all 64 possible sequences (256) is right, with attention that looks two positions back.
• New pieces: 16 sequences stacked into 112 rows with a block-diagonal causal mask; a cross-entropy that counts only scored rows; the attention backward pass through the mask (it needs no extra term); the three-path gradient into the layer norm; and the embedding gradient transpose(x) @ dx0.
• All 16 gradient matrices equal an independent per-position Python calculation at three starting points; 32 weights agree with finite differences (after a repair to the check, reported above).
• Training is chaotic: a 10⁻¹² change in one weight changes the loss at step 100 by a factor of six, so the run agrees with the reference digit for digit only up to about step 50 and is checked by its claims afterwards. The 256 of 256 is the best of four seeds (256, 239, 249, 239).
• Eighteen wrong programs were run through the checker: seventeen caught, one not (softmax without the row-max subtraction, which only matters for large scores).

Self-check questions

Each answer is collapsed; try the question first.

1. Why does the attention backward pass need no extra term for the causal mask?

   Answer

   The mask is a constant added to the scores before the softmax, so it has no gradient of its own. The masked entries of the attention matrix a are exactly 0 (softmax of −10⁹), and ds = a * (da − row_sum(a * da)) is a product with a, so ds is exactly 0 there: nothing flows back to positions that were not seen. The mask is already inside a.

2. The embedding is a lookup, yet its gradient is written transpose(x) @ dx0. Why does that work?

   Answer

   x has one 1 per row, in the column of that row's token. transpose(x) @ dx0 gives row t = the sum of the rows of dx0 over the positions whose token is t: a scatter-add. That is the gradient of a lookup: a token used in several places receives the sum of the gradients from all of them, and a token not used receives zero.

3. Why are the 16 sequences stacked into one 112×112 attention instead of being processed one at a time?

   Answer

   Mountain Goat has only rank-2 matrices, no batch dimension; stacking is the way to process several sequences in one program, and the mask (block-diagonal as well as causal) prevents any row from attending to another sequence. The price is that the score matrix is 112×112, 16 times larger than 16 separate 7×7 matrices would need, mostly entries that the mask zeroes.

4. At step 100 the program's loss is 0.0356 and the Python reference's is 0.0346. Is the program wrong?

   Answer

   No evidence says so. The two agree to six digits up to step 50, all 16 gradient matrices agree at three starting points, and sensitivity.py shows that the same Python code, with one weight changed by 10⁻¹², goes from 0.0346 to 0.2115 at step 100. The training is chaotic after about step 70, so two correct implementations whose floating-point operations are ordered differently will not agree digit for digit there. This is why the check compares claims instead of digits from step 100 on.

5. The sweep shows 256, 239, 249, 239 correct for the four seeds at learning rate 1.0. What does that say about the "256 of 256" result?

   Answer

   That it is the best of the four runs, not typical: the model reliably fits its 16 training sequences (loss below 0.005 in all four) but generalizes to the other 48 sequences only 93% to 100% of the time, depending on the starting weights. "The model learns the task" is true of the training; "the model has found the rule" is true of some runs and only approximately of others.

6. One mutant, softmax without the row-maximum subtraction, was not caught. Does that mean the checks are wrong?

   Answer

   No, it means they have a gap. Without the subtraction softmax is the same function, and for scores of this size the numbers are the same; the subtraction exists to prevent overflow of exp for large scores. None of the checks uses large scores, so none can tell the two apart. A test with scores in the hundreds would catch it; none was written, and the page says so rather than counting the mutant as caught.

7. Why was every intermediate result bound with its own let, instead of writing each gradient as a function as in Chapter 34?

   Answer

   Because Mountain Goat expands a function call in place wherever it is used, so a quantity used by several later gradients is recomputed for each, and the work multiplies with depth; the generated code also never frees memory. The nested version was killed, most likely for lack of memory. Binding each intermediate with let computes it once.