Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add CUequivariance backend #127

Open
CompRhys opened this issue Nov 21, 2024 · 1 comment
Open

Add CUequivariance backend #127

CompRhys opened this issue Nov 21, 2024 · 1 comment

Comments

@CompRhys
Copy link
Contributor

How hard would it be to add a CUequivariance back end to SevenNet?

@YutackPark
Copy link
Member

It was harder than I firstly thought. Now available on cu_euqi branch.

Nearly x3 times faster in training, when the channel is consistent across the L value. If # of channels is different for each L value, like SevenNet-0, I found it becomes slightly slower.
Slightly (~50%?) faster for MD even with an inconsistent number of channels.

Here's the limitations and usage training.

Issues

  1. SevenNet-0 is not available. Due to this issue:
    [Bug] torch.autograd.grad never ends, if ChannelWiseTensorProduct is used with specific configuration NVIDIA/cuEquivariance#33
    Currently, sevennet automatically fallback to the original e3nn implementation if the condition is met. Models with higher Lmax (>2) will work.

  2. LAMMPS is not available. Due to this issue:
    Questions about TorchScript support NVIDIA/cuEquivariance#30
    It is a more unclear than the previous limitation. PyTorch itself may deprecated the torchscript than we have to re-implement LAMMPS part when newer replacement, torch.export came out. I'm not sure whether they gonna support torchscript for cuEquivariance.

Usage:

Installation:

git clone https://github.com/MDIL-SNU/SevenNet.git
cd SevenNet
git checkout cu_equi
pip install .

pip install cuequivariance-torch
pip install cuequivariance-ops-torch-cu12  # choose between cu12 or cu11 based on CUDA version

Training:

sevenn input.yaml --enable_cueq

Calculator:

from sevenn.sevennet_calculator import SevenNetCalculator
calc = SevenNetCalculator(PATH_TO_CHECKPOINT, enable_cueq=True)

The interface is intentionally made similar to that of MACE.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants