Very large dataset

[Mathieu-Istas]

Dear Nequip team,
I have a large dataset (240 000 configurations of 200 atoms). I probably don't need that many configurations to get a good model, but if I can use the full dataset, I might as well do it.
When calling nequip-train, I get an error when processing the dataset. However, it is not an obvious OOM like this discussion, and I can't understand the error (I only copied the the part of the traceback that seemed relevant):

"""
Traceback (most recent call last):
  File "/home/istas/.local/lib/python3.9/site-packages/torch/serialization.py", line 441, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "/home/istas/.local/lib/python3.9/site-packages/torch/serialization.py", line 655, in _save
    zip_file.write_record('data.pkl', data_value, len(data_value))
RuntimeError: [enforce fail at inline_container.cc:471] . PytorchStreamWriter failed writing file data.pkl: file write failed

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/SHARE/App/Python-3.9.2/lib/python3.9/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
  File "/SHARE/App/Python-3.9.2/lib/python3.9/multiprocessing/pool.py", line 48, in mapstar
    return list(map(*args))
  File "/home/istas/.local/lib/python3.9/site-packages/nequip/data/_dataset/_ase_dataset.py", line 67, in _ase_dataset_reader
    torch.save(datas, path)
  File "/home/istas/.local/lib/python3.9/site-packages/torch/serialization.py", line 442, in save
    return
  File "/home/istas/.local/lib/python3.9/site-packages/torch/serialization.py", line 291, in __exit__
    self.file_like.write_end_of_file()
RuntimeError: [enforce fail at inline_container.cc:337] . unexpected pos 64 vs 0
"""

The traceback finishes with

RuntimeError: Failed to build object with prefix `dataset` using builder ASEDataset

This error disappear if I point to a smaller file in dataset_file_name (for instance with 50 000 configurations), but lowering n_train or n_val with the original file with all configurations (around 5 Gb) result in the same error. Do you know if there something I could do to train on the full dataset?

Regards,
Mathieu

[Linux-cpp-lisp]

Hi @Mathieu-Istas ,

Thanks for your interest in our code!

This is... strange. Are you low on storage on your machine's temporary directory? Is your tempdir on a network filesystem?

Can you try to run with the environment variable NEQUIP_NUM_TASKS=1?

[Linux-cpp-lisp]

Sounds like possibly you may not have enough space in your /tmp dir: the neighborlists can be quite a bit of data, if you have a large dataset with a large cutoff. The error you've posted would also agree with this. You could watch the disk usage in /tmp during processing to confirm this. Easiest workaround is to find a larger scratch space for nequip to work with...

[Mathieu-Istas]

Thanks for the answer, I will try to find a bigger scratch. There is one thing that bothers me, if I use the large dataset with a lot of configurations but assign very small values to n_train and n_val, like 10 for each, I still get the same error. I would have thought that the memory from the neighbor list would be small in that case, or is Nequip building lists with every possible configuration, including configurations not in train neither validation?

[Mathieu-Istas]

Last question, would Allegro be better?

[Linux-cpp-lisp]

By default nequip preprocesses the entire dataset, and then subsamples from it, so that you don't have to repeat preprocessing if you change the sampling seed, n_val, etc. (imagine you run a sweep over training set sizes and have to reprocess for every run rather than once ahead of time, etc.) This behavior can be modified with the dataset_include_frames option in the YAML, which indexes out a specific set of frames before preprocessing or train val split.

[Linux-cpp-lisp]

Allegro will not change this behavior, as it also uses the nequip framework training infrastructure.

[Linux-cpp-lisp]

I'll note also that there's a contributed HDF5 dataset option that may avoid this issue: https://github.com/mir-group/nequip/blob/develop/nequip/data/_dataset/_hdf5_dataset.py.

[Mathieu-Istas]

Do you have an example that transforms a xyz or a npz file into a hdf5 file? I am not sure which keys I should put into the hdf5 file. Also the value of dataset in the config should be equal to hdf5?

[Linux-cpp-lisp]

Not entirely sure; see #227

[nitin6753]

Hello,

I have 500 MB of data, comprising approximately 20,000 configurations generated through VASP. When I used a batch size of 5, the terminal showed processing data, after which it closes.
When I lower the batch size to 2, the process does begin, but it has been running for an extended period (over 4 days) and has yet to complete.
It's worth mentioning that I've already converted the VASP output to the extxyz format.

I'm wondering if this behavior is expected, or if there might be an error in my approach.
I'm seeking advice on how to enhance the processing speed and which parameters could benefit from optimization.

Here are the specifications of my system:
CPU: Intel® Xeon® W-2275 (28 cores)
RAM: 16 GB
GPU: NVIDIA T1000

Any assistance or insights into this matter would be greatly appreciated.
Thank you in advance.

[Mathieu-Istas]

Hello, I am not from nequip team, so take what I say with a grain of salt, maybe I can try to give some answers. 4 days seems a lot, but your dataset is rather big. Have you tried to call nequip-train on a CPU with more memory? You will not be able to perform the training, but that could do the processing dataset step for you. This step create a file named something like processed_dataset_01a24e85085f99f527ec9e3925cd5d58392fee40 that you can transfer to the root folder define in your yaml config file. That is kind of a workaround, but it may allow you to do the training which is not as memory intensive and I/O I think.

Mathieu

[Linux-cpp-lisp]

@Mathieu-Istas is exactly right, thank you for contributing! See for further details e.g. #263 or #293. Note also that you can explicitly control the number of worker processes for the preprocessing with the NEQUIP_NUM_TASKS environment variable, which attempts to get a sane value but could be forced to default to 1 if it can't get a useful core count.