[Feature Request]: Larger than memory datasets.

[JonathanSchmidt1]

Problem

Now that multi-gpu training is working we are very interested in training on some larger crystal structure datasets. However, the datasets do not fit into the RAM. It would be great if it would be possible to either have to only load a partial dataset on each ddp node or be able to load the features on the fly to make large scale training possible. I assume the LMDB datasets that the OCP and MATSCIML are using should work for that. Ps: thank you for the pytorch lightning implementation

Proposed Solution

Add an option to save and load data to/from an LMDB database.

Alternatives

Examples would be here
https://github.com/Open-Catalyst-Project/ocp/blob/main/tutorials/lmdb_dataset_creation.ipynb
or here https://github.com/IntelLabs/matsciml/tree/main/matsciml/datasets

Code of Conduct

• I agree to follow this project's Code of Conduct

[JonathanSchmidt1]

A small update to this request. I also asked the matsciml team about this issue as they have an interface to matgl and other models included in their package and they were kind enough to prepare a guide on how to prepare a suitable dataset https://github.com/IntelLabs/matsciml/issues/85 . I will follow that guide for our data and maybe it could be used to extend the training capabilities of matgl as well.

[shyuep]

@JonathanSchmidt1 The dataloaders in matgl already allows you to do a one-time processing of the structures into a graph dataset. Once that graph dataset is done, it is much smaller in memory than the structures. In fact, that is the way we have been training extremely large datasets.

[JonathanSchmidt1]

Thank you for the reply. The preprocessing is definitely useful.
But after preprocessing a decently sized dataset (4.5M structures) takes up 132Gb on disk and 128Gb loaded into RAM and we would like to train on larger datasets in the future.
Maybe I am also doing sth wrong. Right now I am just doing the following to preprocess data:

elem_list = get_element_list(structures)
# setup a graph converter
converter = Structure2Graph(element_types=elem_list, cutoff=6.0)
# convert the raw dataset into MEGNetDataset
dataset = M3GNetDataset(
    threebody_cutoff=4.0, structures=structures, converter=converter, labels={"energies":energies}
)
dataset.process()
dataset.save()

For me, the issue is also the rather old architecture of the gpu partition of the supercomputer I have to use (64Gb Ram per node, 1 Gpu per node). So in-memory datasets are just not an option there.
However even on a modern architecture with e.g. 512 Gb per node and 8 gpus per node the in-memory dataset becomes a problem as I think with ddp each process will load its own copy of the dataset resulting in ~960 Gb if I would use all gpus?

[JonathanSchmidt1]

Are there any updates on this? Even though we have decent nodes now with 200gb ram per gpu the datasets also take up more than 600gb now.

[jun0c129]

I have a similar issue. I'm just trying to train M3GNet on MPtrj, which contains 1.5M structures. The raw dataset consisting of pymatgen structures and their corresponding labels (energy, forces, and stress) is about 7 GB, however, the processed graph data (dgl_graph.bin, dgl_line_graph.bin) is about 230 GB. It looks like this dataset is stored in CPU RAM, and I cannot train using multi-GPUs via SLURM environment where --ntasks-per-node and gres=gpus are both set to the number of GPUs, since srun command run python script on the same number of instances in parallel, may consuming memory of (230 x number of GPUs) GB which is too large even for the modern nodes. Below is a code snippet that I have tried, and I would appreciate if any workaround or solution is available. Please let me know if my understanding is wrong. Thanks.

For processing the dataset:

element_types = DEFAULT_ELEMENTS
converter = Structure2Graph(element_types=element_types, cutoff=5.0)
mgldataset = MGLDataset(
    threebody_cutoff=4.0,
    structures=structures,
    converter=converter,
    labels=labels,
    include_line_graph=True,
    )

For training:

SLURM_NTASKS_PER_NODE = int(os.getenv('SLURM_NTASKS_PER_NODE', '1'))
trainer = pl.Trainer(
    max_epochs=1000000,
    accelerator="cuda",
    devices=SLURM_NTASKS_PER_NODE,
    strategy="ddp" if SLURM_NTASKS_PER_NODE > 1 else "auto",
    logger=csvlogger,
    inference_mode=False,
    callbacks=[model_checkpoint, early_stopping]
    )

The SLURM script:

#!/bin/bash
#SBATCH -J train_m3gnet
#SBATCH -p long
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --cpus-per-task=6
#SBATCH -o %x.o%j
#SBATCH -e %x.e%j
#SBATCH --time=336:00:00 # 14 days
#SBATCH --gres=gpu:4
#SBATCH --mem=500G

srun python ./train.py

[shyuep]

Hi there, sorry for the delayed response. The dataloaders today already allow a clear_processed kwarg that can be set to True. That presumably will delete the structures after processing. Does this help at all? I think if the structures are provided in a generator, that should help immensely in reducing data needs.

If not, I can definitely explore using something like LMDB. We regularly train on million structure datasets and have not encountered issues.