A complete small training run: from raw text to a checkpoint you can trust

The posts so far took one piece of training at a time: one update, the corpus, mixtures, the tokenizer and the architecture. This post puts them together in a run small enough to do on a laptop CPU in about a minute and a quarter, so that every step can be checked. The model is tiny and the text is 1 MB of Shakespeare, but the steps, and the ways they go wrong, are the ones a large run has.

We train the same model five ways, keep everything else fixed, and read the curves.

1. Prepare the data

The corpus is Tiny Shakespeare: 1,115,394 characters of plays. Four decisions come before any training:

• Split before anything else. The last 10% of the text is held out for validation. The split is contiguous rather than random: neighbouring lines of a play are nearly the same text, so a random split would put near-copies of validation text into training.
• Train the tokenizer on the training split only. A byte-level BPE tokenizer with 2,048 entries, trained as in the tokenization post, gives 2.89 characters per token. Training it on validation text too would leak information about that text into the vocabulary.
• Count tokens: 346,827 for training and 43,559 for validation.
• Count passes. Each step trains on 16 sequences of 128 tokens, 2,048 tokens; 600 steps make 1.23 million tokens, about 3.5 passes over the training data. The mixtures post explains why that number matters.

2. Build the model and check its starting point

The model is the modern block from the architecture post at toy size: 4 layers, width 128, 4 heads, a SwiGLU feed-forward network of width 352, rotary positions, RMSNorm before each sub-layer, and input and output embeddings tied. That is 1,066,112 parameters, a quarter of them in the 2,048 × 128 embedding.

Before training, check the loss. A freshly initialized model with small random weights predicts close to uniformly over the vocabulary, so its loss should be about $\ln V$.

Try it: predict the first loss. What loss should the untrained model have on a vocabulary of 2,048 tokens?

Answer. About $\ln 2{,}048 \approx 7.62$ nats per token, the loss of a uniform guess. The measured value was 7.65. A first loss far above $\ln V$, say 30, would mean the initialization makes some logits huge, so the model starts out confidently wrong; one far below would mean the targets leak into the inputs. Both are bugs, and this check catches them for free.

3. The training loop

The loop is the update from the first post, repeated:

for step in range(600):
    lr = schedule(step)                       # warmup, then cosine decay
    for g in opt.param_groups:
        g["lr"] = lr
    x, y = next_batch(train_ids)              # y is x shifted by one token
    loss = F.cross_entropy(model(x).view(-1, V), y.view(-1))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    if step % 25 == 0:
        log(step, loss, evaluate(model), norm)   # held-out loss on fixed batches

The optimizer is AdamW with betas (0.9, 0.95) and weight decay 0.1 on the weight matrices only (norm weights are not decayed), as in many open recipes. Gradients are clipped to a global norm of 1.

4. The learning-rate schedule

The learning rate rises linearly from almost zero to its peak over the first 50 steps, the warmup, then follows a cosine curve down to 10% of the peak at the end. Both halves have reasons.

Warmup. In the first steps Adam’s moment estimates are based on a handful of gradients, and as the first post showed, its first update moves every parameter by about the full learning rate, whatever its gradient. With random weights and a large learning rate, those early steps can push the model somewhere it never recovers from. Starting small gives the moment estimates time to settle.

Decay. Late in training, a high learning rate keeps the weights bouncing around a good region instead of settling into it. Decaying it lets the loss fall further at the end, which is also why a schedule should be set for the length the run will actually have: the Chinchilla review shows how a mismatched schedule distorts comparisons.

5. Precision

Large runs compute in BF16 to halve memory traffic and use fast matrix units, while keeping optimizer state, and often a master copy of the weights, in FP32. The reason is visible in one line: in BF16 the next number above 1.0 is 1.0078, so 1.0 + 0.001 is exactly 1.0. A small weight update applied directly to BF16 weights can vanish; applied to an FP32 copy it survives. (The efficient-execution post covers FP16, which also needs loss scaling, and FP8.)

This run trains in FP32 on CPU, but checks the other direction: evaluating the trained model with BF16 autocast gives a validation loss of 4.1263, identical to FP32 to four decimals. For a model this size, computing the forward pass in BF16 costs nothing measurable; the risk is in accumulating updates, not in computing activations.

6. Evaluate as you go

Every 25 steps the model is evaluated on the same 20 validation batches. Using fixed batches matters: random batches would add sampling noise to every point and hide small differences between runs. The baseline’s validation loss falls from 7.64 to 4.13, a perplexity of $e^{4.13} \approx 62$: the model is as uncertain, on average, as if choosing among 62 equally likely tokens.

The training loss, smoothed over 25 steps, ends near 3.88, about 0.25 below validation. Some of that gap is the model fitting text it has seen three or four times, which is the cost of repetition from the mixtures post; some is that the last tenth of the file is a different stretch of the plays. Watching the gap during training is how you would notice the model starting to memorize.

Loss is not the only evaluation. A sample from the trained model, prompted with “ROMEO:”, reads:

ROMEO:
Ay, even thy tions to drink them, what
To over one poor behe's steeps, if you yet
Benvouned in the field, adrimling thousand

It has learned the layout of a play, line lengths, punctuation and a vocabulary of real and plausible words, but not meaning. Larger models are evaluated with benchmarks; the post-training and application posts return to evaluation in detail.

7. Checkpoints that resume exactly

A checkpoint must contain everything that determines what happens next: the model weights, the optimizer’s moments and step count, the position in the learning-rate schedule, and the state of the data sampler and random number generators. Leave out the optimizer state and the run restarts Adam from zero moments, which behaves like a sudden warmup-free jump; leave out the data state and it re-reads data it has already seen.

The test is simple and worth automating: stop a run halfway, save, rebuild everything from the checkpoint and continue, then compare with a run that never stopped. Stopping this run at step 300 and resuming reproduced the uninterrupted run’s training loss exactly, at every one of the remaining 300 steps. On GPUs, nondeterministic kernels make bitwise identity harder, but the curves should agree within noise.

8. Read the curves

Four runs change one thing each:

• No warmup (peak 0.003 from step 0): validation loss ends at 4.44, 0.31 nats worse than the baseline. Interestingly, it is ahead for the first 50 steps (6.01 against 6.18 at step 25), because a full learning rate learns fast at first; it falls behind from about step 100 and never catches up. Early speed is not the same as a good trajectory.
• Peak 0.03, ten times higher: 4.77. The loss falls quickly at first and then flattens.
• Peak 0.1: 5.40, a plateau far above the baseline.
• Peak 0.1 without clipping: 5.46. In the baseline the gradient norm has a median of 0.57 after warmup and exceeds 1 in under 1% of steps; here it exceeds 1 in 7% of steps and spikes as high as 15.

Notice what did not happen: no run diverged to infinity or NaN. A small model with tied embeddings and normalization before every sub-layer is hard to blow up; too high a learning rate shows up here as a stall. Large models fail more dramatically, with sudden loss spikes that may or may not recover, which is why large runs monitor gradient norms, keep clipping on, and save checkpoints they can roll back to.

Try it: diagnose a curve. A colleague’s run starts at a loss of 7.6, falls to 6 in 50 steps, and then sits at 6 for the rest of training while gradient norms jump between 0.5 and 10. What would you change first?

Answer. The starting loss is right and early learning works, so the data and model are probably fine. A plateau with spiking gradient norms is the signature of a learning rate that is too high for the later phase, as in the peak-0.1 runs above. Lower the peak (or check that the decay schedule is actually being applied), keep clipping on, and compare validation curves over the first few hundred steps before committing to a long run.

What to remember

• Split data before training anything, contiguously, and train the tokenizer on the training split only. Count tokens and passes.
• Check the initial loss against $\ln V$; it catches initialization and data bugs immediately.
• Warmup protects the first steps, when Adam’s estimates are poor; cosine decay lets the loss settle at the end. In this run, skipping warmup cost 0.31 nats.
• Keep updates in FP32 (BF16 cannot hold $1 + 0.001$); computing activations in BF16 changed nothing measurable here.
• Evaluate on fixed held-out batches and watch the train–validation gap.
• A checkpoint is weights, optimizer state, schedule position and data/RNG state; test that resuming reproduces the run.
• Too high a learning rate may stall rather than diverge; gradient-norm spikes are the early warning.

Sources

• Tiny Shakespeare, from A. Karpathy’s char-rnn repository (data/tinyshakespeare/input.txt, commit 6f9487a).
• PyTorch 2.6 and Hugging Face Tokenizers 0.21: scaled_dot_product_attention, AdamW, clip_grad_norm_, autocast, byte-level BPE.
• Stanford CS336, assignment 1, “Basics”: implementing the tokenizer, model, optimizer and training loop from scratch, the natural next step after this post.