Inference and training times

This tutorial shows how to measure inference and training times for a saved FlashCart model on a fixed atomic structure. It adapts the CUDA-event timing procedure used for the FlashCart paper to a checkpoint and structure of your choice.

Inference includes energy and force evaluation. The training measurement includes prediction, loss evaluation, and differentiation with respect to the model parameters. The graph is built before timing, and both the graph and model are transferred to the GPU. Data loading, neighbor-list construction, device transfers, and optimizer updates are excluded.

Run the example

The example requires an NVIDIA GPU, a checkpoint in the run’s best/ directory, and a structure containing elements supported by the checkpoint. Training measurements also require the run’s config.yaml. You can reuse the run and test.extxyz from Quickstart. The example uses models trained with energies in eV and positions in Å. Its synthetic training targets support loss functions on energies and forces.

Download measure.py into the directory containing your training run and structure file, then run:

CUDA_VISIBLE_DEVICES=0 python measure.py models/my_run test.extxyz

The script is also available at examples/timing/measure.py in the FlashCart repository. Paths are resolved from the directory where you run the command.

Replace 0 with the GPU index for your run. The script reads the first structure in the file, loads models/my_run/best, and measures both inference and training. Use --phase inference or --phase training to measure only one phase. The training loss comes from models/my_run/config.yaml.

The script uses float32 with TF32 disabled and keeps the checkpoint’s tensor backend. Predictions are evaluated without torch.compile by default. To measure compiled prediction, run the script again with:

CUDA_VISIBLE_DEVICES=0 python measure.py models/my_run test.extxyz \
    --compile-mode reduce-overhead

Compilation applies to prediction. The training measurement also includes the loss and its backward pass. Each phase has its own warmup. See Kernel generation and compilation for the distinction between generated Triton kernels and torch.compile. Use --compile-mode default to match the setting described under Compiled evaluation in LAMMPS interface. The LAMMPS interface does not support CUDA-graph modes. The inference measurement excludes neighbor-list construction, communication, and integration, which also contribute to the time required for a complete simulation.

What is timed

The same graph is reused for every evaluation. Each measured call performs the following work:

Phase	Computation
Inference	Evaluate the energy and differentiate it to obtain forces.
Training	Clear parameter gradients, evaluate the properties required by the loss, compute the loss, and differentiate it with respect to the model parameters.

Inference disables parameter gradients but keeps autograd enabled to compute forces. Training keeps the derivative graph with create_graph=True so that the force loss can be differentiated with respect to model parameters. The weights remain unchanged because no optimizer update is performed.

Before timing training, the script creates synthetic targets from the model’s predictions plus fixed Gaussian noise. The standard deviation is \(10^{-3}N\) eV for the total energy of an \(N\)-atom structure and \(10^{-2}\) eV/Å for each force component. The random seed is zero. These targets allow the loss and its parameter gradients to be computed without reference calculations. They are used only for timing, not for assessing prediction accuracy. They replace any reference labels only in the in-memory graph.

How time is measured

By default, each phase runs 10 untimed warmup evaluations followed by 50 measured evaluations. CUDA events surround each measured call. Synchronization before and after the call ensures that earlier GPU work has finished and that the measurement is complete before its elapsed time is read.

Warmup allows initial compilation and autotuning to take place before timing. Increase --warmup if initial work still affects the measurements, and increase --repeat to collect more samples.

The script reports the mean and sample standard deviation in milliseconds per evaluation and microseconds per atom. Time per atom is the measured time divided by the number of atoms in the structure. The standard deviation describes variation between repeated evaluations of this graph, not variation between training runs.

The script also prints the checkpoint, GPU, software versions, tensor backend, compilation setting, cutoff, and numbers of atoms and edges. Keep these settings with the measurements. Time per atom depends on system size and the number of neighbors. Use the same structure when comparing models.

Timing script

"""Measure CUDA inference and training computations on one atomic structure.

    python examples/timing/measure.py models/my_run test.extxyz

The example uses the best checkpoint and synthetic training targets. Timings exclude
structure reading, neighbor-list construction, device transfers, and optimizer updates.
"""

import argparse
import os
from importlib.metadata import version
from pathlib import Path
from statistics import mean, stdev

os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
os.environ["TORCH_ALLOW_TF32_CUBLAS_OVERRIDE"] = "0"

import torch
import yaml
from ase.io import read

from flashcart.data.graph import graph_from_ase
from flashcart.model.flashcart import FlashCartPotential
from flashcart.training.loss import loss_from_config


def prepare_step(model, graph, phase, train_loss=None, compile_mode=None):
    """Prepare a repeated prediction or training computation.

    Training targets are detached predictions plus seeded Gaussian noise, with
    standard deviations of 0.001 eV per atom for energy and 0.01 eV/Å for forces.
    The supplied model and graph are updated in place.

    Args:
        model (FlashCartPotential): Model already on the selected device.
        graph (AtomicData): Graph already on the same device as the model.
        phase (str): ``"inference"`` or ``"training"``.
        train_loss (LossFunction, optional): Required for training. Only energy and
            force losses are supported. Default: None.
        compile_mode (str, optional): Prediction compilation mode. Default: None,
            which leaves prediction uncompiled.

    Returns:
        Callable: No-argument function returning predictions or a scalar training
            loss. Training clears parameter gradients before each computation.
    """
    if phase not in ("inference", "training"):
        raise ValueError("phase must be 'inference' or 'training'.")
    training = phase == "training"
    if training:
        if train_loss is None:
            raise ValueError("Training requires a loss function.")
        keys = train_loss.compute_keys
        if keys["compute_stress"] or keys["compute_virials"]:
            raise ValueError("This example supports energy and force losses, not stress or virial losses.")
    model.train(training)
    model.requires_grad_(training)
    model.zero_grad(set_to_none=True)
    options = {
        "compute_forces": train_loss.compute_keys["compute_forces"] if training else True,
        "use_compile": compile_mode is not None,
        "compile_mode": compile_mode or "reduce-overhead",
    }

    if training:
        with torch.enable_grad():
            reference = model.predict(graph, compute_forces=options["compute_forces"])
        generator = torch.Generator(device=graph.positions.device).manual_seed(0)
        energy = reference["energy"].detach()
        graph.energy = energy + 0.001 * graph.n_atoms.to(energy) * torch.randn(
            energy.shape, generator=generator, device=energy.device, dtype=energy.dtype
        )
        if options["compute_forces"]:
            forces = reference["forces"].detach()
            graph.forces = forces + 0.01 * torch.randn(
                forces.shape, generator=generator, device=forces.device, dtype=forces.dtype
            )
        del reference

        def step():
            model.zero_grad(set_to_none=True)
            prediction = model.predict(graph, create_graph=True, **options)
            loss = train_loss.training_loss(prediction, graph, world_size=1)
            loss.backward()
            return loss

    else:

        def step():
            return model.predict(graph, **options)

    return step


def measure(step, device, warmup=10, repeat=50):
    """Measure individual calls with CUDA events after untimed warmup.

    Args:
        step (Callable): No-argument function performing the measured computation.
        device (torch.device): CUDA device used by the computation.
        warmup (int, optional): Number of untimed calls. Default: 10.
        repeat (int, optional): Number of measured calls. Default: 50.

    Returns:
        list[float]: Elapsed CUDA-event time in milliseconds for each measured call.
    """
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    times = []
    with torch.cuda.device(device), torch.enable_grad():
        for _ in range(warmup):
            output = step()
            del output
        torch.cuda.synchronize(device)
        for _ in range(repeat):
            torch.cuda.synchronize(device)
            start.record()
            output = step()
            end.record()
            torch.cuda.synchronize(device)
            times.append(start.elapsed_time(end))
            del output
    return times


def main():
    """Load a checkpoint and report inference or training timings on CUDA."""
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("run_dir", type=Path, help="training output directory containing best/ and config.yaml")
    parser.add_argument("structure", type=Path, help="structure file readable by ASE (only the first frame is used)")
    parser.add_argument("--phase", choices=("inference", "training", "both"), default="both")
    parser.add_argument("--warmup", type=int, default=10)
    parser.add_argument("--repeat", type=int, default=50)
    parser.add_argument(
        "--compile-mode",
        choices=("default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"),
        default=None,
        help="enable torch.compile for prediction using this mode",
    )
    args = parser.parse_args()
    if args.warmup < 0 or args.repeat < 2:
        parser.error("--warmup must be nonnegative and --repeat must be at least two.")
    if not torch.cuda.is_available():
        parser.error("CUDA is required for this example. Select a GPU with CUDA_VISIBLE_DEVICES.")

    torch.manual_seed(0)
    torch.set_default_dtype(torch.float32)
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False
    torch.set_float32_matmul_precision("highest")
    device = torch.device("cuda:0")
    checkpoint = args.run_dir / "best"
    model = FlashCartPotential.from_checkpoint(checkpoint, device=device).float()
    atoms = read(args.structure, index=0)
    graph = graph_from_ase(atoms, model.elements, model.r_max, device=device)
    if len(atoms) == 0 or graph.edge_index.shape[1] == 0:
        parser.error("The structure must contain atoms and neighbors within the model cutoff.")
    train_loss = None
    if args.phase != "inference":
        config = yaml.safe_load((args.run_dir / "config.yaml").read_text())
        train_loss = loss_from_config(config["train_loss"])
        if train_loss.compute_keys["compute_stress"] or train_loss.compute_keys["compute_virials"]:
            parser.error("This example supports energy and force losses, not stress or virial losses.")

    print(f"checkpoint: {checkpoint.resolve()}")
    print(f"structure: {args.structure.resolve()} (first frame)")
    print(f"GPU: {torch.cuda.get_device_name(device)}")
    print(f"software: FlashCart {version('flashcart')}, PyTorch {torch.__version__}, CUDA {torch.version.cuda}")
    print(f"dtype: {next(model.parameters()).dtype}, TF32: disabled")
    print(f"backend: {'Triton' if model.use_triton else 'PyTorch'}, compile: {args.compile_mode or 'disabled'}")
    print(f"system: {len(atoms)} atoms, {graph.edge_index.shape[1]} edges, cutoff {model.r_max:g} Å")
    print(f"calls: {args.warmup} warmup, {args.repeat} measured per phase")
    phases = ("inference", "training") if args.phase == "both" else (args.phase,)
    for phase in phases:
        step = prepare_step(model, graph, phase, train_loss, args.compile_mode)
        samples = measure(step, device, args.warmup, args.repeat)
        average = mean(samples)
        deviation = stdev(samples)
        print(
            f"{phase}: {average:.6f} ± {deviation:.6f} ms/call, "
            f"{average * 1000 / len(atoms):.6f} ± {deviation * 1000 / len(atoms):.6f} µs/atom "
            "(mean ± sample standard deviation)"
        )


if __name__ == "__main__":
    main()

The diamond benchmark

The model benchmarks in the FlashCart paper use periodic diamond cells with a lattice constant of 3.567 Å. For a checkpoint that supports carbon, prepare a 1000-atom cell with:

from ase.build import bulk
from ase.io import write

atoms = bulk("C", "diamond", a=3.567, cubic=True).repeat((5, 5, 5))
write("diamond.extxyz", atoms)

Pass diamond.extxyz in place of test.extxyz. A (4, 4, 4) repeat produces 512 atoms, the size used for the training comparison between tensor backends. Other benchmarks use different system sizes and compilation settings. To reproduce a published result, also match its checkpoint, backend, hardware, and software versions.

Profiling dataset batches

For operation tables and traces, use Performance and profiling. The flashcart-profile command initializes a new model and measures elapsed wall time over selected dataset batches. This tutorial instead loads a saved checkpoint and measures individual evaluations of one graph with CUDA events. Both exclude optimizer updates, but their reported times describe different workloads and measurement procedures.