Inference Speed

class mlipaudit.benchmarks.inference_speed.inference_speed.InferenceSpeedBenchmark(force_field: ForceField | Calculator, data_input_dir: str | PathLike = './data', run_mode: RunMode | Literal['dev', 'fast', 'standard'] = RunMode.STANDARD)

Benchmark measuring model and MD throughput and how they scale with size.

For each structure in the size-stratified protein dataset it measures two complementary speeds: the model forward-pass time (energy + forces, engine-independent) and the MD step time (end-to-end, including neighbour lists, the integrator and the simulation engine). The gap between them reflects simulation overhead. The model is scored with a Hill function on its forward-pass time relative to a reference hardware cost curve, so that faster models score higher; the score is wall-clock based and only comparable across models run on the same hardware.

name

The unique benchmark name (inference_speed), which also determines the input data directory and the HuggingFace dataset archive name.

Type:

str

category

The benchmark category, used for grouping in the UI.

Type:

str

result_class

The InferenceSpeedResult type returned by analyze.

Type:

type[mlipaudit.benchmark.BenchmarkResult] | None

model_output_class

The InferenceSpeedModelOutput type.

Type:

type[mlipaudit.benchmark.ModelOutput] | None

required_elements

The element types present in the input files.

Type:

set[str] | None

__init__(force_field: ForceField | Calculator, data_input_dir: str | PathLike = './data', run_mode: RunMode | Literal['dev', 'fast', 'standard'] = RunMode.STANDARD) → None

Initializes the benchmark.

Parameters:

• force_field – The force field model to be benchmarked.

• data_input_dir – The local input data directory. Defaults to “./data”. If the subdirectory “{data_input_dir}/{benchmark_name}” exists, the benchmark expects the relevant data to be in there, otherwise it will download it from HuggingFace.

• run_mode – Whether to run the standard benchmark length, a faster version, or a very fast development version. Subclasses should ensure that when RunMode.DEV, their benchmark runs in a much shorter timeframe, by running on a reduced number of test cases, for instance. Implementing RunMode.FAST being different from RunMode.STANDARD is optional and only recommended for very long-running benchmarks. This argument can also be passed as a string “dev”, “fast”, or “standard”.

Raises:

• ChemicalElementsMissingError – If initialization is attempted with a force field that cannot perform inference on the required elements.

• ValueError – If force field type is not compatible.

run_model() → None

For each structure, time the model forward pass (model throughput) and run a short MD simulation on each supported backend (MD throughput). Every measurement fails independently.

The device is warmed up once up front so the first structure is not timed on a cold GPU (see DEVICE_WARMUP_SECONDS).

analyze() → InferenceSpeedResult

Aggregate the timings and compute the throughput score.

Returns:

An InferenceSpeedResult object.

Raises:

RuntimeError – If called before run_model().

class mlipaudit.benchmarks.inference_speed.inference_speed.InferenceSpeedResult(*, failed: bool = False, score: Annotated[float | None, Ge(ge=0), Le(le=1)] = None, structure_names: list[str], structures: list[InferenceSpeedStructureResult], graph_cutoff_angstrom: float | None = None)

Result object for the inference-speed benchmark.

structure_names

The names of the structures.

Type:

list[str]

structures

List of per structure results.

Type:

list[mlipaudit.benchmarks.inference_speed.inference_speed.InferenceSpeedStructureResult]

graph_cutoff_angstrom

The model’s graph (neighbour-list) cutoff in Angstrom, which influences neighbour count and therefore speed. None if unavailable (e.g. some external calculators do not expose it).

Type:

float | None

class mlipaudit.benchmarks.inference_speed.inference_speed.InferenceSpeedStructureResult(*, structure_name: str, num_atoms: Annotated[int, Gt(gt=0)], num_steps: Annotated[int, Gt(gt=0)], num_episodes: Annotated[int, Gt(gt=0)], timestep_fs: float | None = None, average_forward_time: Annotated[float, Ge(ge=0)] | None = None, forward_times: list[float] = [], md: dict[str, MDBackendResult] = {}, failed: bool = False)

Result object for a single structure.

structure_name

The structure name.

Type:

str

num_atoms

The number of atoms in the structure.

Type:

int

num_steps

The number of steps in the simulation.

Type:

int

num_episodes

The number of episodes in the simulation.

Type:

int

timestep_fs

The MD timestep in femtoseconds, used to convert step times into a throughput (ns/day).

Type:

float | None

average_forward_time

The average wall-clock time of a single model forward pass (energy + forces), excluding warm-up and the slowest FORWARD_TRIM_FRACTION of passes. This is the engine-independent model-throughput metric. None if the forward pass failed.

Type:

float | None

forward_times

The kept (trimmed) forward-pass durations, sorted ascending, used to quantify variance. Empty if unavailable.

Type:

list[float]

md

MD throughput per backend, keyed by backend name (ase/jax_md).

Type:

dict[str, mlipaudit.benchmarks.inference_speed.inference_speed.MDBackendResult]

failed

Whether the forward pass and all MD backends failed for this structure.

Type:

bool

class mlipaudit.benchmarks.inference_speed.inference_speed.MDBackendResult(*, average_step_time: Annotated[float, Ge(ge=0)] | None = None, step_time_samples: list[float] = [])

MD throughput for a single backend.

average_step_time

The mean MD step time (s), excluding compilation. None if the run failed.

Type:

float | None

step_time_samples

Per-chunk step times (s/step) between successive logger calls (compilation excluded), used for variance. Empty if unavailable.

Type:

list[float]

class mlipaudit.benchmarks.inference_speed.inference_speed.InferenceSpeedModelOutput(*, structure_names: list[str], forward_times: list[list[float]] = [], md_step_times: list[dict[str, list[float]]] = [])

Model output for the inference-speed benchmark.

structure_names

The names of the structures used.

Type:

list[str]

forward_times

A list, per structure, of the individual timed model forward-pass durations (excluding warm-up), untrimmed and in measurement order so that drift over the run stays visible. Empty for structures whose forward pass failed.

Type:

list[list[float]]

md_step_times

A list, per structure, of a mapping from MD backend name (ase/jax_md) to the per-chunk step times (seconds per step) measured between successive logger calls, with the compilation chunk excluded. Backends that failed or were not run are absent.

Type:

list[dict[str, list[float]]]