Training strategy API

Core training-loop classes and helpers.

See also

• Fine-tuning guide: Fine-Tuning Pretrained Models

• Fine-tuning API: Fine-tuning API

• Losses guide: Losses

• Training update hooks: Training update hooks

Strategies

TrainingStrategy	Pydantic-driven supervised training loop for MLIP models.
default_training_fn	Run a forward pass and prefix output keys with predicted_.

Field	Type	Description
hooks	list[nvalchemi.hooks._protocol.Hook | nvalchemi.training.hooks.update.TrainingUpdateHook | nvalchemi.training.hooks.update.TrainingUpdateOrchestrator]	Hooks to run at training stages. Accepts Hook Protocol instances, bare TrainingUpdateHook instances (auto-wrapped into a single TrainingUpdateOrchestrator), or an explicit TrainingUpdateOrchestrator. Example: hooks=[CheckpointHook(...), MyClipGradHook()].
step_count	<class 'int'>	Runtime optimizer-step counter, excluded from specs. Batches whose optimizer step is skipped by update hooks do not advance this counter.
models	dict[str, nvalchemi.models.base.BaseModelMixin]	Named models visible to training_fn and hooks. Single-model inputs are stored under "main"; torch.nn.ModuleDict inputs are accepted and normalized to a plain dict.
optimizer_configs	dict[str, list[nvalchemi.training.optimizers.OptimizerConfig]]	Optimizer/scheduler configs keyed by model name. Keys may target a subset of models; omitted models are frozen/eval during run.
num_epochs	int | None	Epoch count; mutually exclusive with num_steps. At runtime, epochs are converted into a target step count from the dataloader length and epoch_step_modifier.
num_steps	int | None	Target step count; mutually exclusive with num_epochs.
epoch_step_modifier	<class 'float'>	Positive multiplier applied when converting num_epochs to a target step count. Hooks may inspect this value through ctx.workflow.
training_fn	collections.abc.Callable[..., collections.abc.Mapping[str, torch.Tensor]] | None	Explicit forward-pass callable. Single-model strategies call (model, batch); named-model strategies call (models, batch).
loss_fn	<class 'nvalchemi.training.losses.composition.ComposedLossFunction'>	Composed loss whose components drive target collection. Leaf losses are accepted and normalized to one-component composed losses.
loss_target_assembler	<class 'nvalchemi.training.losses.composition.LossTargetAssemblyProtocol'>	Callable that assembles loss targets from the loss function, training predictions, current batch, and optional workflow.
devices	list[torch.device]	One device shared by all models, or one device per model for helper placement. Named-model run currently supports one device only.
distributed_manager	physicsnemo.distributed.manager.DistributedManager | None	Optional external distributed manager. The strategy passes this through hook contexts for distributed-aware hooks.
global_step_count	<class 'int'>	Runtime optimizer-step counter across all data-parallel workers, excluded from specs. This advances by the distributed world size when an optimizer step runs, so checkpoint restarts can recover sampler progress without assuming the same world size.
batch_count	<class 'int'>	Runtime batch counter, excluded from specs. This advances for every completed batch, including batches whose optimizer step is skipped.
epoch_count	<class 'int'>	Runtime epoch counter, excluded from specs.
epoch_step_count	<class 'int'>	Runtime counter for batches consumed within the current epoch, excluded from specs.
single_model_input	<class 'bool'>	Runtime flag recording whether a single model was supplied (stored under "main") rather than a named mapping. Set automatically during validation and used to pick the training_fn call convention.
last_validation	dict[str, typing.Any] | None	Most recent validation summary dict, or None before the first validation pass. Exposed to hooks via ctx.validation.
inference_model	torch.nn.modules.module.Module | torch.nn.modules.container.ModuleDict | None	Optional inference-time model (e.g. EMA weights) used in place of the live training model for validation when ValidationConfig.use_ema is set.
validation_config	nvalchemi.training._validation.ValidationConfig | None	Validation configuration controlling when and how validation runs. None disables validation.

Optimizer helpers

OptimizerConfig	Declarative optimizer and optional LR-scheduler bundle.
setup_optimizers	Build optimizers and schedulers for configured model names.
zero_gradients	Call zero_grad(set_to_none=True) on each optimizer.
step_optimizers	Call step() on each optimizer.
step_lr_schedulers	Call step() on each non-None time-based scheduler.

Field	Type	Description
optimizer_cls	<class 'type'>	Optimizer class; optimizer_kwargs must match its signature.
optimizer_kwargs	dict[str, typing.Any]	Keyword arguments forwarded to the optimizer constructor; validated against its __init__ signature at construction time.
scheduler_cls	type | None	Optional LR scheduler. Time-based schedulers (StepLR, CosineAnnealingLR, etc.) step every optimizer step. Metric-driven schedulers (ReduceLROnPlateau and subclasses) step only at validation checkpoints via step_metric_schedulers().
scheduler_kwargs	dict[str, typing.Any]	Must be empty unless scheduler_cls is set.
scheduler_metric_adapter	collections.abc.Callable[[dict[str, typing.Any]], float] | str | None	How a metric-driven scheduler (ReduceLROnPlateau) extracts its scalar metric from the validation summary dict. A str is treated as a key lookup into the summary; a callable receives the whole summary dict and returns a float; None uses the default extractor (see _extract_scheduler_metric()).

Serialization and checkpoints

BaseSpec	Base class for JSON-serializable, no-pickle hyperparameter specs.
create_model_spec	Build a BaseSpec instance for target with the given kwargs.
create_model_spec_from_json	Rebuild a BaseSpec from its JSON-dict form.
register_type_serializer	Register JSON (de)serializers for a custom type.
CheckpointManifest	Unified checkpoint manifest and runtime container.
save_checkpoint	Save a checkpoint with a manifest.
load_checkpoint	Load a multi-component checkpoint written by save_checkpoint().