Open MatSci ML Toolkit : A Broad, Multi-Task Benchmark for Solid-State Materials Modeling

This is the implementation of the MatSci ML benchmark, which includes ~1.5 million ground-state materials collected from various datasets, as well as integration of the OpenCatalyst dataset supporting diverse data format (point cloud, DGL graphs, PyG graphs), learning methods (single task, multi-task, multi-data) and deep learning models. Primary project contributors include: Santiago Miret (Intel Labs), Kin Long Kelvin Lee (Intel AXG), Carmelo Gonzales (Intel Labs), Mikhail Galkin (Intel Labs), Marcel Nassar (Intel Labs), Matthew Spellings (Vector Institute).

News

• [2024/08/23] Readthedocs is now online!
• [2023/09/27] Release of pre-packaged lmdb-based datasets from v1.0.0 via Zenodo.
• [2023/08/31] Initial release of the MatSci ML Benchmark with integration of ~1.5 million ground state materials.
• [2023/07/31] The Open MatSci ML Toolkit : A Flexible Framework for Deep Learning on the OpenCatalyst Dataset paper is accepted into TMLR. See previous version for code related to the benchmark.

Introduction

The MatSci ML Benchmark contains diverse sets of tasks (energy prediction, force prediction, property prediction) across a broad range of datasets (OpenCatalyst Project [1], Materials Project [2], LiPS [3], OQMD [4], NOMAD [5], Carolina Materials Database [6]). Most of the data is related to energy prediction task, which is the most common property tracked for most materials systems in the literature. The codebase support single-task learning, as well as multi-task (training one model for multiple tasks within a dataset) and multi-date (training a model across multiple datsets with a common property). Additionally, we provide a generative materials pipeline that applies diffusion models (CDVAE [7]) to generate new unit cells.

The package follows the original design principles of the Open MatSci ML Toolkit, including: - Ease of use for new ML researchers and practitioners that want get started on interacting with the OpenCatalyst dataset. - Scalable computation of experiments leveraging PyTorch Lightning across different computation capabilities (laptop, server, cluster) and hardware platforms (CPU, GPU, XPU) without sacrificing performance in the compute and modeling. - Integrating support for DGL and PyTorch Geometric for rapid GNN development.

The examples outlined in the next section how to get started with Open MatSci ML Toolkit using simple Python scripts, Jupyter notebooks, or the PyTorch Lightning CLI for a simple training on a portable subset of the original dataset (dev-set) that can be run on a laptop. Subsequently, we scale our example python script to large compute systems, including distributed data parallel training (multiple GPU on a single node) and multi-node training (multiple GPUs across multiple nodes) in a computing cluster. Leveraging both PyTorch Lightning and DGL capabilities, we can enable the compute and experiment scaling with minimal additional complexity.

Installation

• Docker: We provide a Dockerfile inside the docker that can be run to install a container using standard docker commands.
• mamba: We have included a mamba specification that provides a complete out-of-the-box installation. Run mamba env create -n matsciml --file conda.yml, and will install all dependencies and matsciml as an editable install.
• pip: In this case, we assume you are bringing your own virtual environment. Depending on what hardware platform you have, you can copy-paste the following commands; because the absolute mess that is modern Python packaging, these commands include the URLs for binary distributions of PyG and DGL graph backends.

For CPU only (good for local laptop development):

console pip install -f https://data.pyg.org/whl/torch-2.4.0+cpu.html -f https://data.dgl.ai/wheels/torch-2.4/repo.html -e './[all]'

For XPU usage, you will need to install PyTorch separately first, followed by matsciml; note that the PyTorch version is lower as 2.3.1 is the latest XPU binary distributed.

console pip install torch==2.3.1+cxx11.abi torchvision==0.18.1+cxx11.abi torchaudio==2.3.1+cxx11.abi intel-extension-for-pytorch==2.3.110+xpu oneccl_bind_pt==2.3.100+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/ pip install -f https://data.pyg.org/whl/torch-2.3.0+cpu.html -f https://data.dgl.ai/wheels/torch-2.3/repo.html -e './[all]'

For CUDA usage, substitute the index links with your particular toolkit version (e.g. 12.1 below):

console pip install -f https://data.dgl.ai/wheels/torch-2.4/cu121/repo.html -f https://data.pyg.org/whl/torch-2.4.0+cu121.html -e './[all]'

Additionally, for a development install, one can specify the extra packages like black and pytest with pip install './[dev]'. These can be added to the commit workflow by running pre-commit install to generate git hooks.

Intel XPU capabilities

[!NOTE] As of PyTorch 2.4+, XPU support has been upstreamed to PyTorch and starting from torch>=2.5.0 onwards, should be available as a pip install. We will update the instructions accordingly when it does. We recommend consulting the PyTorch documentation for updates and instructions on how to get started with XPU use. In the meantime, please consult this page to see how to set up PyTorch on XPUs.

The module matsciml.lightning.xpu implements interfaces for Intel XPU to Lightning abstractions, including the XPUAccelerator and two strategies for deployment (single XPU/tile and distributed data parallel). Because we use PyTorch Lightning, there aren't many marked differences in running on Intel XPU, or GPUs from other vendors. The abstractions we mentioned are registered in the various Lightning registries, and should be accessible simply through pl.Trainer arguments, e.g.:

python trainer = pl.Trainer(accelerator='xpu')

The one major difference is for distributed data parallelism: Intel XPUs use the oneCCL communication backend, which replaces nccl, gloo, or other backends typically passed to torch.distributed. Please see examples/devices for single XPU/tile and DDP use cases.

NOTE: Currently there is a hard-coded torch.cuda.stream context in PyTorch Lightning's DDPStrategy. This issue has been created to see if the maintainers would be happy to patch it so that the cuda.Stream context is only used if a CUDA device is being used. If you encounter a RuntimeError: Tried to instantiate dummy base class Stream, please just set ctx = nullcontext() in the line of code that raises the exception.

Examples

The examples folder contains simple, unit scripts that demonstrate how to use the pipeline in specific ways:

For more advanced use cases:

Data Pipeline

In the scripts folder you will find two scripts needed to download and preprocess datasets: the download_datasets.py can be used to obtain Carolina DB, Materials Project, NOMAD, and OQMD datasets, while the download_ocp_data.py preserves the original Open Catalyst script.

In the current release, we have implemented interfaces to a number of large scale materials science datasets. Under the hood, the data structures pulled from each dataset have been homogenized, and the only real interaction layer for users is through the MatSciMLDataModule, a subclass of LightningDataModule.

```python from matsciml.lightning.data_utils import MatSciMLDataModule

no configuration needed, although one can specify the batch size and number of workers

devsetmodule = MatSciMLDataModule.fromdevset(dataset="MaterialsProjectDataset") ```

This will let you springboard into development without needing to worry about how to wrangle with the datasets; just grab a batch and go! With the exception of Open Catalyst, datasets will typically return point cloud representations; we provide a flexible transform interface to interconvert between representations and frameworks:

Task abstraction

In Open MatSci ML Toolkit, tasks effective form learning objectives: at a high level, a task takes an encoding model/backbone that ingests a structure to predict one or several properties, or classify a material. In the single task case, there may be multiple targets and the neural network architecture may be fluid, but there is only one optimizer. Under this definition, multi-task learning comprises multiple tasks and optimizers operating jointly through a single embedding.

References

• [1] Chanussot, L., Das, A., Goyal, S., Lavril, T., Shuaibi, M., Riviere, M., Tran, K., Heras-Domingo, J., Ho, C., Hu, W. and Palizhati, A., 2021. Open catalyst 2020 (OC20) dataset and community challenges. Acs Catalysis, 11(10), pp.6059-6072.
• [2] Jain, A., Ong, S.P., Hautier, G., Chen, W., Richards, W.D., Dacek, S., Cholia, S., Gunter, D., Skinner, D., Ceder, G. and Persson, K.A., 2013. Commentary: The Materials Project: A materials genome approach to accelerating materials innovation. APL materials, 1(1).
• [3] Batzner, S., Musaelian, A., Sun, L., Geiger, M., Mailoa, J.P., Kornbluth, M., Molinari, N., Smidt, T.E. and Kozinsky, B., 2022. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1), p.2453.
• [4] Kirklin, S., Saal, J.E., Meredig, B., Thompson, A., Doak, J.W., Aykol, M., Rühl, S. and Wolverton, C., 2015. The Open Quantum Materials Database (OQMD): assessing the accuracy of DFT formation energies. npj Computational Materials, 1(1), pp.1-15.
• [5] Draxl, C. and Scheffler, M., 2019. The NOMAD laboratory: from data sharing to artificial intelligence. Journal of Physics: Materials, 2(3), p.036001.
• [6] Zhao, Y., Al‐Fahdi, M., Hu, M., Siriwardane, E.M., Song, Y., Nasiri, A. and Hu, J., 2021. High‐throughput discovery of novel cubic crystal materials using deep generative neural networks. Advanced Science, 8(20), p.2100566.
• [7] Xie, T., Fu, X., Ganea, O.E., Barzilay, R. and Jaakkola, T.S., 2021, October. Crystal Diffusion Variational Autoencoder for Periodic Material Generation. In International Conference on Learning Representations.

Contributing

Please refer to the developers guide for how to contribute the the Open MatSciML Toolkit.

Citations

If you use Open MatSci ML Toolkit in your technical work or publication, we would appreciate it if you cite the Open MatSci ML Toolkit paper in TMLR:

If you use v1.0.0, please cite our paper:

Please cite datasets used in your work as well. You can find additional descriptions and details regarding each dataset here.