Checkpoints & Tests

This page covers saving and loading, the prompt loop in talk.py, and the test suite.

This page has no equations and no backward pass. Saving, loading and the prompt loop compute no loss and no gradient. The mathematics that talk.py reaches is on Autoregressive Generation.

Checkpoint format

What is it?

A checkpoint is three files that share a prefix. With the default prefix, checkpoints/model, they are checkpoints/model.npz, checkpoints/model.json and checkpoints/model.tok.json.

Why does the model need it?

talk.py runs in a separate process from train.py, and it builds nothing itself. It needs the weights, the shapes they fit, and the tokenizer whose ids they were trained on. Each file holds one of the three.

{"vocab_size": 9, "context_len": 8, "emb_dim": 4, "num_heads": 2, "ff_dim": 8, "num_layers": 1, "ln_eps": 1e-05}

{"merges": [{"pair": ["l", "o"], "merged": "lo", "frequency": 4}, {"pair": ["lo", "w"], "merged": "low", "frequency": 4}], "itos": ["<unk>", "l", "o", "w", " ", "e", "r", "lo", "low"]}

Tensor shapes

The arrays in the .npz file have the shapes that init_params gave them, because saving does not change them. The Architecture page lists each one at the values in config.py. V depends on the corpus, so it is not stored in config.py. It is stored in the checkpoint as vocab_size.

Fields of the .json file

asdict(cfg) writes one key per field of ModelConfig, including ln_eps. train.main does not pass ln_eps, so the file holds the default of the dataclass.

Merge records in the .tok.json file

Each merge is a record with the keys pair, merged and frequency, which train_bpe creates (tokenizer.py:59). JSON has no tuples, so pair is written as a list. merge_tokens reads only pair[0] and pair[1], so a list works the same way.

ensure_ascii=False keeps characters outside ASCII readable in the file, and test_non_ascii_characters_survive checks that they come back unchanged.

Saving

16def save_checkpoint(prefix, params, cfg: ModelConfig, tok: Tokenizer):17    folder = os.path.dirname(prefix)18    if folder:19        os.makedirs(folder, exist_ok=True)2021    np.savez(prefix + ".npz", **params) # keys like "b0.wq" are kept as they are2223    with open(prefix + ".json", "w", encoding="utf-8") as f:24        json.dump(asdict(cfg), f)2526    # stoi is not saved, it is just itos turned around, so load rebuilds it27    # json has no tuples, so each merge's "pair" is stored as a list, merge_tokens only28    # uses pair[0] and pair[1], so a list works exactly the same29    with open(prefix + ".tok.json", "w", encoding="utf-8") as f:30        json.dump({"merges": tok.merges, "itos": tok.itos}, f, ensure_ascii=False)

Where is it used?

Called by

train.main

Tested by

tests.test_checkpoint.saved, tests.test_checkpoint.test_non_ascii_characters_survive

Loading

32def load_checkpoint(prefix):33    with np.load(prefix + ".npz") as z:34        params = {k: z[k] for k in z.files} # copy out, the file closes after this block3536    with open(prefix + ".json", encoding="utf-8") as f:37        cfg = ModelConfig(**json.load(f))3839    with open(prefix + ".tok.json", encoding="utf-8") as f:40        d = json.load(f)4142    stoi = {}43    for i in range(len(d["itos"])):44        stoi[d["itos"][i]] = i4546    return params, cfg, Tokenizer(d["merges"], stoi, d["itos"])

Where is it used?

Called by

talk.main

Tested by

tests.test_checkpoint.test_config_round_trips, tests.test_checkpoint.test_loaded_model_gives_identical_logits, tests.test_checkpoint.test_non_ascii_characters_survive, tests.test_checkpoint.test_tokenizer_round_trips, tests.test_checkpoint.test_weights_round_trip_exactly

Only talk.main calls it in the pipeline, so nothing in the repository resumes training from a checkpoint.

What a checkpoint does not hold

save_checkpoint receives params, cfg and tok, and writes only those. So these are not saved: Adam's m, v and t, the random generator's state, the step count, and the training settings in config.py. A checkpoint cannot tell you how it was trained.

Why the files are not in the repository

.gitignore lists checkpoints/ and /data/, so neither a checkpoint nor the corpus is part of a commit. This page shows no output from a trained model, so every number on it is an Illustrative Example.

Interactive inference

talk.py loads a checkpoint and completes text in a loop.

7def main(prefix="checkpoints/model"):8    params, cfg, tok = load_checkpoint(prefix)9    rng = np.random.default_rng(config.seed)10    print(f"loaded {prefix}  (V={cfg.vocab_size})")11    print("Type the start of a text and press Enter. Empty line quits.")1213    while True:14        try:15            prompt = input("> ")16        except EOFError:17            break1819        if prompt.strip() == "":20            break2122        ids = tok.encode(prompt)2324        out = generate(params, cfg, ids, config.max_new_tokens, config.seq_len,25                       config.temperature, config.top_k, config.top_p, rng)2627        print(tok.decode(out[len(ids):])) # only the new tokens, the prompt is already on screen

1. Load. load_checkpoint restores the weights, shapes and tokenizer (talk.py:8). The default prefix is checkpoints/model. Running python talk.py <prefix> passes another one (talk.py:30).
2. Create the generator. It uses config.seed, once, before the loop (talk.py:9). So a new run of talk.py repeats the same random sequence, and a second prompt in one run continues it.
3. Read a prompt. End of input stops the loop (talk.py:14-17), and so does an empty or whitespace-only line (talk.py:19).
4. Encode. tok.encode turns the prompt into ids (talk.py:22). A character the tokenizer never saw becomes id 0, <unk>.
5. Generate. generate appends config.max_new_tokens ids (talk.py:24-25). See Autoregressive Generation.
6. Decode only the new tokens. The prompt is already on screen, so out[len(ids):] skips it (talk.py:27).

The model's shape and tokenizer come from the checkpoint. The window and the sampling settings come from config.py as it is when talk.py runs. If config.seq_len is raised above the checkpoint's context_len, a long enough output fails the assertion in forward (transformer.py:156).

Called by

talk.<module>

Tested by

No test calls it directly

No test calls talk.main. The functions it calls are tested separately.

Reproducing training

These are the commands the code supports. They were not run for this page.

1. Supply a corpus. train.main reads config.corpus_path, which is "data/corpus.csv", using the column "text". A fresh clone does not have this file, because /data/ is git-ignored.
2. Install the dependencies. The repository imports numpy and pandas, and the tests also import pytest. It has no requirements file.
3. Train. Run python train.py from the repository root. The script calls main() with no arguments (train.py:85), so it always saves to checkpoints/model.
4. Talk. Run python talk.py from the repository root.

Training settings are the values in config.py, listed on the Training Loop page. Random numbers start from config.seed.

Test suite

Run the tests with python -m pytest tests from the repository root. The test files import root modules by bare name, and the repository has no conftest.py or pytest configuration.

Tests write only to pytest's temporary folder, so they do not touch checkpoints/.

8 test files hold 64 test functions, which make 67 cases because parametrized tests run once per value, counted from data/generated/tests-index.json. This page does not state a pass count, because it would go out of date.

"Exercises" lists what a test calls directly, or through its fixtures and helpers. A function called from inside one of those is covered too, but it is not listed. Each sentence in "What it verifies" links to its test.

What no test reaches

These functions are not reached by any test, even indirectly, according to the static call graph:

• data.encode_corpus
• talk.main
• train.batch_loss_and_grads
• train.validation_loss
• train.main

The two main functions are the whole pipeline, so an end-to-end run of training or of the prompt loop is not tested. test_overfit_one_window is the closest. It runs loss_and_grads and Adam.step together on one window.

Gradient checks have their own page, Gradient Checking.