Repository: ACEsuit/mace (language: python)

Issue:
Performance and scalability benchmarks for MACE
It would be useful to have some repeatable performance benchmarks showing the expected runtime and scalability of the MACE architecture.  The idea would be to use the PyTorch implementation available in this repo to run some common use cases for molecular dynamics simulations (e.g. inference with a single system).  Ideally the time per evaluation would be converted into an upper bound for the simulation time per day (e.g. ns / day units).  Some possibly interesting axes of investigation:

* Hardware variants: biased toward recent GPUs but still would be interesting to have some coverage from systems available to the MACE community. 
* Data Types: initially comparing float64 to float32
* compilation with PyTorch compared to normal eager mode

It seems reasonable to measure the performance directly from PyTorch as real production MD codes would ideally have minimal overhead over the evaluation through PyTorch.


Discussion / hints:
Indeed, I think we do have some prototype of that on the branch (https://github.com/ACEsuit/mace/blob/high_perf/mace/tools/cuda_tools.py) with the cuda kernels. It is important for the benchmark to include a wide variety of sizes/density to really get an idea (this one is not doing).

Fix patch:
diff --git a/.gitignore b/.gitignore
index 3817d9f3..296776e4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -33,3 +33,5 @@ dist/
 *.xyz
 /checkpoints
 *.model
+
+.benchmarks
diff --git a/mace/tools/torch_tools.py b/mace/tools/torch_tools.py
index e42a74f8..31e837df 100644
--- a/mace/tools/torch_tools.py
+++ b/mace/tools/torch_tools.py
@@ -6,7 +6,7 @@
 
 import logging
 from contextlib import contextmanager
-from typing import Dict
+from typing import Dict, Union
 
 import numpy as np
 import torch
@@ -129,13 +129,18 @@ def init_wandb(project: str, entity: str, name: str, config: dict, directory: st
 
 
 @contextmanager
-def default_dtype(dtype: torch.dtype):
+def default_dtype(dtype: Union[torch.dtype, str]):
     """Context manager for configuring the default_dtype used by torch
 
     Args:
-        dtype (torch.dtype): the default dtype to use within this context manager
+        dtype (torch.dtype|str): the default dtype to use within this context manager
     """
     init = torch.get_default_dtype()
-    torch.set_default_dtype(dtype)
+    if isinstance(dtype, str):
+        set_default_dtype(dtype)
+    else:
+        torch.set_default_dtype(dtype)
+
     yield
+
     torch.set_default_dtype(init)


Test patch:
diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py
new file mode 100644
index 00000000..78b04ccd
--- /dev/null
+++ b/tests/test_benchmark.py
@@ -0,0 +1,122 @@
+import os
+from typing import Optional
+
+import pandas as pd
+import json
+import pytest
+import torch
+from ase import build
+
+from mace import data
+from mace.calculators.foundations_models import mace_mp
+from mace.tools import AtomicNumberTable, torch_geometric, torch_tools
+
+
+def is_mace_full_bench():
+    return os.environ.get("MACE_FULL_BENCH", "0") == "1"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda is not available")
+@pytest.mark.benchmark(warmup=True, warmup_iterations=4, min_rounds=8)
+@pytest.mark.parametrize("size", (3, 5, 7, 9))
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+@pytest.mark.parametrize("compile_mode", [None, "default"])
+def test_inference(
+    benchmark, size: int, dtype: str, compile_mode: Optional[str], device: str = "cuda"
+):
+    if not is_mace_full_bench() and compile_mode is not None:
+        pytest.skip("Skipping long running benchmark, set MACE_FULL_BENCH=1 to execute")
+
+    with torch_tools.default_dtype(dtype):
+        model = load_mace_mp_medium(dtype, compile_mode, device)
+        batch = create_batch(size, model, device)
+        log_bench_info(benchmark, dtype, compile_mode, batch)
+
+        def func():
+            torch.cuda.synchronize()
+            model(batch, training=compile_mode is not None, compute_force=True)
+
+        torch.cuda.empty_cache()
+        benchmark(func)
+
+
+def load_mace_mp_medium(dtype, compile_mode, device):
+    calc = mace_mp(
+        model="medium",
+        default_dtype=dtype,
+        device=device,
+        compile_mode=compile_mode,
+        fullgraph=False,
+    )
+    model = calc.models[0].to(device)
+    return model
+
+
+def create_batch(size: int, model: torch.nn.Module, device: str) -> dict:
+    cutoff = model.r_max.item()
+    z_table = AtomicNumberTable([int(z) for z in model.atomic_numbers])
+    atoms = build.bulk("C", "diamond", a=3.567, cubic=True)
+    atoms = atoms.repeat((size, size, size))
+    config = data.config_from_atoms(atoms)
+    dataset = [data.AtomicData.from_config(config, z_table=z_table, cutoff=cutoff)]
+    data_loader = torch_geometric.dataloader.DataLoader(
+        dataset=dataset,
+        batch_size=1,
+        shuffle=False,
+        drop_last=False,
+    )
+    batch = next(iter(data_loader))
+    batch.to(device)
+    return batch.to_dict()
+
+
+def log_bench_info(benchmark, dtype, compile_mode, batch):
+    benchmark.extra_info["num_atoms"] = int(batch["positions"].shape[0])
+    benchmark.extra_info["num_edges"] = int(batch["edge_index"].shape[1])
+    benchmark.extra_info["dtype"] = dtype
+    benchmark.extra_info["is_compiled"] = compile_mode is not None
+    benchmark.extra_info["device_name"] = torch.cuda.get_device_name()
+
+
+def read_bench_results(files: list[str]) -> pd.DataFrame:
+    def read(file):
+        with open(file, "r") as f:
+            data = json.load(f)
+
+        records = []
+        for bench in data["benchmarks"]:
+            record = {**bench["extra_info"], **bench["stats"]}
+            records.append(record)
+
+        df = pd.DataFrame(records)
+        df["ns/day (1 fs/step)"] = 0.086400 / df["median"]
+        df["Steps per day"] = df["ops"] * 86400
+        columns = [
+            "num_atoms",
+            "num_edges",
+            "dtype",
+            "is_compiled",
+            "device_name",
+            "median",
+            "Steps per day",
+            "ns/day (1 fs/step)",
+        ]
+        return df[columns]
+
+    return pd.concat([read(f) for f in files])
+
+
+if __name__ == "__main__":
+    # Print to stdout a csv of the benchmark metrics
+    import subprocess
+
+    result = subprocess.run(
+        ["pytest-benchmark", "list"], capture_output=True, text=True
+    )
+
+    if result.returncode != 0:
+        raise RuntimeError(f"Command failed with return code {result.returncode}")
+
+    files = result.stdout.strip().split("\n")
+    df = read_bench_results(files)
+    print(df.to_csv(index=False))