kempnerforge.checkpoint.state¶

Training state assembly for checkpointing.

Collects the full training state — model, optimizer, scheduler, dataloader, training metadata, and RNG states — into a single dict for DCP save/load.

RNG state capture ensures exact reproducibility on resume.

Dataloader and RNG state are per-rank. train_state.pt carries every rank’s own payload under PER_RANK_KEY as a versioned envelope:

{"version": 1, "world_size": W, "data_parallel": layout | None,
 "ranks": {global_rank: payload}}

A payload holds that rank’s own "rng" capture and, when its loader is stateful, its "dataloader" state (every rank or none). Every rank must capture the same RNG streams: "torch_cuda" is captured only where CUDA is available, and an entry missing a stream would silently resume that rank with a fresh one. layout maps every global rank to the (dp_rank, dp_size) data partition its loader read, or is None when the saving manager was not given it. The legacy single-copy "rng"/"dataloader" keys (always rank 0’s) are still written beside the envelope, so readers that predate it keep loading, with their documented limitation that they restore rank 0’s copy everywhere.

Functions

`build_per_rank_envelope`(entries, world_size)	Assemble the versioned per-rank envelope from every rank's payload.
`build_train_state`(step, tokens_seen[, ...])	Build the non-distributed portion of the training state.
`capture_rank_state`([dataloader])	Capture this rank's own resume payload: RNG streams and loader cursor.
`dp_coordinates`(value)	`value` as a `(dp_rank, dp_size)` tuple, or None if it is not one.
`get_rng_state`()	Capture all RNG states for reproducibility on resume.
`restore_train_state`(state[, scheduler, ...])	Restore the non-distributed portion of the training state.
`select_rank_state`(envelope, rank, world_size)	Return `rank`'s own payload from a per-rank envelope, or refuse.
`set_rng_state`(state)	Restore all RNG states from a checkpoint.

kempnerforge.checkpoint.state.get_rng_state()[source]¶

Capture all RNG states for reproducibility on resume.

Return type:: dict[str, Any]

kempnerforge.checkpoint.state.set_rng_state(state)[source]¶

Restore all RNG states from a checkpoint.

Parameters:: state (dict[str, Any])
Return type:: None

kempnerforge.checkpoint.state.capture_rank_state(dataloader=None)[source]¶

Capture this rank’s own resume payload: RNG streams and loader cursor.

Parameters:: dataloader (Any | None)
Return type:: dict[str, Any]

kempnerforge.checkpoint.state.dp_coordinates(value)[source]¶

value as a (dp_rank, dp_size) tuple, or None if it is not one.

Valid: a pair of ints with 0 <= dp_rank < dp_size.

Parameters:: value (Any)
Return type:: tuple[int, int] | None

kempnerforge.checkpoint.state.build_per_rank_envelope(entries, world_size, data_parallel=None)[source]¶

Assemble the versioned per-rank envelope from every rank’s payload.

Requires exactly one payload per global rank 0..world_size - 1: a gather that lost a rank would otherwise save an envelope the load side could never restore completely. The ValueError for an incomplete set names world_size, the ranks the gather returned and both ways out, so a caller driving a manager as another rank can fix the world it simulated. data_parallel is the complete global rank -> (dp_rank, dp_size) layout, or None when unknown.

Parameters:

entries (dict[int, dict[str, Any]])
world_size (int)
data_parallel (dict[int, Any] | None)

Return type:

dict[str, Any]

kempnerforge.checkpoint.state.select_rank_state(envelope, rank, world_size, data_parallel=None)[source]¶

Return rank’s own payload from a per-rank envelope, or refuse.

Refuses with ValueError rather than guessing when the envelope is malformed or incomplete, written by a newer schema version, or saved under a different world size or data-parallel layout: a loader/RNG payload restored onto the wrong rank or data partition silently changes the data order and random streams of the resumed run. The whole envelope is checked before rank’s entry is selected, so every rank reaches the same verdict.

Parameters:

envelope (Any) – The broadcast per_rank envelope.
rank (int) – This process’s global rank.
world_size (int) – This run’s world size.
data_parallel (dict[int, tuple[int, int]] | None) – This run’s complete global rank -> (dp_rank, dp_size) layout, or None when unknown (the layout check is then skipped).

Return type:

dict[str, Any]

kempnerforge.checkpoint.state.build_train_state(step, tokens_seen, scheduler=None, dataloader=None, extra=None, per_rank=None)[source]¶

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.state.restore_train_state(state, scheduler=None, dataloader=None, lr_multipliers=None)[source]¶

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]]