From 092ee7806cf227257807549a589f39caddc73768 Mon Sep 17 00:00:00 2001 From: Bowen Deng Date: Mon, 26 Oct 2026 13:51:59 -0700 Subject: [PATCH 1/4] chore(venv): ship cuEquivariance with the mlip CUDA extras MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MACE's e3nn tensor-product kernels have a fused-CUDA replacement upstream ships as cuequivariance, and the warp-speed difference is large: 193 ms -> 32 ms per call for a 32-atom cell, with identical energies. Nothing in the repository asked for it, so every MACE run has been on the slow path. The packages ride in the torch extras because their relationship with the torch build is exact — the ops wheels are compiled against one CUDA major, and cuequivariance_torch imports whichever is installed under the bare module name cuequivariance_ops_torch. Upstream mace-torch declares only the cu12 pair; cu13 exists and is what a driver >= 580 host needs, so cu130 gets it too. About 30 MB per architecture. Installing is not enabling: MaceWrapper turns it on only on a CUDA device, and falls back to the plain path everywhere else. --- .agents/rules/mcp-environments.md | 2 +- venv/README.md | 2 +- venv/mlip/pyproject.toml | 14 +++- venv/mlip/uv.lock | 103 ++++++++++++++++++++++++++++++ 4 files changed, 117 insertions(+), 4 deletions(-) diff --git a/.agents/rules/mcp-environments.md b/.agents/rules/mcp-environments.md index cfef1693..52ad64e8 100644 --- a/.agents/rules/mcp-environments.md +++ b/.agents/rules/mcp-environments.md @@ -11,7 +11,7 @@ AtomisticSkills runs all skill scripts and MCP servers through the unified launc | Environment | Primary Packages | GPU | MCP Servers | | :--- | :--- | :--- | :--- | | `cpu` | ASE, pymatgen, RDKit, pycalphad, MDAnalysis (no torch) | No | `base`, `atomate2`, `drugdisc`, `smol` | -| `mlip` | PyTorch 2.14.1, MACE-torch 0.3.16, MatGL 4.1.2, nvalchemi-toolkit 0.2.0 | Yes | `mace`, `matgl` | +| `mlip` | PyTorch 2.14.1, MACE-torch 0.3.16 + cuEquivariance, MatGL 4.1.2, nvalchemi-toolkit 0.2.0 | Yes | `mace`, `matgl` | | `fairchem` | PyTorch 2.13.0, fairchem-core 2.24.0, nvalchemi-toolkit 0.2.0 | Yes | `fairchem` | ### Research Stacks diff --git a/venv/README.md b/venv/README.md index c8aa633c..b942cc23 100644 --- a/venv/README.md +++ b/venv/README.md @@ -51,7 +51,7 @@ Floors (not strict pins) keep the scientific packages of the shared projects cur | Package | mlip | fairchem | cpu | | :--- | :--- | :--- | :--- | | torch | 2.14.1 (CUDA 12.6 or 13) | 2.13.0 (CUDA 12.6 or 13) | — | -| mace-torch / e3nn | 0.3.16 / 0.4.4 | — | — | +| mace-torch / e3nn | 0.3.16 / 0.4.4 + cuEquivariance (cu126/cu130 extras) | — | — | | matgl | 4.1.2 (PyG only; DGL removed) | — | — | | fairchem-core / e3nn | — | 2.24.0 / 0.6.0 | — | | nvalchemi-toolkit | 0.2.0 | 0.2.0 | — | diff --git a/venv/mlip/pyproject.toml b/venv/mlip/pyproject.toml index d3d2328a..19013a13 100644 --- a/venv/mlip/pyproject.toml +++ b/venv/mlip/pyproject.toml @@ -137,8 +137,18 @@ lammps = ["cython", "mkl-devel ; platform_machine == 'x86_64'"] # torchvision (via timm, physicsnemo) is compiled against one torch build, so it # comes from the same index; a PyPI torchvision with a cu126 torch fails with # "operator torchvision::nms does not exist". -cu126 = ["torch", "torchvision"] -cu130 = ["torch", "torchvision"] +# cuEquivariance ships the same pair of packages for both CUDA majors, so it +# rides along with the torch build it must match: the ops wheels are compiled +# against one CUDA major, and cuequivariance_torch imports them under the bare +# module name cuequivariance_ops_torch either way. Upstream mace-torch declares +# only the cu12 pair (cueq-cuda-11/12); cu13 exists and is what a driver >= 580 +# host needs. Roughly 30 MB per architecture. +# +# Installing it is not the same as using it: MaceWrapper enables CuEquivariance +# only on a CUDA device where the kernels import, and falls back to the plain +# e3nn path otherwise (see _build_calculator). +cu126 = ["torch", "torchvision", "cuequivariance-torch>=0.2.0", "cuequivariance-ops-torch-cu12>=0.2.0"] +cu130 = ["torch", "torchvision", "cuequivariance-torch>=0.2.0", "cuequivariance-ops-torch-cu13>=0.2.0"] [dependency-groups] dev = ["pytest", "pytest-timeout"] diff --git a/venv/mlip/uv.lock b/venv/mlip/uv.lock index 6253698e..522378ca 100644 --- a/venv/mlip/uv.lock +++ b/venv/mlip/uv.lock @@ -345,10 +345,14 @@ dependencies = [ [package.optional-dependencies] cu126 = [ + { name = "cuequivariance-ops-torch-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "cuequivariance-torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "torch", version = "2.14.1+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "torchvision", version = "0.29.1+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cu130 = [ + { name = "cuequivariance-ops-torch-cu13", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "cuequivariance-torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "torch", version = "2.14.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "torchvision", version = "0.29.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] @@ -401,6 +405,10 @@ requires-dist = [ { name = "beautifulsoup4" }, { name = "biopython" }, { name = "cachetools", marker = "extra == 'openmm'" }, + { name = "cuequivariance-ops-torch-cu12", marker = "extra == 'cu126'", specifier = ">=0.2.0" }, + { name = "cuequivariance-ops-torch-cu13", marker = "extra == 'cu130'", specifier = ">=0.2.0" }, + { name = "cuequivariance-torch", marker = "extra == 'cu126'", specifier = ">=0.2.0" }, + { name = "cuequivariance-torch", marker = "extra == 'cu130'", specifier = ">=0.2.0" }, { name = "cvxpy" }, { name = "cython", marker = "extra == 'lammps'" }, { name = "dara-xrd" }, @@ -1123,6 +1131,92 @@ nvtx = [ { name = "nvidia-nvtx", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-25-atomisticskills-venv-mlip-cu126' and extra == 'extra-25-atomisticskills-venv-mlip-cu130') or (platform_machine == 'aarch64' and sys_platform == 'linux' and extra != 'extra-25-atomisticskills-venv-mlip-cu126') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-25-atomisticskills-venv-mlip-cu126') or (platform_machine == 'aarch64' and extra == 'extra-25-atomisticskills-venv-mlip-cu126' and extra == 'extra-25-atomisticskills-venv-mlip-cu130') or (platform_machine == 'x86_64' and extra == 'extra-25-atomisticskills-venv-mlip-cu126' and extra == 'extra-25-atomisticskills-venv-mlip-cu130') or (sys_platform != 'linux' and extra == 'extra-25-atomisticskills-venv-mlip-cu126' and extra == 'extra-25-atomisticskills-venv-mlip-cu130')" }, ] +[[package]] +name = "cuequivariance" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "networkx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "opt-einsum", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "scipy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "sympy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ee/91/c78169c11dfe89c8a200f9eab178b7f637dde39673be6815af3007fa8922/cuequivariance-0.13.0.tar.gz", hash = "sha256:4227705313bdd89650077c5b723416381e76add1bec4aa103a26347946f53efd", size = 219499, upload-time = "2026-10-15T19:53:14.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/e1/40915811975b95dc31dde5c227e3a1e4a6af264cf01b32cd2ea05d2f0e58/cuequivariance-0.13.0-py3-none-any.whl", hash = "sha256:72cc774e2f74bb2666c431fb6b5e6e612317474d38b9114240fc78585f8c1f81", size = 250801, upload-time = "2026-10-15T19:53:13.402Z" }, +] + +[[package]] +name = "cuequivariance-ops-cu12" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-ml-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "platformdirs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/a1/4ab8138424855d55c4e123132d60322987eb0f0177ff4b41dd72fe7ccb2a/cuequivariance_ops_cu12-0.13.0-py3-none-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7be711daa7313f37ed1bff288f4edd718d006ed84429ab62cde90843d639f55", size = 41111805, upload-time = "2026-10-15T19:22:07.153Z" }, + { url = "https://files.pythonhosted.org/packages/5b/30/db90ba304fc07b300c87ff1593acffcd39a17162f5d73359ba8d7171f425/cuequivariance_ops_cu12-0.13.0-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:54fcc31a63b0768946d6ccb83580796dec50de6c795a78022436ca9f47e1bdd2", size = 41168895, upload-time = "2026-10-15T19:22:11.354Z" }, +] + +[[package]] +name = "cuequivariance-ops-cu13" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-ml-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "platformdirs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tqdm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/29/ff/7d956b9bb5cee909f7ac2b653fcf69572bb0658a808b556d2970b16415bc/cuequivariance_ops_cu13-0.13.0-py3-none-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:62d267f500919231f964384762e4eddef2f5bd260608e3932cc3f95b84249fc0", size = 44889148, upload-time = "2026-10-15T19:22:15.7Z" }, + { url = "https://files.pythonhosted.org/packages/7b/07/f5319a6039b9ef5d05ee11125ad2378e4da099117fd35b24e24a787dc960/cuequivariance_ops_cu13-0.13.0-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:12b0d8b736bb195447946771a0092d69f9c597789a917d8609dcaa16220ffa34", size = 44958162, upload-time = "2026-10-15T19:22:20.411Z" }, +] + +[[package]] +name = "cuequivariance-ops-torch-cu12" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuequivariance-ops-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "scipy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", version = "2.14.1+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/0e/dd/c3cbedcaef043456154df7226117c72d1db3c2fa343d9702ab6510250728/cuequivariance_ops_torch_cu12-0.13.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5c4c42f16380639e4282725e204c72643350af3246b5335bb8d42fb5e7426f58", size = 1193257, upload-time = "2026-10-15T19:22:32.672Z" }, + { url = "https://files.pythonhosted.org/packages/f6/c2/bdc3ca4ac9560dca4b70aeaa5f39ef19a03f80a3ce9abcbd39ef2a9aa3ac/cuequivariance_ops_torch_cu12-0.13.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:909d49ac15a4c132b110a20ae755a9f956d2466a19c2f36e0274e03a9558784a", size = 1202213, upload-time = "2026-10-15T19:22:33.875Z" }, +] + +[[package]] +name = "cuequivariance-ops-torch-cu13" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuequivariance-ops-cu13", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "scipy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "torch", version = "2.14.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/bd/7740b6bd8ab680d3ba9edb7d6bc90bcee60aa2f4e2bff4ed4ef1dca53f33/cuequivariance_ops_torch_cu13-0.13.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:28121445776a9a268f55a5a38c9f48f76512cdd051300a7574d9c2eaf82d7507", size = 28512268, upload-time = "2026-10-15T19:23:25.913Z" }, + { url = "https://files.pythonhosted.org/packages/c4/8d/176e83eb4fb411e6b5ff5bda06e5a7799bc6369aea19112277319ced1f8a/cuequivariance_ops_torch_cu13-0.13.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:181e22ab2644b2ac0020aa0654ff5a1468b99d7d6409a8a82269dd6e8941ca09", size = 28530554, upload-time = "2026-10-15T19:23:33.065Z" }, +] + +[[package]] +name = "cuequivariance-torch" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuequivariance", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/96/fd/221132e3f204a5bfe4efbc531242bbbb84a568e8aa834648a371a47f969b/cuequivariance_torch-0.13.0.tar.gz", hash = "sha256:f5ffb513454d64cd7cf29915153ae029a27538ead7672c552dab0f005984141e", size = 170540, upload-time = "2026-10-15T19:53:20.036Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/1b/975e2b9703ead34b24c3b3226074255e0ef02b171ee8ae98921094a57783/cuequivariance_torch-0.13.0-py3-none-any.whl", hash = "sha256:f256fe9ab5c3a8f37e08345ba3c734d37a6ccd1a59b176479d61a4126c56acd0", size = 184761, upload-time = "2026-10-15T19:53:19.001Z" }, +] + [[package]] name = "custodian" version = "2025.12.14" @@ -3441,6 +3535,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/31/83/f3647ce26916c94a6ca4ff1810623e2c405cff2dea6e78d29516b2514df9/nvidia_cusparselt_cu13-0.8.1-py3-none-win_amd64.whl", hash = "sha256:dccbd362f91a7b9024d1f55ee9f548ac065027ff15d8c8b0db889ab3a8f31215", size = 156885108, upload-time = "2025-09-05T18:51:35.958Z" }, ] +[[package]] +name = "nvidia-ml-py" +version = "13.615.71" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/30/b25216758be3d3e2834825d8193609e2d71c770a8bd7984438c058c90268/nvidia_ml_py-13.615.71.tar.gz", hash = "sha256:bebe4e48f51b1dc75028c0815cb7bfa14a31a5bb80be70c9d980c6036953fc3d", size = 57485, upload-time = "2026-09-25T15:15:28.226Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/53/a1/1681dfa1c904d4e3e72e51b55a0ff012d50b766843ef832d584abe2113c6/nvidia_ml_py-13.615.71-py3-none-any.whl", hash = "sha256:959bf4adf6fe1308e4bd739e722236b0d1ec8392e2cefad33ff70c311380b9b6", size = 58132, upload-time = "2026-09-25T15:15:26.54Z" }, +] + [[package]] name = "nvidia-nccl-cu12" version = "2.29.3" From e715dd43f98ef0c52040dcdbbbb18aab6c644bfd Mon Sep 17 00:00:00 2001 From: Bowen Deng Date: Mon, 26 Oct 2026 13:52:03 -0700 Subject: [PATCH 2/4] feat(mace): use cuEquivariance where it works, and never only there MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every MACE calculator was built without an accelerator, so the fused-cueq path was never taken even where it was installed. MaceWrapper now routes all six calculator constructions through one builder that prefers cueq. The constraint that shapes the design is that an optimisation may not cost anyone the ability to run MACE. mace itself raises when cuequivariance is missing and, more subtly, when the installed kernels do not cover the card in use. So: - 'auto' (the default) considers cueq only on a CUDA device with both halves of cueqvariance importable; on a CPU host, in an environment that never synced the cu130 extra, the plain path is used and said so once. - An accelerated attempt that fails — an unsupported card, a missing kernel, a broken install — is logged with its reason and retried once without the accelerator, so MACE loads anyway. - The plain path is deliberately outside that guard: a bad checkpoint or an unreadable model still raises instead of being absorbed as 'no GPU'. - One failure is remembered, so a per-structure rebuild does not retry a conversion that cannot work. - 'none' pins the plain path for timings that can be compared across hosts, and an unknown value is rejected rather than silently ignored. load_model reports what the calculator will attempt without building it: the calculator is lazy, and building it would download and convert the checkpoint just to answer. Measured on one 32-atom Cu cell, MACE-OMAT-0-small, float32: 193 ms/call on the plain e3nn path against 32 ms with cueq, energies agreeing to every printed digit (both -3.744912 eV/atom). No GPU, cueq uninstalled, or the CPU path all give the same energy on the plain path. tests/mace/test_mace_accelerator.py pins the fallbacks with a fake factory, so it needs no GPU, no cueq install and no checkpoint: a simulated unsupported card degrades instead of raising, a genuine error on the plain path still surfaces, and a failure is not retried per structure. --- src/mcp_server/mace_server.py | 24 ++++- src/utils/mlips/mace/mace_wrapper.py | 156 +++++++++++++++++++++++++-- tests/mace/test_mace_accelerator.py | 143 ++++++++++++++++++++++++ 3 files changed, 312 insertions(+), 11 deletions(-) create mode 100644 tests/mace/test_mace_accelerator.py diff --git a/src/mcp_server/mace_server.py b/src/mcp_server/mace_server.py index ff464d01..0b6e5c4d 100644 --- a/src/mcp_server/mace_server.py +++ b/src/mcp_server/mace_server.py @@ -41,6 +41,7 @@ def load_model( model_name: str = "MACE-OMAT-0-small", device: str = "auto", task_name: Optional[str] = None, + accelerator: str = "auto", ) -> str: """ Load a MACE model. @@ -59,9 +60,15 @@ def load_model( task_name: Optional task name that sets the model's head. Supported options for multi-head models (MACE-MH): 'omat_pbe' (default), 'matpes_r2scan', 'omol', 'spice_wB97M', 'oc20_usemppbe'. + accelerator: 'auto' (default) uses CuEquivariance on a CUDA device when + it is installed, and runs the plain e3nn path otherwise — no + GPU, no cuequivariance, or a card its kernels do not cover + all still work, just slower. 'cueq' asks for it and degrades + the same way; 'none' pins the plain path for reproducible + timings. Returns: - Confirmation message. + Confirmation message naming the accelerator actually in use. CRITICAL: This tool must be called before using any other tool to load the model into memory. """ @@ -74,7 +81,12 @@ def load_model( try: from src.utils.mlips.mace.mace_wrapper import MACEWrapper - wrapper = MACEWrapper(model_name=model_name, device=device, head=task_name) + wrapper = MACEWrapper( + model_name=model_name, + device=device, + head=task_name, + accelerator=accelerator, + ) # Load the model - this might take a while and produce output wrapper.load(model_path=model_name if os.path.exists(model_name) else None) @@ -88,7 +100,13 @@ def load_model( f"Model wrapper does not support setting head to {task_name}" ) - return f"Successfully loaded MACE model: {model_name}" + # The calculator, and with it the accelerator, is built lazily on the + # first structure; report what it will attempt rather than forcing a + # build here, which would download and convert the checkpoint. + return ( + f"Successfully loaded MACE model: {model_name}. " + f"Acceleration: {wrapper.describe_accelerator()}." + ) except Exception as e: logger.error(f"Failed to load model {model_name}: {e}") return f"Error loading model: {str(e)}" diff --git a/src/utils/mlips/mace/mace_wrapper.py b/src/utils/mlips/mace/mace_wrapper.py index 39809f80..c49173c8 100644 --- a/src/utils/mlips/mace/mace_wrapper.py +++ b/src/utils/mlips/mace/mace_wrapper.py @@ -12,6 +12,32 @@ logger = logging.getLogger(__name__) +# How a calculator may be accelerated. CuEquivariance ("cueq") swaps MACE's +# e3nn tensor-product kernels for fused CUDA ones, which needs both its Python +# front end and its compiled ops package, and only helps on a GPU. +# +# "auto" is deliberately forgiving. A host with no GPU, no cuequivariance +# installed (the CPU-only environments, or a run that did not sync the cu130 +# extra), or a card that cuequivariance has no kernels for must still be able +# to run MACE, so an accelerator that cannot be used is a log line and a slower +# path, never an exception out of load(). Asking for "cueq" is a request, not a +# guarantee, for the same reason: it is honoured where it works. +ACCELERATORS = ("auto", "cueq", "none") +_CUEQ_MODULES = ("cuequivariance_torch", "cuequivariance_ops_torch") + + +def cueq_installed() -> bool: + """Whether both halves of CuEquivariance can be imported.""" + from importlib.util import find_spec + + return all(find_spec(name) is not None for name in _CUEQ_MODULES) + + +def _is_cuda(device: str) -> bool: + """Whether an ASE/``torch`` device string names a CUDA device.""" + return str(device).split(":")[0] == "cuda" + + # Available MACE checkpoints and models AVAILABLE_MACE_MODELS = { # MACE-MH (Multi-Head) models - MACE-MH-1 is the latest recommended default @@ -76,6 +102,7 @@ def __init__( model_version: str = "latest", device: str = "auto", head: Optional[str] = None, + accelerator: str = "auto", ): """ Initialize MACE wrapper. @@ -85,10 +112,26 @@ def __init__( model_version: Version of the model device: Device to run the model on ("auto", "cpu", "cuda") head: Optional head name for multi-head models (e.g. "omat_pbe") + accelerator: One of ``ACCELERATORS``. "auto" uses CuEquivariance + where it can help and quietly runs the plain e3nn path + elsewhere; "cueq" asks for it but still degrades; "none" pins + the plain path for reproducible timings. """ + if accelerator not in ACCELERATORS: + raise ValueError( + f"Unknown accelerator {accelerator!r}; expected one of " + f"{', '.join(ACCELERATORS)}." + ) super().__init__(model_name, model_version) self.device = device self.head = head + self.accelerator = accelerator + # What actually ended up on the calculator; set by _build_calculator. + self.accelerator_effective = "none" + self._accelerator_warned = False + # Set once the accelerated attempt has failed on this device, so a + # per-structure rebuild does not retry a conversion that cannot work. + self._accelerator_failed = False self.model_class = None self.calculator_class = None self.model = None @@ -131,6 +174,91 @@ def load(self, model_path: Optional[str] = None) -> None: self.load_checkpoint(model_path) logger.info(f"Loaded model from checkpoint: {model_path}") + def _accelerator_kwargs(self) -> Dict[str, Any]: + """ + Calculator kwargs for the accelerated path, or empty when it is out. + + Decides once per call and explains the decision the first time it is + a "no", so a user who expected the speedup can tell whether the host + lacks a GPU, the package, or a compatible card. + """ + if self.accelerator == "none" or self._accelerator_failed: + return {} + if not _is_cuda(self.device): + if not self._accelerator_warned: + logger.info( + "CuEquivariance needs a CUDA device; running %s on %s " + "without it.", + self.model_name, + self.device, + ) + self._accelerator_warned = True + return {} + if not cueq_installed(): + if not self._accelerator_warned: + logger.info( + "CuEquivariance is not installed; running %s without it. " + "It ships with the mlip cu126/cu130 extras.", + self.model_name, + ) + self._accelerator_warned = True + return {} + return {"enable_cueq": True} + + def describe_accelerator(self) -> str: + """ + One line on what the calculator will attempt, and why. + + Cheap and side-effect free: it answers from the device, the installed + packages and the requested mode, without building a calculator or + touching the model. Whether cueq then holds up on the card is only + known when the calculator is built, and _build_calculator logs it. + """ + if self.accelerator == "none": + return "none (disabled by request)" + if not _is_cuda(self.device): + return f"none (device is {self.device}, and cueq needs CUDA)" + if not cueq_installed(): + return "none (cuequivariance is not installed)" + return "cueq (falls back to the plain path if this card has no kernels)" + + def _build_calculator(self, factory, /, **kwargs): + """ + Build a calculator, preferring the accelerated path. + + MACE raises rather than degrading when cueqvariance is missing or has + no kernels for the card in use, and doing that here would take MACE + away from a user whose GPU merely is not one cuequivariance supports. + So the accelerated attempt is the only thing in the try: if it fails, + rebuild on the plain path and keep the reason. The plain path is not + guarded, so a real error (a bad checkpoint, an unreadable model) still + surfaces instead of being absorbed. + """ + extra = self._accelerator_kwargs() + if extra: + try: + calculator = factory(**kwargs, **extra) + except Exception as exc: + logger.warning( + "CuEquivariance could not run %s on %s (%s: %s); " + "continuing on the plain e3nn path, which is slower.", + self.model_name, + self.device, + type(exc).__name__, + exc, + ) + self.accelerator_effective = "none" + self._accelerator_failed = True + else: + self.accelerator_effective = "cueq" + logger.info( + "CuEquivariance acceleration enabled for %s.", self.model_name + ) + return calculator + else: + self.accelerator_effective = "none" + return factory(**kwargs) + def create_calculator(self) -> Calculator: """ Create an ASE calculator from the model. @@ -145,16 +273,22 @@ def create_calculator(self) -> Calculator: if self.model_path and Path(self.model_path).exists(): from mace.calculators import MACECalculator - return MACECalculator( - model_paths=str(self.model_path), device=self.device, head=self.head + return self._build_calculator( + MACECalculator, + model_paths=str(self.model_path), + device=self.device, + head=self.head, ) # 2. Check if model_name is a path (fallback) if Path(self.model_name).exists(): from mace.calculators import MACECalculator - return MACECalculator( - model_paths=str(self.model_name), device=self.device, head=self.head + return self._build_calculator( + MACECalculator, + model_paths=str(self.model_name), + device=self.device, + head=self.head, ) # 3. Standard Pretrained mapping @@ -199,7 +333,9 @@ def create_calculator(self) -> Calculator: ): head = "omat_pbe" - return MACECalculator(model_paths=model_path, device=self.device, head=head) + return self._build_calculator( + MACECalculator, model_paths=model_path, device=self.device, head=head + ) if "OFF" in model_name_upper: # Use MACE-OFF calculator @@ -211,25 +347,29 @@ def create_calculator(self) -> Calculator: "MACE-OFF23-LARGE": "large", } off_model_size = off_model_mapping.get(model_name_upper, "medium") - return mace_off(model=off_model_size, device=self.device) + return self._build_calculator( + mace_off, model=off_model_size, device=self.device + ) elif "ANI-CC" in model_name_upper: # Use MACE-ANI-CC from mace.calculators import mace_anicc + # mace_anicc takes a fixed argument list, so it cannot be handed + # enable_cueq the way the other helpers can. return mace_anicc(device=self.device) elif "OMOL" in model_name_upper: # Use MACE-OMOL from mace.calculators import mace_omol - return mace_omol(model=model_id, device=self.device) + return self._build_calculator(mace_omol, model=model_id, device=self.device) else: # Use MACE-MP calculator for materials models (including MPA, MATPES, OMAT) from mace.calculators import mace_mp - return mace_mp(model=model_id, device=self.device) + return self._build_calculator(mace_mp, model=model_id, device=self.device) def predict_atomic_features(self, structure_data: Any) -> Dict[str, Any]: """ diff --git a/tests/mace/test_mace_accelerator.py b/tests/mace/test_mace_accelerator.py new file mode 100644 index 00000000..ed6a4f84 --- /dev/null +++ b/tests/mace/test_mace_accelerator.py @@ -0,0 +1,143 @@ +""" +Tests for MACE accelerator selection and its fallbacks. + +CuEquivariance is an optimisation, never a requirement: a CPU-only host, an +environment that did not sync the cu130 extra, or a card its kernels do not +cover must all still be able to run MACE. These tests pin that behaviour with +a fake calculator factory, so they need no GPU, no cuequivariance install and +no checkpoint download. + +Run in: venv/run mlip python -m pytest tests/mace/test_mace_accelerator.py +""" + +import logging + +import pytest + +from src.utils.mlips.mace import mace_wrapper as mw + + +def fake_calculator(**kwargs): + """Stand-in for a MACE calculator that records how it was built.""" + return {"kwargs": kwargs} + + +@pytest.fixture +def wrapper_cuda(): + """A wrapper pointed at CUDA, without touching any hardware.""" + w = mw.MACEWrapper(model_name="MACE-MH-1", device="cuda") + w.is_loaded = True + return w + + +@pytest.mark.mace +class TestAcceleratorSelection: + def test_auto_on_cpu_never_requests_cueq(self): + w = mw.MACEWrapper(model_name="MACE-MH-1", device="cpu") + built = w._build_calculator(fake_calculator, device="cpu") + assert "enable_cueq" not in built["kwargs"] + assert w.accelerator_effective == "none" + + def test_auto_without_cueq_installed_degrades(self, wrapper_cuda, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: False) + built = wrapper_cuda._build_calculator(fake_calculator, device="cuda") + assert "enable_cueq" not in built["kwargs"] + assert wrapper_cuda.accelerator_effective == "none" + + def test_auto_with_cueq_requests_it(self, wrapper_cuda, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + built = wrapper_cuda._build_calculator(fake_calculator, device="cuda") + assert built["kwargs"]["enable_cueq"] is True + assert wrapper_cuda.accelerator_effective == "cueq" + + def test_none_never_attempts_even_when_installed(self, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + w = mw.MACEWrapper(model_name="MACE-MH-1", device="cuda", accelerator="none") + built = w._build_calculator(fake_calculator, device="cuda") + assert "enable_cueq" not in built["kwargs"] + assert w.accelerator_effective == "none" + + def test_explicit_cueq_request_still_degrades(self, wrapper_cuda, monkeypatch): + """Asking for cueq is a request, not a promise.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: False) + w = mw.MACEWrapper(model_name="MACE-MH-1", device="cuda", accelerator="cueq") + built = w._build_calculator(fake_calculator, device="cuda") + assert "enable_cueq" not in built["kwargs"] + assert w.accelerator_effective == "none" + + def test_unknown_accelerator_rejected(self): + with pytest.raises(ValueError, match="Unknown accelerator"): + mw.MACEWrapper(model_name="MACE-MH-1", accelerator="fastest") + + +@pytest.mark.mace +class TestAcceleratorFallback: + def test_incompatible_card_degrades_instead_of_raising( + self, wrapper_cuda, monkeypatch, caplog + ): + """A GPU cuequivariance has no kernels for must not take MACE down.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + + def factory(**kwargs): + if kwargs.get("enable_cueq"): + raise RuntimeError("CUDA error: no kernel image is available") + return fake_calculator(**kwargs) + + with caplog.at_level(logging.WARNING): + built = wrapper_cuda._build_calculator(factory, device="cuda") + + assert "enable_cueq" not in built["kwargs"] + assert wrapper_cuda.accelerator_effective == "none" + assert any("no kernel image" in r.getMessage() for r in caplog.records) + + def test_missing_module_degrades(self, wrapper_cuda, monkeypatch): + """mace raises ImportError when cuequivariance is absent; absorb it.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + + def factory(**kwargs): + if kwargs.get("enable_cueq"): + raise ImportError("cuequivariance is not installed") + return fake_calculator(**kwargs) + + built = wrapper_cuda._build_calculator(factory, device="cuda") + assert "enable_cueq" not in built["kwargs"] + + def test_real_error_on_the_plain_path_surfaces(self, wrapper_cuda, monkeypatch): + """The plain path is unguarded, so genuine bugs still raise.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + + def factory(**kwargs): + raise TypeError("model_paths must be a path") + + with pytest.raises(TypeError, match="model_paths"): + wrapper_cuda._build_calculator(factory, device="cuda") + + def test_fallback_is_remembered_for_later_calls(self, wrapper_cuda, monkeypatch): + """One failed attempt must not cost a rebuild on every structure.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + attempts = [] + + def factory(**kwargs): + attempts.append(kwargs.get("enable_cueq", False)) + if kwargs.get("enable_cueq"): + raise RuntimeError("no kernel image") + return fake_calculator(**kwargs) + + wrapper_cuda._build_calculator(factory, device="cuda") + wrapper_cuda._build_calculator(factory, device="cuda") + assert attempts == [True, False, False] + + +@pytest.mark.mace +class TestDescribeAccelerator: + def test_reports_cpu_as_none(self): + w = mw.MACEWrapper(model_name="MACE-MH-1", device="cpu") + assert "CPU" in w.describe_accelerator() or "cpu" in w.describe_accelerator() + + def test_reports_missing_package(self, wrapper_cuda, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: False) + assert "not installed" in wrapper_cuda.describe_accelerator() + + def test_reports_cueq_available(self, wrapper_cuda, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + assert wrapper_cuda.describe_accelerator().startswith("cueq") From 45dcdd342b5879ed54337576e61c858b8bf1a24c Mon Sep 17 00:00:00 2001 From: Bowen Deng Date: Mon, 26 Oct 2026 13:52:32 -0700 Subject: [PATCH 3/4] fix(mace): do not promise cueq to a model that cannot take it MACE-ANI-CC is built by mace_anicc, which fixes its own argument list, so the calculator never receives enable_cueq. describe_accelerator still reported 'cueq (falls back ...)' for it on any CUDA host, and _accelerator_kwargs returned the flag for a path that drops it: the same mistake as announcing a Nitro fallback for a checkpoint that was never eligible. Gate both on whether the model has an accelerated calculator path, and name that reason instead of the device or the package, so the loader reports what this model will do rather than what the host could do in general. --- src/utils/mlips/mace/mace_wrapper.py | 16 ++++++++++++++++ tests/mace/test_mace_accelerator.py | 24 ++++++++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/src/utils/mlips/mace/mace_wrapper.py b/src/utils/mlips/mace/mace_wrapper.py index c49173c8..f5ebf456 100644 --- a/src/utils/mlips/mace/mace_wrapper.py +++ b/src/utils/mlips/mace/mace_wrapper.py @@ -174,6 +174,18 @@ def load(self, model_path: Optional[str] = None) -> None: self.load_checkpoint(model_path) logger.info(f"Loaded model from checkpoint: {model_path}") + def _accepts_accelerator(self) -> bool: + """ + Whether this model's calculator path can take an accelerator at all. + + MACE-ANI-CC is built by mace_anicc, which fixes its own argument list, + so it cannot be handed enable_cueq. Reporting the cuEquivariance + speedup for it would promise a fast path that never runs — the same + mistake as announcing a Nitro fallback for a checkpoint that was never + eligible for Nitro. + """ + return "ANI-CC" not in self.model_name.upper() + def _accelerator_kwargs(self) -> Dict[str, Any]: """ Calculator kwargs for the accelerated path, or empty when it is out. @@ -184,6 +196,8 @@ def _accelerator_kwargs(self) -> Dict[str, Any]: """ if self.accelerator == "none" or self._accelerator_failed: return {} + if not self._accepts_accelerator(): + return {} if not _is_cuda(self.device): if not self._accelerator_warned: logger.info( @@ -216,6 +230,8 @@ def describe_accelerator(self) -> str: """ if self.accelerator == "none": return "none (disabled by request)" + if not self._accepts_accelerator(): + return f"none ({self.model_name} has no accelerated calculator path)" if not _is_cuda(self.device): return f"none (device is {self.device}, and cueq needs CUDA)" if not cueq_installed(): diff --git a/tests/mace/test_mace_accelerator.py b/tests/mace/test_mace_accelerator.py index ed6a4f84..89fa3d0c 100644 --- a/tests/mace/test_mace_accelerator.py +++ b/tests/mace/test_mace_accelerator.py @@ -128,6 +128,30 @@ def factory(**kwargs): assert attempts == [True, False, False] +@pytest.mark.mace +class TestModelsWithoutAcceleratedPath: + def test_ani_cc_never_requests_cueq(self, monkeypatch): + """mace_anicc fixes its own arguments, so cueq cannot be passed.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + w = mw.MACEWrapper(model_name="MACE-ANI-CC", device="cuda") + assert w._accelerator_kwargs() == {} + assert w.accelerator_effective == "none" + + def test_ani_cc_does_not_promise_the_speedup(self, monkeypatch): + """Do not announce an accelerator the model can never use.""" + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + w = mw.MACEWrapper(model_name="MACE-ANI-CC", device="cuda") + described = w.describe_accelerator() + assert described.startswith("none") + assert "no accelerated calculator path" in described + + def test_other_models_do_accept_it(self, monkeypatch): + monkeypatch.setattr(mw, "cueq_installed", lambda: True) + for name in ("MACE-MH-1", "MACE-OFF23-medium", "MACE-OMOL-extra-large"): + w = mw.MACEWrapper(model_name=name, device="cuda") + assert w._accelerator_kwargs() == {"enable_cueq": True}, name + + @pytest.mark.mace class TestDescribeAccelerator: def test_reports_cpu_as_none(self): From d71f40fed8c3a12a6b0b802e86d23b7780541d31 Mon Sep 17 00:00:00 2001 From: Bowen Deng Date: Mon, 26 Oct 2026 15:48:48 -0700 Subject: [PATCH 4/4] test(uv): allow the CUDA-matched cueq ops in the torch extras test_extras_are_the_same_everywhere pinned both torch extras to exactly ['torch', 'torchvision'], which cuEquivariance cannot satisfy: its ops wheel is compiled per CUDA major and published under a different distribution name (...-ops-torch-cu12 / ...-ops-torch-cu13) with no unsuffixed meta-package, so unlike torch it cannot be switched by index. Spell the mlip exception out rather than loosening the assertion: the shared front end must be identical in both flavours, exactly one ops package must be present per flavour, and its suffix must match that flavour. fairchem keeps the original equality check. Anything else appearing in the extras, or the two flavours drifting, still fails. --- tests/test_uv_projects.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/tests/test_uv_projects.py b/tests/test_uv_projects.py index 13285a86..87ea94bf 100644 --- a/tests/test_uv_projects.py +++ b/tests/test_uv_projects.py @@ -125,7 +125,30 @@ def test_extras_are_the_same_everywhere(project): # from each (build support in mlip, the LAMMPS wheel in fairchem); everything # else matches cpu. if project in ("mlip", "fairchem"): - assert extras.pop("cu126") == extras.pop("cu130") == ["torch", "torchvision"] + cu126, cu130 = extras.pop("cu126"), extras.pop("cu130") + # torch and torchvision are the same two names in both flavours, told + # apart by their index (see test_gpu_projects_lock_both_torch_cuda_builds). + assert cu126[:2] == cu130[:2] == ["torch", "torchvision"] + if project == "mlip": + # cuEquivariance is the one dependency whose wheel is compiled per + # CUDA major under a different *distribution name* + # (…-ops-torch-cu12 / …-ops-torch-cu13) with no unsuffixed + # meta-package, so unlike torch it cannot be switched by index. + # Spell both halves out: anything else appearing here, or the two + # flavours drifting apart, still fails. + def split_cueq(extra): + tail = extra[2:] + ops = [e for e in tail if e.startswith("cuequivariance-ops-torch-cu")] + assert len(ops) == 1, extra + return set(tail) - set(ops), ops[0] + + shared126, ops126 = split_cueq(cu126) + shared130, ops130 = split_cueq(cu130) + assert shared126 == shared130 == {"cuequivariance-torch>=0.2.0"} + assert ops126 == "cuequivariance-ops-torch-cu12>=0.2.0" + assert ops130 == "cuequivariance-ops-torch-cu13>=0.2.0" + else: + assert cu126 == cu130 == ["torch", "torchvision"] assert extras.pop("lammps"), project assert extras == load("cpu")["project"]["optional-dependencies"] assert set(extras) == {"openmm", "pymol", "docking", "void", "transport"}