Problem
When you need this. The complete model still has random weights. Train it on batches of stories, watch its development loss and compare its predictions with the counting baseline.
The previous build produced a model and checked its gradients, but every weight is still random. This build trains it on the story corpus and answers two questions with recorded numbers: does it predict the reserved test rows better than the bigram baseline, and can it recall the color that the bigram never could?
The contract is one function. train(tok, text, cfg) splits the corpus exactly as the baseline build did, trains for a fixed number of steps and returns the parameters, the training loss of every step and the development loss at regular intervals. trace_train.py runs the recorded job, scores the 60 test stories once, saves the parameters to params.json and writes training-trace.json.
Success has three checks. The test loss is lower than the bigram’s 1.1702. The model recalls the final color in more test stories than the bigram’s 0 of 60. And a second run with the same seeds gives identical parameters, so that the recorded result can be reproduced: the tests compare every value of two short runs, and the recorded run stores a checksum.
Concepts
One weight update is a step, and one pass over the whole training set is an epoch. Four more terms describe this pipeline:
- Training example
- One whole story. The model reads all its tokens but the last and predicts all its tokens but the first, so one forward pass yields a prediction, and a loss, at every position. The causal mask keeps each prediction from seeing its answer.
- Batch
- 8 stories processed with the same weights. Their gradients are averaged, and the weights change once per batch. Each epoch reshuffles the 280 training stories with a seeded generator.
- Adam
- An optimizer that keeps, for every parameter, a running average of its gradient (m) and of its squared gradient (v). The running averages smooth the update and adjust its scale per parameter. Both averages start at 0, so the update first divides them by their respective 1 − βt corrections. It then divides the corrected gradient average by the square root of the corrected squared-gradient average plus a small ε. The optimizer retains these values between steps; the next lesson must save them when resuming training.
- Development check
- Every 50 steps, the pipeline scores the 60 development stories without updating any weight. The test stories are scored once, after training ends.
Architecture
Read the loop from Training stories. Each step takes a Shuffled batch of them, and the model runs forward and backward on each story in the batch. Mean gradient averages their gradients, Adam update changes every weight, and the path labeled next batch starts the next step. Every 50 steps, the dashed branch to Dev loss check scores the development stories. It reads the weights but never changes them.
train.py holds the pipeline: the split, the batch gradient, Adam, the development and test scoring, the color recall check and a checksum of the parameters. It imports model.py from the previous build unchanged. Nothing runs in parallel, and the only randomness comes from two seeds, one for the starting weights and one for the shuffle.
| Setting | Value |
|---|---|
d | 16 |
hidden | 64 |
max_len | 64 |
init_seed | 3 |
data_seed | 5 |
batch | 8 |
steps | 600 |
learning_rate | 0.01 |
beta1 | 0.9 |
beta2 | 0.999 |
adam_eps | 1e-8 |
eval_every | 50 |
Working code
Runtime: Python 3.10 or newer. Dependencies: Python standard library only.
Download and extract the complete lab using Download the complete lab above. Open build-a-small-language-model-training-pipeline/; keep its Python modules, tests and fixture files together. The bundled README includes setup commands for your platform.
"""Train the one-block model on the story corpus with Adam, standard library only.
Each training example is one whole story: the model reads tokens 0..n-2 and
predicts tokens 1..n-1. A step averages the gradients of a batch of stories,
then Adam updates every parameter. Development loss is measured every
EVAL_EVERY steps; the test stories are scored once, at the end.
"""
import math
import random
from operator import mul
import model as m
from bpe import Tokenizer
CONFIG = {
"d": 16, "hidden": 64, "max_len": 64, "init_seed": 3, "data_seed": 5,
"batch": 8, "steps": 600, "learning_rate": 0.01, "beta1": 0.9, "beta2": 0.999,
"adam_eps": 1e-8, "eval_every": 50,
}
def split_stories(text, train=280, dev=60):
stories = [s + "\n" for s in text.split("\n") if s]
if train + dev >= len(stories):
raise ValueError(f"need more than {train + dev} stories, got {len(stories)}")
return stories[:train], stories[train:train + dev], stories[train + dev:]
def encode_all(tok, stories, max_len):
out = []
for s in stories:
ids = tok.encode(s)
if len(ids) - 1 > max_len:
raise ValueError(f"story of {len(ids)} tokens does not fit {max_len} positions")
out.append(ids)
return out
def zeros_like(p):
return {k: [[0.0] * len(v[0]) for _ in v] if isinstance(v[0], list) else [0.0] * len(v) for k, v in p.items()}
def accumulate(total, g, scale):
for k, v in g.items():
t = total[k]
if isinstance(v[0], list):
for tr, vr in zip(t, v):
for j, x in enumerate(vr):
tr[j] += x * scale
else:
for j, x in enumerate(v):
t[j] += x * scale
def batch_gradient(p, batch):
"""Mean loss and mean gradient over a batch of encoded stories."""
total, losses = zeros_like(p), []
for ids in batch:
logits, cache = m.forward(p, ids[:-1])
loss, d = m.loss_and_dlogits(logits, ids[1:])
losses.append(loss)
accumulate(total, m.backward(p, cache, d), 1.0 / len(batch))
return math.fsum(losses) / len(losses), total
class Adam:
"""Adam with bias correction, as in Kingma and Ba (2015) and torch.optim.Adam without weight decay."""
def __init__(self, p, lr, beta1, beta2, eps):
if not (lr > 0 and math.isfinite(lr)):
raise ValueError(f"learning rate must be a positive finite number, got {lr!r}")
self.lr, self.b1, self.b2, self.eps, self.t = lr, beta1, beta2, eps, 0
self.m, self.v = zeros_like(p), zeros_like(p)
def step(self, p, g):
self.t += 1
c1, c2 = 1 - self.b1 ** self.t, 1 - self.b2 ** self.t
for k in p:
rows = (p[k], g[k], self.m[k], self.v[k])
if not isinstance(p[k][0], list):
rows = ([p[k]], [g[k]], [self.m[k]], [self.v[k]])
for pr, gr, mr, vr in zip(*rows):
for j, gj in enumerate(gr):
if not math.isfinite(gj):
raise ValueError(f"non-finite gradient in {k}")
mr[j] = self.b1 * mr[j] + (1 - self.b1) * gj
vr[j] = self.b2 * vr[j] + (1 - self.b2) * gj * gj
pr[j] -= self.lr * (mr[j] / c1) / (math.sqrt(vr[j] / c2) + self.eps)
def mean_loss(p, seqs):
losses = [m.loss_and_dlogits(m.forward(p, s[:-1])[0], s[1:])[0] * (len(s) - 1) for s in seqs]
return math.fsum(losses) / sum(len(s) - 1 for s in seqs)
def color_recall(p, tok, stories):
"""Predict the color in each story's last sentence from everything before it."""
correct = 0
for story in stories:
color = story[story.index("now "):].split()[4]
cut = story.rindex(" " + color)
ids = tok.encode(story[:cut])
logits, _ = m.forward(p, ids)
guess = max(range(len(logits[-1])), key=lambda i: (logits[-1][i], -i))
correct += tok.vocab[guess] == tok.tokens(" " + color)[0]
return {"correct": correct, "total": len(stories)}
def most_common_color_recall(train_stories, stories):
"""A baseline that ignores the story: always guess the color seen most in training."""
colors = [s.split()[3] for s in train_stories]
top = max(sorted(set(colors)), key=colors.count)
correct = sum(s[s.index("now "):].split()[4] == top for s in stories)
return {"color": top, "correct": correct, "total": len(stories)}
def train(tok, text, cfg=CONFIG, log=None):
train_s, dev_s, _ = split_stories(text)
tr = encode_all(tok, train_s, cfg["max_len"])
dv = encode_all(tok, dev_s, cfg["max_len"])
p = m.init_params(len(tok.vocab), cfg["d"], cfg["hidden"], cfg["max_len"], cfg["init_seed"])
opt = Adam(p, cfg["learning_rate"], cfg["beta1"], cfg["beta2"], cfg["adam_eps"])
rng = random.Random(cfg["data_seed"])
order, history, dev_history = [], [], []
for step in range(1, cfg["steps"] + 1):
if len(order) < cfg["batch"]:
order += rng.sample(range(len(tr)), len(tr)) # a fresh shuffle each epoch
batch = [tr[i] for i in order[:cfg["batch"]]]
del order[:cfg["batch"]]
loss, g = batch_gradient(p, batch)
opt.step(p, g)
history.append(round(loss, 4))
if step % cfg["eval_every"] == 0 or step == cfg["steps"]:
dev_history.append({"step": step, "dev_loss": round(mean_loss(p, dv), 4)})
if log:
log(step, loss, dev_history[-1]["dev_loss"])
return p, history, dev_history
def checksum(p):
"""A short fingerprint of every parameter value, to compare two runs."""
return round(math.fsum(math.fsum(map(mul, v, v)) if not isinstance(v[0], list)
else math.fsum(math.fsum(map(mul, r, r)) for r in v) for v in p.values()), 9)batch_gradient() adds each story’s gradient scaled by 1 / batch size, so the result is the mean. Adam.step() follows the paper’s Algorithm 1 parameter by parameter and stops with an error on a non-finite gradient instead of writing NaN into the weights. mean_loss() weights each story by its number of predictions, so a long story counts more than a short one, the same as the baseline’s scorer.
Read trace_train.py
"""Run the recorded training job and score the result.
Run: python3 trace_train.py (train, print the trace; about two minutes on a laptop CPU)
python3 trace_train.py --write (also rewrite training-trace.json, params.json and training-time.json)
python3 trace_train.py --check (exit 1 if training-trace.json or params.json differ)
training-time.json holds the wall-clock time of the --write run and is not compared,
because timings change from run to run.
"""
import json
import platform
import sys
import time
from pathlib import Path
import train as t
from bpe import Tokenizer
HERE = Path(__file__).parent
SWEEP_RATES, SWEEP_STEPS = (0.001, 0.01, 0.1), 60
def run():
tok = Tokenizer.load(HERE / "tokenizer.json")
text = (HERE / "corpus.txt").read_text()
_, dev, test = t.split_stories(text)
te = t.encode_all(tok, test, t.CONFIG["max_len"])
start = time.perf_counter()
p, history, dev_history = t.train(tok, text)
seconds = time.perf_counter() - start
sweep = []
for lr in SWEEP_RATES:
cfg = dict(t.CONFIG, learning_rate=lr, steps=SWEEP_STEPS, eval_every=SWEEP_STEPS)
_, h, dh = t.train(tok, text, cfg)
sweep.append({"learning_rate": lr, "steps": SWEEP_STEPS, "last_batch_loss": h[-1],
"dev_loss": dh[-1]["dev_loss"]})
record = {
"config": t.CONFIG,
"split": {"train": 280, "dev": len(dev), "test": len(test)},
"train_loss_by_step": history,
"dev_loss": dev_history,
"test_loss": round(t.mean_loss(p, te), 4),
"color_recall_test": t.color_recall(p, tok, test),
"most_common_color_baseline": t.most_common_color_recall(t.split_stories(text)[0], test),
"parameter_checksum": t.checksum(p),
"learning_rate_sweep": sweep,
}
timing = {"seconds_for_recorded_run": round(seconds, 1), "steps": t.CONFIG["steps"],
"python": platform.python_version(), "machine": platform.machine(), "system": platform.system()}
return record, p, timing
if __name__ == "__main__":
record, params, timing = run()
text = json.dumps(record, indent=1) + "\n"
params_text = json.dumps(params) + "\n"
if "--write" in sys.argv:
(HERE / "training-trace.json").write_text(text)
(HERE / "params.json").write_text(params_text)
(HERE / "training-time.json").write_text(json.dumps(timing, indent=1) + "\n")
elif "--check" in sys.argv:
same = (HERE / "training-trace.json").read_text() == text and (HERE / "params.json").read_text() == params_text
sys.exit(0 if same else "training-trace.json or params.json is out of date")
else:
print(text, end="")Run the recorded job:
python3 trace_train.pyThe recorded run of 600 steps took 91.2 seconds, measured on one machine; your numbers will differ. That is one measurement, not a benchmark. python3 trace_train.py --check repeats the job and exits with an error if training-trace.json or params.json differ.
Read the loss curve
The training loss of the first batch is 4.5827, close to ln 96 = 4.5643, and the last is 0.2553. Each training value comes from a different batch, so it moves up and down; the development loss is the steadier signal:
| Step | Development loss |
|---|---|
| 50 | 0.7212 |
| 100 | 0.5114 |
| 150 | 0.4558 |
| 200 | 0.4143 |
| 250 | 0.3454 |
| 300 | 0.3116 |
| 350 | 0.3082 |
| 400 | 0.3040 |
| 450 | 0.2941 |
| 500 | 0.2970 |
| 550 | 0.3141 |
| 600 | 0.3078 |
The development loss falls fast for 300 steps, then flattens. It is lowest at step 450, 0.2941, and ends at 0.3078: a small rise in one run with one seed, not a replicated trend. The run keeps only the final weights.
Score the test stories once
| Model | Test loss | Color recalled in the last sentence |
|---|---|---|
| Bigram baseline | 1.1702 | 0 of 60 |
| Trained one-block transformer | 0.2694 | 39 of 60 |
The transformer’s test loss is 0.2694, against the bigram’s 1.1702. It recalls the color in 39 of 60 test stories, where the bigram recalled none: attention lets the last sentence read the first. A guess that ignores the story, always the color seen most in training (“yellow”), would score 18 of 60. The transformer still misses the color in 21 test stories. Both models use the same tokenizer, split and test rows, so their scores are comparable on this fixture. As the baseline split check shows, 5 of the 60 test rows duplicate training text and the tokenizer saw the whole corpus. These scores describe this synthetic task; they are not a clean estimate on entirely new stories.
Check your understanding
The trace also ran 60 steps at three learning rates, from the same starting weights and the same batches. Before you open the comparison, decide which rate you would pick from these numbers, and what you would still want to know.
| Learning rate | Development loss after 60 steps |
|---|---|
| 0.001 | 3.4556 |
| 0.01 | 0.6196 |
| 0.1 | 0.6608 |
Compare your reasoning
Rate 0.01 has the lowest development loss after 60 steps, 0.6196. Rate 0.001, the paper’s suggested default, is still far behind after the same steps: 3.4556. Rate 0.1 is close behind at 0.6608. A short run favors fast rates; a long run can favor a rate that is slower at first but steadier. Before you commit, run the candidates longer and compare their development loss over the whole curve, not one point.
Read the explanation
Keep the weights with the best development loss, and score the test stories once with them. That requires saving checkpoints during training.
Inspect the recorded trace
{
"config": {
"d": 16,
"hidden": 64,
"max_len": 64,
"init_seed": 3,
"data_seed": 5,
"batch": 8,
"steps": 600,
"learning_rate": 0.01,
"beta1": 0.9,
"beta2": 0.999,
"adam_eps": 1e-08,
"eval_every": 50
},
"split": {
"train": 280,
"dev": 60,
"test": 60
},
"train_loss_by_step": [
4.5827,
4.4245,
4.2887,
4.1223,
3.9969,
3.8289,
3.6895,
3.5444,
3.4556,
3.3367,
3.1635,
3.0186,
2.9194,
2.7607,
2.712,
2.5881,
2.4048,
2.2835,
2.2774,
2.0406,
1.9984,
1.8459,
1.7941,
1.679,
1.6387,
1.6093,
1.5026,
1.4486,
1.2743,
1.3131,
1.2566,
1.2652,
1.1758,
1.1058,
1.0574,
1.0575,
0.9403,
0.9538,
0.8931,
0.8561,
0.8318,
0.8038,
0.8511,
0.791,
0.7698,
0.771,
0.726,
0.6684,
0.7103,
0.6711,
0.6899,
0.6404,
0.6833,
0.6744,
0.6521,
0.6768,
0.663,
0.6191,
0.6451,
0.6068,
0.6248,
0.5868,
0.6067,
0.5952,
0.5789,
0.5679,
0.6024,
0.562,
0.6022,
0.6002,
0.547,
0.56,
0.5513,
0.5119,
0.56,
0.5448,
0.5768,
0.5753,
0.5546,
0.5213,
0.5013,
0.5227,
0.5586,
0.5449,
0.5101,
0.5356,
0.5336,
0.5266,
0.5271,
0.5022,
0.532,
0.5567,
0.5479,
0.5248,
0.5136,
0.5093,
0.5074,
0.494,
0.517,
0.5049,
0.5224,
0.5017,
0.5117,
0.5,
0.4911,
0.4556,
0.496,
0.4786,
0.4922,
0.4837,
0.5094,
0.5045,
0.5015,
0.4926,
0.4926,
0.4576,
0.4516,
0.4649,
0.4519,
0.4863,
0.5065,
0.4974,
0.4739,
0.4479,
0.4821,
0.4931,
0.4637,
0.4506,
0.4641,
0.4179,
0.504,
0.4615,
0.4398,
0.4957,
0.494,
0.4538,
0.4558,
0.461,
0.4968,
0.4328,
0.4386,
0.457,
0.4345,
0.4529,
0.4205,
0.394,
0.4727,
0.4447,
0.4959,
0.4473,
0.4674,
0.3962,
0.483,
0.4403,
0.4372,
0.434,
0.4707,
0.4542,
0.4139,
0.4141,
0.4772,
0.3921,
0.4596,
0.4154,
0.3752,
0.4155,
0.3993,
0.4138,
0.4397,
0.4193,
0.4087,
0.4236,
0.357,
0.4554,
0.4232,
0.3753,
0.4237,
0.4462,
0.3846,
0.4272,
0.43,
0.3672,
0.4269,
0.4459,
0.427,
0.3997,
0.3218,
0.3513,
0.4101,
0.4611,
0.3707,
0.4504,
0.4522,
0.3946,
0.4192,
0.3885,
0.3693,
0.3921,
0.3798,
0.3783,
0.4085,
0.4082,
0.3804,
0.3942,
0.4126,
0.4218,
0.4155,
0.3725,
0.4265,
0.4271,
0.3505,
0.3838,
0.3522,
0.4073,
0.3899,
0.3471,
0.3423,
0.3693,
0.4006,
0.3613,
0.3289,
0.3666,
0.3431,
0.3397,
0.3521,
0.3365,
0.3479,
0.3151,
0.3551,
0.3392,
0.2985,
0.3656,
0.3209,
0.3395,
0.3442,
0.3026,
0.3617,
0.3405,
0.3416,
0.3327,
0.35,
0.3245,
0.3314,
0.3232,
0.3715,
0.3156,
0.3391,
0.3108,
0.2988,
0.3477,
0.3596,
0.3236,
0.3156,
0.2981,
0.3363,
0.3025,
0.322,
0.3159,
0.3019,
0.325,
0.3058,
0.3136,
0.3041,
0.3247,
0.3328,
0.3193,
0.311,
0.3157,
0.3435,
0.3221,
0.2979,
0.3136,
0.3076,
0.3046,
0.3142,
0.2924,
0.3102,
0.2979,
0.2812,
0.3269,
0.3173,
0.2799,
0.2976,
0.3096,
0.2998,
0.3116,
0.3148,
0.3351,
0.3108,
0.3005,
0.3479,
0.33,
0.2866,
0.3067,
0.2882,
0.273,
0.32,
0.3272,
0.3013,
0.2916,
0.286,
0.3176,
0.3144,
0.2887,
0.3207,
0.3058,
0.3059,
0.3023,
0.2639,
0.2977,
0.322,
0.2907,
0.3257,
0.2861,
0.2906,
0.2835,
0.2637,
0.2986,
0.2921,
0.2765,
0.2763,
0.3131,
0.2853,
0.273,
0.2793,
0.3245,
0.2757,
0.2874,
0.2818,
0.2952,
0.2669,
0.2894,
0.3093,
0.2913,
0.3039,
0.3016,
0.2982,
0.2834,
0.309,
0.2799,
0.3103,
0.294,
0.3053,
0.2605,
0.3219,
0.2945,
0.2932,
0.2675,
0.2766,
0.2692,
0.2982,
0.2958,
0.274,
0.2756,
0.2754,
0.3053,
0.3138,
0.25,
0.3425,
0.2904,
0.2964,
0.3018,
0.2781,
0.2837,
0.3088,
0.2752,
0.2671,
0.2796,
0.31,
0.2879,
0.2743,
0.3322,
0.2836,
0.2897,
0.2709,
0.3058,
0.2761,
0.2705,
0.2644,
0.2788,
0.3089,
0.2981,
0.2962,
0.3242,
0.3199,
0.2454,
0.2626,
0.3048,
0.3088,
0.3002,
0.2827,
0.2807,
0.2845,
0.2744,
0.2821,
0.2889,
0.2738,
0.2894,
0.2925,
0.3049,
0.2544,
0.2856,
0.2975,
0.2675,
0.311,
0.3028,
0.2586,
0.2781,
0.2948,
0.3074,
0.2654,
0.2894,
0.3001,
0.2828,
0.2425,
0.2919,
0.3093,
0.2806,
0.3081,
0.2824,
0.2745,
0.2593,
0.2752,
0.2797,
0.2898,
0.2529,
0.2609,
0.321,
0.2783,
0.2669,
0.2385,
0.2765,
0.2718,
0.2681,
0.2615,
0.268,
0.2804,
0.2906,
0.2693,
0.2841,
0.2903,
0.276,
0.3073,
0.259,
0.2792,
0.2787,
0.2603,
0.2825,
0.2857,
0.2822,
0.2829,
0.2674,
0.28,
0.2812,
0.3055,
0.2583,
0.2811,
0.285,
0.31,
0.2502,
0.2807,
0.2753,
0.2801,
0.26,
0.304,
0.2657,
0.3039,
0.2884,
0.2816,
0.2693,
0.2699,
0.2741,
0.2886,
0.2736,
0.2637,
0.2887,
0.2804,
0.289,
0.2664,
0.269,
0.2588,
0.2431,
0.2462,
0.2617,
0.2631,
0.2589,
0.298,
0.2726,
0.2857,
0.2696,
0.2732,
0.2521,
0.2669,
0.263,
0.2775,
0.2943,
0.248,
0.2693,
0.2977,
0.2773,
0.2792,
0.2762,
0.2771,
0.278,
0.2608,
0.2357,
0.2799,
0.3049,
0.27,
0.2777,
0.2697,
0.2818,
0.2657,
0.2686,
0.2765,
0.2774,
0.2664,
0.2745,
0.2734,
0.2583,
0.2672,
0.2732,
0.2851,
0.2851,
0.2875,
0.2561,
0.2552,
0.2509,
0.2484,
0.2431,
0.2821,
0.2363,
0.247,
0.2923,
0.3051,
0.3096,
0.2538,
0.2637,
0.2528,
0.2763,
0.2628,
0.2753,
0.269,
0.284,
0.2648,
0.2715,
0.3263,
0.2971,
0.2758,
0.3103,
0.2772,
0.2638,
0.2711,
0.2895,
0.3468,
0.2567,
0.2996,
0.2932,
0.2995,
0.2627,
0.3093,
0.3126,
0.2831,
0.2749,
0.2883,
0.3093,
0.2586,
0.2705,
0.3041,
0.2884,
0.2687,
0.2889,
0.2812,
0.2686,
0.3013,
0.2665,
0.2898,
0.2995,
0.2647,
0.3196,
0.2884,
0.2958,
0.2741,
0.2798,
0.2428,
0.2801,
0.3229,
0.3049,
0.2866,
0.2902,
0.2768,
0.2765,
0.2875,
0.2808,
0.2737,
0.2686,
0.2557,
0.2552,
0.2778,
0.2553
],
"dev_loss": [
{
"step": 50,
"dev_loss": 0.7212
},
{
"step": 100,
"dev_loss": 0.5114
},
{
"step": 150,
"dev_loss": 0.4558
},
{
"step": 200,
"dev_loss": 0.4143
},
{
"step": 250,
"dev_loss": 0.3454
},
{
"step": 300,
"dev_loss": 0.3116
},
{
"step": 350,
"dev_loss": 0.3082
},
{
"step": 400,
"dev_loss": 0.304
},
{
"step": 450,
"dev_loss": 0.2941
},
{
"step": 500,
"dev_loss": 0.297
},
{
"step": 550,
"dev_loss": 0.3141
},
{
"step": 600,
"dev_loss": 0.3078
}
],
"test_loss": 0.2694,
"color_recall_test": {
"correct": 39,
"total": 60
},
"most_common_color_baseline": {
"color": "yellow",
"correct": 18,
"total": 60
},
"parameter_checksum": 501.972163829,
"learning_rate_sweep": [
{
"learning_rate": 0.001,
"steps": 60,
"last_batch_loss": 3.4624,
"dev_loss": 3.4556
},
{
"learning_rate": 0.01,
"steps": 60,
"last_batch_loss": 0.6068,
"dev_loss": 0.6196
},
{
"learning_rate": 0.1,
"steps": 60,
"last_batch_loss": 0.6755,
"dev_loss": 0.6608
}
]
}Failure cases
Each case has a named test in test_train.py.
| Case | Result with this implementation | What it tells you |
|---|---|---|
| A story longer than the position table | encode_all() raises ValueError before training |
A model cannot read past its last position; check the data against the model before the first step |
| A non-finite gradient | Adam.step() raises ValueError naming the parameter |
One NaN would spread to every weight it touches |
| A learning rate that is zero, negative or NaN | Adam() raises ValueError |
The update would do nothing or corrupt the weights |
| Too few stories to split | split_stories() raises ValueError |
The split sizes are fixed; a shorter corpus must be rejected |
| A larger rate in a 60-step sweep | Development loss after 60 steps is 0.6608, worse than 0.6196 at the chosen rate | Compare rates on development loss, not on how fast the first steps fall |
The complete lab includes these tests and their fixtures. Run the 13 standard-library tests. They run short training jobs and replay the saved parameters; they do not repeat the full recorded run.
python3 -m unittest -v test_train.pyRead the tests
"""Tests for train.py. Run: python3 -m unittest -v test_train.py
These tests run short training jobs and replay the saved parameters. They do not
repeat the full recorded run; `python3 trace_train.py --check` does that.
"""
import json
import math
import unittest
from pathlib import Path
import model as m
import train as t
from bpe import Tokenizer
HERE = Path(__file__).parent
TOK = Tokenizer.load(HERE / "tokenizer.json")
TEXT = (HERE / "corpus.txt").read_text()
RECORD = json.loads((HERE / "training-trace.json").read_text())
def short(steps, **changes):
return dict(t.CONFIG, steps=steps, eval_every=steps, **changes)
class Mechanism(unittest.TestCase):
def test_adam_first_step_matches_the_formula(self):
p = {"w": [0.5, -0.25]}
g = {"w": [0.2, -0.4]}
opt = t.Adam(p, 0.01, 0.9, 0.999, 1e-8)
opt.step(p, g)
# After one step m-hat = g and v-hat = g squared, so each weight moves by lr * g / (|g| + eps).
for before, grad, after in zip([0.5, -0.25], g["w"], p["w"]):
self.assertAlmostEqual(after, before - 0.01 * grad / (abs(grad) + 1e-8), places=15)
def test_adam_bias_correction_matters_early(self):
p, q = {"w": [0.0]}, {"w": [0.0]}
opt = t.Adam(p, 0.01, 0.9, 0.999, 1e-8)
opt.step(p, {"w": [1.0]})
q["w"][0] -= 0.01 * (0.1 * 1.0) / (math.sqrt(0.001 * 1.0) + 1e-8) # uncorrected moments
self.assertAlmostEqual(p["w"][0], -0.01, places=9)
self.assertNotAlmostEqual(q["w"][0], p["w"][0], places=3)
def test_batch_gradient_is_the_mean_of_story_gradients(self):
p = m.init_params(len(TOK.vocab), 16, 64, 64, 3)
stories = t.encode_all(TOK, t.split_stories(TEXT)[0][:2], 64)
_, g = t.batch_gradient(p, stories)
single = []
for ids in stories:
logits, cache = m.forward(p, ids[:-1])
single.append(m.backward(p, cache, m.loss_and_dlogits(logits, ids[1:])[1]))
for name in ("W_out", "Wq", "tok_emb"):
for r, (a, b) in enumerate(zip(single[0][name], single[1][name])):
for j in range(len(a)):
self.assertAlmostEqual(g[name][r][j], (a[j] + b[j]) / 2, places=12)
def test_training_is_deterministic(self):
a = t.train(TOK, TEXT, short(3))[0]
b = t.train(TOK, TEXT, short(3))[0]
self.assertEqual(json.dumps(a), json.dumps(b)) # every value, not only the checksum
def test_short_run_lowers_development_loss(self):
_, _, dev = t.train(TOK, TEXT, short(30))
start = m.init_params(len(TOK.vocab), 16, 64, 64, 3)
_, dv, _ = t.split_stories(TEXT)
self.assertLess(dev[-1]["dev_loss"], t.mean_loss(start, t.encode_all(TOK, dv, 64)))
def test_saved_parameters_reproduce_the_recorded_scores(self):
p = json.loads((HERE / "params.json").read_text())
_, _, test = t.split_stories(TEXT)
self.assertEqual(t.checksum(p), RECORD["parameter_checksum"])
self.assertEqual(round(t.mean_loss(p, t.encode_all(TOK, test, 64)), 4), RECORD["test_loss"])
self.assertEqual(t.color_recall(p, TOK, test), RECORD["color_recall_test"])
def test_trained_parameters_beat_the_bigram_baseline(self):
bigram = json.loads((HERE / "baseline-trace.json").read_text())["test"]["bigram"]["loss"]
p = json.loads((HERE / "params.json").read_text())
_, _, test = t.split_stories(TEXT)
self.assertLess(t.mean_loss(p, t.encode_all(TOK, test, 64)), bigram)
class FailureCases(unittest.TestCase):
def test_story_longer_than_positions(self):
with self.assertRaisesRegex(ValueError, "does not fit 48 positions"):
t.encode_all(TOK, ["ava has a blue cup. ava goes to the garden. at the garden, ava gives "
"the cup to leo. now leo has the blue cup.\n"], 48)
def test_learning_rate_must_be_positive(self):
for bad in (0, -0.01, float("nan")):
with self.subTest(lr=bad), self.assertRaisesRegex(ValueError, "positive finite"):
t.Adam({"w": [0.0]}, bad, 0.9, 0.999, 1e-8)
def test_non_finite_gradient_stops_training(self):
p = {"w": [0.0, 1.0]}
opt = t.Adam(p, 0.01, 0.9, 0.999, 1e-8)
with self.assertRaisesRegex(ValueError, "non-finite gradient in w"):
opt.step(p, {"w": [0.1, float("inf")]})
def test_most_common_color_baseline(self):
train, _, test = t.split_stories(TEXT)
result = t.most_common_color_recall(train, test)
self.assertEqual(result, RECORD["most_common_color_baseline"])
self.assertLess(result["correct"], RECORD["color_recall_test"]["correct"])
def test_too_few_stories(self):
with self.assertRaisesRegex(ValueError, "need more than 340 stories"):
t.split_stories("a.\nb.\n")
def test_larger_rate_in_a_short_sweep_is_worse(self):
dev = {}
for lr in (0.01, 0.1):
dev[lr] = t.train(TOK, TEXT, short(60, learning_rate=lr))[2][-1]["dev_loss"]
self.assertGreater(dev[0.1], dev[0.01])
recorded = {r["learning_rate"]: r["dev_loss"] for r in RECORD["learning_rate_sweep"]}
self.assertEqual(dev, {lr: recorded[lr] for lr in dev})
if __name__ == "__main__":
unittest.main()test_adam_first_step_matches_the_formula and test_adam_bias_correction_matters_early fail if the bias correction is removed. test_batch_gradient_is_the_mean_of_story_gradients fails if the batch sums instead of averaging. test_saved_parameters_reproduce_the_recorded_scores fails if mean_loss() weights stories equally instead of by their predictions.
Deployment
The artifact is params.json together with the tokenizer and model.py; nothing serves the model yet. Keep the configuration and seeds from training-trace.json with it, so that anyone can retrain it and compare checksums.
Scaling
The one measured number is the recorded run’s time above. One step processes 8 stories one at a time; its cost grows with the batch size, with the square of the story length in attention, and with the model width times the vocabulary size in the output layer. Array libraries, batched matrix products and accelerators change every one of those factors; this build does not measure them.
Key takeaways
- One forward pass over a story trains every position at once, because the causal mask hides each answer from its own prediction.
- Adam keeps two running averages per parameter and corrects their bias toward 0 in early steps.
- Score development stories during training and test stories once at the end, and compare with the baseline on the same split and tokenizer.
Best practices
- Seed every source of randomness and record a checksum of the final parameters.
- Stop on a non-finite gradient instead of continuing with damaged weights.
- Report a targeted check, here the recalled color, beside the average loss.
References
Sources checked October 20, 2026.
- Kingma and Ba, Adam: A Method for Stochastic Optimization (ICLR 2015), arXiv 1412.6980v9: Algorithm 1, the bias correction in section 3, and the suggested defaults α = 0.001, β1 = 0.9, β2 = 0.999, ε = 10⁻⁸.
- PyTorch 2.14: torch.optim.Adam: the same update without weight decay, with the same defaults.
- PyTorch tutorials: optimizing model parameters: a training loop of forward pass, loss, backward and optimizer step, followed by an evaluation loop that changes no weights.