Chapter 21 Training, Saving, and Interrogating a Mini GPT

The previous chapter’s model assigned probabilities, but its parameters were still their initial random values. Running that forward computation a thousand times would not teach it anything. We must use the loss to change parameters, then separate two questions: did the program follow the learning rule, and did the resulting model acquire the ability we wanted?

MG-2 completes the same Mini GPT with bounded CPU training, saved state, a resumed run, a fresh-process repeat, and recorded generation. Every number below is either an explicitly derived calculation or a measurement from the shared local run. The small corpus is original teaching data; none of these results comes from a paper benchmark or an industrial language model.

21.1 Turn a loss into one update

Chapter 7 explained a gradient as the local rate at which loss changes with a parameter. Backpropagation follows the forward computation in reverse to calculate those rates. The model structure has not changed: gradients now travel from vocabulary loss through the head, norms, FFN, attention and embeddings. An optimizer uses them to update the stored parameters. Attention weights calculated inside one forward call are temporary activations, not the parameters that the optimizer stores.

Example 21.1 — The first real training step. Keep Chapter 20’s 24 training records in their declared order, numbered 0–23. Select eight indices independently with replacement. The fixed batch generator’s first selection is [15,19,1,22,7,15,5,3]. Record 15 appears twice because a draw does not remove it from the collection. Construct each selected record’s input, valid flags and shifted targets exactly as before. The observed batch mean loss before the update is 5.672425134737334 nats per valid target.

The loop’s essential operations are short because the forward mechanics are already implemented:

indices = torch.randint(len(corpus["train"]), (cfg["batch_size"],), generator=rng)
batch = tuple(array[indices] for array in train_batch)
opt.zero_grad(set_to_none=True)
loss = model(batch[0], batch[1], batch[2])[1]
loss.backward()
opt.step()

Clearing gradients prevents unintended accumulation from the previous step. backward() fills gradients; step() changes parameters. The first batch loss is measured before that change. The curve below instead evaluates the complete fixed splits after the numbered update. Different records, denominators and parameter times explain why these values differ from each other and from MG-1’s two-record loss.

We use Adam, the optimizer already named in Chapter 7’s nonlinear experiment. It keeps running statistics of gradients and squared gradients, so its update depends on history as well as the current gradient. This chapter treats that rule as a declared implementation choice; it does not attribute Adam’s results to the hand-derived plain SGD subtraction. The saved optimizer state will matter when we resume.

Table 21.1 — The run contract. Settings are fixed in data/mini-gpt/config.json; the environment and complete dependency list accompany the measured record.

Item	Value and meaning
Model and data	The same one-block, 13875-parameter model; 24 train and 8 validation records
Context / batch	64 slots per record; 8 records drawn with replacement per update
Updates / measurements	600 updates; complete-split evaluation at steps 0,50,…,600
Adam	Learning rate 0.01; betas 0.9 and 0.999; epsilon 0.00000001; weight decay 0; `foreach=False`
Initialization seed	1729, including the initial matrices
Batch seed	101, in a separate CPU generator
Sampling seed	31415, separate from training
Numeric execution	CPU, float64, one thread, deterministic algorithms enabled; dropout 0
Tested stack	Python 3.12.10, NumPy 2.0.2, PyTorch 2.7.0; macOS 26.5.2 arm64
Selection rule	First minimum validation loss among the 13 measured steps

A seed is an initial state for a particular random-number procedure, not a universal recipe independent of software. Separate generators prevent taking a sample for display from consuming the next batch draw. Source files, ordered data, tokenizer, configuration and dependencies are also part of the run identity.

21.2 Read the entire curve before choosing a snapshot

Evaluation performs model running (inference) with eval() and no_grad(), and performs no optimizer update. Each complete training evaluation scores 521 valid next-token events; validation scores 178. For either split, add the negative log probabilities of the actual targets and divide by that split’s number of valid events. Every event has equal weight; we do not first average each record and give short and long records equal total weight.

For example, a two-event record with losses 1 and 3, and a one-event record with loss 5, give $(1+3+5)/3=3$ nats per event. Averaging the record means instead gives $((1+3)/2+5)/2=3.5$ . Those are different questions. Our denominator follows Chapter 20’s event contract throughout training and evaluation.

Solid line and circles: training. Dashed line and squares: validation. Dotted vertical: selected step 50. Lines join measured points only.

Figure 21.1 — Training fit improves while validation worsens after step 50 Original measured CPU experiment, CC BY-SA 4.0. Source: data/mini-gpt/training-run.json; every plotted point is also in Table 21.2.

Table 21.2 — All measured points, in nats per valid target. Values are rounded to six decimals; training-run.json retains full precision. Lines in Figure 21.1 join these measurements and do not claim intermediate evaluations.

Completed updates	Train loss	Validation loss
0	5.674427	5.669803
50	1.269421	2.405571
100	0.368759	2.703053
150	0.216352	2.993822
200	0.176652	3.324710
250	0.175422	3.522885
300	0.168024	3.595803
350	0.166762	3.748807
400	0.170217	3.741431
450	0.159973	3.896360
500	0.164760	3.901104
550	0.155814	4.009695
600	0.159599	4.068374

Initially both losses are near the uniform baseline from Chapter 20. By step 50, both have fallen. After that, the training loss mostly falls further while validation becomes worse. At step 600 the model scores its training targets much better than at step 50, yet scores the withheld records worse. This divergence is evidence of overfitting under this split and objective. A lower training loss alone would have hidden it.

The declared rule selects step 50, with validation loss 2.405571231866604. It does not select step 600 just because training has finished. We retain both snapshots to compare their behavior. Because no measurements were taken between steps 0 and 50, the record does not identify the best possible earlier stopping time. We also do not report this repeatedly consulted validation set as a final independent test. The related templates and tiny data limit what even its minimum can establish.

21.3 Save the state needed for the next question

A file called a “checkpoint” can contain different things. For fixed model running, parameters plus the matching architecture and tokenizer may be enough. For continuing the same optimization path, Adam’s accumulators, step number and random-generator states also matter. A project milestone is a different kind of object: it includes evidence that the whole pipeline works.

Table 21.3 — Three meanings that should not be confused.

Saved object	What it contains or references	What it supports
Weight snapshot	Parameter state dictionary; separately versioned architecture and tokenizer	Recompute predictions with that model
Training-state checkpoint	Parameters, Adam state, step, batch/global Torch random states, configuration and data/tokenizer identities	Continue from the same next batch and update state
MG-2 milestone bundle	MG-1 trace, corpus/tokenizer/configuration, source/dependency identities, all curve points, checkpoint files, resume/repeat checks and generation traces	Audit and repeat this small end-to-end experiment

The local weights-selected.pt contains step 50’s parameters; weights-final.pt contains step 600’s. The training-state files are resume-step300.pt and resume-final.pt. Their measured sizes total 966540 bytes, about 0.967 decimal MB, including serialization overhead. This is a file-size measurement, not peak training memory. All four are generated from the original synthetic corpus; no pretrained weights are downloaded.

Saving only model.state_dict() would omit Adam’s running statistics. Recreating a batch generator with seed 101 would restart its sequence, not continue after step 300. Our loader verifies the corpus, tokenizer and configuration identities before loading the declared architecture, then restores the model, optimizer and generators. It uses torch.load(..., map_location="cpu", weights_only=True) for these locally produced tensor/container files. The setting does not supply missing configuration or repair a changed vocabulary.

The continuation comparison concerns parameter updates. MG-2 also retains the earlier curve and the selected weight snapshot separately: a step-300 training state alone cannot tell us which earlier checkpoint had the lowest validation loss. Those records are needed to preserve selection over the whole run, rather than choosing again only among the resumed steps.

Example 21.2 — Resume means reproducing the next update. The uninterrupted run’s step 301 selects [12,16,9,15,22,6,3,1]; its pre-update batch loss is 0.17280904760175986. A second run loads the step-300 checkpoint and starts there. The batch indices and loss match, followed by every recorded step through 600.

Table 21.4 — Actual continuation and repeat evidence. Equality here is exact in the declared stack, not merely equal rounded displays.

Comparison	Observed result
Resume: batch indices and pre-update losses at steps 301–600	All equal
Resume: complete-split curve points at steps 300–600	All equal
Resume: final parameters, optimizer and both random states	All equal by deterministic tensor/state digests
Resume: next batch after step 600	Both give `[15,21,7,2,23,16,7,16]`
Fresh process: all 600 steps and all 13 curve points	All equal
Fresh process: selected step, final states and next batch	All equal

The fresh process begins again from initialization rather than loading the trained weights. That checks a different route from continuation. The recorded final parameter digest is:

6de30d8677afce8e39691f446344d4bd67ae0eda5136743e1894152b78c54c01

The code hashes each tensor’s type, shape and bytes in a fixed traversal; checkpoint file hashes separately identify serialized files. Matching digests detect changes in the covered state, but cannot prove that an implementation has the intended mathematics. The independent mechanism, causal-visibility and loss checks from earlier chapters remain necessary.

21.4 A generated sentence is a sequence of fresh decisions

For generation, freeze the chosen parameters. Encode the prompt once, add BOS, run the model, and use the logits at the last real input position. After choosing a token, append its ID and run again. This simple implementation recomputes the whole prefix; it uses no key/value cache. It does not rerun the tokenizer over the growing text after every choice.

Example 21.3 — A complete three-decision continuation. Load the final step-600 snapshot, explicitly chosen here to inspect its memorized pattern. The prompt is The key is red. Answer: , including one space after the colon. Its input IDs, including BOS, are [256,84,265,268,32,277,270,261,65,276,119,262,259]. Use greedy selection: select the highest eligible logit, breaking an exact tie by smaller ID. BOS and PAD are forbidden as outputs, EOS is allowed, and the maximum new-token budget is 24.

Decision 1 · final step-600 weights

Row number13 IDs; last ID 259 = colon-space
An operationForward → last logits → greedy
Row numberAppend 270 = red; p = 0.986202

The model has not yet received a generated red token. BOS/PAD cannot be selected.

Decision 2 · same weights

Row number14 IDs; last ID 270 = red
An operationRecompute with the new prefix
Row numberAppend 46 = period; p = 0.999693

The period probability belongs to a new vocabulary distribution.

Decision 3 · same weights

Row number15 IDs; last ID 46 = period
An operationRecompute → select 257 = EOS
Text piecesStop; ordinary text: red.

p(EOS) = 0.999583. Probabilities shown are raw model probabilities, before excluding BOS/PAD.

Figure 21.2 — The same parameters, three changing prefixes Original teaching diagram, CC BY-SA 4.0. Values and rules are declared in the chapter and shared data.

Table 21.5 — Each appended ID changes the next input. The probability column is the raw model probability over all 291 candidates, before forbidding BOS/PAD.

Decision	Input length before decision	Selected ID / piece	Selected logit	Model probability	Result
1	13	270 / `red`	12.756594	0.986202	Append ID 270
2	14	46 / `.`	17.023494	0.999693	Append ID 46
3	15	257 / EOS	16.481496	0.999583	Stop; display `red.`

The first row’s next-most-probable candidates include blu (ID 274, probability 0.006150) and key (ID 268, 0.001701). The distribution strongly favors red, but the next row is a new computation with red already in its input. Its probability of a period is not a conditional entry read from the first distribution. EOS ends the record and is not printed as ordinary text.

Sampling changes the selection rule. For a declared temperature tau greater than zero, divide eligible logits by tau before exponentiating. With top-k sampling, keep only the k highest eligible logits, then normalize over that retained set. Thus, for retained candidate IDs R:

p_i^{\text{sample}}=\frac{\exp(z_i/\tau)}{\sum_{j\in R}\exp(z_j/\tau)},\qquad i\in R.

Here z_i is the last-position score for candidate i, not an attention score. Excluded candidates get probability zero. Lower temperature sharpens score differences; higher temperature flattens them within the retained set. For logits 2 and 0 with both candidates retained, tau=1 gives the first candidate $e^2/(e^2+1)\approx0.880797$ ; tau=2 gives $e/(e+1)\approx0.731059$ . Neither rule supplies information absent from the model.

Our implementation subtracts the largest retained scaled logit before exponentiation, preserving these probabilities while avoiding large exponentials. It logs both raw model probabilities and selection probabilities. Even greedy generation’s probability after excluding BOS/PAD is a slightly different normalized quantity: in the first row above it is 0.9862021226445513 rather than the raw 0.9862019510381717.

21.5 Change the clue, not just the display of success

The original red diagnostic omitted the final space: The key is red. Answer:. Its last token is the colon byte, ID 58. The training record instead has a single merged : token, ID 259. Exact text round trip does not make the shorter encoding a token prefix of the longer encoding. This boundary mismatch was observed before adding the space-aligned probes; the original failures remain in generation-baseline.json and generation.json.

Table 21.6 — Greedy continuations under the same token budget. Code spans preserve punctuation; “space” means one literal space after the colon. These are measured completions, not corrections of the model’s spelling.

Prompt condition	Selected step 50	Final step 600
`The key is red. Answer:` without space	`of Anspitap the blue.`	`Result: keyred red.`
`The key is red. Answer:` with space	`bluesuersther: blue key.`	`red.`
`The key is blue. Answer:` with space	`blue.`	`blue.`
`The key was red. Now blue. Answer:` with space	`blue key.`	`red.`
`Old limit: 600. New limit: 750. Answer:` with space	`peng.`	`shipedwion: old.`

Aligning the boundary makes the familiar final-model red continuation work. It does not establish robust handling of a changed clue. The fourth prompt explicitly changes the key to blue, yet the final model returns red. The fifth supplies 750 as the new fictional limit, yet neither snapshot returns it. These controlled changes are reasons to reject a broad claim that the model reliably follows updates. They are not a population accuracy estimate.

The original corpus also includes Old limit: 750. New limit: with an authored expected suffix 600.. That prompt never supplies the new value. A mismatch with 600 therefore cannot by itself prove violation of an explicit update. We retain the probe and identify this limitation instead of silently treating its desired suffix as a logically determined answer. The later explicit-value probe tests a better specified question; it is a separately versioned addition, not a replacement of the unfavorable baseline or a new checkpoint-selection set.

Table 21.7 — Changing temperature on the aligned red prompt. Both sampled runs use top-k=5, seed 31415 and at most 24 new tokens.

Rule	Selected step 50	Final step 600
Greedy	`bluesuersther: blue key.`	`red.`
Sample, temperature 0.7	`blueshindiwer: ng.`	`red.`
Sample, temperature 1.2	`blue.`	`red.`

The same final completion at two temperatures is possible: probability distributions can differ while a particular sequence of draws chooses the same IDs. The selected snapshot’s minimum validation loss likewise does not promise success on every individual prompt. Picking only the final red success would conceal both the validation curve and the changed-clue failure.

Every displayed run stops on EOS and decodes as valid UTF-8. The generator also has explicit stops for its new-token budget and a 64-ID total input-buffer limit, including BOS; it does not slide the context window or invent new position rows. Arbitrary generated bytes need not form valid UTF-8. That failure is tested separately: raw bytes are retained, valid_utf8 is false, and a replacement-character display is identified as a display, not an exact round trip.

21.6 Repeat the bounded experiment

After the CPU environment setup in code/part-iv/README.md, run from the book root with that environment’s Python:

python code/mini-gpt/forward.py
python code/mini-gpt/run.py --output /tmp/mini-gpt-reproduction
python -m unittest discover -s code/mini-gpt -p 'test_*.py' -v

The training command performs the full run, continuation check, fresh-process repeat and recorded generation. A separate output directory preserves the accepted reference evidence. code/mini-gpt/README.md maps the files and commands; notebooks/part-iv/03-mini-gpt.ipynb executes the same core functions in a fresh CPU kernel and retains its outputs. A readable notebook is supplementary; the tables above already expose the decisive evidence.

The reference record was produced on September 13, 2026 in the Asia/Shanghai timezone. The main 600-step run took 1.7165 seconds, continuation from 300 took 0.7023 seconds, and the fresh-process training section took 1.7270 seconds. The complete pipeline took 4.7675 seconds. These are measured wall-clock durations on one machine, including different surrounding work; process startup and cache conditions vary. They are not a hardware performance comparison or a promised runtime for a reader. Dependency installation is separate and needs network access once; the experiment itself uses only local data.

Within the recorded stack, exact repeats succeeded. Across operating systems, devices or library versions, a matching seed does not guarantee identical floating-point results or sampled text. Such a repeat should first check the invariant mathematics, split/tokenizer identity, finite loss, lack of future leakage, and the direction of the overfitting evidence, then report numerical differences explicitly. This task has not measured a cross-platform or GPU reproduction.

21.7 Part IV checkpoint: connect the evidence to the claim

We can now trace a head from input coordinates to weighted values, state exactly which keys a query can see, combine independent heads, and place that operation inside a residual/norm/FFN block. Position information and token boundaries have explicit interfaces. Architecture and training objective are separate choices; none of them alone turns a network into a reliable assistant.

The same diagram now corresponds to real CPU arrays and code. MG-1 establishes the path from text to logits/loss. MG-2 establishes parameter updates, saved state, continuation, repeatability under declared conditions and inspectable generation. The curve and changed-clue probes establish limits as well as progress. No live order operation, factual database lookup or general problem-solving reasoning was demonstrated by the generated text.

The next question is what changes when data, model size and the surrounding learning process grow. Chapter 22 takes up that question through the modern LLM lifecycle. Our small model supplies a mechanism to reason from, together with a discipline: state the desired ability, choose evidence that could fail, and keep the failure visible.

21.8 Exercises

A new validation split has records containing 2 and 3 valid targets. Their individual losses are [1,3] and [2,2,7]. Compute the event-weighted mean and the equally weighted record mean. Which matches this project?
A colleague chooses step 600 because its train loss is 0.159599, far below step 50’s 1.269421. Use Table 21.2 to assess the choice under the declared rule. Can the same table locate the best step between 1 and 49?
Two programs restore the same step-300 weights. One recreates Adam and resets seed 101; the other restores Adam and the batch-generator state. Are their next predictions necessarily different? Are their next parameter updates necessarily the same? Explain the distinction.
Two eligible logits are 2 and 0. Calculate the first candidate’s probability at temperature 2. What does top-k=1 do to that probability? Does changing temperature then guarantee a different completion?
The final model says red for both the familiar red prompt and the changed-blue prompt in Table 21.6. State the narrow supported success and the unsupported broader claim. Why is the original unspecified new-limit probe weaker evidence?
A new tokenizer keeps vocabulary size 291 but exchanges the meanings of IDs 270 and 274. Shapes and weight loading still pass. What would break, and which checkpoint check should reject it?
A generation trace ends because the 24-new-token budget is reached, without EOS. A display shows a tidy final sentence. What must its record retain before another reader can interpret that display?
Design an evidence checklist for extending this model to 32-coordinate representations while keeping two heads. Include one shape consequence, one information-flow test, one loss check, and one evaluation choice that avoids selecting only appealing outputs.

21.9 Reference answers

Total loss is 15 over five events, so the project mean is 3. The record means are 2 and 11/3; their equal average is 17/6, approximately 2.833333. The latter gives the shorter record more weight per event and does not follow the declared loss reduction.
The rule chooses step 50 because its validation loss is the minimum measured value, 2.405571. Step 600 has worse validation loss, 4.068374. Training fit is useful evidence but is not the selection criterion. Nothing in the table measures steps 1–49; interpolating a drawn line does not create those evaluations.
Before another update, matching model parameters, inputs and evaluation settings can produce the same predictions without any optimizer state. Subsequent updates need not match: Adam’s accumulators differ, and a restarted generator need not choose the same batch. Restoring both states is necessary for the declared continuation; identities and the remaining execution conditions must match too.
Dividing by 2 gives logits 1 and 0, so the probability is $e/(e+1)\approx0.731059$ . Top-k=1 retains only the larger logit and makes its selection probability 1 for any positive temperature. With larger k, different distributions can still yield the same sampled completion on particular draws.
The model reproduces a familiar training-template suffix when the token boundary aligns. It has not demonstrated reliable use of revised clues; returning red after the explicit change to blue contradicts that claim. The unspecified-limit prompt supplies no new limit, so failure to match an author’s preferred 600 is not by itself a violation of supplied information.
Embedding lookup and output-column meanings would no longer correspond to the saved model, even though all shapes remain legal. The tokenizer identity check should reject the load. A deliberate migration would have to update every affected parameter/data interface and be recorded as a new version.
Retain the prompt and token IDs, selected checkpoint/tokenizer, sampling settings and seed, full generated IDs/bytes, stop reason, and UTF-8 validity. A budget stop is not an EOS decision; cleaned punctuation or replacement decoding must not be silently substituted for the actual output.
Head width becomes 16 and the scale becomes the square root of 16; Q/K/V and residual widths must be updated consistently. Change future IDs and require earlier real logits to remain unchanged. Check that zero logits still give mean loss $\ln(291)$ over valid targets only. Keep the split and step-selection rule fixed, retain the whole validation curve and all original probes, and add separately versioned tests for the ability the wider model is supposed to improve.

21.10 References

PyTorch, Adam, 2.7. Specifies the optimizer’s running gradient statistics, configuration and saved state; the measured curve is this book’s experiment.
PyTorch, Saving and Loading Models and torch.load, 2.7. Support the distinction between model state and continuation state, and the explicit CPU tensor/container loading used here. The implementation and tested versions are pinned in the accompanying records.
PyTorch, Reproducibility, 2.7. Explains controlling random sources and deterministic algorithms, and why a seed does not guarantee equality across software releases and platforms.