Problem with installing DMFoldv1

[DSJ]

I created alphafold2nondocker env with following code:

conda install -y -c conda-forge openmm==7.5.1 cudnn==8.2.1.32 cudatoolkit==11.3.1 cudatoolkit-dev==11.3.1 pdbfixer==1.7
conda install -y -c bioconda hmmer==3.3.2 hhsuite==3.3.0 kalign2==2.04
pip install absl-py==0.13.0 biopython==1.79 chex==0.0.7 dm-haiku==0.0.4 dm-tree==0.1.6 immutabledict==2.0.0 jax==0.2.14 ml-collections==0.1.0 numpy==1.20.1 scipy==1.7.0 tensorflow==2.8.4 pandas==1.3.4 tensorflow-cpu==2.8.4
pip install --upgrade jax==0.2.14 jaxlib==0.1.69+cuda111 -f https://storage.googleapis.com/jax-rele ... eases.html

When running Run_DeepMSA2.py, I encountered the following error:

2024-01-16 13:56:17.671131: W external/org_tensorflow/tensorflow/stream_executor/gpu/asm_compiler.cc:235] Falling back to the CUDA driver for PTX compilation; ptxas does not support CC 8.9
2024-01-16 13:56:17.671152: W external/org_tensorflow/tensorflow/stream_executor/gpu/asm_compiler.cc:238] Used ptxas at ptxas
2024-01-16 13:56:17.671678: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:625] failed to get PTX kernel "shift_right_logical_3" from module: CUDA_ERROR_NOT_FOUND: named symbol not found
2024-01-16 13:56:17.671711: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2040] Execution of replica 0 failed: Internal: Could not find the corresponding function
Traceback (most recent call last):
File "/home/dsj/Software/DMFold1.0/bin/alphafold/run_alphafold_msa_benchmark.py", line 306, in <module>
app.run(main)
File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/home/dsj/Software/DMFold1.0/bin/alphafold/run_alphafold_msa_benchmark.py", line 274, in main
predict_structure(
File "/home/dsj/Software/DMFold1.0/bin/alphafold/run_alphafold_msa_benchmark.py", line 151, in predict_structure
prediction_result = model_runner.predict(processed_feature_dict)
File "/home/dsj/Software/DMFold1.0/bin/alphafold/alphafold/model/model.py", line 134, in predict
result = self.apply(self.params, jax.random.PRNGKey(0), feat)
File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/random.py", line 75, in PRNGKey
k1 = convert(lax.shift_right_logical(seed_arr, lax._const(seed_arr, 32)))
File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/lax.py", line 382, in shift_right_logical
return shift_right_logical_p.bind(x, y)
File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/core.py", line 264, in bind
out = top_trace.process_primitive(self, tracers, params)
File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/core.py", line 603, in process_primitive
return primitive.impl(*tracers, **params)
File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/xla.py", line 249, in apply_primitive
return compiled_fun(*args)
File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/xla.py", line 365, in _execute_compiled_primitive
out_bufs = compiled.execute(input_bufs)
RuntimeError: Internal: Could not find the corresponding function

This one seems to be caused by the wrong version of cuda, then I updata my alphafold2nondocker env with following code:

Code: Select all

conda install -y -c conda-forge openmm==7.5.1 cudatoolkit==11.8 pdbfixer==1.7
conda install -y -c bioconda hmmer==3.3.2 hhsuite==3.3.0 kalign2==2.04
pip install absl-py==0.13.0 biopython==1.79 chex==0.0.7 dm-haiku==0.0.4 dm-tree==0.1.6 immutabledict==2.0.0 jax==0.3.25 ml-collections==0.1.0 numpy==1.20.1 scipy==1.7.0 tensorflow==2.8.4 pandas==1.3.4 tensorflow-cpu==2.8.4
pip install --upgrade jax==0.3.25 jaxlib==0.3.25+cuda11.cudnn805 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Then I encountered a new error:

Code: Select all

Traceback (most recent call last):
  File "/home/dsj/Software/DMFold1.0/bin/alphafold/run_alphafold_msa_benchmark.py", line 306, in <module>
    app.run(main)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "/home/dsj/Software/DMFold1.0/bin/alphafold/run_alphafold_msa_benchmark.py", line 274, in main
    predict_structure(
  File "/home/dsj/Software/DMFold1.0/bin/alphafold/run_alphafold_msa_benchmark.py", line 151, in predict_structure
    prediction_result = model_runner.predict(processed_feature_dict)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold/alphafold/model/model.py", line 134, in predict
    result = self.apply(self.params, jax.random.PRNGKey(0), feat)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/transform.py", line 125, in apply_fn
    out, state = f.apply(params, {}, *args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/transform.py", line 313, in apply_fn
    out = f(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold/alphafold/model/model.py", line 60, in _forward_fn
    return model(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold/alphafold/model/modules.py", line 374, in __call__
    _, prev = hk.while_loop(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/stateful.py", line 610, in while_loop
    val, state = jax.lax.while_loop(pure_cond_fun, pure_body_fun, init_val)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/stateful.py", line 605, in pure_body_fun
    val = body_fun(val)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold/alphafold/model/modules.py", line 367, in <lambda>
    get_prev(do_call(x[1], recycle_idx=x[0],
  File "/home/dsj/Software/DMFold1.0/bin/alphafold/alphafold/model/modules.py", line 335, in do_call
    return impl(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold/alphafold/model/modules.py", line 163, in __call__
    representations = evoformer_module(batch0, is_training)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold/alphafold/model/modules.py", line 1725, in __call__
    msa_activations = jax.ops.index_add(msa_activations, 0,
AttributeError: module 'jax.ops' has no attribute 'index_add'

After google this error, I found index_add were removed in jax version 0.3.2. And my gpu device is 4090. Could you please give me some suggestions for resolving this error? Thank you so much for your help.

Best regards,
Shijie Dong

[jlspzw]

Dear Shijie,

Thank you for using our software, we only test the program on A40, V100, and A100 GPU cards, and never tried RTX4090, so I am not sure I can help with this issue.

But have you ever tried this ticket viewtopic.php?f=3&t=1010&sid=5e1ad01058 ... 8834ced841

From the issue what I can see is a Jax problem, can you down the version of your Jax, and select the version suite able for RTX4090, also the jax problem, I encourage you to check the issue section of Alphafold on github

https://github.com/google-deepmind/alph ... open++4090

see if any solution is suitable for RTX4090. Thank you.

Best Regards
Wei Zheng

[DSJ]

Dear Wei Zheng,

Thanks for your suggestions. I downgraded the version with jax==0.2.14 jaxlib==0.1.69+cuda111 -f https://storage.googleapis.com/jax-rele ... eases.html. It's work, but a new error when running Run_DMFold.py, as follow:

Code: Select all

I0117 07:23:01.527364 22723190105920 run_alphafold_mymsa.py:209] Running model model_1_multimer_v3_pred_0 on seq
I0117 07:23:01.527707 22723190105920 model.py:165] Running predict with shape(feat) = {'aatype': (1778,), 'residue_index': (1778,), 'seq_length': (), 'msa': (2701, 1778), 'num_alignments': (), 'template_aatype': (4, 1778), 'template_all_atom_mask': (4, 1778, 37), 'template_all_atom_positions': (4, 1778, 37, 3), 'asym_id': (1778,), 'sym_id': (1778,), 'entity_id': (1778,), 'deletion_matrix': (2701, 1778), 'deletion_mean': (1778,), 'all_atom_mask': (1778, 37), 'all_atom_positions': (1778, 37, 3), 'assembly_num_chains': (), 'entity_mask': (1778,), 'num_templates': (), 'cluster_bias_mask': (2701,), 'bert_mask': (2701, 1778), 'seq_mask': (1778,), 'msa_mask': (2701, 1778)}
Traceback (most recent call last):
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/run_alphafold_mymsa.py", line 514, in <module>
    app.run(main)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/run_alphafold_mymsa.py", line 490, in main
    predict_structure(
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/run_alphafold_mymsa.py", line 217, in predict_structure
    prediction_result = model_runner.predict(processed_feature_dict,
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/model.py", line 167, in predict
    result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 183, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/api.py", line 424, in cache_miss
    out_flat = xla.xla_call(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/core.py", line 1560, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/core.py", line 1551, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/core.py", line 1563, in process
    return trace.process_call(self, fun, tracers, params)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/core.py", line 606, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/xla.py", line 592, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/linear_util.py", line 262, in memoized_fun
    ans = call(fun, *args)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/xla.py", line 667, in _xla_callable
    jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1284, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1262, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/transform.py", line 125, in apply_fn
    out, state = f.apply(params, {}, *args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/transform.py", line 313, in apply_fn
    out = f(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/model.py", line 77, in _forward_fn
    return model(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 485, in __call__
    prev, safe_key = hk.fori_loop(0, num_iter, recycle_body, (prev, safe_key))
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/stateful.py", line 544, in fori_loop
    state, val = jax.lax.fori_loop(lower, upper, pure_body_fun, init_val)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 183, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 212, in fori_loop
    (_, result), _ = scan(_fori_scan_body_fun(body_fun), (lower_, init_val),
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 183, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 1288, in scan
    init_flat, carry_avals, carry_avals_out, init_tree, *rest = _create_jaxpr(init)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 1274, in _create_jaxpr
    jaxpr, consts, out_tree = _initial_style_jaxpr(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 186, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 179, in cached
    return f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 76, in _initial_style_jaxpr
    jaxpr, consts, out_tree = _initial_style_open_jaxpr(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 186, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 179, in cached
    return f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 70, in _initial_style_open_jaxpr
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1252, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1262, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 143, in scanned_fun
    return (i + 1, body_fun(i, x)), None
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/stateful.py", line 524, in pure_body_fun
    val = body_fun(i, val)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 482, in recycle_body
    ret = apply_network(prev=prev, safe_key=safe_key2)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 449, in apply_network
    return impl(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 321, in __call__
    repr_shape = hk.eval_shape(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/stateful.py", line 665, in eval_shape
    out_shape = jax.eval_shape(fun, *args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/api.py", line 2403, in eval_shape
    out = pe.abstract_eval_fun(wrapped_fun.call_wrapped,
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 410, in abstract_eval_fun
    _, avals_out, _ = trace_to_jaxpr_dynamic(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1252, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1262, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 322, in <lambda>
    lambda: embedding_module(batch, is_training))
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 661, in __call__
    template_act = template_module(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 841, in __call__
    summed_template_embeddings, _ = hk.scan(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/stateful.py", line 504, in scan
    (carry, state), ys = jax.lax.scan(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 183, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 1288, in scan
    init_flat, carry_avals, carry_avals_out, init_tree, *rest = _create_jaxpr(init)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 1274, in _create_jaxpr
    jaxpr, consts, out_tree = _initial_style_jaxpr(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 186, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 179, in cached
    return f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 76, in _initial_style_jaxpr
    jaxpr, consts, out_tree = _initial_style_open_jaxpr(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 186, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 179, in cached
    return f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 70, in _initial_style_open_jaxpr
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1252, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1262, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/stateful.py", line 487, in stateful_fun
    carry, out = f(carry, x)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 837, in scan_fn
    return carry + partial_template_embedder(*x), None
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 824, in partial_template_embedder
    return template_embedder(query_embedding,
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 987, in __call__
    act, safe_key = template_stack((act, safe_subkey))
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/layer_stack.py", line 265, in wrapped
    ret = _LayerStackNoState(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/layer_stack.py", line 156, in __call__
    carry, zs = hk.scan(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/stateful.py", line 504, in scan
    (carry, state), ys = jax.lax.scan(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 183, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 1288, in scan
    init_flat, carry_avals, carry_avals_out, init_tree, *rest = _create_jaxpr(init)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 1274, in _create_jaxpr
    jaxpr, consts, out_tree = _initial_style_jaxpr(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 186, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 179, in cached
    return f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 76, in _initial_style_jaxpr
    jaxpr, consts, out_tree = _initial_style_open_jaxpr(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 186, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/util.py", line 179, in cached
    return f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 70, in _initial_style_open_jaxpr
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1252, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1262, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/stateful.py", line 487, in stateful_fun
    carry, out = f(carry, x)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/layer_stack.py", line 149, in layer
    out_x, z = self._call_wrapped(carry.x, *scanned.args_ys)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/layer_stack.py", line 182, in _call_wrapped
    ret = self._f(*args)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 973, in template_iteration_fn
    act = template_iteration(
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules_multimer.py", line 1034, in __call__
    act = dropout_wrapper_fn(
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules.py", line 76, in dropout_wrapper
    residual = module(input_act, mask, is_training=is_training, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/Software/DMFold1.0/bin/alphafold_multimer/alphafold/model/modules.py", line 1283, in __call__
    act = hk.LayerNorm(axis=[-1], create_scale=True, create_offset=True,
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 428, in wrapped
    out = f(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/module.py", line 279, in run_interceptors
    return bound_method(*args, **kwargs)
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/layer_norm.py", line 128, in __call__
    scale = hk.get_parameter("scale", param_shape, inputs.dtype,
  File "/home/dsj/miniconda3/envs/alphafold2nondocker/lib/python3.8/site-packages/haiku/_src/base.py", line 293, in get_parameter
    raise ValueError(
jax._src.traceback_util.UnfilteredStackTrace: ValueError: Unable to retrieve parameter 'scale' for module 'alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/layer_norm_input'. All parameters must be created as part of `init`.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

Could you please provide some suggestions?

Best regards,
Shijie Dong

[jlspzw]

Hi Shijie,

Based on my knowledge, it should be still the Jax problem with RTX4090 card, since we do not have that card we never tried if AF2 could run with RTX4090. now I think we need some time to see if Deepmind has any solution https://github.com/google-deepmind/alph ... open++4090. I will see if I can find an RTX4090 to test it.

Best
Wei

[DSJ]

Dear Wei Zheng,

I found the DMFold used AlphaFoldv2.0. Have you considered using the latest AlphaFoldv2.3.0? I can run the AlphaFoldv2.3.0 locally. Maybe this will help us run sucessfully with RTX4090.
I have another question. Does the online version of DMFold use the model_multimer_v2? I check the result of the online DMFold. The pTM-score is lower than what I have done with Alphafold2complex. I guess using the latest model_multimer_v3 model might greatly increase the accuracy of the multimeric structure. I sincerely hope to solve my problem locally using DMFold.

Best regards,
Shijie Dong

[jlspzw]

Hi Shijie,

Thank you for the suggestion.

First, we want to clarify, that although pTMscore is overall trustable, it still has some problems, especially when compared to different version AF2 models. For example, version 2.1 pTMscore is higher than version 2.2 in most cases, but the real performance actually is that 2.2 is better than 2.1. So without a native structure to make a comparison, it is not safe to say that 2.3 is better than 2.2, despite we believe it may be right since DeepMind uses more data to train the model.

Second, we actually use the 2.2, not the 2.0.

We also noticed that only 2.3 supports the RTX4090 card, now we are considering implementing the DMFold with 2.3, but since a lot of work needs to be done, it may take some time before it is available. But we will try to release it soon.

Last, before the next version release, based on the suggestion from one of our lab members, you can try to close the GPU cuda and only use CPU mode to run DMFold. You can both set the CUDA device environment to unavailable by system level (before installation), and during the running (either in the shell program by set env or in the python program by set.env function.

Thank you for using our software, we will focus on working on the next version of DMFold to solve the problem completely.

Best Regards
Wei Zheng