From c26f6db8151901bbacf69bb77382d1d3f2701684 Mon Sep 17 00:00:00 2001 From: Thomas Warford Date: Mon, 15 Sep 2025 20:33:52 +0000 Subject: [PATCH 1/2] Add cueq parameterization --- tests/test_benchmark.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index fae104b41..674eae8c3 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -21,17 +21,18 @@ def is_mace_full_bench(): @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("enable_cueq", [False, True]) @pytest.mark.parametrize("compile_mode", [None, "default"]) def test_inference( - benchmark, size: int, dtype: str, compile_mode: Optional[str], device: str = "cuda" + benchmark, size: int, dtype: str, enable_cueq: bool, 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) + model = load_mace_mp_medium(dtype, enable_cueq, compile_mode, device) batch = create_batch(size, model, device) - log_bench_info(benchmark, dtype, compile_mode, batch) + log_bench_info(benchmark, dtype, enable_cueq, compile_mode, batch) def func(): torch.cuda.synchronize() @@ -41,12 +42,13 @@ def func(): benchmark(func) -def load_mace_mp_medium(dtype, compile_mode, device): +def load_mace_mp_medium(dtype, enable_cueq, compile_mode, device): calc = mace_mp( model="medium", default_dtype=dtype, - device=device, + enable_cueq=enable_cueq, compile_mode=compile_mode, + device=device, fullgraph=False, ) model = calc.models[0].to(device) @@ -71,10 +73,11 @@ def create_batch(size: int, model: torch.nn.Module, device: str) -> dict: return batch.to_dict() -def log_bench_info(benchmark, dtype, compile_mode, batch): +def log_bench_info(benchmark, dtype, enable_cueq, 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["cueq_enabled"] = enable_cueq benchmark.extra_info["is_compiled"] = compile_mode is not None benchmark.extra_info["device_name"] = torch.cuda.get_device_name() @@ -95,6 +98,7 @@ def process_benchmark_file(bench_file: Path) -> pd.DataFrame: "num_atoms", "num_edges", "dtype", + "cueq_enabled", "is_compiled", "device_name", "median", From 6e1e502b23357ab1fa24e51862adcaf7dac680c7 Mon Sep 17 00:00:00 2001 From: Thomas Warford Date: Mon, 15 Sep 2025 20:34:01 +0000 Subject: [PATCH 2/2] add header --- tests/test_benchmark.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 674eae8c3..2658267d3 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -1,3 +1,16 @@ +""" +Benchmark inference performance of the MACE medium model on GPU. + +Run benchmarks: + pytest tests/test_benchmark.py --benchmark-save= + +To also include torch.compile benchmarks: + MACE_FULL_BENCH=1 pytest tests/test_benchmark.py --benchmark-save= + +Convert results to CSV: + python tests/test_benchmark.py > results.csv +""" + import json import os from pathlib import Path