# Nvalchemi Training API

> How to configure nvalchemi training workflows with TrainingStrategy, custom training functions, standalone or composed losses, loss-weight schedules, optimizer and scheduler configs, validation, hooks, restartable checkpoints, model-agnostic inputs, and scaling to multiple GPUs or nodes with DistributedManager and DDPHook. Use when training a model from scratch, setting up optimizers, schedulers, validation, or checkpointing, or scaling a run across GPUs or nodes (DDP); for adapting a pretrained model, see nvalchemi-fine-tuning.

- Skill: `nvidia/nvalchemi-training-api` (Agent Skill)
- Install (CLI): `npx skillmds@latest add nvidia/nvalchemi-training-api`
- Raw SKILL.md: https://api.skillmd.com/api/skills/nvidia/nvalchemi-training-api/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: AI & ML
- Author: NVIDIA (https://skillmd.com/u/nvidia)
- Updated: 2026-09-17
- Page: https://skillmd.com/skills/nvidia/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`.

```python
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

```python
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`.

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.

```python
def training_fn(models: dict[str, BaseModelMixin], batch: Batch):
    student = models["student"](batch)
    with torch.no_grad():
        teacher = models["teacher"](batch)
    return {
        "student_energy": student["energy"],
        "teacher_energy": teacher["energy"].detach(),
    }

loss_fn = ComposedLossFunction(
    [EnergyMSELoss(prediction_key="student_energy", target_key="teacher_energy")]
)

strategy = TrainingStrategy(
    models={"student": student_model, "teacher": teacher_model},
    optimizer_configs={
        "student": [
            OptimizerConfig(
                optimizer_cls=torch.optim.AdamW,
                optimizer_kwargs={"lr": 3e-5},
            )
        ]
    },
    training_fn=training_fn,
    loss_fn=loss_fn,
    num_epochs=5,
)
```

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:

```python
loss_fn = EnergyMSELoss() + ForceMSELoss()
loss_fn.dtype_policy = "prediction_to_target"
```

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.

```python
loss_fn = (
    1.0 * EnergyMSELoss()
    + LinearWeight(start=0.0, end=10.0, num_steps=1000) * ForceMSELoss()
    + CosineWeight(start=0.0, end=0.1, num_steps=5000) * StressMSELoss()
)
```

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.

```python
strategy = TrainingStrategy.load_checkpoint("runs/example/checkpoints", map_location="cuda")
strategy.run(train_loader)
```

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.

```python
from nvalchemi.distributed import DistributedManager
from nvalchemi.training.hooks import DDPHook

DistributedManager.initialize()          # also handles single-process runs
manager = DistributedManager()

strategy = TrainingStrategy(
    models=model,
    optimizer_configs=OptimizerConfig(
        optimizer_cls=torch.optim.AdamW, optimizer_kwargs={"lr": 1e-4}
    ),
    loss_fn=EnergyMSELoss() + ForceMSELoss(),
    distributed_manager=manager,
    hooks=[DDPHook(), CheckpointHook("runs/ddp/checkpoints", epoch_interval=1)],
    num_epochs=20,
)
strategy.run(train_loader)
```

`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`:

```bash
torchrun --standalone --nproc_per_node=4 train.py
# runnable example:
uv run --extra cuXX torchrun --standalone --nproc_per_node=2 \
    examples/intermediate/06_ddp_mlp_training.py --backend auto
```

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`.