MACE Training with ALCHEMI Training Utilities

This example walks through a complete model-training lifecycle on the ALCHEMI Toolkit, using a baseline ScaleShiftMACE model trained on the MatPES r2SCAN dataset as the reference workflow. The training script run configuration is loaded with Hydra from 10_vanilla_mace.yaml.

At a high level, the ALCHEMI training workflow has the following structure:

[Graph Data] -> [Model Architecture] -> [Supervised Objective] -> [Runtime Hooks] -> [TrainingStrategy]

Data — MatPES r2SCAN 2025.2 structures, obtained from MatPES, are read from ALCHEMI-compatible Zarr splits. Each sample contains graph inputs (positions, atom types, periodic boundary metadata) and supervised labels (energy, forces, stress).

Model — A 9.06M-parameter ScaleShiftMACE model from ACEsuit is wrapped with MACEWrapper so it can be used by TrainingStrategy. NVIDIA cuEquivariance kernels are enabled by default in the Hydra config (model.cueq.enabled: true).

Loss — Energies, forces, and stresses are fit with a weighted sum of Huber losses. PiecewiseWeight schedules are used to change the loss-term weights at a configured optimizer step for the second training stage.

Runtime — Distributed wrapping, EMA, neighbor-list rebuild, gradient clipping, metrics logging, and checkpointing are attached through runtime hooks rather than being implemented directly in the core trainin loop. Validation is configured separately using ValidationConfig on TrainingStrategy. Validation runs automatically during run().

Dataset-derived metadata (E0s, avg_num_neighbors, atomic_inter_shift / atomic_inter_scale), must be precomputed and set in cfg.model before training. The default YAML includes values computed from the MatPES r2SCAN training split.

/home/kelvin/Repos/nvalchemi-toolkit/.venv/lib/python3.13/site-packages/torch/jit/_script.py:1488: DeprecationWarning: `torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`.
  warnings.warn(
GPU information: NVIDIA GB10, 12, 1, 128, 48, 140, 1700

Loading train and validation data

The source data is the MatPES r2SCAN 2025.2 release on Hugging Face. This example expects separate train, validation, and test Zarr stores. Load each MatPES split file, for example the train, validation, and test JSON/JSONL files, and convert each split independently to an ALCHEMI Zarr store with AtomicDataZarrWriter.

A minimal converter for this workflow should map each pymatgen/MSON structure dictionary to atomic numbers, Cartesian positions, cell, and PBC tensors; write energy as a system label, forces as an atom label, and convert Voigt-6 stress to a 3 x 3 system tensor.

import periodictable as pt
import torch

from nvalchemi.data import AtomicData
from nvalchemi.data.atomic_data import voigt_to_matrix
from nvalchemi.data.datapipes import AtomicDataZarrWriter

def atomic_numbers_from_element_symbols(sites):
    return torch.as_tensor(
        [
            int(pt.elements.symbol(site["species"][0]["element"]).number)
            for site in sites
        ],
        dtype=torch.int32,
    )

writer = AtomicDataZarrWriter("r2scan-2025.2-train.zarr")
chunk_size = 8192
chunk = []
initialized = False
for record in jsonl_records:
    structure = record["structure"]
    chunk.append(
        AtomicData(
            atomic_numbers=atomic_numbers_from_element_symbols(structure["sites"]),
            positions=torch.as_tensor([site["xyz"] for site in structure["sites"]]),
            cell=torch.as_tensor(structure["lattice"]["matrix"]).reshape(1, 3, 3),
            pbc=torch.as_tensor(structure["lattice"].get("pbc", [True] * 3)).reshape(1, 3),
            energy=torch.as_tensor([[record["energy"]]]),
            forces=torch.as_tensor(record["forces"]),
            stress=voigt_to_matrix(torch.as_tensor(record["stress"])).reshape(1, 3, 3),
        )
    )

    if len(chunk) >= chunk_size:
        writer.append(chunk) if initialized else writer.write(chunk)
        initialized = True
        chunk.clear()

if chunk:
    writer.append(chunk) if initialized else writer.write(chunk)

This pipeline reads those Zarr splits with AtomicDataZarrReader. InMemoryDataset materializes each split once as a Batch on the target device, and DataLoader selects shuffled or sequential batches from that in-memory batch.

The default configuration uses a per-process training batch size of 256 and a validation batch size of 512. Given that the structure sizes in this dataset range from 1 atom to 240 atoms, SizeAwareSampler can also be used as an alternative to cap the atom count per batch when memory is tight.

from pathlib import Path

import torch
from nvalchemi.data.datapipes import AtomicDataZarrReader, DataLoader, InMemoryDataset

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

train_dataset = InMemoryDataset(
    reader=AtomicDataZarrReader(Path("/path/to/r2scan-2025.2-train.zarr")),
    device=device,
    skip_validation=True,
)

train_batches = DataLoader(
    train_dataset,
    batch_size=256,
    shuffle=True,
)

val_dataset = InMemoryDataset(
    reader=AtomicDataZarrReader(Path("/path/to/r2scan-2025.2-valid.zarr")),
    device=device,
    skip_validation=True,
)

val_batches = DataLoader(
    val_dataset,
    batch_size=512,
    shuffle=False,
)

The runnable script wraps this pattern in _loader(...) so Hydra can supply paths, batch sizes, and optional stress scaling transforms.

Building the MACE model

The default configuration trains ScaleShiftMACE to predict energy, force, and stress. Any model object passed to TrainingStrategy must follow BaseModelMixin. MACEWrapper handles input adaptation, neighbor-list metadata, and routes model outputs for MACE model variants.

Before building the model, populate the Hydra config with dataset-derived metadata: E0s (from structure-energy regression or isolated-atom DFT), avg_num_neighbors, and the ScaleShiftMACE pair atomic_inter_shift / atomic_inter_scale. The default YAML includes values precomputed from the training split.

import torch
from mace.modules import ScaleShiftMACE

from nvalchemi.models.mace import MACEWrapper

mace_model = ScaleShiftMACE(...)
model = MACEWrapper(mace_model.to(device=device, dtype=torch.float32))
model.model_config.active_outputs = {"energy", "forces", "stress"}

The runnable script reads architecture hyperparameters from Hydra and builds the wrapped model through _build_model(cfg, device), which calls build_training_mace_model() to set active_outputs and attach a checkpointable model spec.

Defining the loss

The default configuration fits energies, forces, and stresses. The loss is a weighted sum of Huber terms composed with + and * into a ComposedLossFunction. Stage-one weights hold until stage_two_start, then switch (for example 1/10/100 to 10/1/10 at step 54,400 of 68,000).

from nvalchemi.training import (
    ComposedLossFunction,
    EnergyHuberLoss,
    ForceHuberLoss,
    PiecewiseWeight,
    StressHuberLoss,
)

stage_two_start = 54_400

loss_fn: ComposedLossFunction = (
    PiecewiseWeight(
        boundaries=(stage_two_start,),
        values=(1.0, 10.0),
        per_epoch=False,
    )
    * EnergyHuberLoss(per_atom=True, delta=0.01)
    + PiecewiseWeight(
        boundaries=(stage_two_start,),
        values=(10.0, 1.0),
        per_epoch=False,
    )
    * ForceHuberLoss(normalize_by_atom_count=False, delta=0.01)
    + PiecewiseWeight(
        boundaries=(stage_two_start,),
        values=(100.0, 10.0),
        per_epoch=False,
    )
    * StressHuberLoss(delta=0.01)
)

loss_fn.normalize_weights = False

The runnable script builds the same composition from cfg.training.loss through _build_mace_huber_loss(cfg.training.loss).

Configuring the optimizer and scheduler

Schedulers are attached through OptimizerConfig. The runnable example uses TwoStageCosineConstantLR — cosine annealing for stage one, then a constant stage-two learning rate; any torch.optim.lr_scheduler.LRScheduler subclass can be passed via scheduler_cls and scheduler_kwargs.

import torch

from _mace_training_helpers import TwoStageCosineConstantLR
from nvalchemi.training import OptimizerConfig

optimizer_config = OptimizerConfig(
    optimizer_cls=torch.optim.AdamW,
    optimizer_kwargs={
        "lr": 5.0e-3,
    },
    scheduler_cls=TwoStageCosineConstantLR,
    scheduler_kwargs={
        "first_stage_steps": 54_400,
        "second_stage_lr": 1.0e-3,
        "eta_min": 1.0e-3,
    },
)

Hydra supplies learning-rate and schedule values; _optimizer(cfg) maps them onto OptimizerConfig.

Adding runtime hooks

Hooks extend the core training loop without embedding that logic in the loop itself. For example, DDPHook wraps the model in DDP at the TrainingStage SETUP stage when training.distributed.enabled is true. EMAHook maintains shadow weights for validation at AFTER_OPTIMIZER_STEP, and NeighborListHook rebuilds the interaction graph at BEFORE_FORWARD before every forward pass.

from pathlib import Path

from _mace_training_helpers import (
    GradientClipHook,
    TrainingMetricsLogger,
)
from nvalchemi.hooks import NeighborListHook
from nvalchemi.training import (
    CheckpointHook,
    DDPHook,
    EMAHook,
    TrainingStage,
)

hooks = [
    DDPHook(backend="nccl", sampler_kwargs={"seed": 42}),
    EMAHook(model_key="main", decay=0.995),
    GradientClipHook(max_norm=2.0),
    NeighborListHook(
        model.model_config.neighbor_config,
        max_neighbors=256,
        method="batch_naive_tile",
        stage=TrainingStage.BEFORE_FORWARD,
    ),
    TrainingMetricsLogger(every=100),
    CheckpointHook(
        checkpoint_dir=Path("outputs/checkpoints"),
        step_interval=10_000,
    ),
]

GradientClipHook and TrainingMetricsLogger are implemented in this example’s helper module. The other hooks shown above are public ALCHEMI training APIs. The runnable script assembles the full hook list from Hydra through _hooks(cfg, model).

Configuring validation

Validation is configured with ValidationConfig. The configuration specifies the validation data, validation function, loss function, evaluation cadence, and whether to use EMA weights. During run, the strategy evaluates validation at this cadence and once more at the end of training. The latest validation summary is stored on strategy.last_validation.

In multi-GPU runs, each rank evaluates a disjoint validation shard through a DistributedSampler. The runnable script builds this configuration with _build_validation_config(...) after the validation loader and loss function have been constructed.

from nvalchemi.training import ValidationConfig, default_training_fn

validation_config = ValidationConfig(
    validation_data=val_batches,
    validation_fn=default_training_fn,
    loss_fn=loss_fn,
    every_n_steps=1000,
    grad_mode="auto",
    use_ema="auto",
    name="validation",
)

Running TrainingStrategy

The final step is to assemble the objects created above and hand them to TrainingStrategy, which runs the training loop. On each step, it calls the training function, steps the optimizer and scheduler, invokes hooks at their registered stages, runs validation when configured, and tracks checkpointable training state.

DistributedManager provides distributed runtime information such as rank, local rank, world size, and device placement. The same code path works for single-GPU and multi-GPU launches. In distributed runs, DDPHook uses the distributed manager to wrap the model and coordinate rank-specific behavior.

from nvalchemi.distributed import DistributedManager
from nvalchemi.training import TrainingStrategy, default_training_fn

DistributedManager.initialize()
manager = DistributedManager()
device = torch.device(manager.device)

strategy = TrainingStrategy(
    models=model,
    optimizer_configs=optimizer_config,
    num_steps=68_000,
    training_fn=default_training_fn,
    loss_fn=loss_fn,
    devices=[device],
    distributed_manager=manager,
    hooks=hooks,
    validation_config=validation_config,
)

strategy.run(train_loader)

Run the Hydra entrypoint on one or more GPUs:

Single GPU:

uv run --extra cu12 --extra mace python examples/advanced/10_mace_training.py

Multi-GPU:

uv run --extra cu12 --extra mace torchrun --standalone --nproc_per_node=8 \
    examples/advanced/10_mace_training.py \
    --config-name=10_vanilla_mace distributed.enabled=true

The commands above use --extra cu12; change uv run --extra cuXX to match your NVIDIA driver / CUDA toolkit availability (for example, cu13 on CUDA 13.x).

Note that training.batch_size is the per-process batch size. The global batch size is therefore training.batch_size * nproc_per_node.

Skipping Hydra training during docs build. Run with:
uv run --extra cu12 --extra mace python examples/advanced/10_mace_training.py
(use --extra cu13 instead depending on your CUDA availability)

Validation curves and reference results

The figure below shows validation Huber losses from a full default-config run on 1× H100 GPU, which took about 80 minutes of wall time. Actual wall time may differ depending on system configuration, hardware, and software stack. The sharp transition near step 54,400 marks the stage-two loss-weight schedule configured by training.loss.stage_two.start_step.

With this default config (68,000 optimizer steps, about 50 epochs on the MatPES r2SCAN train set), the trained model reaches held-out test MAEs of energy 25.5 meV/atom, forces 145 meV/Å, and stress 0.703 GPa. These values are comparable to the MatPES r2SCAN benchmarks reported in the MatPES paper and to training with the MACE CLI.