`nvalchemi.training`.TrainingStrategy#

pydantic model nvalchemi.training.TrainingStrategy[source]#

Pydantic-driven supervised training loop for MLIP models.

TrainingStrategy is the top-level object that owns a supervised training run. You construct it declaratively with the models to train, the optimizer/scheduler recipe, a duration, a loss, and a forward-pass function, then call run() with a dataloader to execute the loop. Because the strategy is a pydantic.BaseModel, construction validates the whole configuration up front — mismatched optimizer keys, an invalid duration, or an ill-typed training_fn surface as a pydantic.ValidationError before any training starts.

The strategy accepts either a single wrapped model (a BaseModelMixin), which is stored internally under the key "main", or a {name: model} mapping for distillation and multi-model workflows. That choice drives the training_fn calling convention: single-model strategies call training_fn(model, batch) and named-model strategies call training_fn(models, batch). Most workflows use the provided default_training_fn(), which runs a forward pass and prefixes output keys with predicted_ for loss-target assembly. The loss may be a leaf loss (auto-normalized to a one-component ComposedLossFunction) or an explicit composition such as EnergyMSELoss() + ForceMSELoss(...).

optimizer_configs accepts a single OptimizerConfig, a list, or a {model_name: [OptimizerConfig, ...]} mapping; unkeyed forms require a single-model input. Keys may target a subset of models — any model without a config is temporarily set to eval mode and frozen during run(), which is what makes teacher/auxiliary networks in distillation work. Named-model training functions that consume those frozen models must still run their forward passes under torch.no_grad() or detach the outputs unless autograd through them is intentional.

Duration is set by exactly one of num_epochs or num_steps (both the default, or setting both, is rejected). Internally the target is always an optimizer-step count; num_epochs is converted from the dataloader length scaled by epoch_step_modifier, so epoch-based runs require a sized dataloader. Behavior is customized through hooks (checkpointing, logging, gradient clipping, EMA, DDP, mixed precision, …) and validation is enabled by attaching a ValidationConfig via validation_config. Runs are restartable: save_checkpoint() and restore_checkpoint() persist the recipe plus runtime counters, while to_spec_dict() / from_spec_dict() handle JSON-based recipe-only save/load.

Examples

Single-model supervised training for a fixed number of epochs:

>>> import torch
>>> from nvalchemi.training import (
...     EnergyMSELoss,
...     ForceMSELoss,
...     OptimizerConfig,
...     TrainingStrategy,
...     default_training_fn,
... )
>>> strategy = TrainingStrategy(
...     models=model,
...     optimizer_configs=OptimizerConfig(
...         optimizer_cls=torch.optim.Adam,
...         optimizer_kwargs={"lr": 1e-3},
...     ),
...     num_epochs=10,
...     training_fn=default_training_fn,
...     loss_fn=EnergyMSELoss() + ForceMSELoss(normalize_by_atom_count=True),
...     devices=[torch.device("cuda")],
... )
>>> strategy.run(train_loader)

Step-based training with periodic validation and a checkpoint hook:

>>> from nvalchemi.training import ValidationConfig
>>> strategy = TrainingStrategy(
...     models=model,
...     optimizer_configs=OptimizerConfig(optimizer_cls=torch.optim.AdamW),
...     num_steps=50_000,
...     training_fn=default_training_fn,
...     loss_fn=EnergyMSELoss(),
...     validation_config=ValidationConfig(
...         validation_data=val_batches,
...         every_n_steps=1_000,
...     ),
...     hooks=[CheckpointHook(checkpoint_dir="runs/exp")],
... )
>>> strategy.run(train_loader)

Named-model (distillation) setup optimizing only the student while the teacher stays frozen because it is absent from optimizer_configs:

>>> strategy = TrainingStrategy(
...     models={"student": student, "teacher": teacher},
...     optimizer_configs={
...         "student": [OptimizerConfig(optimizer_cls=torch.optim.Adam)]
...     },
...     num_steps=10_000,
...     training_fn=distillation_training_fn,
...     loss_fn=EnergyMSELoss(),
... )

Notes

Exactly one of num_epochs and num_steps must be set; num_epochs additionally requires a sized dataloader so it can be converted to a step target. Every optimizer_configs key must name a model present in models, and each entry must contain at least one OptimizerConfig. devices must have length 1 or len(models); named-model run() currently supports a single shared device only.

Use to_spec_dict() / from_spec_dict() for JSON-based save/load. Optimizer configs, loss specs, devices, importable training functions, and best-effort model specs are serialized. Runtime models and training_fn overrides passed to from_spec_dict() take precedence; the serialized model call mode is used only when no runtime model override is supplied. hooks, step_count, global_step_count, batch_count, epoch_count, and epoch_step_count remain runtime-only.

Bare TrainingUpdateHook instances are auto-wrapped into a single TrainingUpdateOrchestrator on registration; the orchestrator owns the zero_gradients / backward / optimizer.step / scheduler.step calls that the strategy otherwise issues by default. Construction-time hook validation errors surface as pydantic.ValidationError; register_hook() raises ValueError directly.

field models: dict[str, BaseModelMixin] [Required]#: 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.

field optimizer_configs: dict[str, list[OptimizerConfig]] [Optional]#: Optimizer/scheduler configs keyed by model name. Keys may target a subset of models; omitted models are frozen/eval during run.

field num_epochs: int | None = 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.

Constraints:

ge = 1

field num_steps: int | None = None#

Target step count; mutually exclusive with num_epochs.

Constraints:

ge = 1

field epoch_step_modifier: float = 1.0#

Positive multiplier applied when converting num_epochs to a target step count. Hooks may inspect this value through ctx.workflow.

Constraints:

gt = 0
allow_inf_nan = False

field hooks: list[Hook | TrainingUpdateHook | TrainingUpdateOrchestrator] [Optional]#: 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()].

field training_fn: Callable[[...], Mapping[str, Tensor]] | None = None#: Explicit forward-pass callable. Single-model strategies call (model, batch); named-model strategies call (models, batch).

field loss_fn: ComposedLossFunction [Required]#: Composed loss whose components drive target collection. Leaf losses are accepted and normalized to one-component composed losses.

field loss_target_assembler: LossTargetAssemblyProtocol = <function assemble_loss_targets>#: Callable that assembles loss targets from the loss function, training predictions, current batch, and optional workflow.

field devices: list[device] [Optional]#: One device shared by all models, or one device per model for helper placement. Named-model run currently supports one device only.

field distributed_manager: DistributedManager | None = None#: Optional external distributed manager. The strategy passes this through hook contexts for distributed-aware hooks.

field step_count: int = 0#

Runtime optimizer-step counter, excluded from specs. Batches whose optimizer step is skipped by update hooks do not advance this counter.

Constraints:

ge = 0

field global_step_count: int = 0#

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.

Constraints:

ge = 0

field batch_count: int = 0#

Runtime batch counter, excluded from specs. This advances for every completed batch, including batches whose optimizer step is skipped.

Constraints:

ge = 0

field epoch_count: int = 0#

Runtime epoch counter, excluded from specs.

Constraints:

ge = 0

field epoch_step_count: int = 0#

Runtime counter for batches consumed within the current epoch, excluded from specs.

Constraints:

ge = 0

field single_model_input: bool = False#: 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.

field last_validation: dict[str, Any] | None = None#: Most recent validation summary dict, or None before the first validation pass. Exposed to hooks via ctx.validation.

field inference_model: Module | ModuleDict | None = 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.

field validation_config: ValidationConfig | None = None#: Validation configuration controlling when and how validation runs. None disables validation.

nvalchemi.training.TrainingStrategy#

`nvalchemi.training`.TrainingStrategy#