GHMC is slow

[kyleabeauchamp]

My crude experiments suggest that, for a 4500 atom system, the GHMC integrator is 3 or 4 times slower than Langevin.

Given this slowdown, most users would be more inclined to just reduce their timestep and stick with Langevin. A faster implementation of GHMC might help with that.

[jchodera]

I've noticed a number of CustomIntegrator versions are pretty slow, so I'm wondering if there is some general optimization that can be done.

Were you using CUDA or OpenCL here?

[kyleabeauchamp]

CUDA

[jchodera]

It certainly shouldn't be 3-4 times slower. Yes, it does require the force and energy be computed each timestep, but this shouldn't be 4x slower.

One issue is that sigma is recomputed each iteration, rather than just once at the beginning:
https://github.com/choderalab/openmmtools/blob/master/OpenMMTools/integrators.py#L553

If we extended the API to add a addInitializePerDof() method that would ensure this is just called once, that would speed things up.

There might be other hidden issues.

[jchodera]

For now, a viable route would be to adaptively tune the timestep for very high acceptance rates (e.g. 99.999%) with GHMC then lock in that timestep for VVVR.

[jchodera]

CUDA

Could you give OpenCL a try?

@peastman is certainly more familiar with the implementation details, but I believe the CUDA version interprets the various steps to launch kernels on the GPU, while the OpenCL version actually compiles the integrator into a GPU kernel. But I may be wrong---that may have been an older implementation.

[peastman]

The CUDA and OpenCL versions are nearly identical.

Looking at the code, I see several reasons for it to be slower. First, it applies constraints five different times (twice for positions and three times for velocities). The Langevin integrator only applies constraints once. Depending on what sort of constraints your system has, that could have a big impact on speed.

Second, it requires two force/energy evaluations per time step. You use the force and energy, so it has to compute them:

integrator.addComputeGlobal("Eold", "ke + energy")
...
integrator.addComputePerDof("v", "v + 0.5*dt*f/m")

Then you modify the positions, so the forces it computed before are no longer valid:

integrator.addComputePerDof("x", "x + v*dt")

Then you use the force and energy again, so it has to recompute them:

integrator.addComputePerDof("v", "v + 0.5*dt*f/m + (x-x1)/dt")
...
integrator.addComputeGlobal("Enew", "ke + energy")

And then you modify the positions yet again, which again invalidates the forces and means they'll have to be recomputed at the start of the next time step:

integrator.addComputePerDof("x", "x*accept + xold*(1-accept)")

Anything that modifies the positions invalidates the forces, either addComputePerDof("x", ...), addConstrainPositions(), or potentially addUpdateContextState().

[kyleabeauchamp]

Is there something like a "stack" that can be used to cache the old forces, energies, and positions without massive recomputation?

[peastman]

No, it only caches one set of forces at a time, then throws them out when anything causes them to become invalid.

[jchodera]

@kyleabeauchamp Did you want me to take a stab at speeding this up? Or are you on that?

[kyleabeauchamp]

I think eventually we should do that. My own work is still at the level of running some basic tests on these things.

[kyleabeauchamp]

I think the easiest way to speed up our GHMC code is to avoid re-calculating the energy with every timestep. This is easily done by doing n_steps iterations of hamiltonian dynamics with every round of GHMC; our current code essentially hard-wires n_steps = 1, which kills us on the energy calculation.

[jchodera]

We can definitely expose the number of steps as a parameter. The acceptance rate falls off with the number of steps, but there is likely an optimal, and there may be other tricks (as you've pointed out) to keep acceptance rates high.

Should I submit a PR?

[kyleabeauchamp]

Let's punt until this project is our top priority. I believe I already have this code written as well.

[jchodera]

Did you experiment with removing the extra velocity/position constraint steps too? I'm really curious why the CustomIntegrator is so slow---is it really just the energy call?

[kyleabeauchamp]

I did not experiment with the constraint steps yet.

[jchodera]

I think the easiest way to speed up our GHMC code is to avoid re-calculating the energy with every timestep. This is easily done by doing n_steps iterations of hamiltonian dynamics with every round of GHMC; our current code essentially hard-wires n_steps = 1, which kills us on the energy calculation.

Oh, there's another way to avoid this that doesn't involve doing extra hamiltonian dynamics inside GHMC (which decreases the acceptance rate): Have your n_steps parameter run a block of GHMC steps where the steps are unrolled in the integrator such that the total energy from the last GHMC step is not recomputed for the next GHMC step.

[kyleabeauchamp]

I think I like the idea of just exposing n_steps for now, as it's a hyperparameter than can be tuned rather easily. It also cleans up the code base