Lavender-2

Appearance

Checkpointing ​

Why it needs to be more than torch.save ​

A training checkpoint has to preserve enough state to resume training, not just to run inference. Concretely:

  • Model weights — model.state_dict().
  • Optimizer state — Adam's momentum and variance buffers. Without these, a resumed run kicks the model into a new optimization trajectory and effectively restarts Adam's adaptive learning rates from scratch.
  • LR scheduler state — so scheduler.step() picks up at the right point on the cosine curve.
  • Tokenizer — serialized to JSON and saved alongside the weights, so inference loads the exact same string-to-ID mapping the model was trained on.
  • RNG state (nice-to-have) — so data ordering and dropout masks are reproducible from that step.

A naive torch.save(model, path) captures only the first bullet.

What SafeCheckpointer does ​

The notebook defines a SafeCheckpointer class that writes all of the above with triple redundancy:

  1. Primary path — checkpoints/latest/.
  2. Fallback A — checkpoints/latest.backup/.
  3. Fallback B — checkpoints/step_<N>/, a rolling step-indexed snapshot.

The write order is:

  1. Write the primary to a temporary directory.
  2. fsync the directory.
  3. Atomically rename into checkpoints/latest/.
  4. Copy to fallback A.
  5. Copy to fallback B.

Renames are atomic on POSIX filesystems, so if the process is killed mid-write, at most one of the three snapshots is in a torn state and the other two recover cleanly.

Why bother ​

Pretraining runs are long and expensive to redo. A single corrupted checkpoint — from an OOM kill, a preemption, a filesystem hiccup — can cost hours of GPU-time to reach again. Triple redundancy is cheap insurance compared to the cost of re-running.

Invariants ​

  • checkpoints/ is git-ignored. Do not commit it; the files are huge and contain model state, not code.
  • Checkpoints include the tokenizer. A checkpoint without its tokenizer is useless — there is no way to know what string each ID originally meant. chat.py expects both together in the checkpoint folder.

Reading the code ​

Search for SafeCheckpointer in full_llm_pipeline.ipynb. The class is short — the interesting part is the ordering of the writes, not the API surface.

Next stage: Inference — loading a checkpoint and generating text.