Now, let's build your very own AI.

The training goal — answers shifted by one

Turning the training data into token IDs gives one long row of IDs. From it, cut a 13-token window and use the first 12 as input and the 12 shifted by one as the answers.

• If the tokens are [a, b, c, d], the answers for input [a, b, c] are [b, c, d].
• Position 1 sees only a and predicts b, position 2 sees a and b and predicts c, position 3 sees a, b and c and predicts d.
• Thanks to causal attention, one forward pass solves 12 problems at the same time.

Even without people adding answer labels, the text itself becomes the questions and the answers. That is why text from the internet can be used as training data as it is (self-supervised learning). Predicting the next token well requires knowing grammar, context, format and facts, so this simple goal draws out language ability as a whole.

Loss — cross-entropy (the negative log of the correct-answer probability)

If p is the probability the model gave to the correct token at a position, the loss at that position is −log(p). The loss of a window is the average of the 12 position losses.

• p = 1 gives loss 0: the model was sure of the answer.
• p = 0.5 gives 0.693, p = 0.1 gives 2.303, p = 0.01 gives 4.605: the smaller the probability given to the answer, the more steeply the loss grows.
• Giving a high probability to a wrong answer is not penalized directly. Probabilities add up to 1, so if wrong answers take a lot, the correct p shrinks and the loss grows.

You can know the starting loss in advance. Before training, the model gives almost the same probability (1/40) to all 40 vocabulary tokens, so the loss is −log(1/40) = ln 40 ≈ 3.689. If the first validation 3.7568 that train.py prints is close to this value, the model was initialized correctly. For GPT-2, with a vocabulary of 50,257, the starting loss is about 10.8. Applying exp to the loss gives perplexity, which you can read as "how many options it is torn between", like exp(3.689) = 40.

For models with vocabularies of tens of thousands, loss is roughly read like this. 10 or more is random level; 5–7 means it is learning grammar, 3–5 format, 2–3 meaning; and below 2 you should suspect memorizing (overfitting). The starting point depends on the vocabulary size, so in XSLM you look at how far it has dropped from the starting point of 3.69.

import math

for p in [1.0, 0.9, 0.5, 0.1, 0.01]:
    print('p =', p, ' loss =', round(0.0 - math.log(p), 4))
for vocab in [40, 50257]:
    print('guess evenly among', vocab, '->', round(math.log(vocab), 4))

cd ~/xslm
mkdir -p lab
micro lab/loss.py
python3 -m lab.loss

p = 1.0  loss = 0.0
p = 0.9  loss = 0.1054
p = 0.5  loss = 0.6931
p = 0.1  loss = 2.3026
p = 0.01  loss = 4.6052
guess evenly among 40 -> 3.6889
guess evenly among 50257 -> 10.8249

Gradient descent and learning rate — small steps against the gradient

The gradient is "how much the loss grows if this number gets a little bigger". To reduce the loss, move the number in the opposite direction of the gradient.

• new value = old value − learning rate × gradient

The learning rate (learning_rate) is how far to move at once. The experiment below starts at 0 and takes 20 steps to find w = 3, the lowest point of f(w) = (w − 3)².

• 0.01: too cautious, it only gets to about 1.0. Learning is slow.
• 0.1: it reaches 2.97, almost there.
• 0.5: for this function, it hits exactly 3 in one step.
• 1.1: it jumps too far, keeps overshooting to the other side of the valley, and flies off to −112 (divergence).

Real models have thousands to billions of numbers, so the shape of the valley is unknown and the learning rate is chosen by experiment. XSLM’s 0.03 is a cautious value that keeps a small model from blowing up. Changing the learning rate changes settings.json, so start with a new checkpoint.

# f(w) = (w - 3) ** 2 is smallest at w = 3. Its gradient is 2 * (w - 3).
for lr in [0.01, 0.1, 0.5, 1.1]:
    w = 0.0
    for step in range(20):
        w -= lr * 2 * (w - 3)
    print('lr', lr, '-> w =', round(w, 4))

cd ~/xslm
mkdir -p lab
micro lab/descent.py
python3 -m lab.descent

lr 0.01 -> w = 0.9972
lr 0.1 -> w = 2.9654
lr 0.5 -> w = 3.0
lr 1.1 -> w = -112.0128

One turn of the training loop — six steps

One iteration of train.py goes in this order.

• Pick a window: choose a random starting point in the row of training tokens and cut 13 tokens.
• Clear gradients: set grad of every parameter to 0. Gradients accumulate with +=, so leftovers from the previous iteration would get mixed in. Embedding rows for tokens that do not appear in this window are not in the computation graph and backward does not touch them, so if you do not clear them here, old gradients stay as they are.
• Forward pass: model.loss runs forward and computes the loss.
• Backward pass: loss.backward() finds the gradients of all 1,248 numbers.
• Update: after limiting the gradient size, move every number with p.data −= learning rate × gradient × scale.
• Save and log: save the checkpoint and print step and loss.

Training repeats this turn the set number of times. Before the start and after the end, it measures the loss on the first window of the validation data to see whether it also improved on text not used for training.

XSLM counts steps instead of how many times it went through the data (epochs). Windows are picked at random, so some parts are picked several times and some less. Big models are also usually counted in steps.

Gradient clipping — don’t go too far in one step

Sometimes the gradient suddenly becomes very large on some window. Updating as it is would make the numbers jump far at once and ruin the training so far (exploding gradients). So before updating, train.py measures the length of the whole gradient.

• length = √(sum of the squares of all gradients)
• scale = min(1, 1 ÷ length)
• If the length is 1 or less, the scale is 1 and the gradient is used as is; if it is greater than 1, every gradient is shrunk by the same ratio so the length becomes exactly 1.

The key is to keep the direction and shrink only the size. In PyTorch, clip_grad_norm_ does the same thing (usually with a maximum length of 1.0). The 1e-8 in the denominator is a tiny number to avoid dividing by 0 when all gradients are 0.

Optimizers — SGD and AdamW

The rule that adjusts numbers using gradients is called an optimizer.

• SGD (XSLM): new value = old value − learning rate × gradient. It looks only at the current gradient and remembers nothing. So there is no need to save optimizer state in the checkpoint.
• Momentum: remembers the flow (velocity) of past gradients and moves faster when they point the same way and less when they wobble.
• Adam and AdamW: remember, for each number, the average of its gradients and the average of their size, and choose a suitable step size for each number on their own. AdamW adds weight decay, which shrinks numbers a little so they don’t grow too large. It is the standard for training large language models, and the learning rate is usually around 0.0003.

Adam remembers two more numbers per parameter, so to resume training, the checkpoint must also save that state (optimizer state). That is why a PyTorch checkpoint usually contains the optimizer state along with the model weights. Training also uses several times the memory needed for the weights alone.

Batches — the average gradient of several windows

XSLM computes the gradient from one window at a time (batch size 1). Each window has different content, so the gradient swings a lot and the loss goes up and down too.

Big models compute several windows at once and use the average gradient. For example, with 16 windows of 256 tokens (batch size 16), each update sees 4,096 tokens.

• Averaging reduces the wobble, so the update direction is stable.
• A GPU computes several windows at the same time, so a bigger batch adds almost no time.
• In exchange, it uses that much more memory. If memory runs out (OOM), reduce the batch or the window length.

When memory is short, you can use gradient accumulation: run backward several times on small batches to collect gradients, then update once. XSLM’s backward clears the gradients in the graph to 0 when it starts, so this method cannot be used as is.

Overfitting and underfitting — watch training loss and validation loss together

Looking at both losses together tells you the state of the model.

• Both go down: it is learning well.
• Only training loss goes down while validation loss goes up: overfitting. It memorized the training data and actually gets worse on text it has never seen.
• Neither goes down much: underfitting. The model is too small, training is too short, or the learning rate is off.

We measured this ourselves with XSLM (desktop Python, default data). Training 1,500 times at the default learning rate of 0.03 brought validation loss down from 3.76 to 3.23. Training 3,000 times at a learning rate of 0.1 brought training loss down to 1.3 on some windows, but validation loss rose to 4.42, and with 0.3 it rose to 5.28. Memorizing 12 training pairs is easy, but there is not enough data to learn rules that work on 3 unseen pairs.

The surest way to reduce overfitting is more data; shrinking the model, dropout and stopping training early (using the checkpoint with the lowest validation loss) are also used. If it is underfitting, make the model bigger or train more.

Checkpoints — what is saved and why

checkpoints/model.json holds everything needed to stop training and resume exactly where you left off.

• weights: the 1,248 numbers of the 9 matrices.
• step: the number of updates done so far.
• random: the state of the random number generator. When resuming, the order of windows continues without a break.
• settings and tokenizer: the model size and the ID rules. generate.py and runner.py read only these to rebuild the model.
• data_hash: the SHA-256 fingerprint of the training and validation data. It detects whether the data changed.

So training 3 steps twice gives exactly the same result, down to every number, as training 6 steps at once. Check it yourself. Move the model.json from running --steps 3 twice to another name, then compare it with the result of a fresh single run of --steps 6.

Saving is done safely. It first writes everything to model.tmp and then renames it with os.replace. Renaming happens in an instant, so even if you press Ctrl+C or the power goes off while writing, model.json is either the previous complete version or the new complete version. No half-written broken file is left.

PyTorch saves the weights, optimizer state, step and settings to a binary file (.pt) with torch.save. XSLM saves at every step, but saving a big model also takes time, so it is evaluated and saved every few hundred steps.

Reproducibility — getting the same result again with a seed

A computer’s random numbers are a sequence made by fixed rules from a seed (a starting number). The same seed always gives the same sequence. XSLM uses seed 7 to fix two things.

• The initial weight values (random.Random(seed) in model.py).
• The order of training windows (rng in train.py).

So with the same data and the same settings, anyone running it anywhere goes through the same training. To compare a change in just the learning rate, the rest of the randomness must be the same, so reproducibility is basic to experiments. Changing the seed starts from different initial values and gives different results. A conclusion that depends on one lucky seed is hard to trust, so repeat important comparisons with a few different seeds.

However, if Python or the operating system’s math library differs, the last decimal places can differ, so in very long training runs the numbers may drift apart little by little.

Reading the training log — what the numbers tell you

train.py prints lines like these.

• parameters: 1248 resume: 0 — the number of parameters and the step it resumed from. 0 the first time.
• validation: 3.7568 — validation loss before training. It should be close to ln(vocabulary size).
• step: 1 loss: 3.6875 — the training loss of that step’s window.
• validation: … — validation loss after training.
• saved: checkpoints/model.json — where it was saved.

The loss at each step goes up and down because the windows differ. Look at the trend over many steps, not one number. With the default 20 steps it only moves a little around 3.7; with hundreds to thousands of steps it drops clearly.

flush=True means send the output to the screen right away. Without it, Python collects output and shows it all at once, so training looks frozen, and when nohup writes to a file the log is also late.

When training goes badly — checks by symptom

Narrow down the cause by the symptom.

• Loss stays almost the same around 3.7: too few training steps. Raise --steps, or with a new checkpoint try a slightly bigger learning rate.
• Loss jumps wildly or only grows: the learning rate is too big. Lower it and start fresh.
• ValueError: math domain error or nan: a probability became 0 and log(0) was computed, or the numbers diverged. Lower the learning rate and check that gradient clipping is in place.
• AssertionError: changed settings/tokenizer or changed data: the checkpoint and the current settings or data differ. Move the old checkpoint aside and start fresh.
• AssertionError: add more text: there are fewer than 13 training or validation tokens. Add lines to raw.txt.
• Too slow: reducing width and context makes it faster, but needs a new checkpoint. First, split the run into smaller --steps.
• Training worked but the Runner’s answer is unchanged: the Runner reads the checkpoint only once when it starts. Restart it.