Segmentation fault when calling `get_model_haiku_params`

[nicoceres]

Hello,

I just downloaded and installed AF_unmasked (main branch of this repo) following the procedure described in the README.

I made a simple test by trying to get a 'standard' AF prediction for a homodimeric protein (600x2 aa), ie. no cross_chain_templates, no separate_homomer_msas, no templates at all.
I tried a prediction with the model_5_multimer_v3 parameter set.

Unfortunately, a segmentation fault error stops the execution quite early in the pipeline, at the moment of loading parameter set, if I understand correctly. The error message looks like this:

`I0129 11:29:29.626704 139818377099072 templates.py:857] Using precomputed obsolete pdbs /localHDD/nceres/data/alphafold_db/pdb_mmcif/obsolete.dat.
I0129 11:29:29.632656 139818377099072 run_alphafold.py:478] Setting max_recycles to 20
I0129 11:29:29.632744 139818377099072 run_alphafold.py:479] Setting early stop tolerance to 0.5
I0129 11:29:29.632779 139818377099072 run_alphafold.py:480] Setting dropout to False
Fatal Python error: Segmentation fault

Current thread 0x00007f2a00b33740 (most recent call first):
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/xla_bridge.py", line 322 in _version_check
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/xla_bridge.py", line 387 in _check_cuda_versions
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/xla_bridge.py", line 428 in make_gpu_client
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/xla_bridge.py", line 935 in _init_backend
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/xla_bridge.py", line 853 in backends
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/xla_bridge.py", line 960 in _get_backend_uncached
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/xla_bridge.py", line 981 in get_backend
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/xla_bridge.py", line 1093 in local_devices
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/interpreters/pxla.py", line 1698 in _get_default_device
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/interpreters/pxla.py", line 1743 in _get_and_check_device_assignment
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/interpreters/pxla.py", line 1707 in call
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/pjit.py", line 1322 in _resolve_in_shardings
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/pjit.py", line 1399 in _pjit_call_impl_python
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/pjit.py", line 1471 in call_impl_cache_miss
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/pjit.py", line 1488 in _pjit_call_impl
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/core.py", line 913 in process_primitive
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/core.py", line 425 in bind_with_trace
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/core.py", line 2788 in bind
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/pjit.py", line 176 in _python_pjit_helper
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/pjit.py", line 298 in cache_miss
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 179 in reraise_with_filtered_traceback
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/dispatch.py", line 87 in apply_primitive
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/core.py", line 913 in process_primitive
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/core.py", line 425 in bind_with_trace
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/core.py", line 422 in bind
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/lax/lax.py", line 558 in _convert_element_type
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 2197 in array
File "/projects/DAMM/AF_unmasked/alphafold/model/utils.py", line 116 in flat_params_to_haiku
File "/projects/DAMM/AF_unmasked/alphafold/model/data.py", line 33 in get_model_haiku_params
File "/projects/DAMM/AF_unmasked/run_alphafold.py", line 481 in main
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/absl/app.py", line 258 in _run_main
File "/projects/DAMM/miniconda3/envs/AF_unmasked/lib/python3.9/site-packages/absl/app.py", line 312 in run
File "/projects/DAMM/AF_unmasked/run_alphafold.py", line 531 in
`.

Any idea about this segmentation fault error?
Many thanks in advance.

[clami66]

Hi,

This is something I have not seen before. Could you try and clone this older commit, then try and install a fresh environment from that?

[nicoceres]

Thank you Claudio, I' ll try this, but are you sure about the link? I end up on a commit with just one change in the README.md file :).

On my side I made some tests, and I think that the problem is related to a library incompatibility when I create the conda environment with jax==0.4.26, jaxlib==0.4.26+cuda12.cudnn89, cudnn=8.9.7.29, in agreement with the requisites.txt file from the main branch.

Note that my system-wide installation of nvcc returns 'Build cuda_11.8.r11.8/compiler.31833905_0', whereas the conda installed nvcc returns 'Build cuda_12.4.r12.4/compiler.34097967_0'.

In this context, I get a segmentation fault by simply importing jax and running jax.devices() through python in the AF_unmasked conda environment.

Current working 'solution':

• turn flash attention off (by commenting out code introduced in this commit)
• downgrading jax and jaxlib to versions from this commit.

What a pity; I cannot profit from flash attention :).

Please do not hesitate to share with me any suggestions you might have in order to be able to profit from flax and flash attention.

[clami66]

The one I linked is one of the latest commits before I added flash attention (which led to bumping Jax version), so that's why it felt like a safe option. But glad you were able to work it out on your own!

Perhaps it is possible to find a newer jax version built for cuda 11, but I couldn't find it when I added the latest features

[clami66]

Perhaps you can try with jax 0.4.25+cuda11?