Unbalanced memory distribution with multiple gpus

[Hongyu-yu]

When doing simulation in Lammps with pair allegro on multiple gpus, the memory distribution is not equal and differs a lot.
Details of the simulation:
NPT from 1K to 1200 K with 0 pressure.
System: LiNbO3 with a 8 * 8 * 8 supercell of 7680 atoms and the box is:

ITEM: BOX BOUNDS xy xz yz pp pp pp
-2.0891960044998164e+01 4.1783920089996343e+01 -2.0891960044998164e+01
0.0000000000000000e+00 3.6185936267635796e+01 0.0000000000000000e+00
0.0000000000000000e+00 5.6516967256636690e+01 0.0000000000000000e+00

It is equally distributed at the early stage of the simulation but then the distribution is weird.
Is that normal? Any suggestions to avoid that? Thanks!

[Linux-cpp-lisp]

[Hongyu-yu]

Hi @Linux-cpp-lisp ,

Sure! However, I don't find a discussion part in pair allegro. May you first build a discussion part there? Thanks!

[Linux-cpp-lisp]

@Hongyu-yu --- entirely my fault, apologies, I hadn't enabled discussions here and that's why I couldn't transfer it.

[Hongyu-yu]

It's okay :) Thanks for the transfer.

[Linux-cpp-lisp]

Regarding the actual issue, I believe LAMMPS does some logging about how many atoms it puts where? Hard to say otherwise whether this is an issue at the LAMMPS or PyTorch level.

balance or fix balance in LAMMPS may resolve it if the actual atom distribution is strange. (Though, in a bulk material, such a difference in atomic density would probably be a sign of something going wrong.)

[Hongyu-yu]

Thanks for acknowledging fix balance in Lammps and I will try it soon.
After a detailed check, something went wrong in the simulation that atoms flew away when temperature was high.
So, it could be a problem with the quality of the potential from a bad dataset, I think.
For sampling, I use DFT NPT to get about 1400 training configurations. The better plan should be active learning maybe, such as DPGEN and the newly developed psiflow. Do you have any suggestions for a better sampling method on this stage? Thanks!

[Linux-cpp-lisp]

That sounds right 😆

Often these things happen when two atoms approach each other too closely and go out of domain of the training data (see, for example, Fig 6 of the excellent https://arxiv.org/abs/2210.07237). The simplest solution to try is ensure that your DFT data is always sampled from a higher temperature than the highest temperature you intend to run MD at. For NPT, it can also be helpful to include a few frames run at a higher pressure than the highest pressure you intend to run at, as well. Often following these two simple rules of thumb should produce a stable potential with our models.

Active learning is definitely a nice approach as well, if more complicated. Psiflow certainly sounds interesting, and they do appear to support Allegro right out of the box: https://github.com/svandenhaute/psiflow/blob/main/psiflow/models/_nequip.py#L136-L160. Another active learning approach you may want to consider is https://github.com/mir-group/flare/ from our group.

[Hongyu-yu]

Thanks for the detailed and quick reply!

Your suggestions are great! Maybe I need to rebuild the dataset.

Flare sounds good! I will try it. Could it work with Allegro to build potential or providing reliable dataset with active learning? Thanks!

[Linux-cpp-lisp]

Yes, Allegro can be used with data from any source / sampling scheme including active learning; there is no requirement that the data come from AIMD.