Ctrl+K

Training checkpoints#

A TrainingStrategy is the runtime object that owns one training job: models, optimizers, schedulers, dataloaders supplied at run time, hooks, counters, and the callable used to turn a batch into a loss. The strategy is the object you run during training; a checkpoint is a serialized snapshot of enough of that strategy state to resume the same job later.

Training checkpoints therefore contain more than model weights. They include model weights, optimizer state, learning-rate scheduler state, strategy runtime counters, checkpointable hook state, and a serializable recipe for rebuilding the strategy components that can be reconstructed from metadata. They are intended for training restarts, not just inference weight export.

Manual save and restart#

Use save_checkpoint() when a live strategy should write its current training state:

from nvalchemi.training import TrainingStrategy

strategy = TrainingStrategy(...)
strategy.run(train_loader)

checkpoint_index = strategy.save_checkpoint("runs/example/checkpoints")

Use load_checkpoint() to reconstruct a new TrainingStrategy from one of those checkpoint snapshots. The method name refers to the checkpoint source; the return value is a strategy ready to keep training after you provide any runtime-only objects the checkpoint cannot serialize:

from nvalchemi.training import TrainingStrategy

strategy = TrainingStrategy.load_checkpoint(
    "runs/example/checkpoints",
    map_location="cpu",
    training_fn=training_fn,
)

strategy.num_steps = 20_000
strategy.run(train_loader)
strategy.save_checkpoint("runs/example/checkpoints")

checkpoint_index=-1 loads the latest checkpoint recorded in manifest.json. Pass an explicit index to restart from an older point:

strategy = TrainingStrategy.load_checkpoint(
    "runs/example/checkpoints",
    checkpoint_index=3,
)

Training functions#

Checkpoint metadata stores the training function only when it can be expressed as an importable dotted path. Importable functions do not need to be passed again:

# my_project/train_fns.py
def supervised_step(model, batch):
    predictions = model(batch)
    return loss_fn(predictions, batch)

strategy = TrainingStrategy(..., training_fn=supervised_step)
strategy.save_checkpoint(checkpoint_dir)

restored = TrainingStrategy.load_checkpoint(checkpoint_dir)

If the original strategy used a local function, a closure, or another non-importable callable, the checkpoint records that the function must be provided by the caller. Pass training_fn=... when loading:

def build_loss(scale):
    def training_fn(model, batch):
        predictions = model(batch)
        return scale * loss_fn(predictions, batch)

    return training_fn

restored = TrainingStrategy.load_checkpoint(
    checkpoint_dir,
    training_fn=build_loss(scale=0.5),
)

Hooks are runtime objects and are intentionally supplied at load time. For example, pass a new CheckpointHook when the restarted job should continue periodic checkpoint writes:

from nvalchemi.training import CheckpointHook, TrainingStrategy

strategy = TrainingStrategy.load_checkpoint(
    "runs/example/checkpoints",
    hooks=[
        CheckpointHook("runs/example/checkpoints", step_interval=1000),
    ],
)

Restartable hook state#

Hooks are still runtime objects and must be supplied when loading a strategy. However, hooks that implement CheckpointableHook have their runtime state stored in strategy checkpoints and restored into the matching hook supplied at load time. This is intended for hooks whose state changes training semantics, such as EMAHook and its averaged weights.

from nvalchemi.training import CheckpointHook, EMAHook, TrainingStrategy

checkpoint_dir = "runs/example/checkpoints"

ema = EMAHook(model_key="main", decay=0.999)
strategy = TrainingStrategy(
    ...,
    hooks=[
        ema,
        CheckpointHook(checkpoint_dir, step_interval=1000),
    ],
)
strategy.run(train_loader)

restored_ema = EMAHook(model_key="main", decay=0.999)
restored = TrainingStrategy.load_checkpoint(
    checkpoint_dir,
    hooks=[
        restored_ema,
        CheckpointHook(checkpoint_dir, step_interval=1000),
    ],
)

When a script already constructs the strategy and its runtime hooks, use restore_checkpoint() to restore checkpoint state into those live objects in place instead of reconstructing the strategy from metadata:

restored = TrainingStrategy(
    ...,
    hooks=[
        restored_ema,
        CheckpointHook(checkpoint_dir, step_interval=1000),
    ],
)
restored.restore_checkpoint(checkpoint_dir)

Checkpointable hooks are matched by class occurrence in the runtime hook list, so load-time hooks should be registered in the same relative order as the hooks that wrote the checkpoint. Non-checkpointable hook state remains the user’s responsibility. Prefer deriving transient state from restored strategy counters or rebuilding caches at setup time when possible.

Periodic checkpoint hook#

Use CheckpointHook for long-running jobs that should save without custom logic in the training loop:

from nvalchemi.training import CheckpointHook, TrainingStrategy

strategy = TrainingStrategy(
    ...,
    hooks=[
        CheckpointHook("runs/example/checkpoints", step_interval=1000),
    ],
)
strategy.run(train_loader)

A checkpoint hook owns one cadence policy. Use step_interval to save every N completed optimizer steps, or epoch_interval to save every N completed epochs. Register separate hooks only when a job intentionally needs separate checkpoint roots or policies.

By default, CheckpointHook captures a CPU snapshot on the training thread and writes that snapshot on a background thread. This avoids racing live model and optimizer tensors while moving filesystem writes off the main training path. Pending async writes are flushed when the strategy exits its hook context.

Model reconstruction specs#

Strategy checkpoints store model weights separately from a small JSON model spec. A model spec records an importable callable plus JSON-serializable keyword arguments. For ordinary modules, this callable is usually the class constructor:

create_model_spec(torch.nn.Linear, in_features=16, out_features=1)

For models that are created by a factory, adapter, monkey patch, or optimized conversion pass, the spec can point at that factory instead:

create_model_spec(
    MACEWrapper.from_checkpoint,
    checkpoint_path="small-0b",
    dtype=torch.float32,
    enable_cueq=True,
)

During load, the checkpoint layer rebuilds the model from the spec and then loads the saved training weights. If the factory accepts device, the loader passes map_location into the factory so device-sensitive conversions, such as MACE cuEquivariance conversion, happen directly on the target device.

Models may provide their own reconstruction spec by implementing checkpoint_spec() and returning a BaseSpec or None. Returning None keeps the default constructor-introspection fallback. This is useful for wrappers whose live module cannot be reconstructed from its transformed __init__ arguments.

Serialization scope#

Model specs make the broader rule explicit: a checkpoint stores data and references, never live Python objects. The manifest and every component recipe are plain JSON, and tensor bundles are reloaded with torch.load(..., weights_only=True). That keeps checkpoints portable and safe to reload without pickle, but it also bounds what a checkpoint can carry.

A checkpoint can embed:

  • Reconstruction specs for models, optimizers, and schedulers — an importable dotted path to a constructor or factory plus its keyword arguments, as described above.

  • Spec keyword arguments that are JSON-serializable (strings, numbers, booleans, None, lists, and dicts) or one of the registered tensor-aware types: torch.Tensor (stored as {dtype, shape, data}), torch.dtype, and torch.device. Register additional types with register_type_serializer from nvalchemi.training.

  • Tensor state — model weights, optimizer state, and scheduler state — written as separate weights-only bundles.

  • Strategy runtime counters, so local/global step, batch, and epoch positions resume.

  • Checkpointable hook state for hooks implementing CheckpointableHook.

  • The training function, but only when it resolves to an importable dotted path.

A checkpoint cannot embed the following; supply them from your training script at load time:

  • Non-importable callables — lambdas, locally defined or nested functions, closures, functools.partial objects, and bound methods of live instances. The training function is the common case: when it is one of these, the manifest records that it must be provided, and you pass training_fn=... to load_checkpoint() (see Training functions).

  • Spec arguments that are neither JSON-serializable nor registered — arbitrary Python objects, open file handles, or live modules. Reduce them to plain configuration, or expose a factory that rebuilds them from serializable arguments and point the spec at that factory.

  • Targets with positional-only parameters — create_model_spec builds a recipe from keyword arguments only.

  • Runtime objects such as hooks, dataloaders, and datasets. You reconstruct hooks in your script and pass them at load time; only checkpointable hook state is restored into them.

The practical takeaway: anything that cannot be reduced to an importable reference plus serializable arguments belongs in your training script, and the checkpoint expects you to re-supply it when you reload the strategy.

MACE checkpoints and cuEquivariance#

When training starts from an existing MACE checkpoint, construct the wrapper with from_checkpoint() and then let TrainingStrategy save and reload the full restart checkpoint:

import torch

from nvalchemi.models.mace import MACEWrapper
from nvalchemi.training import EMAHook, TrainingStrategy

model = MACEWrapper.from_checkpoint(
    "small-0b",
    device=torch.device("cuda"),
    dtype=torch.float32,
    enable_cueq=True,
)

ema = EMAHook(model_key="main", decay=0.999)
strategy = TrainingStrategy(
    models=model,
    ...,
    hooks=[ema],
)
strategy.run(train_loader)
strategy.save_checkpoint(checkpoint_dir)

On restart, reload the strategy checkpoint rather than saving or loading the EMA hook in isolation:

restored_ema = EMAHook(model_key="main", decay=0.999)
restored = TrainingStrategy.load_checkpoint(
    checkpoint_dir,
    map_location=torch.device("cuda"),
    hooks=[restored_ema],
    training_fn=training_fn,
)

The saved model spec calls MACEWrapper.from_checkpoint again with the recorded MACE checkpoint and options, then the strategy loader restores model weights, optimizer state, counters, and checkpointable hook state such as EMA averages.

Distributed training#

Distributed checkpointing follows the same file layout as single-process checkpointing, but only one process should write the shared checkpoint. The default CheckpointHook(rank_zero_only=True) uses the TrainContext global rank and saves only on rank 0. Other ranks continue training and do not write duplicate manifests or state files.

The usual end-to-end pattern is:

from nvalchemi.training import CheckpointHook, TrainingStrategy

checkpoint_dir = "runs/example/checkpoints"

strategy = TrainingStrategy(
    ...,
    hooks=[
        CheckpointHook(checkpoint_dir, step_interval=1000),
    ],
)
strategy.run(train_loader)

On restart, launch the distributed job again and have each process load the same checkpoint path:

from nvalchemi.training import CheckpointHook, TrainingStrategy

checkpoint_dir = "runs/example/checkpoints"

strategy = TrainingStrategy.load_checkpoint(
    checkpoint_dir,
    map_location=local_device,
    training_fn=training_fn,
    hooks=[
        CheckpointHook(checkpoint_dir, step_interval=1000),
    ],
)
strategy.num_steps = 20_000
strategy.run(train_loader)

load_checkpoint is not rank-zero-only: every process reconstructs its local strategy, model, optimizer, scheduler, and counters from the shared checkpoint files. Pass map_location when the restored process should load onto a rank-local device instead of the device recorded in the checkpoint metadata.

The checkpoint directory must be visible to every rank before restart. For periodic hook saves, the async writer is flushed when the strategy exits. For manual save workflows, users should coordinate their distributed script so only one rank calls save_checkpoint(), then ensure all ranks wait until the checkpoint is complete before any rank tries to reload it.

Current checkpoints store replicated strategy and optimizer state. They are intended for the training strategies used by this package and do not provide a separate sharded checkpoint format for distributed optimizers or model shards. Workflows that shard model or optimizer state outside the strategy checkpoint must save and restore those sharded states separately.

DistributedDataParallel wrappers are unwrapped before model specs and model weights are written, so native checkpoints store the underlying model state without module. key prefixes. FSDP and FSDP2 require PyTorch Distributed Checkpoint (DCP) so that each rank can save its shard and reload under a possibly different topology. Native strategy checkpoints currently reject FSDP/FSDP2-wrapped models instead of writing incomplete rank-local state. See the PyTorch Distributed Checkpoint recipe for the DCP workflow.

Lower-level loader#

The module-level save_checkpoint() and load_checkpoint() functions remain available when callers need the full manifest, component dictionaries, validators, model subsets, or adapter loads. TrainingStrategy.load_checkpoint deliberately returns only the restored strategy and rejects component-only checkpoints.

When a native checkpoint should seed a new fine-tuning job rather than resume the same run, use from_pretrained_checkpoint() instead. It loads the checkpoint model weights into a fresh fine-tuning strategy without optimizer state, scheduler state, counters, hooks, or run limits. Source loss and optimizer/scheduler configuration are reused only when explicitly requested by that constructor.

API reference#

TrainingStrategy.save_checkpoint

Save this strategy as a restartable checkpoint.

TrainingStrategy.restore_checkpoint

Restore checkpoint state into this already-constructed strategy.

TrainingStrategy.load_checkpoint

Load a restartable strategy checkpoint.

save_checkpoint

Save a checkpoint with a manifest.

load_checkpoint

Load a multi-component checkpoint written by save_checkpoint().

CheckpointHook

Periodically save restartable training strategy checkpoints.
On this page
Show Source