Training

Model training in NVIDIA ALCHEMI Toolkit is designed to be a highly modular, extensible, and ergonomic set of core functionalities: everything was designed with the explicit intention for users and developers to implement their atomistic modeling intentions with as little friction as possible, whilst keeping production readiness (e.g. traceability and reproducibility) a background constant.

Tip

AI coding assistant? Load the nvalchemi-training-api agent skill for concise instructions on configuring and extending TrainingStrategy workflows.

As with many components in nvalchemi-toolkit, many of the core user-facing interfaces are written with pydantic, making early validation and (de)serialization first-class citizens. The core of the training utilities is organized around TrainingStrategy — a pydantic workflow engine that orchestrates the moving pieces during training, as well as ensuring that the “recipe” for training can be and is persisted, reproducible, and user-friendly. This class is similar to a “trainer” abstraction that many are familiar with, albeit with some key differences in what it comprises for the sake of reproducibility as well as with modern model training concepts in mind such as multi-model distillation, multitask loss weight scheduling, model averaging, etc. that have become recently popular in atomistic AI/ML models.

Minimal structure

Before getting into the details, it helps to see the overall anatomy at once. Nearly every training script, however elaborate it eventually becomes, is built from the same five parts:

1. Build or load one or more models.

2. Create a dataloader that emits Batch objects.

3. Define the loss or training function that turns each batch into a scalar loss.

4. Configure optimizer, scheduler, validation, and hook behavior.

5. Execute the strategy, while leaving a trail of metrics and model weights.

from nvalchemi.training import (
    CheckpointHook,
    ComposedLossFunction,
    EnergyMSELoss,
    ForceMSELoss,
    OptimizerConfig,
    TrainingStrategy,
    ValidationConfig,
)

# configure how the model will be trained
loss_fn = EnergyMSELoss() + ForceMSELoss() * 10.0

# pass a model, the optimizer configuration, and supporting
# functionality as hooks
strategy = TrainingStrategy(
    models=model,
    optimizer_configs=OptimizerConfig(lr=1e-4),
    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)

The rest of this page walks through what happens after run() starts. The key idea is that TrainingStrategy is not only a loop over batches; it is a small workflow engine whose public extension points are the values of TrainingStage.

Lifecycle Overview

run() expands into a fixed sequence of stages. The whole of it fits in the single diagram below, which is useful as a reference when you are trying to understand the orchestration flow as well as when you are trying to build new workflows and components:

The diagram is meant to be read as both execution order and API map. Stages are where hooks enter the workflow; the filled operation boxes are where the strategy itself calls the model, loss, backward pass, optimizer, scheduler, validation, or checkpoint machinery. Most stages are placed as observation points: hooks can inspect the current TrainContext, log metrics, update side state, or modify workflow-owned objects when that stage allows it. The exception is the two replacement stages, DO_BACKWARD and DO_OPTIMIZER_STEP, which are unique to training and are owned either by the strategy default path or by the TrainingUpdateOrchestrator.

The per-batch stages in the inner loop are covered in detail in subsequent sections, but the outer stages are worth describing here since they are where coarser-grained work belongs. BEFORE_TRAINING and AFTER_TRAINING fire exactly once, wrapping the whole run: the former suits one-time setup that needs the resolved runtime state, and the latter is the place for final teardown such as flushing a reporting sink or closing a writer. BEFORE_EPOCH and AFTER_EPOCH bracket each pass over the dataloader; the latter is the natural home for epoch-level summaries, periodic checkpoints, and epoch-cadence validation. The validation stages sit slightly apart from the main flow: a validation pass runs on a step or epoch cadence, and the moment it finishes AFTER_VALIDATION fires while its reduced summary is still in hand — which is exactly where validation logging and metric-driven schedulers such as ReduceLROnPlateau do their work.

Configuring a training strategy

TrainingStrategy is organized around two groups of inputs. The forward path — models, training_fn, loss_fn, and loss_target_assembler — controls what happens to each batch. The loop control layer — optimizer_configs, validation_config, hooks, and num_epochs or num_steps — controls when and how training progresses.

The two components on the forward path are training_fn and loss_target_assembler. training_fn owns everything from model call to prediction mapping; replace it when the default model(batch) call is not enough — multi-model workflows, distillation, or any forward pass that needs more than one model to produce the output. loss_target_assembler controls how those predictions are matched to targets before loss_fn is called; replace it when targets are not in the batch directly, e.g. when they come from a teacher model’s output or must be assembled from other sources. Together they let you replace the entire forward-to-loss path without touching the rest of the training strategy.

A training_fn is a callable that receives either (model, batch) for a single-model strategy or (models, batch) for a named multi-model strategy, and returns a Mapping of model outputs. That mapping is passed into compute_supervised_loss(), which reads targets from the batch by TrainingStrategy.target_keys unless a loss_target_assembler is supplied. A loss_target_assembler must satisfy LossTargetAssemblyProtocol: it receives loss_fn, the prediction mapping, the batch, the current workflow, and optional target_keys, then returns the target mapping.

The distillation example below shows both seams in use. training_fn runs both models and returns their outputs under distinct keys; teacher_targets implements the loss_target_assembler protocol and pulls targets from the prediction mapping rather than the batch:

from collections.abc import Mapping, Sequence

import torch

from nvalchemi.data import Batch
from nvalchemi.models.base import BaseModelMixin
from nvalchemi.training.losses import ComposedLossFunction, EnergyMSELoss

# example that employs a student-teacher workflow; the same
# batch is passed into a student and a teacher model. The
# teacher output is returned in the predictions mapping and
# routed into loss_fn by loss_target_assembler.


def training_fn(
    models: Mapping[str, BaseModelMixin],
    batch: Batch,
) -> Mapping[str, torch.Tensor]:
    """Implements the logic for computing predictions, given a set
    of models and an incoming ``Batch`` object.
    """
    student_out = models["student"](batch)
    # teacher is not part of the autograd graph
    with torch.no_grad():
        teacher_out = models["teacher"](batch)
    return {
        "student_energy": student_out["energy"],
        "teacher_energy": teacher_out["energy"].detach(),
    }


def teacher_targets(
    loss_fn: ComposedLossFunction,
    predictions: Mapping[str, torch.Tensor],
    batch: Batch,
    *,
    workflow: object | None = None,
    target_keys: Sequence[str] | None = None,
    batch_label: str = "Batch",
) -> Mapping[str, torch.Tensor]:
    """This method is used to inform the training workflow how
    to obtain the target values to train against.

    Normally, the values would be grabbed from the ``Batch`` object but in
    this case we retrieve them from the ``predictions`` as they
    were returned as part of ``training_fn``.
    """
    return {"teacher_energy": predictions["teacher_energy"]}


# at runtime, the loss function will rely on `teacher_targets`
# to provide the labels
loss_fn = EnergyMSELoss(
    prediction_key="student_energy",
    target_key="teacher_energy",
    per_atom=True,
  )

In this example, prediction_key="student_energy" is read from the mapping returned by training_fn, while target_key="teacher_energy" names the target returned by teacher_targets. Users opt into that routing by passing loss_target_assembler=teacher_targets to TrainingStrategy. The strategy calls the assembler with the configured loss, predictions, batch, and current workflow, then passes the resulting target mapping into loss_fn.

Warning

Having training_fn and loss_target_assembler as a mere callable that’s passed into TrainingStrategy was intentional for the sake of agility: when writing a script, you could simply embed the function within the same script, or persist it in the package or workflow you are developing. For security reasons, we do not serialize the function as part of checkpointing, as there are no effective ways to guarantee your callable function is safe to execute, and that it hasn’t been replaced in-flight.

For this reason, it is up to the user/developer to ensure that their training function is importable, and to ensure that their checkpoints and training function/recipe is up to date with one another.

The hooks in the loop control layer participate at different lifecycle stages. Hooks that need structural changes before the first optimizer is built — DDP wrapping, reporter initialization, profiler setup — run during SETUP. Per-batch output should wait for the batch stages that carry the relevant data. At run(), the strategy resolves this startup sequence: it moves models to devices, lets setup hooks mutate the workflow, normalizes update hooks into a single orchestrator, and builds optimizers and schedulers. The loop then begins.

Understanding how the strategy tracks progress through that loop is the foundation for writing hooks that fire at the right time.

Training Counters

The training workflow tracks progress using a small set of counters:

• batch_count counts the number of completed batches on this worker,

• step_count counts completed optimizer/scheduler steps on this worker,

• global_step_count counts completed optimizer/scheduler steps across all data-parallel workers,

• epoch_count counts the number of times the dataloader has been exhausted,

• epoch_step_count counts the number of batches consumed in the current epoch.

The distinction between batch_count and step_count is important. A batch can finish without an optimizer step if the training workflow uses gradient accumulation, spike skipping, or any other update policy that defers or vetoes the step. Code that cares about data throughput should usually read batch_count, while code that cares about local optimizer state should usually read step_count. Distributed code that needs aggregate optimizer progress, such as fixed compute budgets (i.e. how many FLOPs have I utilized across all ranks) or world-size-independent sampler restarts, should read global_step_count; under DDP it advances by the current world size when an optimizer step runs and is restored from checkpoints.

Inside hooks, these values are available from the TrainContext passed into the hook call:

from nvalchemi.training import TrainingStage


# this is just to illustrate how a logger hook can access state
class ProgressLogger:
    stage = TrainingStage.AFTER_BATCH
    frequency = 1

    def __call__(self, ctx, stage):
        logger.info(
            "epoch=%s batch=%s step=%s",
            ctx.epoch,
            ctx.batch_count,
            ctx.step_count,
        )

Outside hooks, the same state is available on the strategy object as strategy.epoch_count, strategy.batch_count, strategy.step_count, strategy.global_step_count, and strategy.epoch_step_count. These values are part of the strategy runtime state and are restored by checkpoints.

After setup, BEFORE_TRAINING fires once before the first batch. The epoch loop then starts with BEFORE_EPOCH. At each epoch boundary, the strategy calls set_epoch(...) on distributed samplers when available, so each epoch can use a deterministic but distinct sample order.

Batches: Forward, Loss, Backward, Update

With the counters and epoch loop in view, we can zoom in on what happens to a single batch. Before the batch stages run, the strategy moves the batch onto the primary training device — the device the model was placed on, which under DDP is the current rank’s GPU. Each stage then gives hooks access to a progressively richer TrainContext, following the natural data-availability order of the forward pass:

Stage	Available on ctx	Extension opportunity
BEFORE_BATCH	Batch, counters	Per-batch setup, zero-gradient policy
BEFORE_FORWARD	Batch	Transform inputs before the model call
AFTER_FORWARD	ctx.predictions	Log prediction statistics; redirect outputs before loss
AFTER_LOSS	ctx.loss	Per-component and per-sample loss diagnostics
AFTER_BACKWARD	Gradients	Gradient norm logging, gradient-based monitoring
AFTER_OPTIMIZER_STEP	Updated weights	EMA updates, learning-rate logging, step metrics
AFTER_BATCH	Full context	Throughput logging, reporting, checkpoint cadence

The two replacement stages, DO_BACKWARD and DO_OPTIMIZER_STEP, are not observation points — they are owned by either the strategy default path or a TrainingUpdateHook. Extending those stages requires the update orchestrator; see Optimizer Orchestration.

The default supervised path calls training_fn to produce a prediction mapping, then calls compute_supervised_loss() to retrieve targets and evaluate loss_fn. The resulting structured loss contains total_loss for backpropagation, along with per-component and per-sample diagnostics accessible at AFTER_LOSS. See Losses and Losses — Training Terms for the loss object contract.

Optimizer Orchestration

Once the loss has been computed, the TrainingStrategy then needs to be able to handle it, as well as provide the opportunity for developers to interact with how backpropagation is performed, perform gradient surgery, etc. TrainingUpdateHook is the abstraction for this: it owns the replacement stages DO_BACKWARD and DO_OPTIMIZER_STEP, and is the right tool for any workflow that changes how gradients are computed or applied — mixed precision, gradient accumulation, gradient clipping, spike skipping, and EMA all fit here.

Optimizers and learning-rate schedulers are configured through optimizer_configs. Each entry names the model parameters it owns and the optimizer/scheduler objects that should be constructed for those parameters. During setup, TrainingStrategy builds the configured optimizers and schedulers once, stores them on the runtime context, and exposes them to hooks as ctx.optimizers and ctx.lr_schedulers.

When no specialized update hooks are registered (these are discussed below), the strategy owns the default update sequence, which runs on every batch:

1. zero gradients before the forward pass,

2. call loss.backward() after AFTER_LOSS,

3. call step_optimizers() to apply the parameter update,

4. advance step-based learning-rate schedulers with step_lr_schedulers(),

5. advance step_count, but only when the optimizer-step path actually executes.

That last point is the one worth internalizing: because step_count moves only when an optimizer step is taken, gradient accumulation and similar policies can defer an update without corrupting the step bookkeeping.

Metric-based schedulers, such as ReduceLROnPlateau, are the exception to this fixed cadence. Rather than stepping on every optimizer step, they require a validation quantity to track, and that quantity is only exposed on the TrainContext once training reaches TrainingStage.AFTER_VALIDATION.

A TrainingUpdateHook can participate in four update stages: BEFORE_BATCH, DO_BACKWARD, DO_OPTIMIZER_STEP, and AFTER_OPTIMIZER_STEP. When one or more update hooks are registered, the strategy folds them into a single TrainingUpdateOrchestrator. The orchestrator becomes the owner of the replacement stages, so the strategy does not also call its default backward or optimizer-step implementation.

The update stages have distinct responsibilities:

Stage	Responsibility
BEFORE_BATCH	Zero-gradient policy and per-batch accumulation setup
DO_BACKWARD	Backward pass or a transformed version of it
DO_OPTIMIZER_STEP	Optimizer and scheduler stepping; can veto the step
AFTER_OPTIMIZER_STEP	Post-step updates such as EMA weights

Update hooks can be registered directly in hooks=[...]; the strategy will wrap bare update hooks into one orchestrator. They can also be composed explicitly with hook_a + hook_b when a script wants to make the composition visible.

Note

Only one object may own DO_BACKWARD and only one object may own DO_OPTIMIZER_STEP. The dividing line is ownership: a hook that only observes gradients, learning rates, or counters — logging gradient norms at AFTER_BACKWARD, say — should stay a standard hook, while one that changes whether or how gradients are applied belongs in the update orchestrator.

See Training update hooks for the stage contract and the built-in update hooks.

Validation, Schedulers, And Reporting

Once the update path is handling gradients and optimizer steps, the next question is whether the model is actually improving. Validation, metric-driven scheduling, and reporting all live in the outer lifecycle stages — AFTER_VALIDATION, AFTER_EPOCH, AFTER_BATCH — so they observe fully updated weights without interfering with the update path. Some schedulers also cannot make their decisions without a validation signal.

Validation is configured through ValidationConfig on the strategy. It reuses the same model, training_fn/validation_fn, loss function, and target assembly language as training, but executes under validation semantics: evaluation mode by default, configurable autograd, optional EMA weights, and distributed reduction of summary metrics.

Step-cadence validation is checked after AFTER_OPTIMIZER_STEP, so it observes the latest successfully updated weights. Epoch-cadence validation is checked after AFTER_EPOCH. When training finishes, the strategy runs a final validation pass if validation is configured. Immediately after each validation pass, AFTER_VALIDATION fires while the reduced summary is still available on the strategy.

Tip

When using model averaging (EMA), the hook will automatically use the averaged model weights for computing validation. This will generally result in significantly smoother validation curves than the training counterparts.

Use AFTER_VALIDATION for lifecycle-level validation logging and metric-driven scheduler behavior. Use the per-batch callback on ValidationConfig only when you need a tap into individual validation batches, predictions, or losses for a custom sink or offline error analysis.

Listening to validation results

Register a standard hook on AFTER_VALIDATION to read the reduced summary. The summary is available on every rank; guard external side effects with a rank check:

from nvalchemi.training import TrainingStage

class SummaryLogger:
    stage = TrainingStage.AFTER_VALIDATION
    frequency = 1

    def __call__(self, ctx, stage):
        summary = ctx.validation
        if ctx.global_rank == 0 and summary is not None:
            my_tracker.log(val_loss=float(summary["total_loss"]))

strategy.register_hook(SummaryLogger())

Per-batch validation callback

When you need more than the reduced summary — per-sample predictions, domain-level breakdowns, or a custom sink — configure a batch_callback on ValidationConfig. Any callable with the keyword-only signature (*, batch, predictions, loss, batch_count, step_count, epoch) satisfies the BatchValidationCallback protocol:

from nvalchemi.training import ValidationConfig

class ZarrBatchSink:
    def __init__(self, store):
        self._store = store

    def __call__(self, *, batch, predictions, loss, batch_count, step_count, epoch):
        group = self._store.require_group(f"step_{step_count}")
        group[f"batch_{batch_count}"] = predictions["energy"].cpu().numpy()

config = ValidationConfig(
    validation_data=val_data,
    batch_callback=ZarrBatchSink(my_zarr_store),
)

A plain function with the same keyword-only signature also satisfies the protocol.

Logging And Reporting

Logging and reporting are observer behavior, so — unlike the update hooks above — they belong in standard hooks rather than the update path. The only real design choice is the stage at which a logger runs: late enough that the data it needs already exists, but no later than necessary. The lifecycle offers a natural home for each kind of output:

• AFTER_LOSS for loss components and per-sample loss summaries,

• AFTER_BACKWARD for gradient diagnostics,

• AFTER_OPTIMIZER_STEP for learning rate, step status, EMA state, or any optimizer-step-dependent metric,

• AFTER_BATCH for generic counters, throughput, and final per-batch logging,

• AFTER_EPOCH for epoch summaries,

• AFTER_VALIDATION for reduced validation summaries.

Because the hook receives TrainContext, it can read counters, losses, models, optimizers, schedulers, the latest validation summary, and the owning workflow from whichever stage it picks. For a complete guide to writing hooks, see Hooks; for the built-in reporting stack, which uses exactly these stages to write Rich and TensorBoard output, see Reporting.

Checkpointing

A long run will eventually be interrupted — preemption, a crash, or a deliberate pause — and resuming it faithfully takes more than the latest weights. While this may sound straightforward to do with pickle, it is not recommended to do so for security (arbitrary code execution) and reproducibility (code changes): for these reasons, we designed the checkpointing workflow and abstraction heavily around making use of pydantic, to enable developers and researchers to make reloading/restarting training products as safely and turn-key as possible. We have tried to hide the pydantic abstraction as much as possible for checkpointing, so users do not need to be familiar with the framework.

A checkpoint captures four categories of state. Each has a developer-facing requirement for the loader to reconstruct and restore it correctly:

• Model weights and architecture: the model state_dict and the hyperparameters needed to reconstruct the model class. Models based on the ALCHEMI model base classes expose a spec automatically. A custom architecture not derived from BaseModelMixin must implement the spec protocol for its weights and config to be reloadable.

• Optimizer and scheduler state: the optimizer state_dict and scheduler construction parameters, handled automatically when using OptimizerConfig. Custom LossWeightSchedule instances used in composed losses must implement to_spec() for their state to be included in the checkpoint.

• Training counters: step_count, epoch_count, batch_count, and global_step_count are always saved and restored with no action required.

• Hook state: hooks that own restart-critical state — accumulated diagnostics, step-conditioned buffers, custom EMA weights — must implement CheckpointableHook. Hooks that do not implement the protocol are silently skipped; their state is neither saved nor restored. Logging hooks generally do not need this because their artifacts are already flushed to an external sink.

Note that training_fn and loss_target_assembler are not serialized (see the warning above). The checkpoint cannot reconstruct them; they must be supplied again at load time.

Use CheckpointHook to write checkpoints periodically from AFTER_BATCH or AFTER_EPOCH. Use TrainingStrategy.save_checkpoint(...) to save at an explicit point in a script. See Training checkpoints for strategy reconstruction, hook state, model specs, and distributed checkpoint behavior.

Restart semantics

There are two distinct restart scenarios, and the right API depends on which applies.

Resuming an interrupted run restores the full training state: model weights, optimizer state, scheduler state, training counters, and any checkpointable hook state. The run continues from the step immediately after the checkpoint. Supply the same hook objects the strategy was originally constructed with — the loader maps saved state into those live objects:

from nvalchemi.training import CheckpointHook, TrainingStrategy

strategy = TrainingStrategy.load_checkpoint(
    "runs/example/checkpoints/step_1000",
    # the checkpoint hook itself must be provided again for continuity
    # as it is stateless and is not kept with the checkpoint
    hooks=[CheckpointHook("runs/example/checkpoints"),],
)
strategy.run(train_loader)

Starting fresh from pretrained weights loads only model weights. Optimizer state, training counters, and hook state are not restored — the run starts at step zero with freshly built optimizers and schedulers. This is the right path for fine-tuning a pretrained model on a new dataset or task:

from nvalchemi.training import FineTuningStrategy, OptimizerConfig

strategy = FineTuningStrategy.from_pretrained_checkpoint(
    "runs/pretrained/checkpoints/final",
    loss_fn=loss_fn,
    optimizer_configs=OptimizerConfig(lr=1e-5),
    num_epochs=10,
)
strategy.run(finetune_loader)

The key difference: load_checkpoint resumes exactly where training stopped, counters and all. from_pretrained_checkpoint gives the model its learned weights but otherwise treats the run as new. See Fine-Tuning Pretrained Models for the full fine-tuning API, including parameter freezing and layer-wise learning-rate configuration.

Reproducibility

Everything above assumes a checkpoint can actually rebuild the run. That is not automatic — it is a property your code either has or lacks, and the failure mode is quiet: a run that trains happily for a week and cannot be resumed.

The mechanics are a cross-cutting feature of the toolkit and are documented once in Serialization and Reproducibility: objects are persisted as JSON recipes (an importable path plus constructor keyword arguments) rather than pickles, and rebuilt by importing the target and calling it again. This section covers what that demands of a training run specifically.

A training run is reproducible when five things hold:

1. Every model can produce a spec. Models built on BaseModelMixin do this automatically, by matching __init__ parameters against attributes of the same name. A model that stores a constructor argument under a different attribute name loses it silently, so the safe habit is self.<name> = <name>. When the constructor transforms its arguments, implement checkpoint_spec() and return the spec explicitly.

2. training_fn and loss_target_assembler are importable. They are recorded as dotted paths and never as code (see the warning in Configuring a training strategy); lambdas, closures, and locally-defined functions are rejected. Keep them versioned alongside the checkpoints they belong to, and pass them again at load time.

3. Custom loss weight schedules implement to_spec(). The built-in schedules already do. A schedule that does not is dropped from the recipe.

4. Hooks owning restart-critical state implement CheckpointableHook. Hooks are runtime objects supplied at load time; only checkpointable ones have their state restored into the instances you provide. Logging hooks generally need nothing, since their output already lives in an external sink.

5. Every model can actually be rebuilt from its spec. This one is enforced rather than merely encouraged: when automatic spec derivation fails, the framework warns (Omitting model spec for '<name>') and save_checkpoint() then refuses to write anything at all, raising ValueError: Cannot save strategy checkpoint because model spec generation failed. Since checkpointing usually happens well into a run, it is worth proving the round trip before launching one — see Serialization and Reproducibility.

Tip

The configuration alone — with no tensors — round-trips through TrainingStrategy.to_spec_dict() and from_spec_dict(). Dumping that JSON before a long run gives you a reviewable, diffable, version-controllable record of the experiment, and is the same format the CLI reads.

Two caveats worth stating plainly. First, your training script is part of the reproducible artifact: the checkpoint stores data and references, and the script supplies the code they point at. Second, nvalchemi does not seed RNGs for you — faithful recipe reproduction is not the same as bitwise-identical results, which additionally requires seeding and the usual CUDA determinism caveats.

Training CLI

The nvalchemi-training CLI offers a structured path to launching training experiments — scaffold a JSON configuration, validate it, then execute — without requiring a Python script. The train command group handles training from scratch using the same configuration format as fine-tuning workflows, so the JSON representations of hooks, optimizers, and loss functions map one-to-one onto the Python API described in the rest of this page.

Starting a training run

Initialize a spec scaffold for your training job:

nvalchemi-training train init \
    --dataset data/train.zarr \
    --output-dir runs/my-model \
    --lr 1e-4 \
    --num-steps 5000 \
    --out train.json

Multiple --dataset flags create a MultiDataset-backed dataloader across the provided paths:

nvalchemi-training train init \
    --dataset data/domain-a.zarr \
    --dataset data/domain-b.zarr \
    --output-dir runs/my-model \
    --out train.json

The generated train.json includes a fully populated optimizer, loss function, and output configuration. However, strategy.model_specs is left empty — training from scratch requires you to supply the architecture. Fill it in using the same BaseSpec JSON format used by checkpoints (cls_path plus constructor keyword fields):

{
  "strategy": {
    "model_specs": {
      "main": {
        "cls_path": "your.package.ModelClass",
        "hidden_size": 128,
        "num_layers": 4
      }
    }
  }
}

Use jq to merge a separately authored model spec into the base scaffold:

jq -s 'add' train.json model_spec.json > train_merged.json

Validating and configuring hooks

Before allocating compute, validate the spec and review the training intent:

nvalchemi-training spec report train.json

The report renders an optimizer summary, loss function configuration, hook list, and a learning-rate preview curve. It also surfaces warnings for common configuration mistakes: an empty model spec, a missing checkpoint hook, no validation dataset, or a learning rate that is out of the expected range. Resolve warnings before executing.

Hooks use the same JSON spec format as the Python API objects: a cls_path and constructor kwargs serialized with BaseSpec. A graph-based model typically needs a neighbor list and a checkpoint hook at minimum:

{
  "source": {
    "hooks": [
      {
        "spec": {
          "cls_path": "nvalchemi.hooks.neighbor_list.NeighborListHook",
          "cutoff": 6.0,
          "format": "coo"
        },
        "stages": ["BEFORE_FORWARD"]
      },
      {
        "spec": {
          "cls_path": "nvalchemi.training.hooks.checkpoint.CheckpointHook",
          "checkpoint_dir": "runs/my-model/checkpoints",
          "step_interval": 500
        }
      }
    ]
  }
}

The stages list specifies which TrainingStage values the hook fires at, corresponding to the stages described in Batches: Forward, Loss, Backward, Update. Omit stages to use the hook’s constructor default.

Executing and resuming

Once the report shows no critical warnings, launch the run:

nvalchemi-training spec run train.json

For distributed training, wrap with torchrun:

torchrun --nproc_per_node=4 -m nvalchemi.training.cli \
    spec run train.json --distributed

If the run is interrupted, resume from the latest checkpoint:

nvalchemi-training spec resume runs/my-model/checkpoints --spec train.json

spec resume restores model weights, optimizer state, training counters, and hook state, then continues training exactly where it stopped — the same behavior as calling TrainingStrategy.load_checkpoint(...) directly from the Python API (described in Restart semantics). Specify --checkpoint-index N to resume from a particular checkpoint instead of the latest.