Checkpointing

PhiJAX stores complete functional training state for exact resumption and also supports model-only restoration for transfer learning.

ModelCheckpoint callback

ModelCheckpoint delegates storage to CheckpointIO; the default backend uses Orbax. A Trainer accepts only one checkpoint callback. The callback opens its backend when fitting starts and closes it during teardown. Custom backends must support repeated open() and close() calls and must be able to reopen for later tasks.

Set monitor, mode, and save_top_k to retain the best checkpoints by an available scalar metric. save_last=True keeps the terminal state independently. Public best_model_path, best_model_score, and last_model_path attributes expose the resolved results. Checkpoints also persist callback state, including ranking and learning-rate bookkeeping.

ModelCheckpoint [source]

ModelCheckpoint(
    checkpoint_io: CheckpointIO,
    *,
    every_n_steps: int | None = None,
    save_last: bool = True,
    save_on_exception: bool = False,
    monitor: str | None = None,
    mode: MonitorMode = "min",
    save_top_k: int | None = None,
)

Bases: Callback

Save functional training states at a configured interval and at fit completion.

The callback owns checkpoint scheduling policy. Its checkpoint_io backend owns serialization and retention, so alternative storage implementations can be supplied without changing trainer orchestration.

Attributes:

Name	Type	Description
`best_model_path`	`Path \| None`	Path of the best retained checkpoint, when available.
`best_model_score`	`float \| None`	Best retained monitored value, when available.
`last_model_path`	`Path \| None`	Path of the terminal checkpoint, when available.

Initialize checkpoint lifecycle policy.

Parameters:

Name	Type	Description	Default
`checkpoint_io`	`CheckpointIO`	Storage backend used for checkpoint operations.	required
`every_n_steps`	`int \| None`	Positive periodic save interval, or `None` to disable periodic saves.	`None`
`save_last`	`bool`	Whether to save the final state passed to :meth:`on_fit_end`.	`True`
`save_on_exception`	`bool`	Whether to save the last valid state from :meth:`on_exception`.	`False`
`monitor`	`str \| None`	Optional scalar metric used to rank checkpoints.	`None`
`mode`	`MonitorMode`	`"min"` or `"max"` monitored-metric optimization direction.	`'min'`
`save_top_k`	`int \| None`	Retained best checkpoint count. `None` selects one with monitoring and all without it.	`None`

Raises:

Type	Description
`ValueError`	If checkpoint scheduling or monitoring options are invalid.

connect [source]

connect(trainer: Any) -> None

Connect checkpoint persistence to Trainer callback state.

Parameters:

Name	Type	Description	Default
`trainer`	`Any`	Owning Trainer.	required

setup [source]

setup() -> None

Reset fit-local checkpoint scheduling state.

on_fit_start [source]

on_fit_start(context: TrainerContext) -> None

Open checkpoint resources for the fit stage.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Initial trainer context.	required

on_train_batch_end [source]

on_train_batch_end(context: TrainerContext) -> bool

Save the state when the completed optimizer step matches the configured interval.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Post-update trainer state, step, and metrics.	required

Returns:

Type	Description
`bool`	Always `False`; checkpointing never requests early termination.

on_train_metrics [source]

on_train_metrics(context: TrainerContext) -> None

Evaluate periodic checkpoint policy after all metrics are available.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Complete post-update metric context.	required

on_fit_end [source]

on_fit_end(context: TrainerContext) -> None

Save the terminal state when save_last is enabled.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Final trainer state, step, and metrics.	required

on_exception [source]

on_exception(
    exception: BaseException, context: CallbackContext
) -> None

Optionally save the most recent valid state during exceptional shutdown.

Parameters:

Name	Type	Description	Default
`exception`	`BaseException`	Exception or interruption terminating the fit call.	required
`context`	`CallbackContext`	Most recent internally consistent trainer context.	required

teardown [source]

teardown() -> None

Wait for asynchronous writes and release stage resources.

close [source]

close() -> None

Wait for pending writes and release checkpoint backend resources idempotently.

state_dict [source]

state_dict() -> Mapping[str, Any]

Return checkpoint ranking and path state.

Returns:

Type	Description
`Mapping[str, Any]`	JSON-compatible persistent callback state.

load_state_dict [source]

load_state_dict(state: Mapping[str, Any]) -> None

Restore checkpoint ranking and path state.

Parameters:

Name	Type	Description	Default
`state`	`Mapping[str, Any]`	Persistent state returned by :meth:`state_dict`.	required

Raises:

Type	Description
`ValueError`	If the saved state is malformed.

Checkpoint IO protocol

CheckpointIO [source]

Bases: Protocol

Describe the storage operations required by :class:ModelCheckpoint.

Backends must allow open() after close(). Both methods must be idempotent so each Trainer task can own its resources independently.

latest_step `property`

latest_step: int | None

Return the latest committed checkpoint step.

Returns:

Type	Description
`int \| None`	Latest checkpoint step, or `None` when storage is empty.

steps `property`

steps: tuple[int, ...]

Return committed checkpoint steps in ascending order.

Returns:

Type	Description
`tuple[int, ...]`	Ordered checkpoint steps.

open [source]

open() -> None

Prepare checkpoint backend resources for a new fit stage idempotently.

save [source]

save(
    state: Any,
    step: int,
    metrics: Mapping[str, float] | None = None,
    *,
    force: bool = False,
    callback_states: Mapping[str, Mapping[str, Any]]
    | None = None,
) -> bool

Submit a checkpoint to storage.

Parameters:

Name	Type	Description	Default
`state`	`Any`	Functional state to persist.	required
`step`	`int`	Unique checkpoint step.	required
`metrics`	`Mapping[str, float] \| None`	Optional host scalar metrics.	`None`
`force`	`bool`	Whether to bypass backend scheduling policies.	`False`
`callback_states`	`Mapping[str, Mapping[str, Any]] \| None`	Optional JSON-compatible callback states.	`None`

Returns:

Type	Description
`bool`	Whether the backend initiated a save.

checkpoint_path [source]

checkpoint_path(step: int) -> Path | None

Return a user-facing path for one checkpoint step.

Parameters:

Name	Type	Description	Default
`step`	`int`	Committed checkpoint step.	required

Returns:

Type	Description
`Path \| None`	Checkpoint path, or `None` when the backend has no filesystem path.

delete [source]

delete(step: int) -> None

Delete one committed checkpoint.

Parameters:

Name	Type	Description	Default
`step`	`int`	Checkpoint step to remove.	required

restore [source]

restore(target: Any, step: int | None = None) -> Any

Restore a complete state.

Parameters:

Name	Type	Description	Default
`target`	`Any`	State defining the expected restore structure.	required
`step`	`int \| None`	Checkpoint step, or `None` for the latest checkpoint.	`None`

Returns:

Type	Description
`Any`	Restored complete state.

restore_weights [source]

restore_weights(
    target: Any, step: int | None = None
) -> Any

Restore model weights into a fresh state.

Parameters:

Name	Type	Description	Default
`target`	`Any`	Fresh state whose non-model fields are preserved.	required
`step`	`int \| None`	Checkpoint step, or `None` for the latest checkpoint.	`None`

Returns:

Type	Description
`Any`	State containing restored model weights.

restore_callback_states [source]

restore_callback_states(
    step: int | None = None,
) -> dict[str, Mapping[str, Any]]

Restore callback state metadata.

Parameters:

Name	Type	Description	Default
`step`	`int \| None`	Checkpoint step, or `None` for the latest checkpoint.	`None`

Returns:

Type	Description
`dict[str, Mapping[str, Any]]`	Stable callback identifier mapping.

wait_until_finished [source]

wait_until_finished() -> None

Wait until pending checkpoint writes have committed.

close [source]

close() -> None

Wait for pending writes and release checkpoint backend resources idempotently.

Orbax implementation

OrbaxCheckpointIO [source]

OrbaxCheckpointIO(
    directory: str | Path,
    *,
    max_to_keep: int | None = 3,
    enable_async_checkpointing: bool = True,
)

Manage lazily opened asynchronous full-state checkpoints with Orbax.

Initialize an Orbax checkpoint manager.

Parameters:

Name	Type	Description	Default
`directory`	`str \| Path`	Checkpoint-root directory.	required
`max_to_keep`	`int \| None`	Maximum recent checkpoints retained, or `None` to retain all.	`3`
`enable_async_checkpointing`	`bool`	Whether saves may finish in a background thread.	`True`

Raises:

Type	Description
`ValueError`	If `max_to_keep` is not positive or `None`.

latest_step `property`

latest_step: int | None

Return the latest committed checkpoint step.

Returns:

Type	Description
`int \| None`	Latest step, or `None` when no checkpoint exists.

steps `property`

steps: tuple[int, ...]

Return committed checkpoint steps in ascending order.

Returns:

Type	Description
`tuple[int, ...]`	Ordered committed steps.

open [source]

open() -> None

Open the Orbax manager when it is not already active.

save [source]

save(
    state: TrainState,
    step: int,
    metrics: Mapping[str, float] | None = None,
    *,
    force: bool = False,
    callback_states: Mapping[str, Mapping[str, Any]]
    | None = None,
) -> bool

Submit a complete training state for checkpointing.

Parameters:

Name	Type	Description	Default
`state`	`TrainState`	Full functional training state.	required
`step`	`int`	Unique checkpoint step.	required
`metrics`	`Mapping[str, float] \| None`	Optional host scalar metrics used by Orbax retention policies.	`None`
`force`	`bool`	Whether to save independently of manager policies.	`False`
`callback_states`	`Mapping[str, Mapping[str, Any]] \| None`	Optional callback states restored before fit-start hooks.	`None`

Returns:

Type	Description
`bool`	Whether Orbax initiated a save.

checkpoint_path [source]

checkpoint_path(step: int) -> Path

Return the filesystem path associated with a checkpoint step.

Parameters:

Name	Type	Description	Default
`step`	`int`	Committed checkpoint step.	required

Returns:

Type	Description
`Path`	Expected Orbax step directory.

delete [source]

delete(step: int) -> None

Delete one committed checkpoint and wait for completion.

Parameters:

Name	Type	Description	Default
`step`	`int`	Checkpoint step to remove.	required

restore [source]

restore(
    target: TrainState, step: int | None = None
) -> TrainState

Restore a complete state for exact training resumption.

Parameters:

Name	Type	Description	Default
`target`	`TrainState`	Abstract or concrete state defining the expected structure and sharding.	required
`step`	`int \| None`	Checkpoint step, or `None` for the latest committed step.	`None`

Returns:

Type	Description
`TrainState`	Restored model, optimizer, balancer, PRNG, step, and precision state.

Raises:

Type	Description
`FileNotFoundError`	If no requested checkpoint exists.

restore_weights [source]

restore_weights(
    target: TrainState, step: int | None = None
) -> TrainState

Restore model weights while preserving fresh optimization and run state.

Parameters:

Name	Type	Description	Default
`target`	`TrainState`	Fresh training state defining both restore structure and preserved non-model fields.	required
`step`	`int \| None`	Checkpoint step, or `None` for the latest committed step.	`None`

Returns:

Type	Description
`TrainState`	`target` with only `model_state` replaced from the checkpoint.

restore_callback_states [source]

restore_callback_states(
    step: int | None = None,
) -> dict[str, Mapping[str, Any]]

Restore JSON-compatible callback states from checkpoint metadata.

Parameters:

Name	Type	Description	Default
`step`	`int \| None`	Checkpoint step, or `None` for the latest committed step.	`None`

Returns:

Type	Description
`dict[str, Mapping[str, Any]]`	Stable callback identifier mapping.

Raises:

Type	Description
`FileNotFoundError`	If no requested checkpoint exists.
`ValueError`	If callback metadata is malformed.

wait_until_finished [source]

wait_until_finished() -> None

Block until all asynchronous checkpoint writes have committed.

close [source]

close() -> None

Wait for pending writes and release Orbax resources idempotently.

enter [source]

__enter__() -> OrbaxCheckpointIO

Enter a checkpoint-manager context.

Returns:

Type	Description
`OrbaxCheckpointIO`	This checkpoint manager.

exit [source]

__exit__(
    exception_type: Any, exception: Any, traceback: Any
) -> None

Close the manager when leaving a context.

Parameters:

Name	Type	Description	Default
`exception_type`	`Any`	Active exception type, if any.	required
`exception`	`Any`	Active exception instance, if any.	required
`traceback`	`Any`	Active traceback, if any.	required

With asynchronous saving, teardown waits for pending writes and closes Orbax before Trainer.fit() returns. The backend reopens when the same Trainer starts another fit or restore task. Applications do not need to call Trainer.close(). Set max_to_keep=None to retain every saved step.

Configure saving

Save the latest state at a fixed interval:

from phijax.callbacks import ModelCheckpoint
from phijax.training import OrbaxCheckpointIO

checkpoint_io = OrbaxCheckpointIO(
    output_directory / "checkpoints",
    max_to_keep=3,
    enable_async_checkpointing=True,
)
checkpoint_callback = ModelCheckpoint(
    checkpoint_io,
    every_n_steps=5_000,
    save_last=True,
)

Monitor any scalar metric available after module and callback metric collection:

checkpoint_callback = ModelCheckpoint(
    checkpoint_io,
    every_n_steps=1_000,
    monitor="train/loss",
    mode="min",
    save_top_k=3,
    save_last=True,
)

Equal scores do not displace an already retained checkpoint. When a better checkpoint enters a full top-k set, the deterministic worst entry is removed through the backend. Callback state is stored under stable callback identifiers and restored before fit-start hooks. A changed callback set or malformed state raises clearly instead of applying a partial restoration.

Restore behavior

restore_checkpoint [source]

restore_checkpoint(
    target: TrainState,
    ckpt_path: str | Path | None,
    *,
    weights_only: bool = False,
    step: int | None = None,
) -> TrainState

Optionally restore full training state or model weights from a checkpoint root.

A missing path is the fresh-training case and returns target unchanged. Full restoration resumes model, optimizer, balancer, PRNG, precision, and step state. Weights-only restoration supports transfer learning while retaining every non-model field from target.

Parameters:

Name	Type	Description	Default
`target`	`TrainState`	Fresh state providing the expected checkpoint structure.	required
`ckpt_path`	`str \| Path \| None`	Orbax checkpoint-root directory, or `None` for fresh training.	required
`weights_only`	`bool`	Whether to restore only `model_state` for transfer learning.	`False`
`step`	`int \| None`	Specific checkpoint step, or `None` for the latest committed step.	`None`

Returns:

Type	Description
`TrainState`	Unchanged fresh state, fully resumed state, or weights-initialized fresh state.

Raises:

Type	Description
`ValueError`	If restoration options are supplied without a path, the path is empty, or `step` is invalid.
`FileNotFoundError`	If the requested checkpoint does not exist.

Mode	Restored state
Full resume	Model, optimizer, balancer, all PRNG streams, step, and precision fields
`weights_only=True`	Model state only; all other template state remains fresh
`ckpt_path=None`	Target is returned unchanged

Training accepts ckpt_path=None for fresh initialization. Standalone prediction requires a checkpoint root. An explicit step chooses an exact committed step; None selects the latest.

Trainer.fit(..., ckpt_path=...) restores the full state by default. Set weights_only=True to load only model weights, or set ckpt_path="last" to use the latest checkpoint from the configured callback. predict_state() loads model weights into the supplied state template when a path is provided. The template must have a compatible model structure. Use restore_checkpoint() directly only when managing state outside the Trainer.

PhiJAX 0.2 checkpoints use schema version 2 because TrainState stores separate sampling and balancing keys. PhiJAX 0.1 checkpoints remain readable by 0.1; 0.2 rejects them rather than partially restoring incompatible state.

Checkpointing

ModelCheckpoint callback

ModelCheckpoint [source]

connect [source]

setup [source]

on_fit_start [source]

on_train_batch_end [source]

on_train_metrics [source]

on_fit_end [source]

on_exception [source]

teardown [source]

close [source]

state_dict [source]

load_state_dict [source]

Checkpoint IO protocol

CheckpointIO [source]

latest_step property

steps property

open [source]

save [source]

checkpoint_path [source]

delete [source]

restore [source]

restore_weights [source]

restore_callback_states [source]

wait_until_finished [source]

close [source]

Orbax implementation

OrbaxCheckpointIO [source]

latest_step property

steps property

open [source]

save [source]

checkpoint_path [source]

delete [source]

restore [source]

restore_weights [source]

restore_callback_states [source]

wait_until_finished [source]

close [source]

__enter__ [source]

__exit__ [source]

Configure saving

Restore behavior

restore_checkpoint [source]

latest_step `property`

steps `property`

latest_step `property`

steps `property`

enter [source]

exit [source]