Memory and Performance Issues with JAX solver

[badisa]

I have been generating forward and reverse dG plots (example below) from a large number of U_klns in parallel and noticed that memory consumption balloons over time.

Example of forward and reverse dG plots (Figure 5 from https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4420631/)

The issue is that when JAX jit's a function, the function is cached based on the input shape (https://stackoverflow.com/a/75642593 https://docs.jax.dev/en/latest/export/shape_poly.html). So as the shape of the U_klns change over time the number of cached functions increases, leading to large amounts (> GB) of memory being allocated and not freed.

Below is a comparison of the memory usage of MBAR with and without JAX
Test Code: https://github.com/badisa/pymbar/blob/experiment/jax-vs-numpy-memory-usage/pymbar/tests/test_mbar.py#L550-L568

MBAR Numpy Memory Consumption

MBAR Jax Memory Consumption

The performance of using JAX over Numpy is also significantly worse. The test linked to above takes 5 seconds with Numpy (https://github.com/choderalab/pymbar/actions/runs/17384102513/job/49347464100?pr=563#step:6:228) and 111 seconds with Jax (https://github.com/choderalab/pymbar/actions/runs/17384102513/job/49347464103?pr=563#step:6:222), which is more than a 22x slow down.

It is possible to use jax.clear_caches() to resolve the memory issue, but this would harm the performance of other jax code I have. The best solution would likely be to merge #509 or provide some way to disable jit in Pymbar (there is the force_no_jax variable, but not clear how to cleanly modify that at runtime).

[ijpulidos]

Thanks for the detailed report. If I'm getting this, you used the force_no_jax variable in your fork to test between jax and no-jax (numpy) results? Just making sure I can reproduce your results here. On the other hand, I do think we want a way to control this. What's suggested in #509 makes sense to me in general, but it probably requires some more thinking about the API, because this would change it drastically, as far as I can see.

[badisa]

I had tested this by creating an environment with and without JAX. I also verified that setting force_no_jax=True has the desired effect.

No opinion on how this gets resolved, just want a way to disable the JAX solver at runtime.

[mikemhenry]

Thanks for this! Another workaround is to install pymbar-core which doesn't have the jax acceleration (but depending on other packages you have, they might require pymbar).

[badisa]

Another workaround is to install pymbar-core which doesn't have the jax acceleration

Are the differences between pymbar and pymbar-core documented somewhere? And is there a pip package or only conda, as I am trying to get away from conda.

[mikemhenry]

pip install pymbar will install pymbar from PyPI without jax, pip install pymbar[jax] will install pymbar with jax acceleration from PyPI. This is documented here: https://github.com/choderalab/pymbar?tab=readme-ov-file#installation in our README.

The differences between pymbar-core and pymbar (talking about the conda packages now) is just the inclusion of jax https://github.com/conda-forge/pymbar-feedstock/blob/main/recipe/meta.yaml#L50-L58

Let me know if you have any more questions!

[mikemhenry]

We now have an environmental variable PYMBAR_DISABLE_JAX that you can set to TRUE (export PYMBAR_DISABLE_JAX=TRUE) to disable JAX at runtime. I want to leave this issue open to track the memory and performance bug reported.