Checkpoints and resuming

A checkpoint captures both a model and the training state needed to continue from that point. A run can resume after a deliberate pause, or a saved model can be loaded independently for evaluation. The checkpoint position and the choice of raw or averaged parameters are separate decisions.

Stop and resume

StopRequest.catch_signals() temporarily handles SIGINT/SIGTERM in the main thread. Any rank may also call stop.request(reason). The trainer finishes the current accumulation window, saves last when a checkpoint directory is configured, skips final evaluation, and returns with state.stop_reason. Signal handlers only set a flag; saving and communication run in the training loop.

state = trainer.train(
    TrainConfig(epochs=100),
    checkpoint_dir=run_dir / "checkpoints",
    resume_from=run_dir / "checkpoints" / "last",
)

Checkpoint contents

Checkpoints atomically publish one complete PyTorch file containing named raw/EMA weights, optimizer/scheduler/scaler state, data progress, and per-rank RNG state. last, best, and final are relative symlinks. A requested pause updates last; a normal endpoint updates final. Retention preserves all three pointers' targets.

Conditions for exact continuation

Exact continuation requires deterministic map-style Dataset reads and collation, unchanged data/batching/world size, and the same objective/count definition, optimizer, scheduler recipe, EMA and precision setup. Epoch-addressed shuffle and model randomness such as Dropout are restored. Worker-local random transforms, custom stateful samplers, and streaming data need a separate data-state protocol and are outside this guarantee. Increasing the cumulative epoch/step limits is allowed; using a shorter max_steps as a normal endpoint may add a final evaluation and alter a plateau schedule, so use StopRequest for a transparent pause.

Legacy Accelerate checkpoint directories must be loaded with the training code that wrote them.

The training example shows how to install the signal handlers around a run and provide its checkpoint directory.

StopRequest()

catch_signals() -> Iterator[StopRequest]

Loading and exporting weights

from elfes.engine import export_weights, load_weights

load_weights(model, run_dir / "checkpoints" / "best", weights="ema")
model = accelerator.prepare(model)
results = evaluator.evaluate(model)
export_weights(run_dir / "checkpoints" / "best", output_path, weights="ema")

Pass an unwrapped, uncompiled model to load_weights. No optimizer or training DataLoader is needed. Exports are ordinary model state dictionaries, with no DDP or compilation prefixes. EMA includes checkpoint-time raw model buffers. best selects the training position; weights selects the parameter set at that position.

Here accelerator and evaluator are configured for the intended evaluation; the model architecture must match the saved parameters. Trainer.save_model() exports the current model, while export_weights() reads weights from a saved checkpoint.

load_weights(model: Module, checkpoint: Path, *, weights: Weights = 'raw') -> None

export_weights(checkpoint: Path, path: Path, *, weights: Weights = 'raw') -> None