Performance bottleneck when using ENVIRONMENTSIMILARITY with OpenMM on HPC cluster with limited CPU threads

[KingOfChuXuan]

Hi PLUMED community,

I am running enhanced sampling simulations for ice nucleation using OpenMM + DeepMD-kit + PLUMED (openmmplumed). I am using ENVIRONMENTSIMILARITY combined with OPES_EXPANDED to drive the phase transition in a bulk water system (2592 atoms / 864 water molecules).

Problem Description

On my local workstation with 24 CPU threads assigned (PLUMED_NUM_THREADS=16), the simulation runs efficiently with >90% GPU utilization. However, on our remote HPC cluster where each GPU node is strictly limited to 12 CPU threads per GPU task, the simulation performance collapses dramatically:

• Without PLUMED (Pure OpenMM + DeepMD): ~10 ns/day
• With PLUMED (ENVIRONMENTSIMILARITY): Drops to ~1 ns/day
• GPU Utilization: Drops from >90% down to ~10% (GPU is idling most of the time waiting for CPU force evaluation).

Since I cannot change the hardware constraints on the cluster (12 CPU threads max per task), ENVIRONMENTSIMILARITY has become a severe CPU-bound bottleneck.

Questions / Request for Advice

Without changing the remote hardware or upgrading the node allocation, what actionable options do I have to improve simulation speed and remove this CPU bottleneck?

PLUMED Script

env: ENVIRONMENTSIMILARITY ...
   SPECIES=1-2592:3
   SIGMA=0.05
   CRYSTAL_STRUCTURE=CUSTOM
   REFERENCE_1=ref1.pdb
   REFERENCE_2=ref2.pdb
   REFERENCE_3=ref3.pdb
   REFERENCE_4=ref4.pdb
   MORE_THAN={RATIONAL R_0=0.65 NN=8 MM=16}
   MEAN
...

n_ice: CUSTOM ARG=env.morethan FUNC=x PERIODIC=NO
ecv: ECV_UMBRELLAS_LINE ARG=n_ice TEMP=260.0 CV_MIN=360 CV_MAX=432 SIGMA=1.0 BARRIER=50
opes: OPES_EXPANDED ARG=ecv.* PACE=2000 STRIDE=1

PRINT STRIDE=500 ARG=env.mean,n_ice,ecv.*,opes.bias FILE=colvar_opes.out
FLUSH STRIDE=500

Any suggestions or optimization tips would be greatly appreciated!

[PabloPiaggi]

Hi @KingOfChuXuan , I suggest two things:

1. You can add STRIDE=2 in the OPES_EXPANDED line, to evaluate the collective variables every 2 steps. You are probably using a timestep of 0.5 fs in MD and evaluating CVs that depend on the position of oxygens every 2 steps should be a good approximation.
2. Make sure that the PLUMED parallelization with OpenMP is working as intended. Set explicitly PLUMED_NUM_THREADS=12 if you are requesting 12 CPUs per GPU or MPI process. You can check in the plumed output file if these options were processed correctly. Make sure that your PLUMED compilation support OpenMP.

Also, there were some recent efforts to accelerate some PLUMED actions with GPUs, but I don't know if this reached ENVIRONMENTSIMILARITY. Maybe @Iximiel can comment on this.

[KingOfChuXuan]

Hi @PabloPiaggi , thank you for the suggestions.

You are probably using a timestep of 0.5 fs in MD

My timestep is already 1 fs, so I'm not sure if increasing OPES_EXPANDED STRIDE is safe. Would STRIDE=2 at 1 fs still be okay? Or how can I figure out if it's safe?

Make sure that the PLUMED parallelization with OpenMP is working as intended.

OpenMP is definitely working. PLUMED_NUM_THREADS=12 is faster than 8, and the log shows Number of threads: 12. So that's not the problem.

Besides, I suspect the bottleneck is CPU-GPU communication, not the CV itself. I tested this: ENVIRONMENTSIMILARITY alone runs fine with no noticeable slowdown. The performance collapse only happens when the bias force is applied (i.e. when UPPER_WALLS or OPES_EXPANDED is active). So I think the issue might be in the force transfer between CPU (PLUMED) and GPU (OpenMM), not in computing the CV. Any idea how I can check or profile this?

[PabloPiaggi]

The way to check if STRIDE>1 is safe is described here (Ferrarotti, Bottaro, Perez-Villa, and Bussi, J. Chem. Theory Comput. 11, 139 (2015)). However, I am not sure if you can easily compute the conserved quantity. I think using STRIDE=2 with 1 fs could lead to accurate sampling, but it is best to test it.

If the OpenMP parallelization is working fine, I am afraid I don't have other suggestions. I don't have first hand experience with the OpenMM/PLUMED interface. Perhaps the core developers can comment on the potential communication bottleneck between OpenMM and PLUMED.

[KingOfChuXuan]

Thanks for the reference! I'll test STRIDE=2 and see how it goes.