kempnerforge.checkpoint

Distributed checkpointing for KempnerForge.

Public API:

• CheckpointManager: Save/load/cleanup distributed checkpoints

• AsyncCheckpointer: Non-blocking checkpoint saves

• build_train_state / restore_train_state: State assembly

class kempnerforge.checkpoint.AsyncCheckpointer

Bases: object

Non-blocking checkpoint saver.

Wraps dcp.async_save() and manages the background save future. Each new save waits for the previous async save to complete first.

Parameters:

mode – Checkpoint mode (disabled/async/async_with_pinned_mem).

__init__(mode=AsyncCheckpointMode.disabled)

Parameters:

mode (AsyncCheckpointMode)

Return type:

None

save(state_dict, checkpoint_id, process_group=None, planner=None)

Save distributed state, potentially asynchronously.

Parameters:

• state_dict (dict) – DCP-compatible state dict (model + optimizer).

• checkpoint_id (str) – Checkpoint directory path.

• process_group – Process group for DCP. Required for PP where each stage has a different state dict — pass a group scoped to ranks within the same PP stage. None uses the default global group.

• planner (SavePlanner | None) – DCP save planner. None calls DCP as before, with its default planner.

Return type:

None

wait()

Block until any pending async save completes.

Return type:

None

property is_pending: bool

Check if an async save is still in progress.

class kempnerforge.checkpoint.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

kempnerforge.checkpoint.build_train_state(step, tokens_seen, scheduler=None, dataloader=None, extra=None, per_rank=None)

Build the non-distributed portion of the training state.

Model and optimizer state are handled by DCP directly. This function captures everything else needed for exact resumption.

Parameters:

• step (int) – Current training step.

• tokens_seen (int) – Total tokens processed so far.

• scheduler (Any | None) – LR scheduler (must have state_dict()). When a param group of its optimizer names a module, as build_optimizer names those of optimizer.lr_multipliers, the module of every group is recorded under scheduler_group_modules.

• dataloader (Any | None) – Stateful dataloader (must have state_dict()); writes the legacy single-copy dataloader key (the caller’s own loader, rank 0 in distributed runs).

• extra (dict | None) – Additional metadata to include.

• per_rank (dict[str, Any] | None) – Versioned per-rank envelope (build_per_rank_envelope) holding every rank’s own loader/RNG payload.

Returns:

Dict with training state, scheduler state, dataloader state, and RNG states.

Return type:

dict[str, Any]

kempnerforge.checkpoint.get_rng_state()

Capture all RNG states for reproducibility on resume.

Return type:

dict[str, Any]

kempnerforge.checkpoint.restore_train_state(state, scheduler=None, dataloader=None, lr_multipliers=None)

Restore the non-distributed portion of the training state.

Parameters:

• state (dict[str, Any]) – Training state dict (from build_train_state).

• scheduler (Any | None) – LR scheduler to restore. The saved schedule carries on in every param group, peak included, whatever the rates the optimizer was built with. Its state is laid out by module for groups other than those it was saved with (see build_train_state): each group of a module the checkpoint holds resumes at that module’s rates, and any other at the saved peak times its multiplier. Nothing is restored when that fails. When a group names a module, the peak rate of each module’s groups is logged, with a warning when one differs from the rate the optimizer was built with, as when optimizer.lr changed since the save.

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

• lr_multipliers (Mapping[str, float] | None) – The optimizer.lr_multipliers the scheduler’s groups were built with, which must be those of the save. Needed only for a group of a module the checkpoint holds no group of.

Return type:

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

Restores whatever rng/dataloader keys state carries. On an exact resume the manager hands each rank its own view of the broadcast train state (its per-rank envelope entry overlaid on these keys) before calling this; any other load passes the broadcast state as is, so rank 0’s single copy restores on every rank.

Returns:

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

Raises:

ValueError – The scheduler state cannot be laid out for the scheduler’s param groups.

Parameters:

• state (dict[str, Any])

• scheduler (Any | None)

• dataloader (Any | None)

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

Return type:

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

kempnerforge.checkpoint.set_rng_state(state)

Restore all RNG states from a checkpoint.

Parameters:

state (dict[str, Any])

Return type:

None

Modules

async_save	Async checkpointing for non-blocking saves.
expert_layout	Expert-parallel experts under their global identity in a DCP state dict.
manager	Checkpoint manager for distributed checkpointing.
state	Training state assembly for checkpointing.