Training configurations

This tutorial trains models using the configurations from the FlashCart paper and shows how to generate configurations for the five SPICE models, each with three random seeds. For the configuration format and available settings, see Configuration files.

The examples use files from the cloned FlashCart repository. Follow Installation and Data preparation before starting, then run the commands from the repository root.

Select the GPU before training. The following settings also disable TF32 to match the numerical settings used in the FlashCart paper:

export CUDA_VISIBLE_DEVICES=0
export NVIDIA_TF32_OVERRIDE=0
export TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=0

Replace 0 with the GPU index for your run. These settings apply to later commands in the same shell.

Train a SPICE model

examples/configs/spice.yaml defines a model with approximately 1.1 million parameters, three interactions, correlation order three, and 64 feature channels. Train and evaluate it with:

bash examples/configs/train_and_test.sh examples/configs/spice.yaml data/spice/test_subsets

The script trains one model, evaluates its best checkpoint on the configured test set, and then evaluates each supplied test subset. This configuration writes the run to models/i3c3f64l2_s0. The evaluation produces test.log and a separate log for each subset under test_subsets in that run directory.

# SPICE (MACE-OFF23): 1.1M-parameter model (i3c3f64l2) on the split with seed 0.
# Only keys that differ from flashcart/configs/default.yaml are listed.
# make_configs.py generates the five SPICE models with three seeds each.

# data (see examples/datasets/prepare_spice.py)
train_path: data/spice/train_0.extxyz
valid_path: data/spice/val_0.extxyz
test_path: data/spice/test.extxyz
elements: [H, C, N, O, F, P, S, Cl, Br, I]
train_batch_size: 128
eval_batch_size: 128
data_seed: 0
n_workers: 2

# output
output_path: models/i3c3f64l2_s0

# model
model_seed: 0
r_max: 6.0
n_hidden_feats: 64
l_max_hidden_feats: 2
n_interactions: 3
correlation: 3
path_reduced_product: false
gate_activation: silu
# isolated-atom energies (eV) of the SPICE reference, in `elements` order
atomic_shifts: [-13.571964772646918, -1030.5671648271828, -1486.3750255780376,
                -2043.933693071156, -2715.318528602957, -9287.407133426237,
                -10834.4844708122, -12522.649269035726, -70045.28385080204,
                -8102.524593409054]

# optimizer and schedule: 400 epochs (default), 20 linear warmup, then 380 cosine decay
weight_decay: 3.0e-4
end_factor: 1.0e-2
warmup_epochs: 20
decay_epochs: 380
decay_shape: cosine

# training
save_every: 5
ema_decay: 0.999
devices: 1
predict_compile: true

# energy MAE per atom + mean force-vector norm, equally weighted
train_loss:
  weighted_sum:
    - property: energy
      reduce: mae
      normalization: per_atom
      weight: 1.0
    - property: forces
      reduce: norm_mean
      normalization: none
      weight: 1.0

Training and evaluation script

#!/usr/bin/env bash
# Train one model, then evaluate it on its test set and on every test subset.
#
#   bash train_and_test.sh configs/i3c3f64l2_s0.yaml data/spice/test_subsets
#
# Writes <output_path>/test.log and <output_path>/test_subsets/<subset>.log.
set -euo pipefail

CONFIG=$1
SUBSET_DIR=${2:-}

# Disable TF32 for the float32 training protocol.
export NVIDIA_TF32_OVERRIDE=0
export TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=0

OUT=$(python -c "import sys; from flashcart.utils.config import load_config; print(load_config(sys.argv[1])['output_path'])" "$CONFIG")

flashcart-train "$CONFIG"
flashcart-test "$CONFIG"

if [ -n "$SUBSET_DIR" ]; then
    mkdir -p "$OUT/test_subsets"
    for subset in "$SUBSET_DIR"/*.extxyz; do
        name=$(basename "$subset" .extxyz)
        flashcart-test "$CONFIG" valid_path=null test_path="$subset" 2>&1 | tee "$OUT/test_subsets/$name.log"
    done
fi

Generate model configurations

make_configs.py generates the five SPICE model configurations reported in the FlashCart paper, each with seeds 0, 1, and 2:

Model	Configuration name before the seed suffix	GPUs
FlashCart-69k	i2c3f20l2r	1
FlashCart-0.5M	i3c5f32l2	1
FlashCart-1.1M	i3c3f64l2	1
FlashCart-3.0M	i3c5f88l2	1
FlashCart-5.6M	i4c5f104l2	2

The script updates the architecture, seeds, data paths, output directory, and GPU count for each run. Generate the files with:

python examples/configs/make_configs.py --out configs

This writes 15 configuration files without launching training. To train a model, pass one generated file to the training script. For example, this command selects the same architecture as above with seed 1:

bash examples/configs/train_and_test.sh configs/i3c3f64l2_s1.yaml data/spice/test_subsets

Each generated configuration has a distinct output directory. Repeat the command with the other configurations you want to evaluate.

FlashCart-5.6M requests two GPUs. Make two devices available with a setting such as CUDA_VISIBLE_DEVICES=0,1 before running the command. The global batch target is unchanged, but actual batches can differ under dynamic batching. See Batch sizes and data distribution for how batches are distributed across GPUs.

Configuration generation script

"""Generate configurations for the five SPICE models reported in the FlashCart paper.

    python examples/configs/make_configs.py --out configs

Generate the five SPICE models reported in the FlashCart paper, each with seeds
0, 1, and 2. This writes 15 files named <arch>_s<seed>.yaml. Train each with
``flashcart-train configs/<name>.yaml`` (or train_and_test.sh). Architecture
names read i<interactions>c<correlation>f<features>l<l_max_hidden>, with a
trailing r for the path-reduced product.
"""

import argparse
from pathlib import Path

from flashcart.utils.config import load_yaml, save_yaml

# (n_interactions, correlation, n_hidden_feats, l_max_hidden_feats, path_reduced_product, devices)
MODELS = [
    (2, 3, 20, 2, True, 1),  # FlashCart-69k
    (3, 5, 32, 2, False, 1),  # FlashCart-0.5M
    (3, 3, 64, 2, False, 1),  # FlashCart-1.1M
    (3, 5, 88, 2, False, 1),  # FlashCart-3.0M
    (4, 5, 104, 2, False, 2),  # FlashCart-5.6M
]
SEEDS = [0, 1, 2]


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--base", type=Path, default=Path(__file__).with_name("spice.yaml"))
    parser.add_argument("--out", type=Path, default=Path("configs"))
    parser.add_argument("--data", default="data/spice", help="directory written by prepare_spice.py")
    args = parser.parse_args()

    base = load_yaml(args.base)
    for i, c, f, l, reduced, devices in MODELS:
        arch = f"i{i}c{c}f{f}l{l}" + ("r" if reduced else "")
        for seed in SEEDS:
            name = f"{arch}_s{seed}"
            cfg = dict(base)
            cfg.update(
                train_path=f"{args.data}/train_{seed}.extxyz",
                valid_path=f"{args.data}/val_{seed}.extxyz",
                test_path=f"{args.data}/test.extxyz",
                output_path=f"models/{name}",
                data_seed=seed,
                model_seed=seed,
                n_interactions=i,
                correlation=c,
                n_hidden_feats=f,
                l_max_hidden_feats=l,
                path_reduced_product=reduced,
                devices=devices,
            )
            save_yaml(args.out / f"{name}.yaml", cfg)
    print(f"wrote {len(MODELS) * len(SEEDS)} configs to {args.out}")


if __name__ == "__main__":
    main()

Train a MAD model

examples/configs/mad.yaml uses a smaller cutoff and gives force errors a lower weight in the training loss function than the SPICE configuration. MAD already subtracts an isolated-atom baseline from its reference energies. The supplied configuration keeps atomic_shifts: ~, so FlashCart starts from zero prior offsets and fits the remaining per-element energy shifts.

Train and evaluate the three seeds, each with a separate output directory:

for seed in 0 1 2
do
    flashcart-train examples/configs/mad.yaml model_seed="$seed" data_seed="$seed" \
        output_path="models/mad_i3c3f96l2_s${seed}"
    flashcart-test examples/configs/mad.yaml model_seed="$seed" data_seed="$seed" \
        output_path="models/mad_i3c3f96l2_s${seed}"
done

MAD uses fixed dataset splits. Here, data_seed changes the training shuffle. The configuration uses one GPU, selected from the visible devices.

# MAD: model i3c3f96l2, seed 0.
# Only keys that differ from flashcart/configs/default.yaml are listed.

# data (see examples/datasets/prepare_mad.py)
train_path: data/mad/mad-train.extxyz
valid_path: data/mad/mad-val.extxyz
test_path: data/mad/mad-test.extxyz
elements:
  [H, He, Li, Be, B, C, N, O, F, Ne, Na, Mg, Al, Si, P, S, Cl, Ar, K, Ca,
   Sc, Ti, V, Cr, Mn, Fe, Co, Ni, Cu, Zn, Ga, Ge, As, Se, Br, Kr, Rb, Sr, Y,
   Zr, Nb, Mo, Tc, Ru, Rh, Pd, Ag, Cd, In, Sn, Sb, Te, I, Xe, Cs, Ba, La, Ce,
   Pr, Nd, Pm, Sm, Eu, Gd, Tb, Dy, Ho, Er, Tm, Yb, Lu, Hf, Ta, W, Re, Os, Ir,
   Pt, Au, Hg, Tl, Pb, Bi, Po, Rn]
data_seed: 0

# output
output_path: models/mad_i3c3f96l2_s0

# model (atomic_shifts stays ~: MAD energies already have the isolated atoms subtracted)
model_seed: 0
r_max: 5.0
n_hidden_feats: 96
n_interactions: 3
gate_activation: silu

# optimizer and schedule: 400 epochs (default), 20 linear warmup, then 380 cosine decay
lr: 2.5e-3
weight_decay: 1.2e-2
end_factor: 1.0e-2
warmup_epochs: 20
decay_epochs: 380
decay_shape: cosine

# training
save_every: 5
ema_decay: 0.999
devices: 1
predict_compile: true

# energy MAE per atom + mean force-vector norm, forces down-weighted
train_loss:
  weighted_sum:
    - property: energy
      reduce: mae
      normalization: per_atom
      weight: 1.0
    - property: forces
      reduce: norm_mean
      normalization: none
      weight: 0.1

Evaluate and compare models

To evaluate the seed-zero MAD model on the complete PET-MAD benchmark, run:

flashcart-test examples/configs/mad.yaml valid_path=null \
    test_path=data/mad/mad-bench.extxyz output_path=models/mad_i3c3f96l2_s0

This reports errors over the complete benchmark. To obtain results for individual benchmark datasets, evaluate the corresponding mad-bench-*.extxyz files in data/mad/subsets, passing one filename as test_path for each run. Use the matching output_path when evaluating another model seed.

The FlashCart paper reports SPICE errors averaged equally over the seven test subsets, followed by the mean and standard deviation across three training seeds.