Checkpointing and resuming

What a checkpoint contains

A checkpoint is not only weights. A complete one holds five things:

1. The trainable weights: the adapter, or the full model.
2. The optimiser state: AdamW’s two moment buffers per trainable parameter.
3. The learning-rate scheduler’s position.
4. The random number generator states, so data shuffling resumes identically.
5. The trainer state: global step, epoch, logged metrics.

Only the first is needed to use the model. All five are needed to continue the run.

The size follows: a 42M-parameter LoRA is about 160 MB in bf16, and its optimiser state is another 340 MB, so a resumable checkpoint is roughly half a gigabyte. Keeping three is a gigabyte and a half, which is nothing. Keeping three full 8B checkpoints is 130 GB, which is not.

The settings

SFTConfig(
    save_steps=200,
    save_total_limit=3,
    save_safetensors=True,
    load_best_model_at_end=False,   # see below
    output_dir="out/triage-v3",
)

Set save_steps so that a crash costs you at most ten or fifteen minutes of work. At three seconds per step that is every 200 to 300 steps.

save_total_limit=3 keeps the disk bounded and is enough to step back from a checkpoint that turned out to be past the peak.

load_best_model_at_end=False is a deliberate choice. “Best” is defined by validation loss, and (as the epochs lesson established) validation loss is not what you want to select on. Keep all three and choose on your task metric afterwards, with the canary set beside it.

Resuming correctly

trainer.train(resume_from_checkpoint=True)

This restores all five components and fast-forwards the data loader to where it stopped. Two conditions make that correct, and both are easy to break:

The data must be in the same order. Resuming replays the shuffle from the saved RNG state, so a dataset that has been re-shuffled with a different seed, or had rows added, produces a different sequence. The resumed run will not see the examples it was going to see, and may see some twice.

The world size must match. Resuming a four-GPU run on two GPUs changes the effective batch size, which changes the number of optimiser steps per epoch, which puts the scheduler in the wrong place. It will run and it will not be the run you started.

Interruptible instances

Spot and community instances are roughly half price and can be reclaimed with little warning. With checkpointing in place that is an inconvenience rather than a loss, provided one thing: the checkpoint must not be on the instance’s local disk.

# free Colab or Kaggle
output_dir = "/content/drive/MyDrive/runs/triage-v3"

# or push each checkpoint to a private Hub repo
SFTConfig(push_to_hub=True, hub_model_id="you/triage-v3",
          hub_private_repo=True, hub_strategy="checkpoint")

Private repositories on the Hugging Face Hub are free, and hub_strategy="checkpoint" pushes the resumable checkpoint rather than only the final weights. A reclaimed instance then costs you a restart, not a run.

Saving the adapter alone

For distribution you want the weights without the optimiser state:

trainer.model.save_pretrained("out/adapter")    # ~160 MB
tok.save_pretrained("out/adapter")               # always alongside

Save the tokenizer with it, every time. An adapter without its tokenizer is a file that loads and produces subtly wrong results, because somebody six months later will pair it with a slightly different tokenizer revision and have no way to know.

Evaluating checkpoints without stopping the run

Checkpoints accumulate while the job is still going, and evaluating them afterwards serialises two things that could overlap. Point a second process at the output directory and have it score each new checkpoint as it appears:

while true; do
  for c in out/triage-v3/checkpoint-*; do
    [ -f "$c/.scored" ] || { python evaluate.py "$c" >> results.jsonl; touch "$c/.scored"; }
  done
  sleep 120
done

On a single-GPU machine this competes for the card, so it suits a rented second card or a CPU-servable small model. Where it works, the run finishes with its selection evidence already gathered rather than a queue of checkpoints to get through.

Two things to record inside the directory

Write a small run.json next to the checkpoint:

{"base_model": "meta-llama/Llama-3.1-8B-Instruct",
 "base_revision": "0e9e39f",
 "dataset_sha256": "9f2c1e…",
 "config_sha256": "71ba04…",
 "git_commit": "4f1a09",
 "seed": 0,
 "gpu": "RTX 4090", "transformers": "4.5x.y", "trl": "0.x.y"}

This is the only thing that makes a checkpoint identifiable later. Without it you have a directory of tensors and a memory. The base_revision field in particular matters: model repositories are updated in place, and “the same base model” six months later may not be the same bytes.