Repository navigation
Enable cuEquivariance for MACE, with a plain-path fallback that always works - #55
Merged
Merged
Conversation
Collaborator
Author
Why
|
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.
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.
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.
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.
bowen-bd
force-pushed
the
feat/mace-cueq
branch
from
October 27, 2026 02:37
1d89268 to
d71f40f
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
MACE's library is already current here —
mace-torch0.3.16 is both whatvenv/mlippins and the newest release — and no newer universal checkpoint has appeared since MACE-MH-1 (2025-10-30). What had gone unused is the accelerator:MACECalculatoracceptsenable_cueq, and every one of this repository's seven calculator constructions omitted it, so MACE has been running its plain e3nn path even where cuEquivariance was installed.This adds cuEquivariance to the
mlipCUDA extras and routes calculator construction through one builder that uses it. The constraint that shapes the design is that an optimisation may not cost anybody the ability to run MACE: no GPU, no cuequivariance install, and a card whose kernels are missing all degrade to the plain path rather than raising.Changes
cuequivariance-torchpluscuequivariance-ops-torch-cu12/cu13ride in the existingcu126/cu130extras, because their coupling to torch is exact — the ops wheels are compiled against one CUDA major, andcuequivariance_torchimports whichever is installed under the bare module namecuequivariance_ops_torch. Upstreammace-torchdeclares only the cu12 pair (cueq-cuda-11/12); the cu13 pair exists and is what a driver ≥ 580 host needs. About 30 MB per architecture, so the images barely move.MaceWrappergains anacceleratorknob (auto/cueq/none) and one_build_calculatorthat all six accelerator-capable constructions go through.autois the default.load_modelreports what the calculator will attempt, without building it: the calculator is lazy, and building it would download and convert the checkpoint just to answer. The tool result names the accelerator in effect, or the reason there is none.MACE-ANI-CCis excluded from both the request and the report:mace_aniccfixes its own argument list, so it never receivesenable_cueq. Announcing a speedup for a path that drops the flag would be the same mistake as reporting a Nitro fallback for a checkpoint that was never eligible for Nitro.Fallbacks
This is the part worth reviewing closely, because mace itself raises rather than degrading:
cuequivariancenot installed (CPU-only env, or an extra that was never synced)accelerator="none"ValueError, rather than being silently ignoredMeasurements
One 32-atom Cu
fccsupercell, MACE-OMAT-0-small, float32, same process, warm:6.0x, with energies agreeing to every printed digit (both −3.744912 eV/atom). CPU and cueq-not-installed runs reproduce the same energy on the plain path.
Verification
tests/mace/— 30 passed (including the Sella TS/IRC tests).tests/mace/test_mace_accelerator.py— 16 tests, all against a fake calculator factory, so they need no GPU, no cuequivariance install and no checkpoint download: an unsupported card degrades instead of raising, a missing module degrades, a genuine error on the plain path still surfaces, a failure is not retried per structure, ANI-CC is never offered the accelerator, and the unknown-value path is rejected.platforms.tsv,servers.tsv, the plugin wiring and the version manifests are unchanged by the extra;uv lock --checkandruffare clean.Merge order with the open UMA/MatGL pull request
Both branches edit the
mliprow in.agents/rules/mcp-environments.mdand the package table invenv/README.md— different words on the same lines, so whichever lands second will need that one-line conflict resolved. Nothing else overlaps: thevenv/mliplock andpyproject.tomlchanges in the two branches sit in different regions and merge cleanly. If it is easier, land the UMA/MatGL branch first and this one after a rebase.