Performance comparison of Cuequivariance mode in MACE

[hspark1212]

Hi, I have a question about the speed performance of the MACE models with cuequivariance.

I compared the speed of enabling and disabling the cuequivariance mode by creating NaCl supercells. As shown in the graph, when the number of atoms in the unit-cell exceeds 100, the calculator with cuequivariance mode enabled (enable_cueq=True) exhibits slower speed.

Below is the code I used for testing.

from collections import defaultdict
import time
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

from ase.build import bulk
from mace.calculators import mace_mp

mace_calc_cueq = mace_mp(device="cuda", enable_cueq=True)
mace_calc = mace_mp(device="cuda")

# Create bulk structure
atoms = bulk("NaCl", "rocksalt", a=5.64)

# Prepare a defaultdict to hold results
results = defaultdict(list)

# Loop over increasing supercell sizes
for i in range(1, 10):
    # Create a supercell of size i x i x i
    test_atoms = atoms.repeat([i, i, i])
    n_atoms = len(test_atoms)

    # --- Time and calculate energy for mace_calc_cueq ---
    time.sleep(1)  # small pause to avoid GPU "warm-up" differences
    test_atoms_cueq = test_atoms.copy()
    test_atoms_cueq.calc = mace_calc_cueq
    start_time = time.time()
    energy_cueq = test_atoms_cueq.get_potential_energy()
    end_time = time.time()
    cueq_time = end_time - start_time

    # --- Time and calculate energy for mace_calc ---
    time.sleep(1)
    test_atoms_no_cueq = test_atoms.copy()
    test_atoms_no_cueq.calc = mace_calc
    start_time = time.time()
    energy_no_cueq = test_atoms_no_cueq.get_potential_energy()
    end_time = time.time()
    no_cueq_time = end_time - start_time

    # Store the results in long format
    results["calculator"].append("mace_calc_cueq")
    results["n_atoms"].append(n_atoms)
    results["time_s"].append(cueq_time)
    results["energy"].append(energy_cueq)

    results["calculator"].append("mace_calc")
    results["n_atoms"].append(n_atoms)
    results["time_s"].append(no_cueq_time)
    results["energy"].append(energy_no_cueq)

    # Print immediate comparison
    print(f"Supercell {i}x{i}x{i} -> {n_atoms} atoms")
    print(f"  mace_calc_cueq time: {cueq_time:.4f} s, energy: {energy_cueq:.6f}")
    print(f"  mace_calc time:      {no_cueq_time:.4f} s, energy: {energy_no_cueq:.6f}")


# Convert defaultdict to DataFrame
df = pd.DataFrame(results)

# Plot runtime comparison
sns.set_style("whitegrid")
plt.figure(figsize=(8, 5))
sns.lineplot(data=df, x="n_atoms", y="time_s", hue="calculator", marker="o")
plt.title("Comparison of Runtime: MACE Calculations with vs. without cuEq")
plt.xlabel("Number of Atoms")
plt.ylabel("Time (s)")
plt.tight_layout()
plt.show()

Before execution, I cleared the .cache/mace directory to download the model again, and I ran it using PyTorch 2.4.0. Additionally, I installed the following:

pip install cuequivariance-torch
pip install cuequivariance-ops-torch-cu11

[ilyes319]

@hspark1212 As mentioned in this thread: #740 timing models is delicate, and the way you are timing MACE is incorrect. To have a proper timing of the model only you can have a look at this script here: #740 (reply in thread).

[hspark1212]

Thank you for answering @ilyes319.

I took a look at the script in #740, and I noticed that the inference is being performed using model directly rather than the ASE calculator:

def run_inference():
        out = model(batch,training=True)
        torch.cuda.synchronize()
        return out

I'm wondering if I could compare the speed using MACE calculator as follows.

mace_calc_cueq = mace_mp(enable_cueq=True)
mace_calc_cueq.get_potential_energy(atoms)