Full Model

This page joins the earlier pages into one pass through the model, in each direction. It adds no new mathematics. Each step is one call to a function explained elsewhere, and the order of those calls is what this page documents.

The numbers on this page are labelled Illustrative Example. They come from a tiny model whose parameters are random draws from the real init_params, so its logits and loss say nothing about the trained model.

Forward pass

What is it?

Whole model, forwardImplemented

Whole model

forward takes one sequence of ids. It adds the embeddings, runs every block in turn, applies the final LayerNorm and projects with the head. It returns the logits and a cache, and applies no softmax to the logits.

Why does the model need it?

This is the function that turns token ids into one score per vocabulary token, at every position. Training, validation and generation all call it.

Mathematics

forward first checks the length:

The sequence must fit the position table #

T \le T_{\max}

Next it computes the first X. Embeddings has the equation. Then come the blocks, which Transformer Block explains:

Blocks in sequence #

X^{(l+1)} = \operatorname{Block}_l\!\big(X^{(l)}\big), \qquad l = 0, \dots, L - 1

Final LayerNorm #

X_f = \operatorname{LN}_f\!\big(X^{(L)}\big)

The head #

Z = X_f W_h + \mathbf{1} b_h^\top

The last two together:

The last two steps together #

\mathrm{Logits} = \operatorname{LN}_f\!\big(X^{(L)}\big)\, W_h + b_h

What does each symbol mean?

Symbol	Meaning	In code	Shape
$T$	The length of the sequence	`len(ids)`	scalar
$T_{\max}$	The number of rows in the position table	`cfg.context_len`	scalar
$X^{(l)}$	The residual stream before block l. X^{(0)} is the output of embed_forward	`X`	(T, d)
$L$	The number of blocks. On this page L always means this, and the loss is written as the word loss	`cfg.num_layers`	scalar
$\operatorname{Block}_l$	block_forward with the parameters whose keys start with b{l}.	`block_forward`	(T, d) to (T, d)
$\operatorname{LN}_f$	The final LayerNorm	`layer_norm_forward with lnf_g, lnf_b`	(T, d) to (T, d)
$X_f$	The output of the final LayerNorm	`Xf`	(T, d)
$W_h, b_h$	Weights and bias of the head	`head_w, head_b`	(d, V) and (V,)
$Z, \mathrm{Logits}$	One score per vocabulary token, at each position	`logits`	(T, V)
$V$	The vocabulary size. It depends on the corpus	`cfg.vocab_size`	scalar

Tiny numerical example

The whole forward passIllustrative ExampleA tiny model: the parameters are random draws from the real init_params, not learned values. The other inputs are hand-picked. Sizes here: V = 5, T_max = 4, d = 4, H = 2, L = 2. Logits from unlearned parameters mean nothing.

ids = [2, 0, 2]

X after embed_forward, shape (3, 4)
-0.1603	-0.333	0.6043	1.3778
0.9295	-0.6287	0.5279	-0.4679
-1.2061	0.236	0.213	1.5067

X after block 0, shape (3, 4)
-0.8929	-1.0414	-0.152	1.2843
-0.1405	-0.8955	-1.439	-0.1701
-0.9039	0.5315	-0.4214	1.1834

X after block 1, shape (3, 4)
-0.291	-1.7398	0.9376	1.9181
1.3171	0.0703	-0.4637	-0.5218
0.1232	-0.228	0.5934	1.661

Xf, shape (3, 4)
-0.3631	-1.4212	0.5341	1.2502
1.6455	-0.0408	-0.763	-0.8417
-0.5825	-1.0762	0.0788	1.5799

logits, shape (3, 5)
0.3811	-0.3483	0.5361	1.9487	1.7362
-0.2283	1.5423	0.0637	-1.5603	0.4553
0.5927	-0.4419	0.7434	1.9937	1.3703

Tensor shapes

config.py has no value for V, because it depends on the corpus, so V stays a symbol. T is at most 16 in training and generation, and the assertion allows up to 32.

One sequence through the modeltransformer.py:166

ids(T,)(16,)
X(T, d)(16, 32)
Xf(T, d)(16, 32)
logits(T, V)(16, V)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32. V depends on the corpus. It is not a value in config.py.

Actual implementation

transformer.pyforwardGitHub

155def forward(params, ids, cfg: ModelConfig):156    assert len(ids) <= cfg.context_len, "sequence longer than the position table"157158    X = embed_forward(params, ids)159    block_caches = []160161    for l in range(cfg.num_layers):162        X, c = block_forward(X, params, l, cfg)163        block_caches.append(c)164165    Xf, c_lnf = layer_norm_forward(X, params["lnf_g"], params["lnf_b"], cfg.ln_eps)166    logits = Xf @ params["head_w"] + params["head_b"]167168    return logits, dict(blocks=block_caches, lnf=c_lnf, Xf=Xf)

What it returns. forward returns the logits and a cache (transformer.py:168). Logits are not probabilities, and forward applies no softmax to them. That happens outside, in one of two places:

The loss. cross_entropy_forward computes log probabilities from the logits, and returns their exponential as probs, in training (transformer.py:173) and in validation (train.py:31). It does not call softmax_forward, and computes the same result in log form.
The sampler. sample_next calls softmax_forward on the logits divided by the temperature (generate.py:11).

The only softmax_forward call inside forward turns attention scores into attention weights (transformer.py:83). That is not applied to logits.

Only loss_and_grads uses the cache. Validation (train.py:30) and generation (generate.py:32) throw it away, though forward builds it each time.

Passed as	Holds	Built at	Read by
`cache["Xf"]`	`Xf`: The output of the final LayerNorm, the input of the head.	transformer.py:168	`linear_backward`, at transformer.py:176
`cache["lnf"]`	`Xhat`: The normalised rows.`s`: One over the root of variance plus epsilon, per row.	ops.py:47	`layer_norm_backward`, at transformer.py:181
`cache["blocks"][l]`	`ln1`: The cache of LayerNorm 1.`att`: The cache of attention.`ln2`: The cache of LayerNorm 2.`ff`: The cache of the feed-forward layer.	transformer.py:139	`block_backward`, at transformer.py:187

Code and math connection

Equation	Role	Lines	Code
The sequence must fit the position table	Computes it	transformer.py:156	`assert len(ids) <= cfg.context_len, "sequence longer than the position table"`
Blocks in sequence	Computes it	transformer.py:162	`X, c = block_forward(X, params, l, cfg)`
Final LayerNorm	Computes it	transformer.py:165	`Xf, c_lnf = layer_norm_forward(X, params["lnf_g"], params["lnf_b"], cfg.ln_eps)`
The head	Written inline	transformer.py:166	`logits = Xf @ params["head_w"] + params["head_b"]`
The last two steps together	Computes it	transformer.py:165	`Xf, c_lnf = layer_norm_forward(X, params["lnf_g"], params["lnf_b"], cfg.ln_eps)`
The last two steps together	Written inline	transformer.py:166	`logits = Xf @ params["head_w"] + params["head_b"]`
The last two steps together	Tests it	tests/test_transformer.py:83	`la, _ = forward(params, a, cfg)`

Where is it used?

Source

transformer.py, lines 155-168, forward

PDF reference

Section 4.5, Eq. 54, p. 11

Used by

Called by: generate.generate, train.validation_loss, transformer.loss_and_grads

Tested by: tests.test_checkpoint.test_loaded_model_gives_identical_logits, tests.test_transformer.test_future_tokens_do_not_change_earlier_logits

Tests it: tests/test_transformer.py, line 83 (test_future_tokens_do_not_change_earlier_logits), for The last two steps together

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
The sequence must fit the position table	Computes it	`transformer.py`	`forward`	156	GitHub: transformer.py, line 156
Blocks in sequence	Computes it	`transformer.py`	`forward`	162	GitHub: transformer.py, line 162
Final LayerNorm	Computes it	`transformer.py`	`forward`	165	GitHub: transformer.py, line 165
The head	Written inline	`transformer.py`	`forward`	166	GitHub: transformer.py, line 166
The last two steps together	Computes it	`transformer.py`	`forward`	165	GitHub: transformer.py, line 165
The last two steps together	Written inline	`transformer.py`	`forward`	166	GitHub: transformer.py, line 166

Backward pass

What is it?

Whole model, backwardImplemented

Whole model

loss_and_grads runs forward and the cross-entropy loss on one sequence, then the whole backward pass. It goes through the head, the final LayerNorm, the blocks last to first, and the embeddings. It returns the loss and a dictionary with one gradient per parameter, under the parameter's own key.

Why does the model need it?

The optimizer needs one gradient per parameter, and each backward function needs the gradient produced by the step after it in the forward pass. This function calls them in that order and keeps the results together.

Mathematics

Every derivative here was derived on an earlier page. The chain starts where the loss meets the logits, from Softmax & Cross-Entropy:

The combined backward, where E_y is one-hot at each correct token #

dZ = \frac{P - E_y}{T}

The head is a linear layer, so the formulas from Linear / ReLU / Initialisation apply, with X = Xf, W = head_w and dY = dZ:

Gradient on the weights #

dW = X^{\top} dY

Gradient on the bias, the sum of dY over rows #

db = \sum_{t} dY_{t}

Gradient on the input #

dX = dY\, W^{\top}

The next lines state only the order of the calls. They add no new derivative: the first is the three formulas above, and the others are derived on LayerNorm, Transformer Block and Embeddings.

Order of calls, the head #

dX_f,\ dW_h,\ db_h = \operatorname{Head}^{\text{back}}(dZ)

Order of calls, the final LayerNorm #

dX^{(L)},\ d\gamma_f,\ d\beta_f = \operatorname{LN}_f^{\text{back}}(dX_f)

Order of calls, the blocks in reverse #

dX^{(l)} = \operatorname{Block}_l^{\text{back}}\!\big(dX^{(l+1)}\big), \qquad l = L - 1, \dots, 0

Order of calls, the embeddings #

d\,\mathrm{tok\_emb},\ d\,\mathrm{pos\_emb} = \operatorname{Embed}^{\text{back}}\!\big(dX^{(0)}\big)

What does each symbol mean?

Symbol	Meaning	In code	Shape
$dZ$	Gradient of the loss on the logits	`dZ`	(T, V)
$P, E_y$	The probabilities, and one-hot rows of the targets	`probs, and the -= 1.0 in cross_entropy_backward`	(T, V)
$dX_f$	Gradient on the output of the final LayerNorm	`dXf`	(T, d)
$dW_h, db_h$	Gradients on the head	`dW, db`	(d, V) and (V,)
$d\gamma_f, d\beta_f$	Gradients on the final LayerNorm	`dg, db`	(d,)
$dX^{(l)}$	Gradient on the residual stream before block l. The code reuses the name dX	`dX`	(T, d)
$\operatorname{Head}^{\text{back}}, \operatorname{LN}_f^{\text{back}}$	linear_backward with the head, and layer_norm_backward with the final LayerNorm's cache	`linear_backward, layer_norm_backward`	functions
$\operatorname{Block}_l^{\text{back}}, \operatorname{Embed}^{\text{back}}$	block_backward for block l, and embed_backward	`block_backward, embed_backward`	functions
$X, W, dY$	The linear layer's names on the page that derives it. At the head they are Xf, head_w and dZ	`X, W, dY`	(T, d), (d, V) and (T, V)
$d\,\mathrm{tok\_emb}, d\,\mathrm{pos\_emb}$	Gradients on the two embedding tables	`grads["tok_emb"], grads["pos_emb"]`	(V, d) and (T_max, d)

Tiny numerical example

The gradient on its way backIllustrative ExampleA tiny model: the parameters are random draws from the real init_params, not learned values. The other inputs are hand-picked. ids = [2, 0, 2], targets = [0, 2, 1].

loss = 2.5621

dZ, shape (3, 5)
-0.3039	0.0142	0.0344	0.1412	0.1141
0.0319	0.1872	-0.2907	0.0084	0.0631
0.0381	-0.3198	0.0443	0.1546	0.0829

dXf, shape (3, 4)
0.1688	-0.0364	0.2427	0.2007
0.2094	0.0265	0.1092	-0.022
-0.2382	-0.0953	0.1574	0.2309

dX after the final LayerNorm, shape (3, 4)
0.0427	-0.0354	0.0359	-0.0432
0.0189	-0.0696	0.1104	-0.0597
-0.2276	0.0807	0.1849	-0.0381

dX after block 1, shape (3, 4)
-0.0022	0.005	0.0405	-0.0433
0.0062	-0.0329	0.0706	-0.044
-0.2087	0.0139	0.1012	0.0936

dX after block 0, shape (3, 4)
-0.0893	0.1395	-0.0779	0.0277
0.0071	0.0263	0.025	-0.0584
-0.1968	0.037	0.0569	0.1029

Key	Shape	Largest entry
`tok_emb`	`(5, 4)`	`2.86e-1`
`pos_emb`	`(4, 4)`	`1.97e-1`
`b0.ln1_g`	`(4,)`	`2.00e-1`
`b0.ln1_b`	`(4,)`	`1.94e-1`
`b0.wq`	`(4, 4)`	`1.05e-1`
`b0.bq`	`(4,)`	`9.06e-2`
`b0.wk`	`(4, 4)`	`2.87e-1`
`b0.bk`	`(4,)`	`1.39e-17`
`b0.wv`	`(4, 4)`	`2.10e-1`
`b0.bv`	`(4,)`	`2.36e-1`
`b0.wo`	`(4, 4)`	`3.59e-1`
`b0.bo`	`(4,)`	`2.43e-1`
`b0.ln2_g`	`(4,)`	`7.30e-2`
`b0.ln2_b`	`(4,)`	`4.44e-2`
`b0.ff1_w`	`(4, 6)`	`1.40e-1`
`b0.ff1_b`	`(6,)`	`1.21e-1`
`b0.ff2_w`	`(6, 4)`	`1.56e-1`
`b0.ff2_b`	`(4,)`	`2.12e-1`
`b1.ln1_g`	`(4,)`	`1.07e-1`
`b1.ln1_b`	`(4,)`	`1.86e-1`
`b1.wq`	`(4, 4)`	`2.62e-2`
`b1.bq`	`(4,)`	`2.19e-2`
`b1.wk`	`(4, 4)`	`7.08e-2`
`b1.bk`	`(4,)`	`2.08e-17`
`b1.wv`	`(4, 4)`	`2.55e-1`
`b1.bv`	`(4,)`	`1.80e-1`
`b1.wo`	`(4, 4)`	`3.51e-1`
`b1.bo`	`(4,)`	`2.18e-1`
`b1.ln2_g`	`(4,)`	`1.25e-1`
`b1.ln2_b`	`(4,)`	`1.77e-1`
`b1.ff1_w`	`(4, 6)`	`1.85e-1`
`b1.ff1_b`	`(6,)`	`2.40e-1`
`b1.ff2_w`	`(6, 4)`	`2.57e-1`
`b1.ff2_b`	`(4,)`	`3.31e-1`
`lnf_g`	`(4,)`	`6.34e-1`
`lnf_b`	`(4,)`	`5.09e-1`
`head_w`	`(4, 5)`	`6.45e-1`
`head_b`	`(5,)`	`3.04e-1`

Checking the backward numericallyIllustrative ExampleA test-time idea, not part of the pipeline: gradcheck.numerical_grad is only called by tests. It nudges each number up and down by h = 1e-5 and measures how the loss changes. Tests compare the central differences here with that real function.

Gradient	From the backward formula	From central differences	Largest gap
`grads["tok_emb"]`	`[0.0071, 0.0263, 0.0250, -0.0584][0.0000, 0.0000, 0.0000, 0.0000][-0.2861, 0.1765, -0.0209, 0.1306][0.0000, 0.0000, 0.0000, 0.0000][0.0000, 0.0000, 0.0000, 0.0000]`	`[0.0071, 0.0263, 0.0250, -0.0584][0.0000, 0.0000, 0.0000, 0.0000][-0.2861, 0.1765, -0.0209, 0.1306][0.0000, 0.0000, 0.0000, 0.0000][0.0000, 0.0000, 0.0000, 0.0000]`	`5.9e-10`
`grads["pos_emb"]`	`[-0.0893, 0.1395, -0.0779, 0.0277][0.0071, 0.0263, 0.0250, -0.0584][-0.1968, 0.0370, 0.0569, 0.1029][0.0000, 0.0000, 0.0000, 0.0000]`	`[-0.0893, 0.1395, -0.0779, 0.0277][0.0071, 0.0263, 0.0250, -0.0584][-0.1968, 0.0370, 0.0569, 0.1029][0.0000, 0.0000, 0.0000, 0.0000]`	`5.5e-11`
`grads["b0.wk"]`	`[-0.2875, 0.1618, 0.0006, -0.0032][0.0849, -0.0466, -0.0000, 0.0009][-0.0604, 0.0338, 0.0001, -0.0007][0.2630, -0.1489, -0.0007, 0.0030]`	`[-0.2875, 0.1618, 0.0006, -0.0032][0.0849, -0.0466, -0.0000, 0.0009][-0.0604, 0.0338, 0.0001, -0.0007][0.2630, -0.1489, -0.0007, 0.0030]`	`2.0e-10`
`grads["b0.ff1_w"]`	`[0.0020, -0.0642, 0.1397, 0.0000, 0.0000, -0.0112][-0.0014, 0.1045, -0.1138, 0.0000, 0.0000, 0.0075][-0.0015, 0.0791, 0.0038, 0.0000, 0.0000, 0.0082][0.0008, -0.1194, -0.0296, 0.0000, 0.0000, -0.0046]`	`[0.0020, -0.0642, 0.1397, 0.0000, 0.0000, -0.0112][-0.0014, 0.1045, -0.1138, 0.0000, 0.0000, 0.0075][-0.0015, 0.0791, 0.0038, 0.0000, 0.0000, 0.0082][0.0008, -0.1194, -0.0296, 0.0000, 0.0000, -0.0046]`	`4.5e-11`
`grads["b1.ln2_g"]`	`[-0.0470, 0.0010, -0.1254, -0.0685]`	`[-0.0470, 0.0010, -0.1254, -0.0685]`	`1.9e-11`
`grads["b1.wq"]`	`[0.0004, -0.0009, -0.0111, 0.0164][-0.0002, 0.0003, 0.0047, -0.0073][0.0017, -0.0040, -0.0198, -0.0127][-0.0019, 0.0045, 0.0262, 0.0036]`	`[0.0004, -0.0009, -0.0111, 0.0164][-0.0002, 0.0003, 0.0047, -0.0073][0.0017, -0.0040, -0.0198, -0.0127][-0.0019, 0.0045, 0.0262, 0.0036]`	`4.0e-11`
`grads["lnf_g"]`	`[0.4220, 0.1531, 0.0587, 0.6343]`	`[0.4220, 0.1531, 0.0587, 0.6343]`	`1.9e-11`
`grads["head_w"]`	`[0.1406, 0.4892, -0.5165, -0.1275, 0.0142][0.3896, 0.3164, -0.0846, -0.3673, -0.2540][-0.1836, -0.1605, 0.2436, 0.0812, 0.0193][-0.3466, -0.6451, 0.3576, 0.4136, 0.2205]`	`[0.1406, 0.4892, -0.5165, -0.1275, 0.0142][0.3896, 0.3164, -0.0846, -0.3673, -0.2540][-0.1836, -0.1605, 0.2436, 0.0812, 0.0193][-0.3466, -0.6451, 0.3576, 0.4136, 0.2205]`	`3.9e-11`
`grads["head_b"]`	`[-0.2339, -0.1184, -0.2120, 0.3042, 0.2602]`	`[-0.2339, -0.1184, -0.2120, 0.3042, 0.2602]`	`1.9e-11`

Two rows of that table are worth a closer look. Each .bk entry is round-off around 0, because softmax ignores a constant added to a row, as Attention explains. The tiny pos_emb has more rows than this sequence has tokens, so its last row has a gradient of exactly 0, as Embeddings explains.

Tensor shapes

The gradient from the logits to the first block's inputtransformer.py:187

dZ(T, V)(16, V)
dXf(T, d)(16, 32)
dX(T, d)(16, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. V depends on the corpus. It is not a value in config.py. d = emb_dim = 32.

Parameter gradients of the headtransformer.py:176

head_w(d, V)(32, V)
head_b(V,)

d = emb_dim = 32. V depends on the corpus. It is not a value in config.py.

Parameter gradients of the embedding tablestransformer.py:188

tok_emb(V, d)(V, 32)
pos_emb(Tmax, d)(32, 32)

V depends on the corpus. It is not a value in config.py. d = emb_dim = 32. Tmax = context_len = 32.

Actual implementation

transformer.pyloss_and_gradsGitHub

171def loss_and_grads(params, ids, targets, cfg: ModelConfig):172    logits, cache = forward(params, ids, cfg)173    loss, probs = cross_entropy_forward(logits, targets)174    grads = {k: np.zeros_like(v) for k, v in params.items()}175    dZ = cross_entropy_backward(probs, targets)176    dXf, dW, db = linear_backward(cache["Xf"], params["head_w"], dZ)177178    grads["head_w"] += dW179    grads["head_b"] += db180181    dX, dg, db = layer_norm_backward(dXf, cache["lnf"], params["lnf_g"])182183    grads["lnf_g"] += dg184    grads["lnf_b"] += db185186    for l in reversed(range(cfg.num_layers)):187        dX = block_backward(dX, cache["blocks"][l], params, l, cfg, grads)188    embed_backward(dX, ids, grads)189190    return loss, grads

The order is forward, cross_entropy_forward, cross_entropy_backward, linear_backward for the head, layer_norm_backward for the final LayerNorm, block_backward for each block with the last block first (transformer.py:186), and embed_backward.

Code and math connection

Equation	Role	Lines	Code
Order of calls, the head	Computes it	transformer.py:176	`dXf, dW, db = linear_backward(cache["Xf"], params["head_w"], dZ)`
Order of calls, the final LayerNorm	Computes it	transformer.py:181	`dX, dg, db = layer_norm_backward(dXf, cache["lnf"], params["lnf_g"])`
Order of calls, the blocks in reverse	Computes it	transformer.py:186	`for l in reversed(range(cfg.num_layers)):`
Order of calls, the blocks in reverse	Computes it	transformer.py:187	`dX = block_backward(dX, cache["blocks"][l], params, l, cfg, grads)`
Order of calls, the embeddings	Computes it	transformer.py:188	`embed_backward(dX, ids, grads)`
Order of calls, the embeddings	Tests it	tests/test_transformer.py:94	`assert grads.keys() == params.keys()`

Where is it used?

test_every_tensor_passes compares its gradients with central differences, and test_overfit_one_window checks that its gradients can drive the loss down on one window (tests/test_transformer.py:110).

Source

transformer.py, lines 171-190, loss_and_grads

PDF reference

Section 4.5, p. 11

Used by

Called by: train.batch_loss_and_grads

Tested by: tests.test_gradcheck.test_every_tensor_passes, tests.test_gradcheck.test_key_bias_gradient_is_zero, tests.test_transformer.test_grads_have_the_shape_of_params, tests.test_transformer.test_initial_loss_is_close_to_ln_vocab, tests.test_transformer.test_overfit_one_window

Tests it: tests/test_transformer.py, line 94 (test_grads_have_the_shape_of_params), for Order of calls, the embeddings

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Order of calls, the head	Computes it	`transformer.py`	`loss_and_grads`	176	GitHub: transformer.py, line 176
Order of calls, the final LayerNorm	Computes it	`transformer.py`	`loss_and_grads`	181	GitHub: transformer.py, line 181
Order of calls, the blocks in reverse	Computes it	`transformer.py`	`loss_and_grads`	186	GitHub: transformer.py, line 186
Order of calls, the blocks in reverse	Computes it	`transformer.py`	`loss_and_grads`	187	GitHub: transformer.py, line 187
Order of calls, the embeddings	Computes it	`transformer.py`	`loss_and_grads`	188	GitHub: transformer.py, line 188

How this reaches a parameter update. batch_loss_and_grads calls this once per window and averages the dictionaries (train.py:22). train.main then clips the average (train.py:62) and passes it to opt.step (train.py:66). AdamW / Learning Rate / Clipping covers those steps.

Gradient flow

Each arrow below is one statement of loss_and_grads, with its text taken from the code at the line shown. Select a box to see the function that produced it, and the lines in loss_and_grads. Transformer Block has the same diagram for the inside of a block.

Gradient flow through the whole model, in the order of loss_and_grads

logits, cache = forward(params, ids, cfg)transformer.py:172
Also returns the cache that the backward steps read.
loss, probs = cross_entropy_forward(logits, targets)transformer.py:173
The shape is that of probs. loss is one number.
grads = {k: np.zeros_like(v) for k, v in params.items()}transformer.py:174No function call
Not a function call. One array of zeros per parameter, under the same key.
dZ = cross_entropy_backward(probs, targets)transformer.py:175
dXf, dW, db = linear_backward(cache["Xf"], params["head_w"], dZ)transformer.py:176
Fills grads: head_w, head_b
dX, dg, db = layer_norm_backward(dXf, cache["lnf"], params["lnf_g"])transformer.py:181
Fills grads: lnf_g, lnf_b
dX = block_backward(dX, cache["blocks"][l], params, l, cfg, grads)transformer.py:187
Fills grads: all 16 entries named b1.*
dX = block_backward(dX, cache["blocks"][l], params, l, cfg, grads)transformer.py:187
Fills grads: all 16 entries named b0.*
embed_backward(dX, ids, grads)transformer.py:188
Fills grads: tok_emb, pos_emb

Selected: logits, Forward pass, at transformer.py:172

transformer.pyforwardGitHub

155def forward(params, ids, cfg: ModelConfig):156    assert len(ids) <= cfg.context_len, "sequence longer than the position table"157158    X = embed_forward(params, ids)159    block_caches = []160161    for l in range(cfg.num_layers):162        X, c = block_forward(X, params, l, cfg)163        block_caches.append(c)164165    Xf, c_lnf = layer_norm_forward(X, params["lnf_g"], params["lnf_b"], cfg.ln_eps)166    logits = Xf @ params["head_w"] + params["head_b"]167168    return logits, dict(blocks=block_caches, lnf=c_lnf, Xf=Xf)

transformer.pyloss_and_gradsGitHub

171def loss_and_grads(params, ids, targets, cfg: ModelConfig):172    logits, cache = forward(params, ids, cfg)173    loss, probs = cross_entropy_forward(logits, targets)174    grads = {k: np.zeros_like(v) for k, v in params.items()}175    dZ = cross_entropy_backward(probs, targets)176    dXf, dW, db = linear_backward(cache["Xf"], params["head_w"], dZ)177178    grads["head_w"] += dW179    grads["head_b"] += db180181    dX, dg, db = layer_norm_backward(dXf, cache["lnf"], params["lnf_g"])182183    grads["lnf_g"] += dg184    grads["lnf_b"] += db185186    for l in reversed(range(cfg.num_layers)):187        dX = block_backward(dX, cache["blocks"][l], params, l, cfg, grads)188    embed_backward(dX, ids, grads)189190    return loss, grads

The gradient dictionary

loss_and_grads starts by making one array of zeros per parameter, under the same key and with the same shape:

transformer.pyloss_and_gradsGitHub

174    grads = {k: np.zeros_like(v) for k, v in params.items()}

Every backward step then adds to that dictionary. It adds in place, so block_backward and embed_backward return without handing grads back: their caller's dictionary already holds the result. The head and the final LayerNorm are added by loss_and_grads itself, and the gradient that goes on to the next step travels separately, as dX.

Within one call each key is written at one place in the code. Only tok_emb can receive several additions there: embed_backward adds one row per position, so a repeated id accumulates.

Because the keys match, Adam.step can look up the gradient of each parameter by name. A test checks that the keys and shapes match (tests/test_transformer.py:94).

Parameter keys

The dictionary is flat: one string key per array. The keys follow three rules:

A suffix per array. The feed-forward layers and the head end in _w for weights and _b for biases. A LayerNorm ends in _g for its scale and _b for its shift.
Attention's four projections come from a loop over q, k, v and o, so "w" + n gives wq and "b" + n gives bq, and so on (transformer.py:31-33).
Block entries start with a prefix that holds the block index, counted from 0 (transformer.py:28). block_forward and block_backward build the same prefix (transformer.py:132), so b0.wq is the query weights of the first block. The PDF writes this prefix bl. (p. 9). This site follows the code.

The table lists each name once, with b{l}. standing for every block. Names and shapes come from calling init_params. head_w is a separate matrix, not tied to tok_emb.

Name	Shape	With config	Scalars	Role
`tok_emb`transformer.py:21	`(V, d)`	`(V, 32)`	`V × d`32V	Token table: one row per vocabulary id.Starts as Xavier uniform. Weight decay.
`pos_emb`transformer.py:22	`(T_max, d)`	`(32, 32)`	`T_max × d`1,024	Position table: one row per position.Starts as Xavier uniform. Weight decay.
`b{l}.ln1_g`transformer.py:29	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 1 scale, before attention.Starts as ones. No decay.
`b{l}.ln1_b`transformer.py:29	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 1 shift, before attention.Starts as zeros. No decay.
`b{l}.wq`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Query projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bq`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Query projection bias.Starts as zeros. No decay.
`b{l}.wk`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Key projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bk`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Key projection bias.Starts as zeros. No decay.
`b{l}.wv`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Value projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bv`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Value projection bias.Starts as zeros. No decay.
`b{l}.wo`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Attention output projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bo`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Attention output projection bias.Starts as zeros. No decay.
`b{l}.ln2_g`transformer.py:35	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 2 scale, before feed-forward.Starts as ones. No decay.
`b{l}.ln2_b`transformer.py:35	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 2 shift, before feed-forward.Starts as zeros. No decay.
`b{l}.ff1_w`transformer.py:36	`(d, F)`	`(32, 128)`	`d × F`4,096 per block	First feed-forward layer weights.Starts as Xavier uniform. Weight decay.
`b{l}.ff1_b`transformer.py:36	`(F,)`	`(128,)`	`F`128 per block	First feed-forward layer bias.Starts as zeros. No decay.
`b{l}.ff2_w`transformer.py:37	`(F, d)`	`(128, 32)`	`d × F`4,096 per block	Second feed-forward layer weights.Starts as Xavier uniform. Weight decay.
`b{l}.ff2_b`transformer.py:37	`(d,)`	`(32,)`	`d`32 per block	Second feed-forward layer bias.Starts as zeros. No decay.
`lnf_g`transformer.py:39	`(d,)`	`(32,)`	`d`32	Final LayerNorm scale.Starts as ones. No decay.
`lnf_b`transformer.py:39	`(d,)`	`(32,)`	`d`32	Final LayerNorm shift.Starts as zeros. No decay.
`head_w`transformer.py:40	`(d, V)`	`(32, V)`	`V × d`32V	Linear head weights. A separate matrix, not tied to tok_emb.Starts as Xavier uniform. Weight decay.
`head_b`transformer.py:40	`(V,)`	`(V,)`	`V`V	Linear head bias.Starts as zeros. No decay.

config.py has num_layers = 2, so the b{l}. rows appear once per block in the real dictionary:

Every key as it appears in the dictionary: 38 entries at num_layers = 2

Key	Shape with config
`tok_emb`	`(V, 32)`
`pos_emb`	`(32, 32)`
`b0.ln1_g`	`(32,)`
`b0.ln1_b`	`(32,)`
`b0.wq`	`(32, 32)`
`b0.bq`	`(32,)`
`b0.wk`	`(32, 32)`
`b0.bk`	`(32,)`
`b0.wv`	`(32, 32)`
`b0.bv`	`(32,)`
`b0.wo`	`(32, 32)`
`b0.bo`	`(32,)`
`b0.ln2_g`	`(32,)`
`b0.ln2_b`	`(32,)`
`b0.ff1_w`	`(32, 128)`
`b0.ff1_b`	`(128,)`
`b0.ff2_w`	`(128, 32)`
`b0.ff2_b`	`(32,)`
`b1.ln1_g`	`(32,)`
`b1.ln1_b`	`(32,)`
`b1.wq`	`(32, 32)`
`b1.bq`	`(32,)`
`b1.wk`	`(32, 32)`
`b1.bk`	`(32,)`
`b1.wv`	`(32, 32)`
`b1.bv`	`(32,)`
`b1.wo`	`(32, 32)`
`b1.bo`	`(32,)`
`b1.ln2_g`	`(32,)`
`b1.ln2_b`	`(32,)`
`b1.ff1_w`	`(32, 128)`
`b1.ff1_b`	`(128,)`
`b1.ff2_w`	`(128, 32)`
`b1.ff2_b`	`(32,)`
`lnf_g`	`(32,)`
`lnf_b`	`(32,)`
`head_w`	`(32, V)`
`head_b`	`(V,)`

The number of scalars, worked out from these shapes, with V left as a symbol:

Implemented

N = 2V d + T_{\max} d + 2d + V + L\left(4d^{2} + 2d F + 9d + F\right)

N = 65\,V + 26{,}496