How the MACE implement the multi-GPU lammps inference?

[SyntaxSmith]

For my best known, there are two modern MLIP could inference use multi-GPU, the one is Allegro which is not MPNN, don't need consider the message passing method across different GPUs. the other is SevenNet, which implement foward communication and backward communication methods during lammps inference. I can't find similar inplementation in MACE's code. Could you explain how you make multi-GPU inference usable?

[ilyes319]

Communicating internal messages between domains is not a requirement it just makes things more efficient by reducing the impact of ghost atoms. To be clear, this impact of ghost atoms is not a product of message passing per se just of a large receptive field. If you make an Allegro model with 10A receptive field, you will have a similar problem. Just with message passing model you can do better.

Work is being done on an interface that includes internal inter GPUs communication for MACE in lammps.

[SyntaxSmith]

If I have a wrong concept? Allegro use spherical harmonics to generate a environment embedding, then update a atomic embedding through its environment embedding, in this process, neighbors' feature doesn't influence center atom's feature,which is different from MPNN scheme. When I use a MPNN to inference, I must exchange ghost atoms' features between interaction layers through MPI communication to ensure the center atom's feature can update correctly, which lammps couldn't handle it automatically. I want to know in MACE, this problem's solution.

[ilyes319]

What you need to understand is that MPI is NOT required to evaluate an MPNN across multiple GPUs, it just makes things more efficient.
Obviously Allegro will also have ghost atoms and their number will grow with the receptive field. A 2 layers MACE with 5A cutoff will have the same number of ghost atoms as a 10A Allegro. MACE is a local model too so you can evaluate it as a local model.... and that is what is happening now. We just consider MACE as a local model with twice the cutoff.

Now the second concept that is probably confusing you, is that with MPNNs you can in principle do BETTER than the corresponding local model if you MPI calls. So in principle a 2 layers MPNN with 5A cutoff could be evaluated more efficiently with respect to ghost atoms than a 10A local model!

[SyntaxSmith]

Thanks,I 'll rethink this problem