sevenn 0.9.3

Scalable EquiVariance Enabled Neural Network

SevenNet

SevenNet (Scalable EquiVariance Enabled Neural Network) is a graph neural network interatomic potential package that supports parallel molecular dynamics simulations with LAMMPS. Its underlying GNN model is based on nequip.

The project provides parallel molecular dynamics simulations using graph neural network interatomic potentials, which enable large-scale MD simulations or faster MD simulations.

The installation and usage of SevenNet are split into two parts: training + command-line interface + ASE calculator (handled by Python) and molecular dynamics (handled by LAMMPS).

• SevenNet
  • Installation
  • Usage
    • SevenNet-0
    • SevenNet Calculator for ASE
    • Training sevenn
      • Multi-GPU training
    • sevenn_graph build
    • sevenn_inference
    • sevenn_get_model
  • Installation for LAMMPS
  • Usage for LAMMPS
    • To check installation
    • For serial model
    • For parallel model
  • Future Works
  • Citation

Installation

• Python >= 3.8
• PyTorch >= 1.12.0

Please install PyTorch from PyTorch official before installing the SevenNet. Note that for SevenNet, torchvision and torchaudio are redundant. You can omit them from the command provided in the installation guide.

Matching PyTorch + CUDA versions may cause problems, especially when compiling SevenNet with LAMMPS. Here are the recommended versions we've been using internally without an issue.

• PyTorch/2.2.2 + CUDA/12.1.0
• PyTorch/1.13.1 + CUDA/12.1.0
• PyTorch/1.12.0 + CUDA/11.6.2

Using the newer versions of CUDA with PyTorch is usually not an issue. For example, you can compile and use PyTorch/1.13.1+cu117 with CUDA/12.1.0.

PLEASE NOTE: You must install PyTorch before installing SevenNet. They are not marked as dependencies since it is coupled with the CUDA version, therefore manual installation is safer.

After the PyTorch installation, simply run

pip install sevenn

Usage

SevenNet-0

SevenNet-0 is a general-purpose interatomic potential trained on the MPF dataset of M3GNet or MPtrj dataset of CHGNet. You can try SevenNet-0 to your application without any training. If the accuracy is unsatisfactory, SevenNet-0 can be fine-tuned.

SevenNet-0 (11July2024)

This model was trained on MPtrj. We suggest starting with this model as we found that it performs better than the previous SevenNet-0 (22May2024). Check Matbench Discovery leaderborad for this model's performance on materials discovery.

Whenever the checkpoint path is the input, this model can be loaded via 7net-0 | SevenNet-0 | 7net-0_11July2024 | SevenNet-0_11July2024 keywords.

Acknowledgments: This work was supported by the Neural Processing Research Center program of Samsung Advanced Institute of Technology, Samsung Electronics Co., Ltd. The computations for training models were carried out using the Samsung SSC-21 cluster.

SevenNet-0 (22May2024)

This model was trained on MPF.2021.2.8. This is the model used in our paper.

Whenever the checkpoint path is the input, this model can be loaded via 7net-0_22May2024 | SevenNet-0_22May2024 keywords.

SevenNet Calculator for ASE

ASE (Atomic Simulation Environment) is a set of tools and Python modules for atomistic simulations. SevenNet-0 and SevenNet-trained potentials can be used with ASE for its use in python.

For pre-trained models,

from sevenn.sevennet_calculator import SevenNetCalculator
sevenet_0_cal = SevenNetCalculator("7net-0", device='cpu')  # 7net-0, SevenNet-0, 7net-0_22May2024, 7net-0_11July2024 ...

For user trained models,

from sevenn.sevennet_calculator import SevenNetCalculator
checkpoint_path = ### PATH TO CHECKPOINT ###
sevenet_cal = SevenNetCalculator(checkpoint_path, device='cpu')

Training

sevenn_preset base > input.yaml
sevenn input.yaml -s

Other valid preset options are: base, fine_tune, and sevennet-0. Check comments of base yaml for explanations.

To reuse a preprocessed training set, you can specify ${dataset_name}.sevenn_data to the load_dataset_path: in the input.yaml.

Multi-GPU training

We support multi-GPU training features using PyTorch DDP (distributed data parallel). We use one process (or a CPU core) per GPU.

torchrun --standalone --nnodes {number of nodes} --nproc_per_node {number of GPUs} --no_python sevenn input.yaml -d

Please note that batch_size in input.yaml indicates batch_size per GPU.

sevenn_graph_build

sevenn_graph_build -f ase my_train_data.extxyz 5.0

You can preprocess the dataset with sevenn_graph_build to obtain *.sevenn_data files. The cutoff length should be provided. See sevenn_graph_build --help for more information.

sevenn_inference

sevenn_inference checkpoint_best.pt path_to_my_structures/*

This will create dir sevenn_infer_result. It includes .csv files that enumerate prediction/reference results of energy and force. See sevenn_inference --help for more information.

sevenn_get_model

This command is for deploying lammps potentials from checkpoints. The argument is either the path to checkpoint or the name of pre-trained potential.

sevenn_get_model 7net-0

This will create deployed_serial.pt, which can be used as lammps potential under e3gnn pair_style.

The parallel model can be obtained in a similar way

sevenn_get_model 7net-0 -p

This will create multiple deployed_parallel_*.pt files. The number of deployed models equals the number of message-passing layers. These models can be used as lammps potential to run parallel MD simulations with GNN potential using multiple GPU cards.

Installation for LAMMPS

• PyTorch (same version as used for training)
• LAMMPS version of 'stable_2Aug2023_update3' LAMMPS
• (Optional) CUDA-aware OpenMPI for parallel MD

As system-wise requirements, you also need CUDA and Intel MKL. Currently, this installation guide is dedicated to Linux.

PLEASE NOTE: CUDA-aware OpenMPI is optional, but recommended for parallel MD. If it is not available, in parallel mode, GPUs will communicate via CPU. It is still faster than using only one GPU, but its efficiency is low.

PLEASE NOTE: CUDA-aware OpenMPI does not support NVIDIA Gaming GPUs. Given that the software is closely tied to hardware specifications, please consult with your server administrator if unavailable.

Ensure the LAMMPS version (stable_2Aug2023_update3). You can easily switch the version using git. After switching the version, run sevenn_patch_lammps with the lammps directory path as an argument.

git clone https://github.com/lammps/lammps.git lammps_dir
cd lammps_dir
git checkout stable_2Aug2023_update3
sevenn_patch_lammps ./

You can refer to sevenn/pair_e3gnn/patch_lammps.sh for the detailed patch process.

Build LAMMPS with cmake (example):

$ cd {path_to_lammps_dir}
$ mkdir build
$ cd build
$ cmake ../cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'`
$ make -j4

If the compilation is successful, you will find the executable at {path_to_lammps_dir}/build/lmp. To use this binary easily, for example, create a soft link in your bin directory (which should be included in your $PATH).

ln -s {absolute_path_to_lammps_dir}/build/lmp $HOME/.local/bin/lmp

This will allow you to run the binary using lmp -in my_lammps_script.lmp.

If you're having trouble with the above procedure and it is related to MKL, check the notes below.

Note for MKL

You may encounter MKL_INCLUDE_DIR NOT-FOUND during cmake or see hundreds of undefined reference to XXX errors from libtorch_cpu.so at the end of the compilation. This typically indicates that either Intel MKL is not installed or its environment variables are not correctly set.

First, check if there are Intel MKL modules available on your system (these may have names like intel-oneapi, intel-mkl, intel-compiler, etc.). Load the relevant module to automatically configure the necessary settings.

If the module is not available, I recommend using Conda. Install mkl and mkl-include with the following command:

conda install -c intel mkl mkl-include

If you encounter an error, remove -c intel. This is a known bug in the recent Conda version.

Next, configure the LD_LIBRARY_PATH for MKL libraries:

export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH

Finally, append the following to your cmake command:

-DMKL_INCLUDE_DIR="$CONDA_PREFIX/include"

Before running the command, I recommend you clean the files generated from the previous cmake command.

Usage for LAMMPS

To check installation

{lmp_binary} -help | grep e3gnn

You will see e3gnn and e3gnn/parallel as pair_style.

For serial model

pair_style e3gnn
pair_coeff * * {path to serial model} {space separated chemical species}

For parallel model

pair_style e3gnn/parallel
pair_coeff * * {number of segmented parallel models} {space separated paths of segmented parallel models} {space separated chemical species}

For example,

pair_style e3gnn/parallel
pair_coeff * * 4 ./deployed_parallel_0.pt ./deployed_parallel_1.pt ./deployed_parallel_2.pt ./deployed_parallel_3.pt Hf O

The number of segmented *.pt files is the same as the number of message-passing layers of the model.

Check sevenn_get_model for deploying lammps models from checkpoint for both serial and parallel.

One GPU per MPI process is expected. If the available GPUs are fewer than the MPI processes, the simulation may run inefficiently.

PLEASE NOTE: Currently, the parallel version raises an error when there are no atoms in one of the subdomain cells. This issue can be addressed using the processors command and, more optimally, the fix balance command in LAMMPS. This will be patched in the future.

Future Works

• Notebook examples and improved interface for non-command line usage
• Implementation of pressure output in parallel MD simulations.
• Development of support for a tiled communication style (also known as recursive coordinate bisection, RCB) in LAMMPS.
• Easy use of parallel models

Citation

If you use SevenNet, please cite (1) parallel GNN-IP MD simulation by SevenNet or its pre-trained model SevenNet-0, (2) underlying GNN-IP architecture NequIP

(1) Y. Park, J. Kim, S. Hwang, and S. Han, "Scalable Parallel Algorithm for Graph Neural Network Interatomic Potentials in Molecular Dynamics Simulations". J. Chem. Theory Comput., 20(11), 4857 (2024) (https://pubs.acs.org/doi/10.1021/acs.jctc.4c00190)

(2) S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. Smidt, and B. Kozinsky, "E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials". Nat. Commun., 13, 2453. (2022) (https://www.nature.com/articles/s41467-022-29939-5)