nvalchemi.training.hooks.CheckpointHook

pydantic model nvalchemi.training.hooks.CheckpointHook

Periodically save restartable training strategy checkpoints.

The hook observes completed training counters and saves TrainingStrategy checkpoints through the same manifest layout as nvalchemi.training.save_checkpoint(). It fires either every step_interval completed optimizer steps or every epoch_interval completed epochs. The two cadences are mutually exclusive so each hook owns one clear checkpoint policy.

With async_save=True (default), the hook first captures an immutable CPU snapshot of model, optimizer, scheduler, and strategy metadata on the training thread, then writes that snapshot on a single background thread. This avoids racing against live training tensors while still moving the filesystem work off the critical path. If a later checkpoint is due while the previous background write is still running, the hook waits for the previous write before capturing the next snapshot so manifest indices stay ordered.

Raises:

• ValueError – If neither interval is provided, or an interval is not positive.

• RuntimeError – If the hook is called without a strategy workflow in TrainContext.

Examples

>>> from nvalchemi.training import CheckpointHook, TrainingStrategy
>>> hook = CheckpointHook("runs/example/checkpoints", step_interval=1000)
>>> strategy = TrainingStrategy(..., hooks=[hook])
>>> strategy.run(train_loader)

field checkpoint_dir: Path [Required]

Root directory for restartable training checkpoints.

field step_interval: int | None = None

Completed-step save interval.

Constraints:

• gt = 0

field epoch_interval: int | None = None

Completed-epoch save interval.

Constraints:

• gt = 0

field async_save: bool = True

Write checkpoint snapshots on a background thread.

field rank_zero_only: bool = True

Restrict checkpoint writes to distributed rank 0.

field last_checkpoint_index: int | None = None

Most recent checkpoint index known to have been written. In async mode, this updates when the background future completes.

Constraints:

• ge = 0