Time consumption before the train starts

[ssrokyz]

Hello!
It takes too much time when start training like an hour or more.
I guess it is because of the large training dataset size (60000 structures).
What task is it that consumes the time?
And is there any way to go around this task?

--------- Update 2022-07-01 -----------
It's fully automated from the 0.5.5 build.
If it's not working, try setting the environment variables OMP_NUM_THREADS=the number of cores and NEQUIP_NUM_TASKS=the number of cores
Replace the "the number of cores" part with your case.

[simonbatzner]

Hi @ssrokyz the very first time you train NequIP, it parses the entire data set and creates neighborlists. It will then for subsequent runs re-load that dataset, so you only have to do this once. However, depending on data set size and the size of the structures, this can take some time. How large are the structures?

[ssrokyz]

I see. So it's a necessary process. Then it's okay since it will speed up the entire training.
Each structure datum has a size of 25~200 atoms. 60000 structures in total. Thx for the answer!

[Linux-cpp-lisp]

right, that's quite a bit of data, and ASE's neighborlist is not the fastest in the world 😄

[ssrokyz]

Is it possible to parallelize the data parsing step?
OpenMP will help significantly! 😀 Or subprocess if it's a python script?

[Linux-cpp-lisp]

Parallelizing this has been on the TODO list for a while but has not yet gotten high enough in priority to get done. It's something I will probably do in the near future here with multiprocessing.

If this is something you'd be willing to help test once I've written it, that would be great!

[ssrokyz]

@Linux-cpp-lisp Thank you so much for your support! I will happily test it!

[Linux-cpp-lisp]

@ssrokyz can you try the ase-multiprocess branch here: https://github.com/mir-group/nequip/tree/ase-multiprocess.

In particular:

• Does it work?
• Is it faster?
• Do you get the same results in your training?

Make sure to use a different setting for root in your YAML to ensure that it reprocesses the dataset rather than loading the cached one from your previous runs. (You should see Processing... in the log before you see Loaded dataset...)

In theory, the generated dataset files should actually be bit-identical... I think 😄 . That also should be easy to check using something like this cmp command on the .pth files in the processed_dataset_... folders in the different roots.

[Linux-cpp-lisp]

I should mention the caveat that this is for the ASE dataset, and not the NPZ dataset--- I assume ASE is what you are using? What is the original on-disk format of your data, extxyz?

[ssrokyz]

@Linux-cpp-lisp It's an ASE dataset (trajectory file), no worries. 😁 I'm compiling the ASE multiprocess branch. I'll let you know once it is done. Thx for your fast implementation!

[ssrokyz]

@Linux-cpp-lisp
I installed ase-multiprocess branch, and I've got an error message (attached below).

I'm using ASE version 3.22.1 which is the latest in conda repository. My dataset has 60000 structures and is in ASE trajectory format. I think this is not a memory issue since I got the same error with a reduced dataset of 600 structures.
My system node has 64 cores, which is quite many, so I tested with export NEQUIP_NUM_TASKS=6 but still got the same error (I checked that the process used only 6 cores).
Reason: 'OSError(24, 'Too many open files')' This part of error message seems responsible for the failure, maybe it's related to the Linux system?
ulimit -Sn command gives 1024 which seems enough for the dataset with 600 structures...
c.f. I think ulimit command gives the number of resources that a process can require.

(Below, error with 600 structures)

Torch device: cuda
Processing...
Traceback (most recent call last):
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/utils/auto_init.py", line 232, in instantiate
    instance = builder(**positional_args, **final_optional_args)
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/data/dataset.py", line 839, in __init__
    super().__init__(
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/data/dataset.py", line 129, in __init__
    super().__init__(root=root, transform=type_mapper)
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/utils/torch_geometric/dataset.py", line 90, in __init__
    self._process()
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/utils/torch_geometric/dataset.py", line 175, in _process
    self.process()
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/data/dataset.py", line 212, in process
    data = self.get_data()
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/data/dataset.py", line 914, in get_data
    datas = p.map(reader, range(n_proc))
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/multiprocessing/pool.py", line 364, in map
    return self._map_async(func, iterable, mapstar, chunksize).get()
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/multiprocessing/pool.py", line 771, in get
    raise self._value
multiprocessing.pool.MaybeEncodingError: Error sending result: '[[(1, AtomicData(atomic_numbers=[64, 1], cell=[3, 3], edge_cell_shift=[4984, 3], edge_index=[2, 4984], forces=[64, 3], pbc=[3], pos=[64, 3], stress=[6], total_energy=[1])), (7, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3756, 3], edge_index=[2, 3756], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (13, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3866, 3], edge_index=[2, 3866], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (19, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3538, 3], edge_index=[2, 3538], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (25, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3628, 3], edge_index=[2, 3628], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (31, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3670, 3], edge_index=[2, 3670], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (37, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3594, 3], edge_index=[2, 3594], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (43, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3826, 3], edge_index=[2, 3826], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (49, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3626, 3], edge_index=[2, 3626], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (55, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3640, 3], edge_index=[2, 3640], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (61, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3912, 3], edge_index=[2, 3912], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (67, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3784, 3], edge_index=[2, 3784], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (73, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3714, 3], edge_index=[2, 3714], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (79, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3652, 3], edge_index=[2, 3652], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (85, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3850, 3], edge_index=[2, 3850], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (91, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3748, 3], edge_index=[2, 3748], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (97, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3732, 3], edge_index=[2, 3732], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (103, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3440, 3], edge_index=[2, 3440], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (109, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3636, 3], edge_index=[2, 3636], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (115, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3666, 3], edge_index=[2, 3666], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (121, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4060, 3], edge_index=[2, 4060], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (127, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3708, 3], edge_index=[2, 3708], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (133, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3376, 3], edge_index=[2, 3376], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (139, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3740, 3], edge_index=[2, 3740], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (145, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3896, 3], edge_index=[2, 3896], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (151, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3558, 3], edge_index=[2, 3558], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (157, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3682, 3], edge_index=[2, 3682], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (163, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4086, 3], edge_index=[2, 4086], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (169, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3342, 3], edge_index=[2, 3342], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (175, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3770, 3], edge_index=[2, 3770], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (181, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3670, 3], edge_index=[2, 3670], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (187, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3756, 3], edge_index=[2, 3756], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (193, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4102, 3], edge_index=[2, 4102], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (199, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3678, 3], edge_index=[2, 3678], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (205, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3744, 3], edge_index=[2, 3744], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (211, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3590, 3], edge_index=[2, 3590], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (217, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3574, 3], edge_index=[2, 3574], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (223, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3840, 3], edge_index=[2, 3840], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (229, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3988, 3], edge_index=[2, 3988], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (235, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4096, 3], edge_index=[2, 4096], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (241, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3682, 3], edge_index=[2, 3682], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (247, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3668, 3], edge_index=[2, 3668], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (253, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3530, 3], edge_index=[2, 3530], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (259, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3690, 3], edge_index=[2, 3690], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (265, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3830, 3], edge_index=[2, 3830], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (271, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3740, 3], edge_index=[2, 3740], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (277, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3854, 3], edge_index=[2, 3854], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (283, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3694, 3], edge_index=[2, 3694], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (289, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3800, 3], edge_index=[2, 3800], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (295, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3894, 3], edge_index=[2, 3894], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (301, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4056, 3], edge_index=[2, 4056], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (307, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3822, 3], edge_index=[2, 3822], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (313, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3826, 3], edge_index=[2, 3826], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (319, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3476, 3], edge_index=[2, 3476], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (325, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4002, 3], edge_index=[2, 4002], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (331, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3810, 3], edge_index=[2, 3810], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (337, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3946, 3], edge_index=[2, 3946], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (343, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3386, 3], edge_index=[2, 3386], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (349, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[3726, 3], edge_index=[2, 3726], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (355, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4118, 3], edge_index=[2, 4118], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (361, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4018, 3], edge_index=[2, 4018], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (367, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4036, 3], edge_index=[2, 4036], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (373, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4072, 3], edge_index=[2, 4072], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (379, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4134, 3], edge_index=[2, 4134], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (385, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4062, 3], edge_index=[2, 4062], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (391, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4032, 3], edge_index=[2, 4032], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (397, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4052, 3], edge_index=[2, 4052], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (403, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4072, 3], edge_index=[2, 4072], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (409, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4106, 3], edge_index=[2, 4106], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (415, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4060, 3], edge_index=[2, 4060], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (421, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4128, 3], edge_index=[2, 4128], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (427, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4058, 3], edge_index=[2, 4058], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (433, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4092, 3], edge_index=[2, 4092], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (439, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4092, 3], edge_index=[2, 4092], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (445, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4068, 3], edge_index=[2, 4068], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (451, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4134, 3], edge_index=[2, 4134], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (457, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4058, 3], edge_index=[2, 4058], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (463, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4046, 3], edge_index=[2, 4046], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (469, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4050, 3], edge_index=[2, 4050], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (475, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4042, 3], edge_index=[2, 4042], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (481, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4134, 3], edge_index=[2, 4134], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (487, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4078, 3], edge_index=[2, 4078], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (493, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4068, 3], edge_index=[2, 4068], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (499, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4078, 3], edge_index=[2, 4078], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (505, AtomicData(atomic_numbers=[57, 1], cell=[3, 3], edge_cell_shift=[4088, 3], edge_index=[2, 4088], forces=[57, 3], pbc=[3], pos=[57, 3], stress=[6], total_energy=[1])), (511, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6156, 3], edge_index=[2, 6156], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (517, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6290, 3], edge_index=[2, 6290], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (523, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6008, 3], edge_index=[2, 6008], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (529, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6392, 3], edge_index=[2, 6392], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (535, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6314, 3], edge_index=[2, 6314], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (541, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6118, 3], edge_index=[2, 6118], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (547, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6106, 3], edge_index=[2, 6106], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (553, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6176, 3], edge_index=[2, 6176], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (559, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6238, 3], edge_index=[2, 6238], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (565, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6220, 3], edge_index=[2, 6220], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (571, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6350, 3], edge_index=[2, 6350], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (577, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6262, 3], edge_index=[2, 6262], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (583, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[5870, 3], edge_index=[2, 5870], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (589, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6124, 3], edge_index=[2, 6124], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (595, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6102, 3], edge_index=[2, 6102], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (601, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6400, 3], edge_index=[2, 6400], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (607, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6322, 3], edge_index=[2, 6322], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (613, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[5878, 3], edge_index=[2, 5878], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (619, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[5902, 3], edge_index=[2, 5902], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (625, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[5626, 3], edge_index=[2, 5626], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (631, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[6158, 3], edge_index=[2, 6158], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (637, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[5808, 3], edge_index=[2, 5808], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (643, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[5666, 3], edge_index=[2, 5666], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (649, AtomicData(atomic_numbers=[81, 1], cell=[3, 3], edge_cell_shift=[5440, 3], edge_index=[2, 5440], forces=[81, 3], pbc=[3], pos=[81, 3], stress=[6], total_energy=[1])), (655, AtomicData(atomic_numbers=[9, 1], cell=[3, 3], edge_cell_shift=[696, 3], edge_index=[2, 696], forces=[9, 3], pbc=[3], pos=[9, 3], stress=[6], total_energy=[1])), (661, AtomicData(atomic_numbers=[54, 1], cell=[3, 3], edge_cell_shift=[3068, 3], edge_index=[2, 3068], forces=[54, 3], pbc=[3], pos=[54, 3], stress=[6], total_energy=[1]))]]'. Reason: 'OSError(24, 'Too many open files')'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/bin/nequip-train", line 8, in <module>
    sys.exit(main())
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/scripts/train.py", line 74, in main
    trainer = fresh_start(config)
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/scripts/train.py", line 177, in fresh_start
    dataset = dataset_from_config(config, prefix="dataset")
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/data/_build.py", line 78, in dataset_from_config
    instance, _ = instantiate(
  File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/utils/auto_init.py", line 234, in instantiate
    raise RuntimeError(
RuntimeError: Failed to build object with prefix `dataset` using builder `ASEDataset`

[Linux-cpp-lisp]

Hm... sounds like this issue from PyTorch:

https://discuss.pytorch.org/t/too-many-open-files-when-using-dataloader/9476/6

If you add

import torch.multiprocessing
torch.multiprocessing.set_sharing_strategy('file_system')

to the top of dataset.py, does it solve the problem? If it does I can add some kind of environment variable to control this strategy, since whether you ought to use it is a system dependent thing (see https://pytorch.org/docs/stable/multiprocessing.html#sharing-strategies).

[Linux-cpp-lisp]

It's also possible that this is just too many tensors to share, and that I need to batch the frames on the workers first so that fewer bigger tensors are getting moved...

[ssrokyz]

@Linux-cpp-lisp
It worked for a small dataset (600 structures)!
I think the solution you suggested was it! (Change sharing strategy from 'file descriptor' to 'file system'.)
But it fails with a large one (60000 structures).

With a large NEQUIP_NUM_TASKS (32 or 64 tested and failed), it gives the error message below.

Traceback (most recent call last):
File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/threading.py", line 932, in _bootstrap_inner
self.run()
File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/threading.py", line 870, in run
self._target(*self._args, **self._kwargs)
File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/multiprocessing/pool.py", line 576, in _handle_results
task = get()
File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/multiprocessing/connection.py", line 251, in recv
return _ForkingPickler.loads(buf.getbuffer())
File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/torch/multiprocessing/reductions.py", line 305, in rebuild_storage_filename
storage = cls._new_shared_filename(manager, handle, size)
RuntimeError: falseINTERNAL ASSERT FAILED at "/home/conda/feedstock_root/build_artifacts/pytorch-recipe_1644986876908/work/aten/src/ATen/MapAllocator.cpp":323, please report a bug to PyTorch. unable to mmap 41848 bytes from file </torch_171653_4588>: Cannot allocate memory (12)

With a small number of NEQUIP_NUM_TASKS=8 or =2, it also fails but gives a different error message like below.

....(omitted)....(66448, AtomicData(atomic_numbers=[64, 1], cell=[3, 3], edge_cell_shift=[3014, 3], edge_index=[2, 3014], forces=[64, 3], pbc=[3], pos=[64, 3], stress=[6], total_energy=[1])), (66456, AtomicData(atomic_numbers=[64, 1], cell=[3, 3], edge_cell_shift=[3812, 3], edge_index=[2, 3812], forces=[64, 3], pbc=[3], pos=[64, 3], stress=[6], total_energy=[1]))]]'. Reason: 'RuntimeError('falseINTERNAL ASSERT FAILED at "/home/conda/feedstock_root/build_artifacts/pytorch-recipe_1644986876908/work/aten/src/ATen/MapAllocator.cpp":323, please report a bug to PyTorch. unable to mmap 88 bytes from file </torch_178007_64057>: Cannot allocate memory (12)')'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/grad/ssrokyz/.conda/envs/nequip-multi/bin/nequip-train", line 8, in
sys.exit(main())
File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/scripts/train.py", line 74, in main
trainer = fresh_start(config)
File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/scripts/train.py", line 177, in fresh_start
dataset = dataset_from_config(config, prefix="dataset")
File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/data/_build.py", line 78, in dataset_from_config
instance, _ = instantiate(
File "/home/grad/ssrokyz/.conda/envs/nequip-multi/lib/python3.8/site-packages/nequip/utils/auto_init.py", line 234, in instantiate
raise RuntimeError(
RuntimeError: Failed to build object with prefix dataset using builder ASEDataset

All seem like memory issues.
The strange thing about it is that my system has 256GB size memory and the dataset is only 262MB in size.
Another one is that no matter what NEQUIP_NUM_TASKS I set, a single subprocess always reserves 12GB of memory.
Does it mean 12GB times NEQUIP_NUM_TASKS amount of memory is required? 😀

[Linux-cpp-lisp]

That is definitely way too much memory; I think using the shared memory approach for this kind of large, one-time transfer was a mistake. I've pushed a fix that roundtrips the processed data through a tempfile to completely avoid these kinds of problems.

Does the update work for you?

[ssrokyz]

@Linux-cpp-lisp
First, IT WORKED 😀 for data parsing, and the size of the data.pth is EXACTLY SAME! (8282203428 bytes)
However, the cmp command you mentioned(#182 (reply in thread)) says that data.pth files are differ. Maybe it's the sequence of data?

One more thing, the training process seems weird.
Each epoch trains only 120 batches, but it should be (60000 structures)/(5 per batch)= 12000 batches.
It works like my dataset has only 600 structures. (I double-checked that I didn't use the small dataset. Also checked params.yaml file.)
I set the NEQUIP_NUM_TASKS=32, batch size=5, and validation dataset has 600 structures.

Below are the last few dozen lines of the log file.

 training
 # Epoch batch         loss       loss_f       loss_e    Ge_f_rmse    Sb_f_rmse    Te_f_rmse psavg_f_rmse     e/N_rmse
      43   100         60.7       0.0607       0.0166        0.141        0.152        0.154        0.149       0.0789
     43   120         51.3       0.0513        0.014        0.136        0.147        0.136         0.14       0.0724

validation
# Epoch batch         loss       loss_f       loss_e    Ge_f_rmse    Sb_f_rmse    Te_f_rmse psavg_f_rmse     e/N_rmse
     43     2         19.9       0.0199       0.0205       0.0829        0.101       0.0814       0.0884       0.0876


  Train      #    Epoch      wal       LR       loss_f       loss_e         loss    Ge_f_rmse    Sb_f_rmse    Te_f_rmse psavg_f_rmse     e/N_rmse
! Train              43  473.920    0.001       0.0631       0.0607         63.2        0.144        0.157        0.155        0.152        0.151
! Validation         43  473.920    0.001       0.0391        0.291         39.4        0.154        0.139         0.13        0.141         0.42
Wall time: 473.92103092884645
! Best model       43   39.415

training
# Epoch batch         loss       loss_f       loss_e    Ge_f_rmse    Sb_f_rmse    Te_f_rmse psavg_f_rmse     e/N_rmse
     44   100         62.7       0.0627       0.0189        0.149        0.154        0.155        0.152       0.0842
     44   120         79.3       0.0793       0.0192        0.155        0.172        0.179        0.169       0.0847

validation
# Epoch batch         loss       loss_f       loss_e    Ge_f_rmse    Sb_f_rmse    Te_f_rmse psavg_f_rmse     e/N_rmse
     44     2         19.5       0.0195       0.0197       0.0812       0.0996       0.0813       0.0874        0.086


  Train      #    Epoch      wal       LR       loss_f       loss_e         loss    Ge_f_rmse    Sb_f_rmse    Te_f_rmse psavg_f_rmse     e/N_rmse
! Train              44  484.528    0.001       0.0628       0.0595         62.8        0.143        0.156        0.155        0.151        0.149
! Validation         44  484.528    0.001       0.0387        0.289           39        0.153        0.138         0.13        0.141        0.418
Wall time: 484.528886647895
! Best model       44   38.995

And I found that multi-processes were remaining in the training process, even the data parsing was finished.

ssrokyz  184229 184225 62 14:23 ?        00:06:28 /home/grad/ssrokyz/.conda/envs/nequip-multi/bin/python /home/grad/ssrokyz/.conda/envs/nequip-multi/bin/nequip-train input.yaml
ssrokyz  184237 184229  0 14:23 ?        00:00:00 /home/grad/ssrokyz/.conda/envs/nequip-multi/bin/python -c from multiprocessing.resource_tracker import main;main(17)
ssrokyz  184238 184229  0 14:23 ?        00:00:00 /home/grad/ssrokyz/.conda/envs/nequip-multi/bin/python -c from multiprocessing.forkserver import main; main(22, 23, ['__main__'], *

One last thing,
When I kill the training process and restart the training with the processed dataset, it still works like the dataset has only 600 structures.

[ssrokyz]

@Linux-cpp-lisp
It is way faster! I think it nearly linearly scaling. A few minutes for data parsing of 60000 structures.
I didn't separate the training and validation dataset but used one trajectory file with n_train: 60000 & n_val: 600. The actual dataset size is 66457.
I'm attaching four files FYI.

4files.zip

[Linux-cpp-lisp]

Great! Glad that it actually works.

Re the "600 structures", this is because in your config you specify:

# training
n_train: 600                                                                     # number of training data
n_val: 6                                                                         # number of validation data
train_val_split: random

which means that it is randomly sampling 600 of the 66457 configurations to make up the training set, and sampling 6 others from the 66457 for the validation set. If you want to use 60000 frames for training, you need to set n_train: 60000.

[ssrokyz]

@Linux-cpp-lisp
Oh.....!! I got it! I'm so sorry for the confusion. It is my bad. 😭
I tested the training process, and the RMSEs are EXACTLY IDENTICAL due to the same seed values!

[Linux-cpp-lisp]

Ah fantastic, you mean it is identical between the main branch and the ase-multiprocessing branch?

[ssrokyz]

@Linux-cpp-lisp
Yes. Between main and ase-multiprocess branches.

[ssrokyz]

With a 64-core CPU, it took only 4 minutes before the training process starts.
3 minutes: multiprocess
1 minute: (I don't know)
Without multiprocess, It took 1-2 hours!
And no issue has been found yet.
Thank you @Linux-cpp-lisp for your good support!

[Linux-cpp-lisp]

This is great @ssrokyz !

I've pushed an update as well to make it handle float64, but I want to caution that this approach (using multiprocessing) is a little broken because of how it handles global options. (This applies both to global PyTorch and global NequIP options.) The complications that I'm realizing result from that mean that you ought to be a little careful with this, and that I will likely not be merging this approach to speeding it up into NequIP— I'm just not confident enough in the edge cases.

That said, if it works for you, definitely keep using it! Just be a little careful.

[ssrokyz]

Hi @Linux-cpp-lisp !
Your work on multiprocess is neat and good.
I'm using it every time and nothing weird happened.
I'm sorry that other users cannot use it always but also glad those who need it can use it!
Cheers! 😊

[Linux-cpp-lisp]

Thanks @ssrokyz --- I tried a few other things yesterday and they didn't work, so I may still merge this 😆

Good to hear!

[ssrokyz]

@Linux-cpp-lisp
I just found that the multiprocess branch was merged into the develop branch!
Congrats and thanks for your efforts! 😊

[Linux-cpp-lisp]

Yep— I managed to fix most of what was broken in a way that shouldn't break unexpectedly... and the speed-ups were too big to ignore 😄

[AhmedBourouis]

Hi @ssrokyz
Can you please show me how you were able to use multiprocess to parallelize the data parsing step? what changes in the code and YAML file should I make?
I have a relatively large dataset with 30000 structures (with 100 ~ 350 atoms per structure) and it's taking more than one hour to process the data before the training starts

[ssrokyz]

Hello, @AhmedBourouis !
As far as I know, multiprocess is now enabled by default in the latest build (0.5.5). (But I'm using the previous beta version.)
So you don't need to do anything. (As far as I know)
I guess that NequIP is not getting the number of cores for multiprocess.
Please try using the option OMP_NUM_THREADS=the number of cores and NEQUIP_NUM_TASKS=the number of cores
Of course, replace the "the number of cores" part with your case.

[AhmedBourouis]

Thanks, @ssrokyz! It's much faster now after setting OMP_NUM_THREADS and NEQUIP_NUM_TASKS