kempnerforge.checkpoint.manager

Checkpoint manager for distributed checkpointing.

Uses PyTorch Distributed Checkpoint (DCP) for model and optimizer state, which supports automatic resharding (save with N GPUs, load with M GPUs).

Non-distributed state (scheduler, dataloader, training meta, RNG) is saved separately as a torch file written by rank 0 and broadcast from rank 0 on load. Scheduler and training meta restore from that single copy on every rank. Dataloader and RNG state are per-rank: every rank’s payload is gathered into a versioned envelope at save time, and on an exact resume (load(exact_resume=True)) each rank restores its own entry from the broadcast (see kempnerforge.checkpoint.state). Any other load (a load_path warm start, evaluation) restores rank 0’s single copy on every rank. A weights-only warm start (exclude_keys=["optimizer", "train_state"]) reads no train_state.pt, so it restores neither and checks no topology.

Resharding model and optimizer state is therefore free, but an exact resume is not: a rank’s loader cursor and RNG streams are valid only on the world size and data-parallel layout that saved them, so an exact resume under a changed topology raises on every rank. Setting checkpoint.ignore_rank_state_mismatch discards the saved loader/RNG state on every rank instead, leaving step, model, optimizer and scheduler to resume onto the new topology while data order and random streams restart fresh.

Functions

missing_train_state(ckpt_dir)

Whether checkpoint ckpt_dir holds optimizer state but no train_state.pt.

Classes

CheckpointManager	Manages save/load/cleanup of distributed checkpoints.
CheckpointRegistrar	Records checkpoints with the logging backends, which this module never imports.

class kempnerforge.checkpoint.manager.CheckpointRegistrar

Bases: Protocol

Records checkpoints with the logging backends, which this module never imports.

checkpoint_fields(step)

Fields each backend adds to its record of the checkpoint saved at step, by backend.

Called on rank 0 inside save(), so it must answer from what it holds, never a server.

Parameters:

step (int)

Return type:

dict[str, dict[str, str]]

register_checkpoint(step, directory, record)

Register the checkpoint at directory once latest points at it.

Called on rank 0 when the checkpoint is committed, so it must not block. record adds fields to a backend’s record in that checkpoint’s metadata.json, from any thread, later.

Parameters:

• step (int)

• directory (Path)

• record (Callable[[str, dict[str, str]], None])

Return type:

None

__init__(*args, **kwargs)

kempnerforge.checkpoint.manager.missing_train_state(ckpt_dir)

Whether checkpoint ckpt_dir holds optimizer state but no train_state.pt.

save() writes train_state.pt with the DCP shards (after them in a synchronous save, while an async save’s shards still flush), so such a checkpoint is a save that stopped before both were written, or one that lost the file, and load() refuses to resume it. A checkpoint without optimizer state, such as an hf-to-dcp import, needs no train_state.pt. The DCP .metadata of ckpt_dir, or of each pp{k} stage, is read only when the file is absent. One that cannot be read is logged and counted as holding no optimizer state, which leaves the load to report it.

Parameters:

ckpt_dir (Path)

Return type:

bool

class kempnerforge.checkpoint.manager.CheckpointManager

Bases: object

Manages save/load/cleanup of distributed checkpoints.

Each checkpoint is stored in a subdirectory: {dir}/step_{N}/ containing DCP shards and a non-distributed training state file.

A latest symlink always points to the most recent checkpoint.

Parameters:

• config – Checkpoint configuration.

• model – The model (FSDP-wrapped or plain).

• optimizer – The optimizer.

• dp_info – This rank’s (dp_rank, dp_size), the data partition its loader reads. Recorded with the per-rank loader/RNG state and checked on load, so a resume under a different data-parallel layout is refused even at the same world size. None records no layout and checks only the world size.

• weight_decay_scope – The run’s optimizer.weight_decay_scope. When set, every save records it in metadata.json, and a load that restores optimizer state on any rank raises ValueError on every rank, before anything is loaded, when the checkpoint was saved under a different scope. A checkpoint that holds no optimizer state, such as a converted one or one saved before the first step, loads under any scope with a fresh optimizer. A checkpoint that does not record a scope counts as "exclude_norms_and_biases".

• lr_multipliers – The run’s optimizer.lr_multipliers. When set, every save records a non-empty one in metadata.json, with the patterns of its modules and, without pipeline parallelism, the module each optimizer param group names. A load of a checkpoint saved with other multipliers or patterns raises ValueError on every rank when it restores the checkpoint’s optimizer state on any rank, before anything is loaded, or the scheduler state in its train_state.pt, before anything from that file is restored. A load that restores neither, such as a weights-only warm start (exclude_keys=["optimizer", "train_state"]), goes ahead without reading the record. A checkpoint that records none counts as saved without multipliers. One whose metadata.json cannot be read, or holds an invalid record, is refused the same way when multipliers are set; without them it loads as one saved without multipliers, with a warning, and only a record of some is refused. With multipliers set, a load that restores optimizer state into param groups laid out otherwise than at the save, a group of another module at a position included, raises on every rank before anything is loaded. A restored scheduler state carries the saved schedule on in every group, a group of a module the checkpoint holds no group of at the saved peak times its multiplier (see restore_train_state).

• module_patterns – Module name to fnmatch patterns, such as VLMConfig.module_patterns. Required for every module lr_multipliers lists.

__init__(config, model, optimizer, process_group=None, pp_rank=None, dp_info=None, weight_decay_scope=None, lr_multipliers=None, module_patterns=None)

Parameters:

• config (CheckpointConfig)

• model (torch.nn.Module)

• optimizer (torch.optim.Optimizer)

• pp_rank (int | None)

• dp_info (tuple[int, int] | None)

• weight_decay_scope (str | None)

• lr_multipliers (Mapping[str, float] | None)

• module_patterns (Mapping[str, Sequence[str]] | None)

Return type:

None

agree(error, what='Checkpoint save')

Agree on the manager’s dedicated gloo group that every rank’s what succeeded.

See agree_across_ranks. The group is not the default process group that DCP’s async-save background thread issues collectives on.

Parameters:

• error (BaseException | None)

• what (str)

Return type:

None

set_tracking_source(source)

Record the logging runs source reports in every later metadata.json.

source returns one JSON-ready record per live run, and is called on rank 0 only, inside save(), so it must return promptly. The tracking key is omitted when it reports none, raises, or returns what JSON cannot write. None detaches it.

Parameters:

source (Callable[[], list[dict[str, str]]] | None)

Return type:

None

set_checkpoint_registrar(registrar)

Record every later checkpoint with the logging backends through registrar.

Both of its methods run on rank 0 only. checkpoint_fields is asked inside save() when a backend has a live run, and its fields join that save’s tracking records. register_checkpoint gets each checkpoint once it is committed, with a record(backend, fields) that later adds fields to that backend’s record, atomically and from any thread. A registrar that raises, or answers what a record cannot hold, costs only those fields or that registration, never the save or the commit. None detaches it.

Parameters:

registrar (CheckpointRegistrar | None)

Return type:

None

save(step, tokens_seen=0, scheduler=None, dataloader=None, extra=None)

Save a checkpoint at the given step.

Under expert parallelism each EP rank holds a copy of every weight outside the routed experts, and the copies can drift apart on GPU (#482). DCP writes each chunk that several ranks hold from one of them, by default the one with the least to write so far, which could join chunks of different copies into one tensor. An EP save writes each such chunk from its lowest rank instead, so the checkpoint holds one EP rank’s copy whole.

Parameters:

• step (int) – Current training step.

• tokens_seen (int) – Total tokens processed.

• scheduler (Any | None) – LR scheduler to save.

• dataloader (Any | None) – Stateful dataloader to save.

• extra (dict | None) – Additional metadata.

Return type:

None

wait()

Block until any pending async checkpoint save completes.

Once the flush is durable, commit its deferred latest symlink + cleanup. The training loop calls this after the loop exits, so the final checkpoint’s latest is committed before process teardown. A flush that failed on any rank raises on every rank, and nothing is committed.

Return type:

None

flush_pending_save()

Drain any in-flight async save before mutating model state.

Called from the FreezeStage transition hook in the training loop: when a transition fires at step S, any save started at step S-1 must have written metadata.json with the pre-transition spec before the transition flips requires_grad. Otherwise metadata.json lands with the post-transition spec attached to the pre-transition shards.

Also commits the deferred latest symlink for that save, so a transition (or any caller draining the queue) leaves latest pointed at the now-durable checkpoint. A flush that failed on any rank raises on every rank, and nothing is committed.

Return type:

None

peek_saved_step(path=None)

Read step from a candidate checkpoint’s metadata.json.

Returns None if no checkpoint resolves or the metadata is missing/unreadable. Used by the training loop on resume to compute the expected freeze list (which depends on saved_step) before calling load.

Parameters:

path (str | None)

Return type:

int | None

load(path=None, scheduler=None, dataloader=None, exclude_keys=None, vlm_freeze_expected=None, exact_resume=False)

Load a checkpoint and restore all state.

Parameters:

• path (str | None) – Checkpoint path. If None, loads from config.load_path or the latest symlink.

• scheduler (Any | None) – LR scheduler to restore.

• dataloader (Any | None) – Stateful dataloader to restore.

• exclude_keys (list[str] | None) – State to skip: the DCP keys "model" and "optimizer" (e.g. ["optimizer"] for fine-tuning), or ["optimizer", "train_state"] for a weights-only warm start that loads the model alone and ignores train_state.pt (step, token count, scheduler, dataloader position, RNG) and the saved freeze metadata, so the new run starts at step 0 with its own schedule, data order and freeze. "train_state" in any other combination raises ValueError.

• vlm_freeze_expected (list[dict[str, Any]] | None) – Canonical freeze metadata (output of canonical_freeze_meta) for the current run’s VLMConfig. When both the saved metadata and this argument are set, a mismatch raises ValueError unless the checkpoint config has ignore_freeze_mismatch=True, in which case the load proceeds with a warning.

• exact_resume (bool) – True only when the run resumes from its own checkpoint (auto-resume). Each rank then restores its own dataloader/RNG state from the per-rank envelope, which must match this run’s world size and data-parallel layout. Every rank must pass the same value. False (a load_path warm start, evaluation, any other load) restores the single-copy rng/dataloader keys (rank 0’s) on every rank with no topology check. The weights-only exclude_keys pair reads no train state: no rank restores loader/RNG state or enters the train-state broadcasts or the layout gather, and with exact_resume it raises.

Returns:

Tuple of (step, tokens_seen, extra) where extra contains any additional keys saved via build_train_state(extra=...).

Raises:

• FileNotFoundError – No checkpoint resolves although one was named, or the checkpoint holds optimizer state that this load restores but has no train_state.pt. Raised before any state is loaded.

• ValueError – exclude_keys is not a list of distinct DCP keys, names "train_state" outside ["optimizer", "train_state"], or is that pair with exact_resume, raised before the path is resolved; the checkpoint was saved under another weight-decay scope, LR multipliers or patterns of their modules (see __init__); or, on an exact resume, its per-rank state does not fit this run’s topology and checkpoint.ignore_rank_state_mismatch is off (see exact_resume).

• RuntimeError – The saved optimizer state does not fit this run’s optimizers.

Return type:

tuple[int, int, dict[str, Any]]

A load that restores optimizer state refuses it on every rank together, before any state is loaded: each rank checks its own part of the checkpoint, and one all-reduce on the manager’s gloo group, which every rank enters, has the other ranks raise a RuntimeError naming the ranks that refused.

apply_dataloader_state(dataloader)

Apply any dataloader state stashed during load().

Training loops call load() before constructing the dataloader (since the dataloader depends on phase/annealing state that load() restores). This method applies the stashed state once the loader exists.

No-op if no state is pending, or if the loader does not support load_state_dict (e.g., plain torch DataLoader for HF streaming).

Parameters:

dataloader (Any)

Return type:

None