Training Loop

train.main is the whole training run. How the Model Learns follows one step from text to changed parameters, and AdamW / Learning Rate / Clipping documents the update. This page documents the rest of the loop: the batches, the step counts, what runs after each epoch, and what is printed and saved.

train.pymainGitHub

58    step = 059    for epoch in range(config.epochs):60        for xb, yb in iterate_batches(tr_x, tr_y, config.batch_size, rng):61            loss, grads = batch_loss_and_grads(params, xb, yb, cfg)62            norm = clip_grad_norm(grads, config.grad_clip)63            step += 164            lr = lr_schedule(step, config.warmup_steps, total_steps,65                             config.lr_max, config.lr_min)66            opt.step(params, grads, lr)67            if step == 1 or step % 50 == 0:68                print(f"step {step:5d}  loss {loss:.3f}  |g| {norm:.2f}  lr {lr:.5f}")6970        v = validation_loss(params, cfg, va_x, va_y)71        print(f"epoch {epoch + 1}: val loss {v:.3f}  ppl {perplexity(v):.1f}")72        seed_ids = [int(i) for i in va_x.iloc[0][:3]]73        out = generate(params, cfg, seed_ids, 30, config.seq_len,74                       config.temperature, config.top_k, config.top_p, rng)75        print("  sample:", repr(tok.decode(out)))7677    # 4. keep the trained model so talk.py can load it later78    save_checkpoint(prefix, params, cfg, tok)79    print(f"saved checkpoint to {prefix}.*")

Mini-batch gradient

What is it?

Mini-batch loss and gradientDerived from implementation

Data

batch_loss_and_grads turns one batch of windows into one loss and one gradient per parameter. The model has no batch dimension, so it loops over the windows and calls loss_and_grads on each. It divides each result by the batch size as it adds it, so what it returns is the mean.

Why does the model need it?

The optimizer takes one gradient per parameter for each update. The comment above iterate_batches in data.py gives the reason for batches: using all windows at once makes each update slow, and using one window at a time makes updates noisy.

In general, one window's gradient depends a lot on which tokens happen to be in it. Averaging over B windows keeps what those gradients share, and differences between windows partly cancel, so the update changes less from one step to the next. That is general reasoning, not something this project measures.

Mathematics

B comes from the batch itself, not from config.py:

The batch size, read from this batch #

B = \text{number of rows of } x_b

The division by B happens inside the loop, on each window's result, as the code adds it to the total:

How the code accumulates the loss #

L \leftarrow L + \frac{L_i}{B} \qquad i = 1, \dots, B, \quad \text{starting from } L = 0

How the code accumulates each gradient #

\bar g_k \leftarrow \bar g_k + \frac{g_{k,i}}{B} \qquad \text{for every key } k, \quad \text{starting from zeros}

When the loop ends, those totals are the means:

The result, the batch mean loss #

L = \frac{1}{B} \sum_{i=1}^{B} L_i

The result, the mean gradient #

\bar g_k = \frac{1}{B} \sum_{i=1}^{B} g_{k,i}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$B$	Windows in this batch. A short last batch uses its own size	`B = len(xb)`	scalar
$x_b, y_b$	The batch's input and target windows	`xb, yb`	(B, T)
$L_i$	The mean cross-entropy of window i	`loss`	scalar
$L$	The running total, and the batch mean loss when the loop ends	`total_loss`	scalar
$g_{k,i}$	Window i's gradient for parameter k	`g[k]`	same as the parameter
$\bar g_k$	The running total, and the mean gradient when the loop ends	`grads[k]`	same as the parameter

Tiny numerical example

A batch of 3 windowsIllustrative ExampleA tiny model: the parameters are random draws from the real init_params, not learned values, and the ids are hand-picked. Each row is one pass of the loop.

Window	`x`	`y`	`loss`	`loss / B`	`total_loss`
0	`[2, 0, 2]`	`[0, 2, 1]`	`2.56207484`	`0.85402495`	`0.85402495`
1	`[1, 3, 4]`	`[3, 4, 0]`	`1.81800983`	`0.60600328`	`1.46002822`
2	`[0, 0, 1]`	`[0, 1, 2]`	`1.45529321`	`0.48509774`	`1.94512596`

The function returns 1.94512596.

Tensor shapes

B is 16 and T is 16, except that the last batch of an epoch can be shorter. V depends on the corpus.

A batch goes through the model one window at a timetrain.py:17

xb(B, T)(16, 16)
x(T,)(16,)
logits inside loss_and_grads(T, V)(16, V)

B = batch_size = 16. A final short batch has fewer rows. 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.

One key as an example. Averaging does not change a gradient's shapetrain.py:22

g["head_w"](d, V)(32, V)
grads["head_w"](d, V)(32, V)

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

Actual implementation

train.pybatch_loss_and_gradsGitHub

12def batch_loss_and_grads(params, xb, yb, cfg: ModelConfig):13    B = len(xb)14    grads = {k: np.zeros_like(v) for k, v in params.items()}15    total_loss = 0.01617    for x, y in zip(xb, yb):18        loss, g = loss_and_grads(params, x, y, cfg)19        total_loss += loss / B2021        for k in grads:22            grads[k] += g[k] / B2324    return total_loss, grads

A short last batch divides by its own size, because B is len(xb). So each batch counts the same in the update, whatever its size.

Code and math connection

Equation	Role	Lines	Code
The batch size, read from this batch	Computes it	train.py:13	`B = len(xb)`
How the code accumulates the loss	Computes it	train.py:15	`total_loss = 0.0`
How the code accumulates the loss	Computes it	train.py:19	`total_loss += loss / B`
How the code accumulates each gradient	Computes it	train.py:14	`grads = {k: np.zeros_like(v) for k, v in params.items()}`
How the code accumulates each gradient	Computes it	train.py:22	`grads[k] += g[k] / B`
The result, the batch mean loss	Computes it	train.py:24	`return total_loss, grads`
The result, the mean gradient	Computes it	train.py:24	`return total_loss, grads`
The result, the mean gradient	Calls it	train.py:61	`loss, grads = batch_loss_and_grads(params, xb, yb, cfg)`

Where is it used?

Source: train.py, lines 12-24, batch_loss_and_grads
PDF reference: Not stated in the PDF
Used by: Called by: train.main
Calls it: train.py, line 61 (main), for The result, the mean gradient
Status: Derived from implementation

Implemented in

Equation	Role	File	Function	Lines	GitHub
The batch size, read from this batch	Computes it	`train.py`	`batch_loss_and_grads`	13	GitHub: train.py, line 13
How the code accumulates the loss	Computes it	`train.py`	`batch_loss_and_grads`	15	GitHub: train.py, line 15
How the code accumulates the loss	Computes it	`train.py`	`batch_loss_and_grads`	19	GitHub: train.py, line 19
How the code accumulates each gradient	Computes it	`train.py`	`batch_loss_and_grads`	14	GitHub: train.py, line 14
How the code accumulates each gradient	Computes it	`train.py`	`batch_loss_and_grads`	22	GitHub: train.py, line 22
The result, the batch mean loss	Computes it	`train.py`	`batch_loss_and_grads`	24	GitHub: train.py, line 24
The result, the mean gradient	Computes it	`train.py`	`batch_loss_and_grads`	24	GitHub: train.py, line 24

The block says "Not stated in the PDF" because the PDF has no equation for this. It describes the idea in prose on p. 14, so the equations here are derived from the code.

How this reaches a parameter update. train.main receives the mean loss and the mean gradients (train.py:61). It clips those gradients (train.py:62) and passes them to opt.step (train.py:66). The loss is only printed, so it plays no part in the update. The backward functions themselves are on Full Model.

Step counts

What is it?

Step countsDerived from implementation

Data

train.main works out how many updates training will make before the loop starts. One epoch has one step per batch, and total_steps is that number times the number of epochs.

Why does the model need it?

The learning-rate schedule needs to know where training ends, so that the cosine reaches lr_min on the last step. Nothing else reads total_steps.

Mathematics

Steps per epoch #

S_{\text{epoch}} = \left\lceil \frac{W}{B} \right\rceil

Total steps #

S_{\text{total}} = E \cdot S_{\text{epoch}}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$W$	Training windows	`len(tr_x)`	scalar
$B$	Windows per batch	`config.batch_size`	scalar
$E$	Passes over the training windows	`config.epochs`	scalar
$S_{\text{epoch}}$	Steps in one epoch. The ceiling counts a short last batch as a step	`steps_per_epoch`	scalar
$S_{\text{total}}$	Steps in the whole run, passed to the schedule	`total_steps`	scalar

Tiny numerical example

Steps from windows, batch size and epochsIllustrative ExampleThe number of windows is hand-picked, because the real number depends on the corpus. batch_size and epochs start at their config.py values.

Windows, len(tr_x)

batch_size

epochs

steps_per_epoch = ceil(1000 / 16) = 63
total_steps = 20 * 63 = 1260
The last batch of each epoch has 8 windows.

Tensor shapes

These are plain numbers, so there are no tensors.

Actual implementation

train.pymainGitHub

54    steps_per_epoch = math.ceil(len(tr_x) / config.batch_size)55    total_steps = config.epochs * steps_per_epoch

iterate_batches hands out exactly steps_per_epoch batches, because it also keeps a short last batch (Data Pipeline). So the counter step ends at exactly total_steps, and the last step gets lr_min.

Code and math connection

Equation	Role	Lines	Code
Steps per epoch	Computes it	train.py:54	`steps_per_epoch = math.ceil(len(tr_x) / config.batch_size)`
Total steps	Computes it	train.py:55	`total_steps = config.epochs * steps_per_epoch`
Total steps	Calls it	train.py:64	`lr = lr_schedule(step, config.warmup_steps, total_steps,`

Where is it used?

Source

train.py, lines 36-81, main

PDF reference

Not stated in the PDF

Used by

Called by: train.<module>

Written inline in train.py, line 54-55. These two lines are the whole calculation. lr_schedule receives total_steps.

Calls it: train.py, line 64 (main), for Total steps

Status

Derived from implementation

Implemented in

Equation	Role	File	Function	Lines	GitHub
Steps per epoch	Computes it	`train.py`	`main`	54	GitHub: train.py, line 54
Total steps	Computes it	`train.py`	`main`	55	GitHub: train.py, line 55

The PDF has no equation for this. These lines appear only in its code listing, on p. 16.

Validation

What is it?

Validation lossDerived from implementation

Data

validation_loss runs the model on every validation window, one at a time, and returns the mean of each window's mean cross-entropy. train.main calls it once at the end of each epoch and prints the result next to its perplexity. It computes no gradients and changes nothing.

Why does the model need it?

The training loss comes from text the model is being trained on, one batch at a time. The validation windows come from the last part of the ids, which training never updates on, so this number shows how the model does on that text.

Mathematics

Each L_j is the cross-entropy from Softmax & Cross-Entropy, and it is already a mean over the window's positions.

Mean cross-entropy over validation windows #

v = \frac{1}{W_v} \sum_{j=1}^{W_v} L_j

train.main passes that result to perplexity, so the L in this equation is v:

Perplexity #

\text{perplexity} = e^{L}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$v$	The validation loss	`np.mean(losses), v in train.main`	scalar
$W_v$	Validation windows	`len(va_x)`	scalar
$L_j$	The mean cross-entropy of window j	`loss`	scalar

Tiny numerical example

3 validation windowsIllustrative ExampleThe same tiny model and windows as the batch example, so the three losses are the same. Validation takes a mean of all windows, and those parameters are not learned values.

Value	Code	Result
Window 0	`loss`	`2.56207484`
Window 1	`loss`	`1.81800983`
Window 2	`loss`	`1.45529321`
Mean	`float(np.mean(losses))`	`1.94512596`
Perplexity	`perplexity(v)`	`6.99451281`

Tensor shapes

Validation windows go through the model one at a timetrain.py:30

inputs(W, T)(W, 16)
x(T,)(16,)
logits(T, V)(16, V)

W is the number of windows. It depends on the corpus. 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.

W here is the number of validation windows, which depends on the corpus.

Actual implementation

train.pyvalidation_lossGitHub

26def validation_loss(params, cfg, inputs, targets):27    losses = []2829    for x, y in zip(inputs.to_numpy(), targets.to_numpy()):30        logits, _ = forward(params, x, cfg)31        loss, _ = cross_entropy_forward(logits, y)32        losses.append(loss)3334    return float(np.mean(losses))

train.pymainGitHub

70        v = validation_loss(params, cfg, va_x, va_y)71        print(f"epoch {epoch + 1}: val loss {v:.3f}  ppl {perplexity(v):.1f}")

It runs at the end of every epoch, on all validation windows, one at a time. It calls forward and discards the cache (train.py:30), so no backward pass runs and the parameters do not change.

Code and math connection

Equation	Role	Lines	Code
Mean cross-entropy over validation windows	Computes it	train.py:32	`losses.append(loss)`
Mean cross-entropy over validation windows	Computes it	train.py:34	`return float(np.mean(losses))`
Mean cross-entropy over validation windows	Calls it	train.py:70	`v = validation_loss(params, cfg, va_x, va_y)`

Where is it used?

Source: train.py, lines 26-34, validation_loss
PDF reference: Not stated in the PDF
Used by: Called by: train.main
Calls it: train.py, line 70 (main), for Mean cross-entropy over validation windows
Status: Derived from implementation

Implemented in

Equation	Role	File	Function	Lines	GitHub
Mean cross-entropy over validation windows	Computes it	`train.py`	`validation_loss`	32	GitHub: train.py, line 32
Mean cross-entropy over validation windows	Computes it	`train.py`	`validation_loss`	34	GitHub: train.py, line 34

The PDF has no equation for this function. It appears only in the code listing, on p. 15-16, under the name train.evaluate. The PDF's Eq. 62 is a different quantity, over the whole corpus with its full history, and the code does not compute it.

The sample printed each epoch

After validation, train.main generates a short text so that a person can look at it. Nothing reads the result.

train.pymainGitHub

72        seed_ids = [int(i) for i in va_x.iloc[0][:3]]73        out = generate(params, cfg, seed_ids, 30, config.seq_len,74                       config.temperature, config.top_k, config.top_p, rng)75        print("  sample:", repr(tok.decode(out)))

The seed tokens are the first 3 tokens of the first validation window, va_x.iloc[0][:3]. So the seed is the same on every epoch.
It generates 30 new tokens. This is a literal in the call, and max_new_tokens in config.py is not used here. It applies only to talk.py.
The window, temperature, top_k and top_p come from config.py: seq_len is 16, temperature is 0.8, top_k is 0 and top_p is 0.9.
The printed text includes the seed. generate returns the seed ids followed by the new ids, and tok.decode(out) decodes them all. repr shows newlines as \n and adds quotes.
It uses the same rng as the batch shuffling (train.py:60). So drawing the sample changes the order of the next epoch's batches.

Sampling and Autoregressive Generation explain how the tokens are chosen.

Logging lines

train.py prints these lines, and nothing else. The output column comes from evaluating the file's own format strings with the values of one real step on a hand-picked text and a tiny model. The sampled ids are stored from one earlier run of generate. It is not output from a training run.

Every line train.py printsIllustrative Example

Once, before training train.py:52
```
V=20  ln V=2.996
```
V is the vocabulary size. ln V is the loss of a model that gives every token the same probability, which is a first number to compare loss with.
At step 1, and at every step divisible by 50 train.py:68
```
step     1  loss 2.930  |g| 1.94  lr 0.00002
```
step counts updates from 1. loss is this batch's mean cross-entropy, before this step's update. |g| is the gradient's length before clipping. lr is this step's learning rate.
After each epoch train.py:71
```
epoch 1: val loss 3.007  ppl 20.2
```
The epoch counts from 1, because the code prints epoch + 1. val loss is the mean cross-entropy over all validation windows. ppl is exp of that value.
After each epoch, on the line after validation train.py:75
```
  sample: 'the dth sathogatdtheatsathdtd ddmdgatd\nthe   gnat'
```
No numbers: the decoded seed and generated tokens.
Once, after the last epoch train.py:79
```
saved checkpoint to checkpoints/model.*
```
No numbers: the prefix of the three files.

Checkpoint

After the last epoch, train.main saves the model once. It saves nothing during training, so stopping a run early leaves no checkpoint from that run.

train.pymainGitHub

78    save_checkpoint(prefix, params, cfg, tok)

save_checkpoint writes three files with the same prefix: the weights, the model shape and the tokenizer. When train.py runs as a script, the prefix is always the default, checkpoints/model, because the file calls main() with no arguments. Checkpoints & Tests documents the files.