Checkpointing
A training checkpoint saves model weights, optimizer state, the step counter, the data position and the RNG state so a run can resume exactly where it died. It is not the model artifact, and a resume that does not restore all of it silently trains a different run.
The problem, the obvious approach, and why it breaks
Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.
A long run will die before it finishes — what must be saved, how often, and what does resuming have to reproduce for the second half to be the same run as the first?
A recommendation team's embedding model trains for four days on preemptible GPUs because on-demand ones cost three times as much. Preemptions arrive about daily. After the third restart from scratch, the team lead asks why "just saving the model" every hour is not enough — they tried, and the resumed run behaved differently.
Save the model weights to object storage every hour. On restart, load the latest weights, and continue training from there.
The resumed run starts the learning-rate schedule from step zero: a full warmup at a point where the model wanted the low end of a cosine decay. The loss spikes and the model spends the next hour recovering ground it had already covered.
- The resumed run starts the learning-rate schedule from step zero: a full warmup at a point where the model wanted the low end of a cosine decay. The loss spikes and the model spends the next hour recovering ground it had already covered.
- Adam's moments start from zero on resume. For the first few hundred steps the optimizer's bias correction treats the run as brand new and takes large, poorly scaled steps; the effective trajectory is that of a run with an optimizer reset in the middle.
- The data loader restarts its shuffle from the epoch's beginning, so the examples between the checkpoint and the preemption are skipped and the examples before the checkpoint are seen twice in this epoch. The epoch's sample is no longer uniform.
- The checkpoint was written by worker 0 while worker 3 was one step ahead. The restored weights and the restored step count disagree by one, and the sharded optimizer state, saved per worker, is from two different steps. It loads. It is not any state the run was ever in.
What is being predicted, and from what data
This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.
- The surrounding model learns user and item embeddings from interaction logs (Embedding Training). This lesson's target is a resumed run whose second half is the run the first half would have continued into — same schedule, same data order, same optimizer momentum.
- That is a stronger requirement than "the weights are restored". The weights are one of five things the run's state consists of.
- The run processes a sharded interaction log in a shuffled order determined by a seeded RNG at the start of each epoch; the position within the epoch is a step counter.
- The optimizer is Adam, holding two moment vectors per parameter — together twice the size of the weights — that encode a running estimate of the gradient's mean and variance.
- The learning rate follows a warmup-then-cosine schedule keyed on the global step, and the data loader's workers hold their own RNG state for augmentation.
How it actually works
Precisely enough to predict its behaviour — not a framework API.
- A training run's state is: the parameters, the optimizer state (moments, step-dependent bias-correction counters), the global step (which drives the schedule), the data-loader position (epoch, shuffle seed, offset within the epoch) and the RNG states (model-side for dropout, loader-side for augmentation). Resuming reproduces the run only if all five are restored together, from the same instant.
- A checkpoint is therefore a *consistent snapshot* of that state, taken at a step boundary when every worker holds the same parameters and the same step count. Under data parallelism the parameters are identical on every worker so one copy suffices; under sharded optimizer state or model parallelism each worker holds a different piece, and the snapshot must gather or index all of them from the same step. This is the Distributed Systems problem of a consistent snapshot, appearing in a training loop (Recovered State Is a Checkpoint Plus the Log After It on the data side is the same idea for a stream).
- The write has to be atomic from the reader's point of view — write to a temporary key, then rename or publish a manifest — because a preemption during the write leaves a half-written file that a naive loader will try to resume from.
- Frequency is a trade between the cost of a checkpoint (the time the run pauses to serialise, the storage, the egress) and the expected loss on failure (the interval times the failure rate). With a daily preemption rate and a checkpoint that takes a minute, hourly checkpoints lose at most an hour per preemption and cost twenty-four minutes a day; every ten minutes would lose six minutes per preemption and cost two and a half hours a day.
Five pieces of state, one instant
Weights are what a checkpoint is usually imagined to be, and they are the piece least likely to be forgotten. The optimizer state — Adam's two moment vectors and its step-dependent bias correction — is twice the size and decides how large the next steps will be. The global step drives the learning-rate schedule. The loader position decides which examples are next. The RNG states decide dropout masks and augmentations.
A resume that restores four of the five produces a run that continues from the right weights under the wrong conditions. The most common omissions — optimizer moments and loader position — both look like a loss spike and a slightly different final model, which is to say they look like nothing at all.
1def save_checkpoint(path, step, model, optimizer, loader, rngs):2 tmp = f"{path}/step-{step}.tmp"3 write(f"{tmp}/model.bin", model.state())4 write(f"{tmp}/optimizer.bin", optimizer.state()) # moments + counters5 write(f"{tmp}/loader.json", {6 "epoch": loader.epoch, "seed": loader.seed, "offset": loader.offset,7 })8 write(f"{tmp}/rng.bin", {"model": rngs.model.state(), "loader": rngs.loader.state()})9 write(f"{tmp}/manifest.json", {"step": step, "files": [...]})10 rename(tmp, f"{path}/step-{step}") # the rename is the commit11 12def resume(path, model, optimizer, loader, rngs):13 ckpt = latest_complete(path) # only directories with a manifest14 model.load(ckpt.model); optimizer.load(ckpt.optimizer)15 loader.seek(ckpt.epoch, ckpt.seed, ckpt.offset) # not restart the epoch16 rngs.load(ckpt.rng)17 return ckpt.step # schedule continues from hereTwo lines carry the design: the rename that makes the write atomic, and loader.seek, which is the line most save helpers do not have. A loader that cannot seek to an offset within a seeded shuffle cannot resume an epoch.
How often, and what it costs
Each checkpoint costs a pause to serialise, storage for the state, and egress if the storage is remote. Each failure costs the work since the last checkpoint. With a failure rate of once a day and a checkpoint that pauses the run for a minute, the expected loss at hourly intervals is half an hour per day and the cost is twenty-four minutes; at ten-minute intervals the loss is five minutes and the cost two and a half hours. The interval that minimises the sum sits between, and it moves with the failure rate.
Asynchronous checkpointing changes the cost side: copy the state to host memory in seconds, resume compute, and upload in the background. That makes frequent checkpoints cheap in run time and leaves storage and egress as the remaining cost, which is why very long runs on preemptible hardware checkpoint far more often than a synchronous scheme would justify.
The most recent checkpoint is complete, consistent and restorable, and its storage outlives the compute that wrote it.
holds when Writes are atomic via rename or manifest; every worker's shard is from the same step; several recent checkpoints are retained; the storage is remote from the preemptible node; resumption is exercised on a dry run before each long run.
breaks when A preemption lands during the write; a worker writes a shard from a different step; the local disk is reclaimed with the node; the layout changes and the shards no longer match the worker count.
respond Fall back to the previous complete checkpoint, and fix the writer before the next long run.
| Interval | Expected loss per failure | Checkpoint cost per day (1-min pause) | Total at 1 failure/day |
|---|---|---|---|
| 6 hours | ~3 hours | 4 min | ~3 h |
| 1 hour | ~30 min | 24 min | ~54 min |
| 20 min | ~10 min | 72 min | ~82 min |
| 10 min | ~5 min | 144 min | ~149 min |
A checkpoint is not an artifact
A checkpoint is training state: weights in training precision, optimizer moments, a step counter, a loader position. Its consumer is the training loop. An artifact is the exported inference function: weights in serving precision, the preprocessing, the signature, the metadata that the registry and the serving system need (What a Model Artifact Contains). Its consumer is production.
The two are produced from the same weights and should be produced by different steps. Exporting an artifact from a checkpoint is a deliberate transformation — strip the optimizer state, cast, attach preprocessing, sign — not a copy. A team that ships checkpoints ships three times the bytes, in the wrong precision, without the contract the serving path needs, and calls the resulting confusion "the model".
Training-precision weights with optimizer moments, no preprocessing, no signature, no version metadata; the serving path guesses at the rest.
Optimizer state stripped, weights cast for inference, preprocessing bundled, signature and lineage attached, registered ([[model-registry]]).
The two objects have different consumers with different requirements. The checkpoint has to resume a run; the artifact has to serve a prediction under a contract. Conflating them makes both jobs harder.
How to build it
Most important first.
- Save all five pieces of state, keyed by the global step, in one manifest: parameters, optimizer state, step, data-loader position (epoch, seed, offset) and every RNG state. Make the loader restore the loader position and skip to the offset rather than restarting the epoch.
- Take the snapshot at a step boundary after a synchronisation, so every worker's piece is from the same step; under sharded state, have every worker write its shard and one worker write a manifest that names all of them — the manifest's presence is the commit.
- Write to a temporary location and publish atomically. Keep the last few checkpoints, not only the latest, so a corrupted or diverged one can be skipped (Atomic Publish on the data side).
- Choose the interval from the failure rate and the checkpoint cost, and write it down. Checkpoint asynchronously — copy state to host memory, resume compute, upload in the background — so the pause is the copy, not the upload.
- Test resumption before the long run: train to step k, checkpoint, kill, resume, train to step k+m, and compare against an uninterrupted run to step k+m. The loss curves should agree to reduction-order tolerance (Gradient Synchronisation).
What to measure
Which number actually maps to the decision — and which numbers look relevant and are not.
- The loss curve after a resume compared with the curve before it and with an uninterrupted reference. A spike or a step change at the resume point is a missing piece of state.
- Checkpoint cost as a fraction of run time — the pause plus the storage and egress bill — against the expected loss per failure at the chosen interval. The interval should minimise the sum.
- Time from preemption to first resumed step, which is the part of the recovery that the checkpoint format and the loader control.
- Do not measure "checkpoints written" as success. A checkpoint that was never resumed from has not been tested.
What must stay true after deployment
The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.
- Every checkpoint is a consistent snapshot: parameters, optimizer state, step, loader position and RNG states all from the same step boundary, and the manifest that names them is the only thing a loader trusts.
- The resumed run reproduces the schedule and the data order the uninterrupted run would have followed — the learning rate at step k+1 and the batch at step k+1 are the same on both.
- The storage the checkpoints live on outlives the compute that wrote them, and the latest checkpoint is readable — the write completed and was not corrupted by the preemption that followed it.
- Offline: the kill-and-resume test on a short run, compared against an uninterrupted run at the same final step, with the loss curves and the consumed example ids diffed.
- Before every long run: resume from the most recent real checkpoint in a dry run and confirm the step, the learning rate and the next batch's example ids match the manifest.
- Over time: an alert on any resume whose loss in the first hundred steps departs from the pre-preemption trend by more than the measured seed variance.
What can go wrong
- The checkpoint restores everything except the loader offset, and the epoch's sample is silently non-uniform on every resume — a bias that compounds with the number of preemptions.
- A checkpoint from a run that had already begun to diverge is the latest one, and resumption faithfully continues the divergence. Keeping several and checking the loss at save time is the remedy.
- The checkpoint format is tied to the parallel layout. Resuming a sixteen-worker sharded run on eight workers requires re-sharding the optimizer state, and the tool for that was never written.
- The checkpoint is treated as the deployable model. It carries optimizer state, is in training precision, lacks the preprocessing and the signature, and is three times the size of the artifact it should have been exported to (What a Model Artifact Contains).
- Full-state checkpoints are large — parameters plus twice that again in optimizer moments — and frequent ones cost storage, egress and a pause per save.
- Consistent snapshots under model or sharded parallelism need a coordination step and a manifest protocol; the simplicity of "save the weights" is exactly what makes it wrong.
- Tying the checkpoint to the parallel layout makes resumption fast and makes changing the layout mid-run an engineering project.
- "We save the model every hour, so we can resume." The model is a fifth of the state. Resuming without the optimizer moments, the step and the loader position starts a new run from old weights, and it behaves like one.
- "The checkpoint is the artifact — ship it." A checkpoint is training state in training precision with optimizer moments attached. An artifact is the exported inference function with its preprocessing and signature (What a Model Artifact Contains). They are different objects with different consumers.
- "Checkpoint as often as possible." Every checkpoint pauses the run and costs storage and egress; the interval that minimises total cost depends on the failure rate, and for a reliable cluster it is long.
Where this applies
ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.
- GENERALThat a resumed run needs the optimizer state, the step, the data position and the RNG state as well as the weights holds for every gradient-based training loop, whatever the framework; frameworks differ only in how much of it their default save helper includes.
- SIMPLIFIEDThe interval arithmetic in the section below assumes a constant failure rate and a fixed checkpoint cost for the shape of the argument; real preemption rates vary by time of day and instance type, and asynchronous checkpointing changes the cost side substantially.
- SCALE-SPECIFICFor a run that finishes in an hour on one reliable machine, checkpointing is a convenience and restart-from-scratch is an acceptable recovery; it becomes a requirement when the run is long enough, or the hardware unreliable enough, that a failure is expected before the end.
Where the depth lives
This domain teaches the model and hands the rest off by name.
- — Testing & Reliability Engineering — the kill-and-resume test is a fault-injection test, and the discipline of running restore drills on schedule rather than trusting a write is a reliability practice this lesson assumes rather than teaches.