A Complete GPT from Scratch
Every piece assembled into one program that trains
37.1 What You Have Built
So we’ve finally reached the pinnacle that we’ve been building toward. Thirty six chapters ago this book began with a single perceptron reading a two input truth table. Since then you have built everything from backpropagation and embeddings, to tokenizers and the transformer block. You also learned about things like the decoder only stack, the key value cache, reinforcement learning low rank adaptation and quantization. That’s more than half a century worth of material and if you reached this far, congratulations, you now know more about how AI truly works than 99% of the planet. Insofar as we built, each of those arrived as a separate program that demonstrated one idea and then stopped at the edge of it. This chapter takes them apart and reassembles them into a single self contained C file that trains on text and generates new text from a prompt, essentially your own GPT a “mini” LLM.
Complete means complete here. The forward pass implements a real architecture with nothing standing in for anything else, and the backward pass computes an exact gradient for every one of the parameters in the model, of which ours will have 27680 of them. The program verifies all of the derivatives we covered thus far against numerical differences before it trains. It then prints the comparison where you can read it. Running the check first is not caution for its own sake, since a backward pass that is wrong in one term still produces a loss curve that falls and looks broadly reasonable. On the machine used to prepare this chapter the verification, the training and the generation together finish in roughly twelve seconds.
After this chapter is done when many of the self described “experts” say “no one understands how AI works” you can choose to be the guy that says “well actually…” or you can be silent, for as the ancient Chinese philosopher Lao Tzu said “He who knows does not speak. He who speaks does not know”.
37.2 The Architecture
We begin by looking at the architecture we will build. What you see presented in the table below is that architecture, and every stage in the pipeline below came out of a chapter you have already worked through, and the program implements them in exactly this order. Our flow is what you can intuitively follow by this point, a token index enters, becomes a vector by lookup, gains a position vector by addition, and then passes through the transformer blocks. Each block normalizes and attends causally over everything to its left, adds the result back into the residual stream. It then normalizes again, applies a position wise feed forward network, and adds that back too. A final normalization prepares the activations for the output head, which produces one logit for every vocabulary entry. Softmax over those logits gives the distribution the next token is drawn from. And that’s your LLM.
| Stage | What it does | Chapter |
|---|---|---|
| Token embedding | Index becomes a vector | 10 |
| Position embedding | Added, not concatenated | 22 |
| RMSNorm | Scale by root mean square | 23 |
| Causal attention | Attend left, 4 heads | 21, 25 |
| Residual | Add the sublayer output | 24 |
| RMSNorm | Again, before the FFN | 23 |
| FFN with GELU | Widen to 128 and back | 24 |
| Residual | Add the sublayer output | 24 |
| Final RMSNorm | Once after the last block | 23 |
| Output head | Tied to the embedding | 26 |
| Softmax | Logits become a distribution | 4 |
When designing this architecture, the choice was rather difficult because there is always a trade off between too small to be useful or illustrate the principle or too large to train properly on the slowest machine. Thus the tradeoff I decided on was somewhere in between, the configuration is small enough to train in seconds and large enough to be a real transformer. The model width is 32, split across 4 attention heads of 8 dimensions each, with a feed forward width of 128 and two stacked blocks. The vocabulary holds 28 entries covering the lowercase alphabet, the space and the period, and the maximum sequence length is 64. If we scale the width to 7168, the heads to 128 and the layers to 61, then replace the feed forward network with the mixture of experts from the previous chapter, and the same code describes DeepSeek V3′s operation very closely, which we cover in the bonus material.
Figure 37-1 is the table above drawn twice, once as the whole stack and once with a single block opened out. Every box on the left is a row of that table, and the two transformer blocks in the middle of it expand into everything on the right.
If we read the right hand panel from the bottom then we will see that the two sublayers are the same shape twice. In a nutshell we normalize, do the work, add the input back. The four heads run in parallel inside the first one and their outputs are laid end to end and projected, while the second widens to 128 and comes back to 32. The two residual lines on the right are the reason a stack of these trains at all.
37.3 The Complete Program
Now we get to the interesting part, the implementation. The program does three things in order.
- It checks every gradient against a numerical estimate
- It trains all 27680 parameters with Adam,
- It generates text from the trained model.
The forward pass caches the activations it will need on the way back, which is why there is a Cache structure holding the normalized inputs, the per head queries, keys and values, the attention weights after softmax, and the feed forward activations before and after GELU. Caching costs memory that a production implementation would manage far more carefully, but at these dimensions the whole cache is a few hundred kilobytes. Every quantity stored there is one the backward pass genuinely needs, and working out which those are is most of the difficulty in writing a backward pass by hand.
/* 178_Capstone.c */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stddef.h>
#define VOCAB 28
#define D_MODEL 32
#define N_HEADS 4
#define D_HEAD 8
#define D_FF 128
#define N_LAYERS 2
#define MAX_SEQ 64
static float randf(void)
{
return (float)rand() / RAND_MAX;
}
static float gelu(float x)
{
float u = 0.7978846f * (x + 0.044715f * x * x * x);
return 0.5f * x * (1.0f + tanhf(u));
}
static float gelu_grad(float x)
{
float a = 0.044715f, c = 0.7978846f;
float u = c * (x + a * x * x * x);
float t = tanhf(u);
return 0.5f * (1.0f + t)
+ 0.5f * x * (1.0f - t * t) * c
* (1.0f + 3.0f * a * x * x);
}
static int char_to_id(char c)
{
if (c >= 'a' && c <= 'z') return c - 'a';
if (c == ' ') return 26;
return 27;
}
static char id_to_char(int id)
{
if (id < 26) return 'a' + id;
if (id == 26) return ' ';
return '.';
}
static void softmax(float *x, int n)
{
int i;
float mx = x[0], s = 0;
for (i = 1; i < n; i++) if (x[i] > mx) mx = x[i];
for (i = 0; i < n; i++) {
x[i] = expf(x[i] - mx);
s += x[i];
}
for (i = 0; i < n; i++) x[i] /= s;
}
/* ---- Parameters ---- */
typedef struct {
float tok_emb[VOCAB][D_MODEL];
float pos_emb[MAX_SEQ][D_MODEL];
float W_Q[N_LAYERS][N_HEADS][D_HEAD][D_MODEL];
float W_K[N_LAYERS][N_HEADS][D_HEAD][D_MODEL];
float W_V[N_LAYERS][N_HEADS][D_HEAD][D_MODEL];
float W_O[N_LAYERS][D_MODEL][D_MODEL];
float W_up[N_LAYERS][D_FF][D_MODEL];
float W_down[N_LAYERS][D_MODEL][D_FF];
float norm1_g[N_LAYERS][D_MODEL];
float norm2_g[N_LAYERS][D_MODEL];
float norm_f_g[D_MODEL];
}
GPT;
#define NPARAM (sizeof(GPT) / sizeof(float))
/* ---- Activation cache, one per forward pass ---- */
typedef struct {
int T;
int tok[MAX_SEQ];
/* embeddings */
float X0[MAX_SEQ][D_MODEL];
/* block input */
float Xa[N_LAYERS][MAX_SEQ][D_MODEL];
/* normed */
float n1[N_LAYERS][MAX_SEQ][D_MODEL];
/* 1/rms */
float s1[N_LAYERS][MAX_SEQ];
float q[N_LAYERS][N_HEADS][MAX_SEQ][D_HEAD];
float k[N_LAYERS][N_HEADS][MAX_SEQ][D_HEAD];
float v[N_LAYERS][N_HEADS][MAX_SEQ][D_HEAD];
/* post softmax */
float p[N_LAYERS][N_HEADS][MAX_SEQ][MAX_SEQ];
/* concat heads */
float cc[N_LAYERS][MAX_SEQ][D_MODEL];
/* after attn add */
float Xb[N_LAYERS][MAX_SEQ][D_MODEL];
float n2[N_LAYERS][MAX_SEQ][D_MODEL];
float s2[N_LAYERS][MAX_SEQ];
/* pre GELU */
float hf[N_LAYERS][MAX_SEQ][D_FF];
/* post GELU */
float gf[N_LAYERS][MAX_SEQ][D_FF];
/* after last block */
float Xc[MAX_SEQ][D_MODEL];
float sf[MAX_SEQ];
/* final normed */
float Xf[MAX_SEQ][D_MODEL];
float logit[MAX_SEQ][VOCAB];
}
Cache;
static void gpt_init(GPT *m)
{
int i, l, d;
float *w = (float *)m;
for (i = 0; i < (int)NPARAM; i++)
w[i] = (randf() * 2 - 1) * 0.08f;
for (l = 0; l < N_LAYERS; l++)
for (d = 0; d < D_MODEL; d++) {
m->norm1_g[l][d] = 1.0f;
m->norm2_g[l][d] = 1.0f;
}
for (d = 0; d < D_MODEL; d++)
m->norm_f_g[d] = 1.0f;
}
/* y = g * x / rms, returns 1/rms */
static float rmsnorm(const float *x, const float *g,
float *y, int n)
{
int i;
float ss = 0;
for (i = 0; i < n; i++) ss += x[i] * x[i];
float s = 1.0f / sqrtf(ss / n + 1e-5f);
for (i = 0; i < n; i++) y[i] = g[i] * x[i] * s;
return s;
}
/* dx from dy, accumulating dg */
static void rmsnorm_back(const float *x,
const float *g,
float s, const float *dy,
float *dx, float *dg, int n)
{
int i;
float dot = 0;
for (i = 0; i < n; i++) {
float xhat = x[i] * s;
dg[i] += dy[i] * xhat;
dot += dy[i] * g[i] * xhat;
}
dot /= n;
for (i = 0; i < n; i++)
dx[i] += s * (dy[i] * g[i] - x[i] * s * dot);
}
static void forward(const GPT *m,
const int *tok, int T,
Cache *c)
{
int i, j, l, h, d, e;
float X[MAX_SEQ][D_MODEL];
c->T = T;
for (i = 0; i < T; i++) c->tok[i] = tok[i];
for (i = 0; i < T; i++)
for (d = 0; d < D_MODEL; d++)
X[i][d] = c->X0[i][d] =
m->tok_emb[tok[i]][d] + m
->pos_emb[i][d];
for (l = 0; l < N_LAYERS; l++) {
memcpy(c->Xa[l], X, sizeof(X));
for (i = 0; i < T; i++)
c->s1[l][i] = rmsnorm(X[i], m->norm1_g[l],
c->n1[l][i], D_MODEL);
for (h = 0; h < N_HEADS; h++) {
for (i = 0; i < T; i++)
for (d = 0; d < D_HEAD; d++) {
float qa = 0, ka = 0, va = 0;
for (e = 0; e < D_MODEL; e++) {
float u = c->n1[l][i][e];
qa += m->W_Q[l][h][d][e] * u;
ka += m->W_K[l][h][d][e] * u;
va += m->W_V[l][h][d][e] * u;
}
c->q[l][h][i][d] = qa;
c->k[l][h][i][d] = ka;
c->v[l][h][i][d] = va;
}
float scale = 1.0f / sqrtf((float)D_HEAD);
for (i = 0; i < T; i++) {
float row[MAX_SEQ];
for (j = 0; j <= i; j++) {
float s = 0;
for (d = 0; d < D_HEAD; d++)
s += c->q[l][h][i][d]
* c->k[l][h][j][d];
row[j] = s * scale;
}
softmax(row, i + 1);
for (j = 0; j <= i; j++)
c->p[l][h][i][j] = row[j];
for (j = i + 1; j < T; j++)
c->p[l][h][i][j] = 0;
for (d = 0; d < D_HEAD; d++) {
float ctx = 0;
for (j = 0; j <= i; j++)
ctx += row[j] * c
->v[l][h][j][d];
c->cc[l][i][h * D_HEAD + d] = ctx;
}
}
}
for (i = 0; i < T; i++)
for (d = 0; d < D_MODEL; d++) {
float o = 0;
for (e = 0; e < D_MODEL; e++)
o += m->W_O[l][d][e] * c
->cc[l][i][e];
X[i][d] += o;
}
memcpy(c->Xb[l], X, sizeof(X));
for (i = 0; i < T; i++)
c->s2[l][i] = rmsnorm(X[i], m->norm2_g[l],
c->n2[l][i], D_MODEL);
for (i = 0; i < T; i++) {
for (d = 0; d < D_FF; d++) {
float a = 0;
for (e = 0; e < D_MODEL; e++)
a += m->W_up[l][d][e] * c
->n2[l][i][e];
c->hf[l][i][d] = a;
c->gf[l][i][d] = gelu(a);
}
for (d = 0; d < D_MODEL; d++) {
float o = 0;
for (e = 0; e < D_FF; e++)
o += m->W_down[l][d][e]
* c->gf[l][i][e];
X[i][d] += o;
}
}
}
memcpy(c->Xc, X, sizeof(X));
for (i = 0; i < T; i++)
c->sf[i] = rmsnorm(X[i], m->norm_f_g, c->Xf[i],
D_MODEL);
for (i = 0; i < T; i++)
for (j = 0; j < VOCAB; j++) {
float z = 0;
for (d = 0; d < D_MODEL; d++)
z += c->Xf[i][d] * m->tok_emb[j][d];
c->logit[i][j] = z;
}
}
static float loss_and_backward(const GPT *m, Cache *c,
GPT *gr)
{
int i, j, l, h, d, e;
int T = c->T;
float total = 0;
static float dX[MAX_SEQ][D_MODEL];
static float dXf[MAX_SEQ][D_MODEL];
memset(dX, 0, sizeof(dX));
memset(dXf, 0, sizeof(dXf));
float inv = 1.0f / (T - 1);
for (i = 0; i < T - 1; i++) {
float pr[VOCAB];
int tgt = c->tok[i + 1];
memcpy(pr, c->logit[i], sizeof(pr));
softmax(pr, VOCAB);
total += -logf(pr[tgt] + 1e-9f);
/* The loss is a mean, so the gradient is too */
for (j = 0; j < VOCAB; j++) {
float hot = (j == tgt) ? 1.0f : 0.0f;
float dz = (pr[j] - hot) * inv;
for (d = 0; d < D_MODEL; d++) {
gr->tok_emb[j][d] += dz * c->Xf[i][d];
dXf[i][d] += dz * m->tok_emb[j][d];
}
}
}
for (i = 0; i < T; i++)
rmsnorm_back(c->Xc[i], m->norm_f_g, c->sf[i],
dXf[i], dX[i], gr->norm_f_g,
D_MODEL);
for (l = N_LAYERS - 1; l >= 0; l--) {
static float dn2[MAX_SEQ][D_MODEL];
static float dn1[MAX_SEQ][D_MODEL];
memset(dn2, 0, sizeof(dn2));
memset(dn1, 0, sizeof(dn1));
/* FFN */
for (i = 0; i < T; i++) {
float dg[D_FF] = {0};
for (d = 0; d < D_MODEL; d++)
for (e = 0; e < D_FF; e++) {
gr->W_down[l][d][e] += dX[i][d]
* c->gf[l][i][e];
dg[e] += m
->W_down[l][d][e] * dX[i][d];
}
for (e = 0; e < D_FF; e++) {
float da = dg[e]
* gelu_grad(c->hf[l][i][e]);
for (d = 0; d < D_MODEL; d++) {
gr->W_up[l][e][d] += da
* c->n2[l][i][d];
dn2[i][d] += m->W_up[l][e][d] * da;
}
}
}
for (i = 0; i < T; i++)
rmsnorm_back(c->Xb[l][i], m->norm2_g[l],
c->s2[l][i], dn2[i], dX[i],
gr->norm2_g[l], D_MODEL);
/* Attention */
static float dcc[MAX_SEQ][D_MODEL];
memset(dcc, 0, sizeof(dcc));
for (i = 0; i < T; i++)
for (d = 0; d < D_MODEL; d++)
for (e = 0; e < D_MODEL; e++) {
gr->W_O[l][d][e] += dX[i][d]
* c->cc[l][i][e];
dcc[i][e] += m
->W_O[l][d][e] * dX[i][d];
}
float scale = 1.0f / sqrtf((float)D_HEAD);
for (h = 0; h < N_HEADS; h++) {
static float dq[MAX_SEQ][D_HEAD];
static float dk[MAX_SEQ][D_HEAD];
static float dv[MAX_SEQ][D_HEAD];
memset(dq, 0, sizeof(dq));
memset(dk, 0, sizeof(dk));
memset(dv, 0, sizeof(dv));
for (i = 0; i < T; i++) {
float dp[MAX_SEQ] = {0}, ds[MAX_SEQ];
for (j = 0; j <= i; j++) {
float acc = 0;
for (d = 0; d < D_HEAD; d++) {
acc += dcc[i][h * D_HEAD + d]
* c->v[l][h][j][d];
dv[j][d] += c->p[l][h][i][j]
* dcc[i][h * D_HEAD + d];
}
dp[j] = acc;
}
float dot = 0;
for (j = 0; j <= i; j++)
dot += c->p[l][h][i][j] * dp[j];
for (j = 0; j <= i; j++)
ds[j] = c->p[l][h][i][j]
* (dp[j] - dot) * scale;
for (j = 0; j <= i; j++)
for (d = 0; d < D_HEAD; d++) {
dq[i][d] += ds[j]
* c->k[l][h][j][d];
dk[j][d] += ds[j]
* c->q[l][h][i][d];
}
}
for (i = 0; i < T; i++)
for (d = 0; d < D_HEAD; d++)
for (e = 0; e < D_MODEL; e++) {
float u = c->n1[l][i][e];
gr->W_Q[l][h][d][e]
+= dq[i][d] * u;
gr->W_K[l][h][d][e]
+= dk[i][d] * u;
gr->W_V[l][h][d][e]
+= dv[i][d] * u;
dn1[i][e] +=
m->W_Q[l][h][d][e] * dq[i][d]
+ m->W_K[l][h][d][e] * dk[i][d]
+ m->W_V[l][h][d][e] * dv[i][d];
}
}
for (i = 0; i < T; i++)
rmsnorm_back(c->Xa[l][i], m->norm1_g[l],
c->s1[l][i], dn1[i], dX[i],
gr->norm1_g[l], D_MODEL);
}
for (i = 0; i < T; i++)
for (d = 0; d < D_MODEL; d++) {
gr->tok_emb[c->tok[i]][d] += dX[i][d];
gr->pos_emb[i][d] += dX[i][d];
}
return total / (T - 1);
}
static double loss_only(const GPT *m, const int *tok,
int T, Cache *c)
{
int i;
double total = 0;
forward(m, tok, T, c);
for (i = 0; i < T - 1; i++) {
float pr[VOCAB];
memcpy(pr, c->logit[i], sizeof(pr));
softmax(pr, VOCAB);
total += -logf(pr[c->tok[i + 1]] + 1e-9f);
}
return total / (T - 1);
}
/* ---- Adam ---- */
static GPT Grad, Mom, Vel;
static void adam_step(GPT *m, float lr, int t)
{
float *w = (float *)m, *g = (float *)&Grad;
float *u = (float *)&Mom, *v = (float *)&Vel;
float b1 = 0.9f, b2 = 0.999f;
float c1 = 1.0f - powf(b1, t);
float c2 = 1.0f - powf(b2, t);
int i;
for (i = 0; i < (int)NPARAM; i++) {
u[i] = b1 * u[i] + (1 - b1) * g[i];
v[i] = b2 * v[i] + (1 - b2) * g[i] * g[i];
w[i] -= lr * (u[i] / c1)
/ (sqrtf(v[i] / c2) + 1e-8f);
}
}
/* ---- Gradient check ----
Perturb one whole block along a random direction and
compare the measured change in loss against the
analytic
gradient projected onto that direction. Averaging
over a
direction removes the float32 noise that ruins single
element checks. */
static double check_block(GPT *m,
const int *tok, int T,
Cache *c, int lo, int hi,
double *out_num)
{
static float dir[NPARAM];
float *w = (float *)m, *g = (float *)&Grad;
double ana = 0, eps = 1e-3, Lp, Lm;
int i;
for (i = lo; i < hi; i++) {
dir[i] = randf() * 2 - 1;
ana += (double)g[i] * dir[i];
}
for (i = lo; i < hi; i++) w[i] += eps * dir[i];
Lp = loss_only(m, tok, T, c);
for (i = lo; i < hi; i++) w[i] -= 2 * eps * dir[i];
Lm = loss_only(m, tok, T, c);
for (i = lo; i < hi; i++) w[i] += eps * dir[i];
*out_num = (Lp - Lm) / (2 * eps);
return ana;
}
/* ---- Corpus and training ---- */
static float unigram_entropy(const int *tok, int n)
{
int count[VOCAB] = {0};
float h = 0;
int i;
for (i = 0; i < n; i++) count[tok[i]]++;
for (i = 0; i < VOCAB; i++)
if (count[i]) {
float p = (float)count[i] / n;
h -= p * logf(p);
}
return h;
}
#define WINDOW 48
#define STRIDE 16
int main(void)
{
static GPT model;
static Cache c;
static int tok[512];
const char *names[] = {
"tok_emb", "pos_emb", "W_Q", "W_K",
"W_V", "W_O",
"W_up", "W_down", "norm1_g", "norm2_g",
"norm_f_g" };
size_t off[] = {
offsetof(GPT, tok_emb), offsetof(GPT, pos_emb),
offsetof(GPT, W_Q), offsetof(GPT, W_K),
offsetof(GPT, W_V), offsetof(GPT, W_O),
offsetof(GPT, W_up), offsetof(GPT, W_down),
offsetof(GPT, norm1_g), offsetof(GPT, norm2_g),
offsetof(GPT, norm_f_g), sizeof(GPT) };
const char *corpus =
"the cat sat on the mat. "
"the dog ran in the park. "
"a big cat sat on a big mat. the small dog "
"ran fast. "
"the cat and the dog sat on the mat. "
"a cat ran. "
"the big dog sat. a small cat ran in the park. "
"the dog and the cat ran fast. a big dog sat "
"on a mat. ";
int n = 0, i, b, e, s, step = 0, g;
double worst = 0;
for (i = 0; corpus[i] && n < 512; i++)
tok[n++] = char_to_id(corpus[i]);
srand(42);
gpt_init(&model);
printf("=== Capstone: a complete GPT ===\n\n");
printf(" vocab %d d_model %d heads %d\n",
VOCAB, D_MODEL, N_HEADS);
printf(" d_ff %d layers %d\n", D_FF, N_LAYERS);
printf(" parameters %d corpus %d characters\n",
(int)NPARAM, n);
printf(" uniform guess scores %.3f\n",
logf(VOCAB));
printf(" character frequency scores %.3f\n\n",
unigram_entropy(tok, n));
/* Verify every gradient before trusting any of
them */
memset(&Grad, 0, sizeof(Grad));
forward(&model, tok, WINDOW, &c);
loss_and_backward(&model, &c, &Grad);
printf(" Gradient check, "
"analytic vs numerical:\n");
printf(" block analytic numerical "
"err\n");
for (b = 0; b < 11; b++) {
double num;
double ana = check_block(&model,
tok, WINDOW, &c,
off[b]
/ sizeof(float),
off[b + 1]
/ sizeof(float),
&num);
double rel = fabs(num - ana)
/ (fabs(num) + fabs(ana) + 1e-12);
if (rel > worst) worst = rel;
printf(" %-9s %11.6f %11.6f %8.1e\n",
names[b], ana, num, rel);
}
printf(" worst relative error %.1e against a\n",
worst);
printf(" tolerance of 5e-3, the limit of "
"float32\n");
printf(" forward passes allow. %s\n\n",
worst < 5e-3 ? "Every gradient agrees."
: "A gradient is wrong.");
/* Train every weight in the model */
memset(&Mom, 0, sizeof(Mom));
memset(&Vel, 0, sizeof(Vel));
printf(" Training all %d parameters:\n",
(int)NPARAM);
for (e = 1; e <= 100; e++) {
double total = 0;
int nb = 0;
for (s = 0; s + WINDOW < n; s += STRIDE) {
memset(&Grad, 0, sizeof(Grad));
forward(&model, tok + s, WINDOW, &c);
total
+= loss_and_backward(&model, &c,
&Grad);
adam_step(&model, 0.01f, ++step);
nb++;
}
if (e == 1 || e % 20 == 0)
printf(" epoch %3d loss %.4f\n",
e, total / nb);
}
printf("\n Generation (temperature 0.8):\n");
for (g = 0; g < 3; g++) {
int seq[MAX_SEQ], len = 4;
const char *p = (g == 1) ? "a bi" : "the ";
for (i = 0; i < 4; i++)
seq[i] = char_to_id(p[i]);
printf(" %s", p);
for (i = 0; i < WINDOW - 4; i++) {
float pr[VOCAB], r, acc = 0;
int pick = VOCAB - 1, j;
forward(&model, seq, len, &c);
for (j = 0; j < VOCAB; j++)
pr[j] = c.logit[len - 1][j] / 0.8f;
softmax(pr, VOCAB);
r = randf();
for (j = 0; j < VOCAB; j++) {
acc += pr[j];
if (r < acc) {
pick = j;
break;
}
}
putchar(id_to_char(pick));
seq[len++] = pick;
}
printf("\n");
}
printf("\n Every component, and its chapter:\n");
printf(" Token embedding Ch. 10\n");
printf(" Position embedding Ch. 22\n");
printf(" RMSNorm and its backward Ch. 23\n");
printf(" Multi-head attention Ch. 21\n");
printf(" Causal mask Ch. 25\n");
printf(" GELU and its derivative Ch. 24\n");
printf(" Feed-forward network Ch. 24\n");
printf(" Residual connections Ch. 24\n");
printf(" Weight tying Ch. 26\n");
printf(" Softmax, cross entropy Ch. 4\n");
printf(" Backpropagation Ch. 5\n");
printf(" Adam Ch. 8\n");
printf(" Temperature sampling Ch. 26\n");
return 0;
}
37.4 Checking the Gradients First
Figure 37-2 has the whole run in it, the gradient check across all eleven parameter blocks, the training from the frequency baseline down to memorization, and the text generated from the trained model. A backward pass that is subtly wrong does not announce itself in any obvious way. It produces a loss that falls slowly and unconvincingly, and it leaves you wondering whether the model wants more epochs or a different step size when what it actually wants is a corrected sign somewhere in the softmax Jacobian. What I’m trying to say is that there is really no way for us to know outright. The only reliable defense against this is numerical differentiation, which treats your own code as the thing under test. Move one weight by a small epsilon in both directions, measure how the loss responds, and compare that response against the derivative your code claims to have computed. If the two disagree then the code is wrong, and no quantity of additional training will ever repair it.
Doing this one weight at a time runs into float32. The loss sits near 3.3 and a single perturbed weight changes it by something near 0.00001, which is close to the resolution a 32-bit float on most machines has left at that magnitude. The difference of two nearly equal numbers throws away most of the significant digits, and the estimate comes back with a few percent of noise attached. That noise is indistinguishable from a genuinely wrong gradient in exactly the parameters whose gradients are smallest, which in this model are the query and key projections. Those two are three orders of magnitude below the largest gradients in the file, so they are the last place you want a measurement you cannot trust. We already discussed similar problems in earlier chapters so you should have no trouble understanding the limitations.
We do however have a fix. The fix is to perturb an entire block along a single random direction and compare against the analytic gradient projected onto that same direction. Thousands of weights move together, so the change in loss is large enough to measure cleanly and the per element noise averages away rather than dominating. Eleven directional checks cover the whole model, one for each parameter block, and together they run in under a second. The check is weaker than a per element test in principle, because a direction could in theory miss a compensating pair of errors, but in practice a random direction through several thousand dimensions catches anything real. To satisfy the itch of those who would rather see this expressed mathematically, take a look at this equation:
The vector d is a random direction with one component per weight in the block, and epsilon is the step size, small enough to approximate a derivative and large enough that the change in loss clears the noise floor. The left side is the directional derivative, meaning the analytic gradient dotted with d, and the right side measures the same quantity by moving the whole block a little each way and differencing. What this essentially tells us is that agreement between the two is evidence about every weight in the block at once rather than about one of them.
In our configuration every block agrees. The largest disagreement across the eleven is 2.3 parts in a thousand, appearing in the final normalization gain whose expected analytic gradient of 0.001285 sits against a measured 0.001291. The token embedding carries the largest gradient in the model at 0.728081 and agrees to nine parts in a hundred thousand, which is about as close as float32 arithmetic can bring two independently computed numbers. A tolerance of five parts per thousand is the right bar to set, because that is what float32 forward passes permit rather than what the mathematics demands. So essentially, everything clears it, and the model that follows can be trusted to be training on live gradients.
37.5 What the Model Learns
What the model learns brings us to a discussion on how the model is trained. Training starts from 2.5934, which is close to where any untrained model starts and slightly above the 2.558 that character frequencies alone would score. By epoch 20 the loss is 0.4683, well below the 1.075 a bigram model of this corpus would reach, so the model is already using context rather than frequency. By epoch 40 it is at 0.0323 and by epoch 100 at 0.0179. Those last figures are below the trigram entropy of the corpus, which means the model has stopped generalizing and started memorizing, and on 248 characters that is exactly what a model with 27680 parameters should be expected to do and is not a problem for us.
We can confirm what the model learned when we reach generation. Prompted with four characters the model gives the corpus back, correctly punctuated and correctly spaced across all 48 positions. The sequences it emits are recombinations of sentences it trained on rather than the character soup an untrained model produces. One of the three samples wanders after a few words, which is what temperature sampling does as soon as it steps off a path the model has actually seen. The point is not that the text is good, because 248 characters of cats and mats cannot teach anything worth saying to anybody. The point is that every weight in the network moved in order to make that text possible, and that the attention heads and the feed forward layers moved along with the embeddings.
There is one detail in the training setup that took a working model to find. The sliding window is 48 tokens wide and generation also runs to 48, and those two numbers have to match. While I was building an earlier version, it trained on 32 token windows, which meant the position embeddings from 32 upward never received a gradient and stayed at their random initialization. Generation was clean for the first thirty characters and then collapsed into repeated fragments, because the model had walked off the end of the positions it had been trained on. The symptom looked exactly like a broken attention gradient and was nothing of the kind, which is some of the difficulty I explained earlier about knowing exactly what breaks. A gradient check would never have found it either, because the gradients were all correct and the training was simply never visiting those parameters.
37.6 Reading the Backward Pass
We can now move on to reading the backward pass. The chain starts at the output head and cross entropy against a softmax gives the gradient of the loss with respect to the logits. This is as the predicted distribution minus the one hot target, which is the result from when we looked at loss functions way back at the beginning of the book, and the loss is a mean over positions so the gradient carries the same divisor. Because the head is tied to the embedding table, that gradient lands on the token embeddings directly and also flows back into the final activations. Tying works here precisely because the input path is differentiated too, so the two contributions to the embedding table arrive as one consistent gradient rather than fighting each other.
Normalization comes next, and its backward form is the one you derived already when we did our deep dive into normalization. The gain accumulates the product of the upstream gradient with the normalized activation, which is the straightforward half. The gradient passed downward is less obvious, because the normalizer depends on every component of the input at once, so moving any one component changes the scaling applied to all of them. That coupling is the second term below, and leaving it out produces a gradient close enough to look plausible and wrong enough to train badly. It is the single most common mistake in a hand written normalization backward pass.
The first term is what an elementwise operation would give, the upstream gradient scaled by the gain and divided by the root mean square. The second term is the correction that makes this a normalization rather than a scaling, because sigma depends on every component, so changing one input moves the divisor and therefore every output. That sum is the same for all j and is computed once per position.
Then the blocks run in reverse order. The feed forward network takes its gradient back through the down projection, through the GELU derivative, and then through the up projection, updating both weight matrices as it goes. The derivative of the tanh approximation to GELU is worth writing out in full, because it is one of the few places in the whole file where a moment of ordinary calculus does real work. The product rule applies twice and the inner cubic contributes the final factor, which is easy to drop and hard to notice once it has been dropped. Writing it as a separate function keeps it verifiable on its own, and the check confirms it:
Here u is the inner expression c times x plus a x cubed, with c near 0.7979 and a equal to 0.044715, the same constants the forward pass uses. The first term is the derivative of the x factor and the second is the derivative of the tanh factor, so the two come from the product rule applied to the forward definition. Appendix A derives both and notes what happens if the second term is dropped, for those of you who want to explore the mathematics a bit further.
As the infamous paper would say “Attention is All you Need”, but in a practical implementation from scratch, attention is in my perspective the hardest piece. The gradient enters through the output projection, splits into per head slices, and passes back through the weighted sum of values. These weighted sums then send gradient to the attention weights and to the values at once. From the weights it goes through the softmax Jacobian, then through the scaled dot product to reach the queries and keys. Finally we go through the three input projections. The causal mask matters here because a position masked in the forward pass must receive no gradient in the backward one, which the code enforces by only ever looping as far as the diagonal. Getting that wrong leaks information from the future into the past and produces a model that scores beautifully during training and fails completely the moment it has to generate, ask me how I know, take a look at this formula here:
This is the softmax backward pass written as a vector operation instead of a Jacobian. The sum is a single number for the whole row, the average of the incoming gradient weighted by the probabilities, and subtracting it is what keeps the result consistent with probabilities that must continue to sum to one. Computing it this way costs one pass over the row rather than building an n by n matrix and it’s a bit easier to read as well.
Each residual connection passes its gradient through unchanged, so the sublayer gradient adds to what is already flowing down the stream rather than replacing it. That property is why deep residual stacks train at all, and it makes the bookkeeping simpler than the diagrams suggest. The gradient finally reaches the embedding lookup, where it accumulates into the token table for whichever token appeared at each position, and into the position table for the position itself. A token appearing five times in the window accumulates five contributions, which is correct and is the one place where the flat gradient buffer earns its keep. Adam then consumes the whole buffer in a single loop over every float in the model, because the parameters live in one contiguous structure.
37.7 What You Have Accomplished
Take a step back to realize what you accomplished. You started with a single perceptron and you now have a transformer that trains itself and writes text. The model built here holds 27680 parameters. For comparison GPT-2 Small holds 124 million, LLaMA-7B holds seven billion, and DeepSeek-V3 holds 671 billion of them. The architecture underneath all four is very similar to the one printed in this chapter, and the differences are scale and engineering. Scale means a wider model, more heads, more layers, more data and more compute. Engineering means the mixture of experts we looked at earlier for capacity, latent attention for memory, low rank adaptation (which we covered in Chapter 34) for cheap fine tuning. If you start working with these models you’ll realize the quantization concepts you learned are crucial for deployment of real world models. You are now well equipped to start dissecting how these models work, what tradeoffs the designers make and understand the constraints of deployment.
There is also something to take from the order in which this program does its work. It verifies before it trains and prints the verification where you can read it. A loss curve on its own cannot distinguish a correct gradient from a merely plausible one, and that difference is the difference between a model which learns and a model which only appears to. Every chapter in this book has argued for measuring rather than assuming, and the final program is where that habit stops being a style preference and becomes the thing that made the model work. The first version of this backward pass disagreed with the numerical check by a factor of 23 in every block simultaneously. That factor was the sequence length, because the loss was averaged over positions while the gradient was not, and no amount of tuning would ever have found it. I said that to say that building, training and using LLMs is part science and part art that you can only master by running, experimenting and examining the output yourself than trusting theory alone.
37.8 Exercises
Break one gradient on purpose. Remove the causal restriction from the attention backward pass so it loops over all positions rather than stopping at the diagonal. Which blocks does the check flag, and how large does the disagreement get?
Replace Adam with plain stochastic gradient descent and find the largest step size that still converges. How many epochs does it need to reach the loss Adam reaches in 100?
Raise the model to a width of 64 with 8 heads, a feed forward width of 256 and 4 layers. Does the gradient check still pass, and how far does the loss fall?
Train on a corpus ten times larger and watch the final loss stop short of memorization. At what corpus size does 27680 parameters become the binding constraint rather than the data?
Add a key value cache from Chapter 27 to the generation loop and measure the speedup across 100 generated tokens. The cache changes only the forward pass, so the gradient check should be unaffected.
Add low rank adapters from Chapter 34 to the attention weights, freeze the base model, and fine tune only the adapters on a different corpus. Verify the adapter gradients with the same directional check.
Quantize the trained weights to int8 using Chapter 35 and compare the generated text before and after. How much loss does the quantization add?
Replace the feed forward network with a small mixture of experts from Chapter 36 using 4 experts and top-2 routing. You will need a gradient through the router, which the check will tell you whether you got right.