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 

Minimal Pattern

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.

For multiple models, pass a named mapping. optimizer_configs must use the same model keys for trainable models. Models absent from optimizer_configs may be used in the forward path but are frozen during training.

If targets do not come directly from the batch, also provide a loss_target_assembler; see docs/userguide/training.md.

Losses And Scheduling

A standalone leaf loss such as EnergyMSELoss() can be used when the objective has one target. Use ComposedLossFunction or operator sugar for multi-target objectives. Leaf losses consume unweighted tensors; weights and schedules live on the composition. Built-in schedules include ConstantWeight, LinearWeight, CosineWeight, and PiecewiseWeight.

Built-in losses default to dtype_policy="strict" and raise when prediction and target dtypes differ. When building or reviewing workflows, check likely label/model dtype alignment, such as float64 dataset labels with float32 model outputs. If the mismatch is intentional, tell the user they can set dtype_policy="prediction_to_target" to cast outputs to labels or dtype_policy="target_to_prediction" to cast labels to outputs. Set the policy on an explicit ComposedLossFunction(...), on a leaf loss, or after operator-sugar construction:

A leaf loss with its own explicit dtype_policy overrides the composed-level policy. The setting is included in serializable loss specs for restartable training workflows. For CLI scaffolds, pass --loss-dtype-policy strict, --loss-dtype-policy prediction_to_target, or --loss-dtype-policy target_to_prediction to nvalchemi-training train init or nvalchemi-training finetune init ...; spec report shows the selected policy.

Caveats:

• normalize_weights=True is the default; set False for raw coefficient sums.
• per_epoch=True schedules require epoch during loss calls.
• Custom schedules must implement per_epoch, __call__(step, epoch), and to_spec() if they are used in restartable strategy checkpoints.
• For custom leaf-loss internals, use nvalchemi-loss-api and docs/userguide/losses.md.

Optimizers And Schedulers

Use OptimizerConfig(optimizer_cls=..., optimizer_kwargs=...); add scheduler_cls and scheduler_kwargs when needed. Keyword arguments are validated against class constructors before training starts.

Time-based schedulers step after optimizer steps. ReduceLROnPlateau-style metric schedulers step after validation; set scheduler_metric_adapter to a validation-summary key or callable when the default "total_loss" is not right.

Checkpoints And Reproducibility

Training workflows should be fully checkpointable and reproducible:

• Use deterministic model/wrapper constructors or create_model_spec(...).
• Keep loss functions, schedules, optimizer configs, and restart-critical hooks serializable; implement to_spec() where protocols require it.
• Use CheckpointHook for periodic checkpoints and save early enough for preempted jobs, including Slurm-style cluster runs.
• Make data splits, sampler state, seeds, units, dtype/device choices, and config files explicit in the run directory.
• For multi-GPU or multi-node runs (DDP, rank-safe checkpointing), see the Scaling to multiple GPUs section below.

Strategy checkpoints are restart packages: model weights, optimizer and scheduler state, strategy counters, checkpointable hook state, and reconstruction metadata.

Resume Training

Use resume when continuing the same run after interruption. This is different from fine-tuning, which imports weights into a new objective or dataset.

Resume only from native TrainingStrategy checkpoints when optimizer, scheduler, hook state, and counters matter. Plain pretrained weight files are not sufficient for faithful continuation. To start a fresh fine-tuning run from native checkpoint weights, use FineTuningStrategy.from_pretrained_checkpoint(...) from nvalchemi-fine-tuning; opt into source loss or optimizer classes with use_original_loss=True or use_original_opt_class=True when those defaults are desired. See docs/modules/training/checkpoints.rst.

Scaling to multiple GPUs (DDP)

Data-parallel training routes through DistributedManager (re-exported from PhysicsNeMo as nvalchemi.distributed.DistributedManager); prefer it as the single entry point. It owns rank, device, and process-group state, and passing it to TrainingStrategy alongside a DDPHook gives every hook the same runtime view, so one script runs unchanged on one process or many (with world size one, DDPHook is a no-op). See docs/userguide/distributed_training.md for the full guide.

DDPHook (during strategy setup) wraps the trainable models in DistributedDataParallel, selects the rank-local device, and injects a distributed sampler into the active dataloader, so no manual sampler wiring is needed. Rank-safety is handled for you: validation all-reduces its metrics across ranks (so never rank-gate the validation call), and CheckpointHook writes from global rank 0 only, unwrapping DDP so checkpoints store plain weights. Reporting is rank-aware too (see nvalchemi-reporting).

Launch one process per GPU with torchrun:

For multi-node launches (torchrun --nnodes/--rdzv_endpoint or Slurm srun), the rank helpers (get_rank, get_world_size, barrier, all_reduce in nvalchemi/training/distributed.py), and sampler/backend tuning, see docs/userguide/distributed_training.md.