nvalchemi-training-api

nvalchemi Training API

Overview

Use TrainingStrategy as the owner of one training job: model(s), dataloaders, loss, optimizer/scheduler config, validation, hooks, runtime counters, and checkpoints. For full details, see docs/userguide/training.md, docs/userguide/losses.md, and docs/modules/training/checkpoints.rst.

import torch

from nvalchemi.data import Batch
from nvalchemi.models.base import BaseModelMixin
from nvalchemi.training import (
    CheckpointHook,
    ComposedLossFunction,
    CosineWeight,
    EnergyMSELoss,
    ForceMSELoss,
    LinearWeight,
    OptimizerConfig,
    StressMSELoss,
    TrainingStrategy,
    ValidationConfig,
    create_model_spec,
)

---

Minimal Pattern

loss_fn = ComposedLossFunction(
    [EnergyMSELoss(), ForceMSELoss()],
    weights=[1.0, 10.0],
    normalize_weights=False,
)

strategy = TrainingStrategy(
    models=model,
    optimizer_configs=OptimizerConfig(
        optimizer_cls=torch.optim.AdamW,
        optimizer_kwargs={"lr": 1e-4, "weight_decay": 1e-5},
    ),
    loss_fn=loss_fn,
    validation_config=ValidationConfig(validation_data=val_loader, every_n_epochs=1),
    hooks=[CheckpointHook("runs/example/checkpoints", epoch_interval=1)],
    num_epochs=20,
)
strategy.run(train_loader)

---

Model-Agnostic Inputs

Accept any torch.nn.Module that works with the selected training_fn. Prefer wrapped BaseModelMixin models for standard AtomicData/Batch contracts; see the nvalchemi-model-wrapping skill or docs/userguide/models.md when adapting arbitrary MLIPs.

Make model construction reproducible when possible. Use native checkpoint constructors that carry a spec, or store a create_model_spec(...) for custom wrappers so strategy checkpoints can rebuild the model before loading weights. Treat foreign checkpoints as imported weights until a fresh TrainingStrategy checkpoint has been saved.

---

Custom Training Functions

Use training_fn when the batch needs custom routing, multiple models, teacher outputs, auxiliary predictions, or non-standard model outputs. It receives (model, batch) for a single model or (models, batch) for named models and returns the prediction mapping consumed by loss_fn.