The slow rate of OC_10M.pb for MD simulation

[Guancred]

Hi, I run the MD simulation for sulfides (about 600 atoms) with the OC_10M.pb, the rate is shown in the following：

which exhibits the slow speed (0.332 ns/day) by adopting 40 cores.
CPU: 2 x Intel Xeon ICX Platinum 8358 (2.6GHz, 32 cores)

And my question is how to improve it?

[iProzd]

Hi @Guancred , the speed you're experiencing (~250ms/step) with 600 atoms using OC_10M.pb (a pretrained DPA1 model with 2 attention layers) on two CPUs is expected. This model is memory-intensive and optimized for use on a GPU.

To improve performance, consider the following options:

1. Switch to GPU: Using a GPU can significantly reduce processing time.
2. Use a Smaller Model: You might also try a smaller model like se_atten_v2, which has zero attention layers and is compressible. Note that you may need to train this model from scratch on your data.
3. Do model compression: Model compression can significantly boost the speed of a trained model. It's supported for limited models (such as se_atten_v2 with zero attention layers or se_a), see here for details.

By switching to a GPU, using a model with zero attention layers, and applying model compression, you can achieve speeds of around ~2.5ms/step, resulting in a 100x performance boost.

I hope this helps!