openfe icon indicating copy to clipboard operation
openfe copied to clipboard
verified publisher icon OpenFreeEnergy

Inquiry regarding ABFE calculation time on RTX4090 and JAX CPU fallback

Open ORCAaAaA-ui opened this issue 2 months ago • 0 comments

Describe the bug

I recently ran the ABFE tutorial (Toluene complex) using OpenFE 1.7. Despite running on a machine with an NVIDIA RTX 4090, the calculation took approximately 15 hours and 45 minutes to complete.

I noticed a warning log indicating that JAX failed to initialize the CUDA backend and fell back to CPU. I am not entirely sure if JAX is heavily utilized for the core computation in this specific tutorial, but I suspect this might be related to the unexpected performance drop.

I would like to ask:

  1. Is a runtime of ~15 hours expected for the Toluene ABFE tutorial on a single RTX 4090 (especially with protocol_repeats = 1)?

  2. Crucially, while monitoring the simulation, I observed that the GPU utilization (RTX 4090) frequently dropped to single-digit percentages (<10%), Is this pattern expected during OpenFE/ABFE calculations?

  3. Does the JAX CPU fallback significantly impact the performance of OpenFE/OpenMM simulations in this context?

To Reproduce

Install OpenFE on a Linux machine with an NVIDIA GPU using conda-forge (default installation).

Run the ABFE tutorial notebook (Toluene complex).

Observe the JAX initialization warning in the logs.

Observe the total simulation time.

Software versions

Operating System: Ubuntu 22.04

Install Method: conda-forge

GPU: NVIDIA GeForce RTX 4090

Conda List Output:

Name Version Build Channel

_libgcc_mutex 0.1 conda_forge conda-forge _openmp_mutex 4.5 2_gnu conda-forge _python_abi3_support 1.0 hd8ed1ab_2 conda-forge adjusttext 1.3.0 pyhd8ed1ab_0 conda-forge adwaita-icon-theme 49.0 unix_0 conda-forge ambertools 24.8 cuda_None_nompi_py313hb0f039c_101 conda-forge amberutils 21.0 pypi_0 pypi annotated-types 0.7.0 pyhd8ed1ab_1 conda-forge anyio 4.11.0 pyhcf101f3_0 conda-forge argon2-cffi 25.1.0 pyhd8ed1ab_0 conda-forge argon2-cffi-bindings 25.1.0 py313h07c4f96_1 conda-forge arpack 3.9.1 nompi_hf03ea27_102 conda-forge arrow 1.4.0 pyhcf101f3_0 conda-forge asttokens 3.0.0 pyhd8ed1ab_1 conda-forge async-lru 2.0.5 pyh29332c3_0 conda-forge at-spi2-atk 2.38.0 h0630a04_3 conda-forge at-spi2-core 2.40.3 h0630a04_0 conda-forge atk-1.0 2.38.0 h04ea711_2 conda-forge attrs 25.4.0 pyh71513ae_0 conda-forge babel 2.17.0 pyhd8ed1ab_0 conda-forge beautifulsoup4 4.14.2 pyha770c72_0 conda-forge blas 2.138 openblas conda-forge blas-devel 3.9.0 38_h1ea3ea9_openblas conda-forge bleach 6.2.0 pyh29332c3_4 conda-forge bleach-with-css 6.2.0 h82add2a_4 conda-forge blosc 1.21.6 he440d0b_1 conda-forge brotli 1.2.0 h41a2e66_0 conda-forge brotli-bin 1.2.0 hf2c8021_0 conda-forge brotli-python 1.2.0 py313h09d1b84_0 conda-forge bson 0.5.10 pyhd8ed1ab_0 conda-forge bzip2 1.0.8 hda65f42_8 conda-forge c-ares 1.34.5 hb9d3cd8_0 conda-forge c-blosc2 2.22.0 h4cfbee9_0 conda-forge ca-certificates 2025.10.5 hbd8a1cb_0 conda-forge cached-property 1.5.2 hd8ed1ab_1 conda-forge cached_property 1.5.2 pyha770c72_1 conda-forge cachetools 6.2.1 pyhd8ed1ab_0 conda-forge cairo 1.18.4 h3394656_0 conda-forge certifi 2025.10.5 pyhd8ed1ab_0 conda-forge cffi 2.0.0 py313hf46b229_1 conda-forge cftime 1.6.4 py313h29aa505_2 conda-forge charset-normalizer 3.4.4 pyhd8ed1ab_0 conda-forge cinnabar 0.5.0 pyhd8ed1ab_0 conda-forge click 8.3.0 pyh707e725_0 conda-forge click-option-group 0.5.6 pyhd8ed1ab_0 conda-forge colorama 0.4.6 pyhd8ed1ab_1 conda-forge comm 0.2.3 pyhe01879c_0 conda-forge contourpy 1.3.3 py313h7037e92_2 conda-forge cpython 3.13.9 py313hd8ed1ab_101 conda-forge cudatoolkit 11.8.0 h4ba93d1_13 conda-forge cycler 0.12.1 pyhd8ed1ab_1 conda-forge cyrus-sasl 2.1.28 hd9c7081_0 conda-forge dbus 1.16.2 h3c4dab8_0 conda-forge debugpy 1.8.17 py313h5d5ffb9_0 conda-forge decorator 5.2.1 pyhd8ed1ab_0 conda-forge defusedxml 0.7.1 pyhd8ed1ab_0 conda-forge dill 0.4.0 pyhd8ed1ab_0 conda-forge edgembar 3.0 pypi_0 pypi epoxy 1.5.10 hb03c661_2 conda-forge exceptiongroup 1.3.0 pyhd8ed1ab_0 conda-forge executing 2.2.1 pyhd8ed1ab_0 conda-forge fftw 3.3.10 nompi_hf1063bd_110 conda-forge filelock 3.20.0 pyhd8ed1ab_0 conda-forge flexcache 0.3 pyhd8ed1ab_1 conda-forge flexparser 0.4 pyhd8ed1ab_1 conda-forge font-ttf-dejavu-sans-mono 2.37 hab24e00_0 conda-forge font-ttf-inconsolata 3.000 h77eed37_0 conda-forge font-ttf-source-code-pro 2.038 h77eed37_0 conda-forge font-ttf-ubuntu 0.83 h77eed37_3 conda-forge fontconfig 2.15.0 h7e30c49_1 conda-forge fonts-conda-ecosystem 1 0 conda-forge fonts-conda-forge 1 0 conda-forge fonttools 4.60.1 py313h3dea7bd_0 conda-forge fqdn 1.5.1 pyhd8ed1ab_1 conda-forge freetype 2.14.1 ha770c72_0 conda-forge freetype-py 2.3.0 pyhd8ed1ab_0 conda-forge fribidi 1.0.16 hb03c661_0 conda-forge fsspec 2025.10.0 pyhd8ed1ab_0 conda-forge gdk-pixbuf 2.44.4 h2b0a6b4_0 conda-forge glib-tools 2.86.1 hf516916_1 conda-forge gmp 6.3.0 hac33072_2 conda-forge gmpy2 2.2.1 py313h86d8783_1 conda-forge graphite2 1.3.14 hecca717_2 conda-forge graphviz 13.1.2 h87b6fe6_0 conda-forge greenlet 3.2.4 py313h7033f15_1 conda-forge griddataformats 1.0.2 pyhd8ed1ab_1 conda-forge gsd 4.2.0 h31e983c_0 conda-forge gtk3 3.24.43 h021d004_4 conda-forge gts 0.7.6 h977cf35_4 conda-forge gufe 1.7.0 pyhd8ed1ab_1 conda-forge h11 0.16.0 pyhd8ed1ab_0 conda-forge h2 4.3.0 pyhcf101f3_0 conda-forge h5py 3.15.1 nompi_py313h253c126_100 conda-forge harfbuzz 12.1.0 h15599e2_0 conda-forge hdf4 4.2.15 h2a13503_7 conda-forge hdf5 1.14.6 nompi_h6e4c0c1_103 conda-forge hicolor-icon-theme 0.17 ha770c72_2 conda-forge hpack 4.1.0 pyhd8ed1ab_0 conda-forge httpcore 1.0.9 pyh29332c3_0 conda-forge httpx 0.28.1 pyhd8ed1ab_0 conda-forge hyperframe 6.1.0 pyhd8ed1ab_0 conda-forge icu 75.1 he02047a_0 conda-forge idna 3.11 pyhd8ed1ab_0 conda-forge importlib-metadata 8.7.0 pyhe01879c_1 conda-forge importlib_resources 6.5.2 pyhd8ed1ab_0 conda-forge iniconfig 2.3.0 pyhd8ed1ab_0 conda-forge intel-openmp 2025.0.0 h06a4308_1171
ipycytoscape 1.3.3 pyhd8ed1ab_1 conda-forge ipykernel 7.1.0 pyha191276_0 conda-forge ipython 9.6.0 pyhfa0c392_0 conda-forge ipython_pygments_lexers 1.1.1 pyhd8ed1ab_0 conda-forge ipywidgets 8.1.8 pyhd8ed1ab_0 conda-forge isoduration 20.11.0 pyhd8ed1ab_1 conda-forge jax 0.5.2 pyhd8ed1ab_0 conda-forge jaxlib 0.5.2 cpu_py313ha650dea_0 conda-forge jedi 0.19.2 pyhd8ed1ab_1 conda-forge jinja2 3.1.6 pyhd8ed1ab_0 conda-forge joblib 1.5.2 pyhd8ed1ab_0 conda-forge jq 1.8.1 h73b1eb8_0 conda-forge json5 0.12.1 pyhd8ed1ab_0 conda-forge jsonpointer 3.0.0 py313h78bf25f_2 conda-forge jsonschema 4.25.1 pyhe01879c_0 conda-forge jsonschema-specifications 2025.9.1 pyhcf101f3_0 conda-forge jsonschema-with-format-nongpl 4.25.1 he01879c_0 conda-forge jupyter-lsp 2.3.0 pyhcf101f3_0 conda-forge jupyter_client 8.6.3 pyhd8ed1ab_1 conda-forge jupyter_core 5.9.1 pyhc90fa1f_0 conda-forge jupyter_events 0.12.0 pyh29332c3_0 conda-forge jupyter_server 2.17.0 pyhcf101f3_0 conda-forge jupyter_server_terminals 0.5.3 pyhd8ed1ab_1 conda-forge jupyterlab 4.4.10 pyhd8ed1ab_0 conda-forge jupyterlab_pygments 0.3.0 pyhd8ed1ab_2 conda-forge jupyterlab_server 2.28.0 pyhcf101f3_0 conda-forge jupyterlab_widgets 3.0.16 pyhcf101f3_1 conda-forge kartograf 1.2.0 pyhd8ed1ab_0 conda-forge kernel-headers_linux-64 5.14.0 he073ed8_2 conda-forge keyutils 1.6.3 hb9d3cd8_0 conda-forge kiwisolver 1.4.9 py313hc8edb43_1 conda-forge konnektor 0.2.0 pyhe01879c_0 conda-forge krb5 1.21.3 h659f571_0 conda-forge lark 1.3.1 pyhd8ed1ab_0 conda-forge lcms2 2.17 h717163a_0 conda-forge ld_impl_linux-64 2.44 h1aa0949_4 conda-forge lerc 4.0.0 h0aef613_1 conda-forge libabseil 20240722.0 cxx17_hbbce691_4 conda-forge libaec 1.1.4 h3f801dc_0 conda-forge libblas 3.9.0 38_h4a7cf45_openblas conda-forge libboost 1.86.0 hed09d94_4 conda-forge libboost-python 1.86.0 py313hfaae9d9_4 conda-forge libbrotlicommon 1.2.0 h09219d5_0 conda-forge libbrotlidec 1.2.0 hd53d788_0 conda-forge libbrotlienc 1.2.0 h02bd7ab_0 conda-forge libcblas 3.9.0 38_h0358290_openblas conda-forge libcups 2.3.3 hb8b1518_5 conda-forge libcurl 8.16.0 h4e3cde8_0 conda-forge libdeflate 1.25 h17f619e_0 conda-forge libdrm 2.4.125 hb03c661_1 conda-forge libedit 3.1.20250104 pl5321h7949ede_0 conda-forge libegl 1.7.0 ha4b6fd6_2 conda-forge libegl-devel 1.7.0 ha4b6fd6_2 conda-forge libev 4.33 hd590300_2 conda-forge libexpat 2.7.1 hecca717_0 conda-forge libffi 3.5.2 h9ec8514_0 conda-forge libfreetype 2.14.1 ha770c72_0 conda-forge libfreetype6 2.14.1 h73754d4_0 conda-forge libgcc 15.2.0 h767d61c_7 conda-forge libgcc-ng 15.2.0 h69a702a_7 conda-forge libgd 2.3.3 h6f5c62b_11 conda-forge libgfortran 15.2.0 h69a702a_7 conda-forge libgfortran-ng 15.2.0 h69a702a_7 conda-forge libgfortran5 15.2.0 hcd61629_7 conda-forge libgl 1.7.0 ha4b6fd6_2 conda-forge libgl-devel 1.7.0 ha4b6fd6_2 conda-forge libglib 2.86.1 h32235b2_1 conda-forge libglvnd 1.7.0 ha4b6fd6_2 conda-forge libglx 1.7.0 ha4b6fd6_2 conda-forge libglx-devel 1.7.0 ha4b6fd6_2 conda-forge libgomp 15.2.0 h767d61c_7 conda-forge libgrpc 1.67.1 h25350d4_2 conda-forge libhwloc 2.12.1 default_h3d81e11_1000 conda-forge libiconv 1.18 h3b78370_2 conda-forge libjpeg-turbo 3.1.2 hb03c661_0 conda-forge liblapack 3.9.0 38_h47877c9_openblas conda-forge liblapacke 3.9.0 38_h6ae95b6_openblas conda-forge liblzma 5.8.1 hb9d3cd8_2 conda-forge libmpdec 4.0.0 hb9d3cd8_0 conda-forge libnetcdf 4.9.2 nompi_h21f7587_118 conda-forge libnghttp2 1.67.0 had1ee68_0 conda-forge libntlm 1.8 hb9d3cd8_0 conda-forge libopenblas 0.3.30 pthreads_h94d23a6_3 conda-forge libpciaccess 0.18 hb9d3cd8_0 conda-forge libpng 1.6.50 h421ea60_1 conda-forge libpq 18.0 h3675c94_0 conda-forge libprotobuf 5.28.3 h6128344_1 conda-forge librdkit 2025.09.1 h3c5c181_1 conda-forge libre2-11 2024.07.02 hbbce691_2 conda-forge librsvg 2.58.4 h49af25d_2 conda-forge libsodium 1.0.20 h4ab18f5_0 conda-forge libsqlite 3.51.0 hee844dc_0 conda-forge libssh2 1.11.1 hcf80075_0 conda-forge libstdcxx 15.2.0 h8f9b012_7 conda-forge libstdcxx-ng 15.2.0 h4852527_7 conda-forge libtiff 4.7.1 h9d88235_1 conda-forge libtorch 2.5.1 cpu_generic_h213959a_8 conda-forge libuuid 2.41.2 he9a06e4_0 conda-forge libuv 1.51.0 hb03c661_1 conda-forge libwebp-base 1.6.0 hd42ef1d_0 conda-forge libxcb 1.17.0 h8a09558_0 conda-forge libxcrypt 4.4.36 hd590300_1 conda-forge libxkbcommon 1.11.0 he8b52b9_0 conda-forge libxml2 2.13.9 h04c0eec_0 conda-forge libxslt 1.1.43 h7a3aeb2_0 conda-forge libzip 1.11.2 h6991a6a_0 conda-forge libzlib 1.3.1 hb9d3cd8_2 conda-forge lightning-utilities 0.15.2 pyhd8ed1ab_0 conda-forge llvmlite 0.45.1 py313hdd307be_0 conda-forge lomap2 3.2.1 pyhd8ed1ab_3 conda-forge lxml 6.0.2 py313h6b0e12e_0 conda-forge lz4-c 1.10.0 h5888daf_1 conda-forge markdown-it-py 4.0.0 pyhd8ed1ab_0 conda-forge markupsafe 3.0.3 py313h3dea7bd_0 conda-forge matplotlib-base 3.10.7 py313h683a580_0 conda-forge matplotlib-inline 0.2.1 pyhd8ed1ab_0 conda-forge mda-xdrlib 0.2.0 pyhd8ed1ab_1 conda-forge mdanalysis 2.10.0 py313h08cd8bf_1 conda-forge mdtraj 1.11.0 np2py313h48ae260_3 conda-forge mdurl 0.1.2 pyhd8ed1ab_1 conda-forge mistune 3.1.4 pyhcf101f3_0 conda-forge mkl 2025.0.0 hacee8c2_941
ml_dtypes 0.5.1 py313h08cd8bf_1 conda-forge mmpbsa-py 16.0 pypi_0 pypi mmtf-python 1.1.3 pyhd8ed1ab_0 conda-forge mpc 1.3.1 h24ddda3_1 conda-forge mpfr 4.2.1 h90cbb55_3 conda-forge mpiplus v0.0.2 pyhd8ed1ab_0 conda-forge mpmath 1.3.0 pyhd8ed1ab_1 conda-forge mrcfile 1.5.4 pyhd8ed1ab_0 conda-forge msgpack-python 1.1.2 py313h7037e92_0 conda-forge munkres 1.1.4 pyhd8ed1ab_1 conda-forge narwhals 2.10.2 pyhcf101f3_0 conda-forge nbclient 0.10.2 pyhd8ed1ab_0 conda-forge nbconvert-core 7.16.6 pyhcf101f3_1 conda-forge nbformat 5.10.4 pyhd8ed1ab_1 conda-forge ncurses 6.5 h2d0b736_3 conda-forge nest-asyncio 1.6.0 pyhd8ed1ab_1 conda-forge netcdf-fortran 4.6.2 nompi_h5aa5643_101 conda-forge netcdf4 1.7.2 nompi_py313ha67fc0a_103 conda-forge networkx 3.5 pyhe01879c_0 conda-forge nomkl 3.0 0
notebook 7.4.7 pyhd8ed1ab_0 conda-forge notebook-shim 0.2.4 pyhd8ed1ab_1 conda-forge numba 0.62.1 py313hd8e3f9f_0 conda-forge numexpr 2.14.1 py313h24ae7f9_100 conda-forge numpy 2.2.6 py313h17eae1a_0 conda-forge ocl-icd 2.3.3 hb9d3cd8_0 conda-forge ocl-icd-system 1.0.0 1 conda-forge oniguruma 6.9.10 hb9d3cd8_0 conda-forge openblas 0.3.30 pthreads_h6ec200e_3 conda-forge opencl-headers 2025.06.13 h5888daf_0 conda-forge openfe 1.7.0 pyhd8ed1ab_2 conda-forge openfe-analysis 0.3.1 pyhd8ed1ab_0 conda-forge openfe-base 1.7.0 pyhd8ed1ab_2 conda-forge openff-amber-ff-ports 2025.09.0 pyhd8ed1ab_0 conda-forge openff-forcefields 2025.10.1 pyhd8ed1ab_0 conda-forge openff-interchange 0.4.8 pyhd8ed1ab_2 conda-forge openff-interchange-base 0.4.8 pyhd8ed1ab_2 conda-forge openff-nagl 0.5.4 pyhd8ed1ab_1 conda-forge openff-nagl-base 0.5.4 pyhd8ed1ab_1 conda-forge openff-nagl-models 2025.9.0 pyhd8ed1ab_0 conda-forge openff-toolkit 0.17.1 pyhd8ed1ab_0 conda-forge openff-toolkit-base 0.17.1 pyhd8ed1ab_0 conda-forge openff-units 0.3.1 pyhd8ed1ab_2 conda-forge openff-utilities 0.1.16 pyhd8ed1ab_0 conda-forge openjpeg 2.5.4 h55fea9a_0 conda-forge openldap 2.6.10 he970967_0 conda-forge openmm 8.2.0 py313h489b4b4_3 conda-forge openmmforcefields 0.15.0 pyhd8ed1ab_1 conda-forge openmmtools 0.25.3 pyhd8ed1ab_0 conda-forge openssl 3.5.4 h26f9b46_0 conda-forge opt_einsum 3.4.0 pyhd8ed1ab_1 conda-forge overrides 7.7.0 pyhd8ed1ab_1 conda-forge packaging 25.0 pyh29332c3_1 conda-forge packmol-memgen 2025.1.29 pypi_0 pypi pandas 2.3.3 py313h08cd8bf_1 conda-forge pandocfilters 1.5.0 pyhd8ed1ab_0 conda-forge pango 1.56.4 hadf4263_0 conda-forge parmed 4.3.0 py313hb83d11f_2 conda-forge parso 0.8.5 pyhcf101f3_0 conda-forge patsy 1.0.2 pyhcf101f3_0 conda-forge pcre2 10.46 h1321c63_0 conda-forge pdb4amber 22.0 pypi_0 pypi pdbfixer 1.12 pyhd8ed1ab_0 conda-forge perl 5.32.1 7_hd590300_perl5 conda-forge pexpect 4.9.0 pyhd8ed1ab_1 conda-forge pickleshare 0.7.5 pyhd8ed1ab_1004 conda-forge pillow 12.0.0 py313h50355cd_0 conda-forge pint 0.24.4 pyhe01879c_2 conda-forge pip 25.2 pyh145f28c_0 conda-forge pixman 0.46.4 h54a6638_1 conda-forge platformdirs 4.5.0 pyhcf101f3_0 conda-forge plotly 6.4.0 pyhd8ed1ab_0 conda-forge plugcli 0.2.1 pyhd8ed1ab_0 conda-forge pluggy 1.6.0 pyhd8ed1ab_0 conda-forge pooch 1.8.2 pyhd8ed1ab_3 conda-forge prometheus_client 0.23.1 pyhd8ed1ab_0 conda-forge prompt-toolkit 3.0.52 pyha770c72_0 conda-forge psutil 7.1.3 py313h54dd161_0 conda-forge pthread-stubs 0.4 hb9d3cd8_1002 conda-forge ptyprocess 0.7.0 pyhd8ed1ab_1 conda-forge pure_eval 0.2.3 pyhd8ed1ab_1 conda-forge py-cpuinfo 9.0.0 pyhd8ed1ab_1 conda-forge py3dmol 2.5.3 pyhd8ed1ab_0 conda-forge pycairo 1.28.0 py313h3f29d12_1 conda-forge pycparser 2.22 pyh29332c3_1 conda-forge pydantic 2.11.10 pyh3cfb1c2_0 conda-forge pydantic-core 2.33.2 py313h4b2b08d_0 conda-forge pyedr 0.8.0 pyhd8ed1ab_1 conda-forge pygments 2.19.2 pyhd8ed1ab_0 conda-forge pygraphviz 1.14 py313h45fd7b0_2 conda-forge pymbar 4.2.0 pyha770c72_2 conda-forge pymbar-core 4.2.0 pyhd8ed1ab_2 conda-forge pymsmt 22.0 pypi_0 pypi pyparsing 3.2.5 pyhcf101f3_0 conda-forge pysocks 1.7.1 pyha55dd90_7 conda-forge pytables 3.10.2 py313h0546747_10 conda-forge pytest 8.4.2 pyhd8ed1ab_0 conda-forge pytest-datadir 1.8.0 pyhd8ed1ab_0 conda-forge pytest-regressions 2.8.3 pyhc455866_0 conda-forge pytest-rerunfailures 16.1 pyhd8ed1ab_0 conda-forge python 3.13.9 hc97d973_101_cp313 conda-forge python-constraint 1.4.0 pyhff2d567_1 conda-forge python-dateutil 2.9.0.post0 pyhe01879c_2 conda-forge python-fastjsonschema 2.21.2 pyhe01879c_0 conda-forge python-gil 3.13.9 h4df99d1_101 conda-forge python-json-logger 2.0.7 pyhd8ed1ab_0 conda-forge python-tzdata 2025.2 pyhd8ed1ab_0 conda-forge python_abi 3.13 8_cp313 conda-forge pytng 0.3.4 py313h75f1e36_0 conda-forge pytorch 2.5.1 cpu_generic_py313_h6691dd5_8 conda-forge pytorch-lightning 2.5.5 pyh2a12c56_0 conda-forge pytraj 2.0.6 pypi_0 pypi pytz 2025.2 pyhd8ed1ab_0 conda-forge pyyaml 6.0.3 py313h3dea7bd_0 conda-forge pyzmq 27.1.0 py312hfb55c3c_0 conda-forge qhull 2020.2 h434a139_5 conda-forge rdkit 2025.09.1 py313h66a1ddb_1 conda-forge re2 2024.07.02 h9925aae_2 conda-forge readline 8.2 h8c095d6_2 conda-forge referencing 0.37.0 pyhcf101f3_0 conda-forge reportlab 4.4.4 py313h07c4f96_0 conda-forge requests 2.32.5 pyhd8ed1ab_0 conda-forge rfc3339-validator 0.1.4 pyhd8ed1ab_1 conda-forge rfc3986-validator 0.1.1 pyh9f0ad1d_0 conda-forge rfc3987-syntax 1.1.0 pyhe01879c_1 conda-forge rich 14.2.0 pyhcf101f3_0 conda-forge rlpycairo 0.4.0 pyh6c17108_0 conda-forge rpds-py 0.28.0 py313h843e2db_1 conda-forge sander 22.0 pypi_0 pypi scikit-learn 1.7.2 py313h06d4379_0 conda-forge scikit-mol 0.6.1 pyhd8ed1ab_0 conda-forge scipy 1.16.3 py313h11c21cd_0 conda-forge seaborn 0.13.2 hd8ed1ab_3 conda-forge seaborn-base 0.13.2 pyhd8ed1ab_3 conda-forge send2trash 1.8.3 pyh0d859eb_1 conda-forge setuptools 80.9.0 pyhff2d567_0 conda-forge six 1.17.0 pyhe01879c_1 conda-forge sleef 3.9.0 ha0421bc_0 conda-forge snappy 1.2.2 h03e3b7b_0 conda-forge sniffio 1.3.1 pyhd8ed1ab_1 conda-forge soupsieve 2.8 pyhd8ed1ab_0 conda-forge spectate 1.0.1 pyhd8ed1ab_0 conda-forge sqlalchemy 2.0.44 py313h07c4f96_0 conda-forge stack_data 0.6.3 pyhd8ed1ab_1 conda-forge statsmodels 0.14.5 py313h29aa505_1 conda-forge sympy 1.14.0 pyh2585a3b_105 conda-forge sysroot_linux-64 2.34 h087de78_2 conda-forge tbb 2022.3.0 h8d10470_0 conda-forge tbb-devel 2022.3.0 h74b38a2_0 conda-forge terminado 0.18.1 pyh0d859eb_0 conda-forge threadpoolctl 3.6.0 pyhecae5ae_0 conda-forge tidynamics 1.1.2 pyhd8ed1ab_0 conda-forge tinycss2 1.4.0 pyhd8ed1ab_0 conda-forge tinydb 4.8.2 pyhd8ed1ab_1 conda-forge tk 8.6.13 noxft_hd72426e_102 conda-forge tomli 2.3.0 pyhcf101f3_0 conda-forge torchmetrics 1.8.2 pyhd8ed1ab_0 conda-forge tornado 6.5.2 py313h07c4f96_1 conda-forge tqdm 4.67.1 pyhd8ed1ab_1 conda-forge traitlets 5.14.3 pyhd8ed1ab_1 conda-forge typing-extensions 4.15.0 h396c80c_0 conda-forge typing-inspection 0.4.2 pyhd8ed1ab_0 conda-forge typing_extensions 4.15.0 pyhcf101f3_0 conda-forge typing_utils 0.1.0 pyhd8ed1ab_1 conda-forge tzdata 2025b h78e105d_0 conda-forge uri-template 1.3.0 pyhd8ed1ab_1 conda-forge urllib3 2.5.0 pyhd8ed1ab_0 conda-forge validators 0.35.0 pyhd8ed1ab_0 conda-forge wayland 1.24.0 hd6090a7_1 conda-forge wcwidth 0.2.14 pyhd8ed1ab_0 conda-forge webcolors 25.10.0 pyhd8ed1ab_0 conda-forge webencodings 0.5.1 pyhd8ed1ab_3 conda-forge websocket-client 1.9.0 pyhd8ed1ab_0 conda-forge widgetsnbextension 4.0.15 pyhd8ed1ab_0 conda-forge xkeyboard-config 2.46 hb03c661_0 conda-forge xmltodict 1.0.2 pyhcf101f3_0 conda-forge xorg-libice 1.1.2 hb9d3cd8_0 conda-forge xorg-libsm 1.2.6 he73a12e_0 conda-forge xorg-libx11 1.8.12 h4f16b4b_0 conda-forge xorg-libxau 1.0.12 hb9d3cd8_0 conda-forge xorg-libxcomposite 0.4.6 hb9d3cd8_2 conda-forge xorg-libxcursor 1.2.3 hb9d3cd8_0 conda-forge xorg-libxdamage 1.1.6 hb9d3cd8_0 conda-forge xorg-libxdmcp 1.1.5 hb9d3cd8_0 conda-forge xorg-libxext 1.3.6 hb9d3cd8_0 conda-forge xorg-libxfixes 6.0.2 hb03c661_0 conda-forge xorg-libxi 1.8.2 hb9d3cd8_0 conda-forge xorg-libxinerama 1.1.5 h5888daf_1 conda-forge xorg-libxrandr 1.5.4 hb9d3cd8_0 conda-forge xorg-libxrender 0.9.12 hb9d3cd8_0 conda-forge xorg-libxt 1.3.1 hb9d3cd8_0 conda-forge xorg-libxtst 1.2.5 hb9d3cd8_3 conda-forge xorg-libxxf86vm 1.1.6 hb9d3cd8_0 conda-forge xorg-xorgproto 2024.1 hb9d3cd8_1 conda-forge yaml 0.2.5 h280c20c_3 conda-forge zeromq 4.3.5 h387f397_9 conda-forge zipp 3.23.0 pyhd8ed1ab_0 conda-forge zlib 1.3.1 hb9d3cd8_2 conda-forge zlib-ng 2.2.5 hde8ca8f_0 conda-forge zstandard 0.25.0 py313h54dd161_0 conda-forge zstd 1.5.7 hb8e6e7a_2 conda-forge

Output

I received the following log messages at the start of the execution:

INFO:2025-11-18 22:18:58,945:jax._src.xla_bridge:924: Unable to initialize backend 'cuda': INFO:2025-11-18 22:18:58,945:jax._src.xla_bridge:924: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig' INFO:2025-11-18 22:18:58,954:jax._src.xla_bridge:924: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory WARNING:2025-11-18 22:18:58,954:jax._src.xla_bridge:966: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

Expected behavior

Given the hardware (RTX 4090), I expected the simulation to complete much faster than 15 hours. I also assumed that conda install openfe would pull the CUDA-enabled version of jaxlib if a GPU is present, or at least compatible with the installed cudatoolkit.

Additional context As seen in the conda list, jaxlib version 0.5.2 is installed with the build string cpu_py313.... I am curious if I should manually enforce the installation of jax[cuda] to resolve this performance issue.

ORCAaAaA-ui avatar Nov 20 '25 07:11 ORCAaAaA-ui