LM8-T5 — Choosing a Base Model and Managing Checkpoints

Learning objectives

After this lesson, you should be able to:

• List the criteria for choosing a base model and apply them to a task with a hardware budget.
• Estimate whether a candidate model fits training and serving budgets.
• Describe what a checkpoint or adapter must record to be reproducible.
• Set up an experiment log and an evaluation protocol that supports honest comparison and rollback.

The idea: the base sets the ceiling, the record sets the trust

Choosing a base model is a search under constraints. Fine-tuning elicits and shapes abilities that the base already has (Lesson LM7-T1), so a base that lacks the language, domain or reasoning ability you need cannot be fixed by tuning, while an oversized one may be unaffordable to train or serve. The practical method: write down the requirements, shortlist models that satisfy the hard constraints, test a few on your own examples before training, and choose the smallest that is good enough.

Criteria

Managing checkpoints

A checkpoint is a snapshot of training state; an adapter or a final model is an output. Each should be saved with enough information to be understood and reproduced: which base, which data, which settings, which code, how it scored. Save checkpoints on a schedule, keep the best by validation loss and the last for resuming, and delete the rest by a stated retention rule. Never overwrite the only copy of a good result.

Experiments

Keep a table: one row per run, with the settings that varied, the seed, the validation metrics on the same evaluation set, cost and time, and a one-line conclusion. Change one thing at a time. Compare runs on identical data and evaluation, and repeat promising runs with another seed before believing a small difference (Lesson LM6-T5).

The mathematics: sizing and comparing

Will it fit? For a candidate with P parameters, check training memory (Lesson LM8-T3) and serving memory:

A 7B model in 4 bits is 3.4 GB plus a KV cache of 0.5 MB per token in 16-bit: with 4,096 tokens and 8 concurrent requests that cache is 17 GB, larger than the weights. Serving budgets are therefore set by cache and concurrency as much as by the model.

Cost per token. A forward pass costs about 2P FLOPs per generated token. At 40% utilisation of a 10¹⁴ FLOP/s device, a 7B model generates at most 4 · 10¹³ / (1.4 · 10¹⁰) ≈ 2,900 tokens per second in compute-bound batches; single-stream generation is limited by memory bandwidth (reading the weights for each token), typically far lower. Halving P roughly doubles throughput on the same hardware.

Comparing candidates. With n evaluation items, an accuracy has standard error √(p(1 − p)/n) (Lesson LM6-T5). Candidates whose scores differ by less than about two standard errors cannot be separated: on 100 items at 80%, that is 8 points. Evaluate candidates on a few hundred items before choosing, and prefer the cheaper one when they tie.

Identity of a checkpoint. Let H(·) be a cryptographic hash (SHA-256) of the weights. Store H(base) in the adapter’s metadata, and H(adapter) in the evaluation report. Equality of hashes proves that the thing evaluated is the thing deployed; a mismatch proves nothing was reproduced.

Retention. A run saving a full checkpoint of P parameters in float32 plus optimizer state (16 bytes per parameter in mixed precision) every k steps for S steps consumes P · 16 · (S/k) bytes unless old ones are deleted: 124 M parameters, 10,000 steps, every 500 steps: 20 checkpoints × 2.0 GB = 40 GB. Keep the last and the best, and an adapter-only artifact for each reported result.

A worked example: a checkpoint that carries its identity

The code saves a small checkpoint together with metadata (base identifier and hash, data version, seed, step, learning rate, code revision), loads it back, verifies the weights against the recorded hash and shows that a tiny change is detected.

import hashlib
import io
import json
import torch
import torch.nn as nn

def sha256_of_state(state):
    h = hashlib.sha256()
    for k in sorted(state):
        h.update(k.encode())
        h.update(state[k].detach().cpu().numpy().tobytes())
    return h.hexdigest()

torch.manual_seed(0)
model = nn.Linear(4, 2)
meta = {"base": "toy-linear-v1", "base_sha256": sha256_of_state(model.state_dict())[:12],
        "data": "toy-data-v3", "seed": 0, "step": 100, "lr": 3e-4, "git": "abc1234"}

buf = io.BytesIO()                                      # a file in memory; use a real path in practice
torch.save({"model": model.state_dict(), "meta": meta}, buf)
buf.seek(0)
ck = torch.load(buf)
print(sorted(ck["meta"]))                               # ['base', 'base_sha256', 'data', 'git', 'lr', 'seed', 'step']
print(sha256_of_state(ck["model"])[:12] == meta["base_sha256"])      # True: the weights are the ones that were recorded

model2 = nn.Linear(4, 2)
model2.load_state_dict(ck["model"])
print(all(torch.equal(a, b) for a, b in zip(model.parameters(), model2.parameters())))   # True: an exact round trip

ck["model"]["weight"][0, 0] += 1e-3                     # a tiny corruption
print(sha256_of_state(ck["model"])[:12] == meta["base_sha256"])      # False: the hash detects it
print(json.dumps(meta)[:40])                            # the metadata is plain JSON and can be kept in a run table

In a real project the metadata also holds the evaluation scores and the path of the data; the principle is that the artifact explains itself and can be checked.

A decision process

1. Write the requirements: task, language, context length, latency, hardware, budget, licence constraints.
2. Shortlist three or four models that meet the hard constraints (licence, fit, context).
3. Evaluate each zero-shot and few-shot on a few hundred of your own examples, with a fixed protocol (Lesson LM7-T5).
4. Estimate fine-tuning and serving cost for the two best (Lesson LM8-T3).
5. Fine-tune the smallest that has headroom, compare with the larger one if budget allows, and choose on validation results and cost.
6. Record the decision and the alternatives considered, with the date, so that it can be revisited when new models appear.

Engineering implications and limitations

Prefer safe formats. A torch.save checkpoint is a pickle and can execute code when loaded. Load only files you trust, or use a tensor-only format such as safetensors for weights you receive from others.

Pin versions. Record library versions (PyTorch, tokenizer, tiktoken or equivalent), the CUDA version and the random seeds. Identical code on a different library version can produce different numbers.

Verify downloads. Check the hash or signature of base weights against the publisher’s. A silently updated or tampered base invalidates every adapter built on it.

Name consistently. A name like base-v1__data-v3__lora-r16__seed0__step2000 carries the main choices; the metadata carries the rest.

Keep an evaluation set that never changes. If the evaluation set changes, scores across time are not comparable. Version it and record which version each score used.

Plan rollback. Keep the previous production model and its evaluation until the new one has proven itself; deployment should be switchable by changing a pointer, not by rebuilding.

Document intended use and limits. A short model card: what it was trained on, what it was evaluated on, known failures, and what it must not be used for.

For robots. A model on a robot is part of a safety case. Tie each deployed model to its evaluation report, the simulator tests it passed and the hash that is running on the device, and log that hash with every mission.

Common misconceptions

Engineering takeaways

1. Choose the base by licence, fit, capability on your own examples, context length, tokenizer, data provenance and ecosystem; pick the smallest that is good enough.
2. Serving memory is weights plus KV cache; the cache can exceed the weights with long contexts and concurrency.
3. A checkpoint or adapter must record its base (id and hash), data version, settings, seed, code revision and evaluation.
4. Hash the weights, pin library versions, version the evaluation set and keep a run table; keep the best, the last and the reported checkpoints.
5. Plan rollback, load only trusted files, and document intended use and limits.

References and further reading

See Chapter 5 of the book (loading and saving model weights in PyTorch) for the mechanics of checkpoints; the checklist and the process in this lesson are general engineering practice rather than a statement about any particular model.

• Mitchell, M. et al. (2019). Model Cards for Model Reporting. arXiv:1810.03993.
• Build a Large Language Model (From Scratch), Appendix E: Parameter-efficient Finetuning with LoRA (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

That completes Module 8. Next, Module 9 turns to evaluation, inference and reliability of whole systems.

Related content

Nothing below is required to finish this lesson.

• Before
  LM8-T4: Quantization Concepts and Trade-Offs

  Map real numbers to low-bit integers, measure the error and know when it hurts.

• Lab
  LM8-L3: LoRA Fine-Tuning Experiment

  Add LoRA layers to a small GPT, freeze the base weights and compare quality, parameters and memory with full fine-tuning.