Training Loop and Evaluation

A training loop that only prints the loss fails in three ways. It cannot tell you whether the loss is falling for the right reason, because a model can overfit the training set and never learn anything general. It cannot tell you when a divergence is starting: the loss can spike for one step and recover, or spike for one step and crash. And it cannot tell you what the model has learned, since loss is a scalar and a generated sample is a paragraph. All three stay hidden unless the loop measures.

The loop in this lesson measures three things. It logs the loss on the training batch every step. It computes the loss on held out batches every K steps and on the final step. And it generates a continuation from a fixed prompt every K steps. Everything lands in a JSONL log, so the artifact is the loop's testimony.

A training loop repeats one cycle: forward pass, loss, backward pass, optimizer step. The loop here adds measurement around that cycle. Two pieces are not obvious, the loss alignment and the AdamW decay split, so they come first.

Loss alignment

Loss alignment means the target at each position is the token that follows the input token. If the input batch is tokens [t0, t1, t2, t3], the target batch is [t1, t2, t3, t4]. Cross entropy is computed on the flat logits of shape (batch * seq, vocab) against the flat targets of shape (batch * seq,). Skip the shift and you train the model to predict its own input, which converges to zero loss while learning nothing useful.

AdamW decay split

The decay split gives weight decay to weight tensors and withholds it from normalization scales and biases. Decay on a LayerNorm scale slowly drives the scale toward zero and breaks normalization. Decay on a bias is mathematically harmless but wasted work. So matrix shaped tensors (linear weights, embedding tables) get decay, and anything that acts as a scale or shift does not.

Warmup plus cosine schedule

Warmup ramps the learning rate linearly up to the target over the first few hundred steps (10 in the demo). It starts at max_lr / warmup_steps, not at zero, so step 0 runs at 3e-4 in the demo. The slow start gives the optimizer state time to populate. Cosine decay then lowers the learning rate toward a floor, min_lr, over the remaining steps (3e-4 in the demo, 10 percent of the 3e-3 max_lr), so the final phase adjusts the weights with small steps. This pairing is the most common schedule in open weights LLM training because it removes most of the brittle moments in the first thousand steps and the last thousand steps.

Held out evaluation

evaluate_model runs a fixed number of batches from the validation split, accumulates the loss, divides by the batch count, and returns the mean. It uses no gradient and no dropout. The training loop rebuilds the eval batches at every probe with make_batches on the validation tokens and a fixed seed (cfg.seed + 1), so the same split and seed give the same number across runs. Read the held out loss next to the training loss and overfitting shows up as a gap.

Qualitative sampling as an early signal

A model whose training loss falls nicely but whose samples are all the same token is broken. A model whose loss curve looks flat but whose samples sharpen into coherent words is learning. The probe is quicker to read than the full curve, and it catches failure modes the scalar misses.

code/main.py implements:

• make_batches(token_ids, batch_size, context_length, seed) which draws random start positions with a seeded generator, so runs reproduce, and yields input and target pairs (target is the input shifted by one). It is an infinite generator, so callers stop after a fixed number of batches.
• calc_loss_batch(model, inputs, targets) which runs the forward pass, flattens, and returns the scalar cross entropy.
• evaluate_model(model, val_loader, max_batches) which iterates a fixed number of validation batches under no grad, stops after max_batches, and returns the mean loss.
• generate_and_print_sample(model, prompt, max_new_tokens, temperature, top_k, seed) which runs its own sampling loop on a fixed prompt rather than calling lesson 35's function, and prints the generated token list. train calls it with temperature 0.8, top_k 20, and seed=step.
• build_param_groups(model, weight_decay) which produces the two-group AdamW parameter list.
• cosine_with_warmup(step, warmup_steps, total_steps, max_lr, min_lr) which returns the learning rate at a given step.
• train(...) which runs the loop, persists outputs/losses.jsonl, and at every probe step (every eval_every steps and the last) prints the train loss, the eval loss, and a sample.
• A demo that trains a tiny model on synthetic data for a small number of steps, writes the JSONL log, and prints the eval loss and a sample at the probe points. On CPU it runs in well under a minute.

Run it:

python3 code/main.py

Output: at every probe step, a line with the train loss and eval loss plus a generated token list, and a final outputs/losses.jsonl you can load with json.loads per line. Per step loss goes to the JSONL log, not the console. Each record is {step, train_loss, lr}, plus val_loss at probe steps.

• torch for autograd, the optimizer, and modules.
• main.py reimplements the lesson 35 GPTModel and its supporting modules locally.

Three patterns turn the textbook loop into something you can leave running overnight.

Clip the gradient norm. A bad batch (anomalous data, an LR spike, a numerical edge case) can produce a huge gradient that wipes out hours of training. Call torch.nn.utils.clip_grad_norm_(params, max_norm=1.0) after backward and before step to keep the optimizer in a safe range. The clipping value is a free parameter, and one is the default that survives most setups.

Log to resumable JSONL, not pickled state. Per step records shaped like {"step": int, "train_loss": float, "lr": float} are durable. A crash leaves a readable artifact, you can grep it, you can plot it with thirty lines of Python, and you can resume training by reading the last step. Pickled state ties you to the exact module layout that produced the file, which is brittle across refactors.

Draw eval batches from a fixed seed. Build the eval batches from the validation tokens with the same seed at every probe (the demo uses cfg.seed + 1). The windows are random start positions, but a fixed seed makes them identical from probe to probe. Comparing eval loss between two runs only works if the eval batches are identical from run to run. Otherwise the comparison measures the batch shuffle as much as the model.

• This loop is the same skeleton that trains a 124M model on real data. Swap the synthetic token tensor for a datasets-style loader and the loop runs unchanged.
• The JSONL log is the deliverable that turns a training run into evidence. The next lesson uses one to compare a freshly trained checkpoint with a pretrained one.
• The qualitative sample probe covers what scalar loss cannot.

1. Add weight_decay_groups() unit tests that confirm scale and bias parameters land in the no decay group, and that linear and embedding weights land in the decay group.
2. Replace the synthetic random tokens with bytes from a small text file so the demo trains on something legible. Verify that the generated sample uses characters present in the file.
3. Change the existing min_lr floor in cosine_with_warmup (the demo default is 3e-4, 10 percent of max_lr) and re-plot the schedule.
4. Save a checkpoint every eval_every steps in addition to the JSONL log. Add a resume_from flag that reloads the model state and the optimizer state.
5. Log per step throughput (tokens per second) next to the loss and confirm it stays in a steady band.

• Phase 19 lesson 35 for the model the loop drives.
• Phase 19 lesson 37 for loading pretrained weights into the same model.
• Phase 10 lesson 04 (pre training mini GPT) for the procedure on real data.
• Phase 10 lesson 10 (evaluation) for the broader eval surface beyond cross entropy loss.