Skip to content
LaunchHelm
Documentation — Checkpoints and resume

Checkpoints and resume

tr.save_checkpoint() writes a bundle — model, optimizer, scheduler, scaler, step, RNG state (python, torch CPU, all CUDA devices, numpy when importable), tunables and control state — uploads it as an artifact and records it for crash resume. In DDP a barrier precedes the save, every rank captures its own RNG at the same safe point (format launchhelm.torch.checkpoint/2 carries one rng_by_rank record per rank; single-rank saves stay on checkpoint/1), and rank 0 owns the upload.

tr.load_checkpoint("latest") fetches the newest checkpoint and restores everything — model, optimizer, scheduler, scaler, step, tunables — and each rank restores its own RNG stream. Loads default to map_location="cpu" so peers do not inherit the coordinator's device. Pass a path instead of "latest" to restore a specific file. The payload is torch.load(weights_only=True)-safe — no arbitrary pickles.

The resume half of the example:

# A relaunched process resumes the run and restores the latest checkpoint.
model2, optimizer2 = build()
tr2 = lht.init("quickstart", resume=run_id)
tr2.watch(model2, optimizer2)
extra = tr2.load_checkpoint("latest")
print("restored at step:", tr2.global_step)
print("extra payload:", extra)

python-torch_checkpoint.py · verbatim from sdks/python/examples/torch_checkpoint.py

lht.init("project", resume=<run id>) reopens the run instead of creating a new one. tr.load_checkpoint("resume") restores the checkpoint someone chose for the run, else the newest: when a run stops without ending its session, its page says so and "Resume…" lists its checkpoints, with the steps each leaves to redo and which ones were saved before the last control change. Choosing one there makes it the one "resume" restores, once; "latest" always means the newest. extra is whatever you passed to save_checkpoint(extra=...). Resume requires the run to still be open — a crashed process leaves it open; tr.finish() closes it permanently. End a session deliberately while keeping the run resumable with tr.finish(close_run=False).

What resume means for metrics

A crash between checkpoints re-executes the steps since the last save — saving at step 64 and crashing at 352 replays steps 65–352 after resume. When load_checkpoint restores an earlier step than the run had reached, it records a rewind on the run: the step it continued from, the furthest step reached, and the checkpoint. The run's charts mark it ("Resumed from step 64 after reaching 352") and draw the redone steps' new values as the run's line; "Show the abandoned attempt" on the Charts tab draws the earlier values over those steps, dimmed. Queries return both values at a redone step unless they ask for duplicates: "last". Keep checkpoints frequent enough that replay is cheap.

tr.on_rewind(callback) runs on every rank after such a restore, with the Rewind (to_step, reached, cause), for state the checkpoint does not hold, such as a data position kept in extra.

Evaluate each checkpoint

On a run's Details tab, "Evaluate every checkpoint…" adds a rule to the project: each time a checkpoint finishes uploading, one evaluation of it is queued. A worker you run takes them one at a time:

@launchhelm.evaluate(project="quickstart")
def score(ctx: launchhelm.EvaluationContext) -> dict[str, float]:
    model = load(ctx.download_to("checkpoints"))
    return {"accuracy": accuracy(model), "loss": loss(model)}


if __name__ == "__main__":
    score.serve()

load, accuracy and loss are yours; the whole worker is sdks/python/examples/evaluate_checkpoints.py. Besides the file, ctx.checkpoint, ctx.step and ctx.run name what is being scored.

Each number the function returns is logged to the training run as eval/<name> at the checkpoint's step, so it charts beside the run's own metrics. An exception marks that evaluation failed with its message, and the worker goes on to the next. A worker that dies mid-evaluation leaves it to the next worker once its two-minute lease lapses, and a score already logged at that step is not logged twice. The checkpoint list shows each evaluation's state; "Evaluate this checkpoint" queues one for a checkpoint saved before the rule existed, and "Try again" retries a failed one. The worker's key is a "Checkpoint evaluation" key from Settings → API keys.

Details

  • For FSDP models, configure a full state dict before saving.
  • save_checkpoint returns the artifact id on rank 0 and None on other ranks.
  • Restore a distributed checkpoint with the same world size and rank mapping as the save.
  • The adapter captures process-global Python/Torch RNGs and NumPy's global MT19937 state. To also restore custom torch.Generator or numpy.random.Generator instances, sampler positions, or dataloader-worker state, gather them into extra before saving and apply them after loading — restoring weights and global RNG alone does not replay a shuffled or multi-worker data pipeline.
  • If any rank fails while reading, validating, or applying a checkpoint, every rank raises CheckpointError together, and the next step raises the stored error instead of training on with mixed state.
  • Non-torch runs can still run.checkpoint(path) — a plain artifact upload recorded under the checkpoint metric name.