OpenMP

[terryfrankcombe]

(Apologies in advance for the noob question!)

I'm running nequip on CPU. Can I convince the computationally expensive bits of the code to run with OpenMP? Simply setting OMP_NUM_THREADS has no obvious effect.

[Linux-cpp-lisp]

Hi @terryfrankcombe,

I've tried to play around with this before... basically: any use of OpenMP will be inherited from PyTorch, so your best bet is to look through the PyTorch discussion forums to find out the best way to enable this for your hardware. It is very picky to get this working, in my experience, but these things also change a lot between PyTorch versions and I last looked at this a while ago.

Pay particular attention to the combination of whether or not PyTorch is using MKL, whether you have AMD or Intel, etc. --- you may for example need to set MKL_NUM_THREADS.

[tgmaxson]

Just to chime in with some of my experience, using NumPy from Conda (not using MKL), I am able to get a small speedup of 1.5-2X by setting OMP_NUM_THREADS=8, but only for CPU only training on an AMD EPYC processor. It's worth a try if you don't have a GPU available.

[Linux-cpp-lisp]

Good to know, thanks for adding this. Worth noting I've in the past used MKL with the "secret" option to make it work on AMD, and that may or may not give you more CPU performance with PyTorch... unclear.

[terryfrankcombe]

Since writing this I have discovered that at least the ASE calculator takes device="CPU" and device="CUDA" args, which seem to work as expected. Using the CPU exclusively does indeed use all available physical cores (at least on one of my machines... on another cluster it seems to stick to one thread).

I have a small model, so which choice is faster depends on what task I'm performing.

Incidentally, when running trajectories I've found for my small system that running multiple copies (multiple trajectories) in a big grid (large separations) is more efficient, at least until I get beyond dozens of simultaneous trajectories.

[Linux-cpp-lisp]

@terryfrankcombe 👍

Incidentally, when running trajectories I've found for my small system that running multiple copies (multiple trajectories) in a big grid (large separations) is more efficient, at least until I get beyond dozens of simultaneous trajectories.

This is definitely expected; worth noting that the models themselves have full support for batched calculations over multiple independent systems in a single call without this hack, including independent systems with PBC and different cells. This just isn't exposed in the ASE calculator or LAMMPS since, to the best of my knowledge, they don't support that kind of multi-MD thing. (I guess you could set something up where you have a bunch of fake calculators that just store the inputs and block until all a received, run a batched computation, and then return--- maybe something like this is good for your application so that you can correctly run separate thermostats on each individual system?)

[Linux-cpp-lisp]

Latest discussion on CPUs can be found here: #303 (comment)

(Apparently building from source makes the speed much better)