Warning on slow compilation of model

[marcmk6]

Hi,

While running ColabFold locally, I get warnings like the following every time the model is (re)compiled, which is slow.

2022-02-15 16:22:05.982180: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:55]
********************************
Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
Compiling module jit_apply_fn.99150
********************************

This seems due to that the model is compiled on CPU rather than GPU. So I guess this is because something is not properly configured. But I didn't find any related configuration that I can change or set within ColabFold. Is there any idea?

(I have tried setting the env var as the warning suggests, but the corresponding output is a bunch of files describing JAX behaviors, which I have no idea with.)

Thanks!

[chris-kafka]

Same thing here. With an older version of jax and jaxlib (like stated in the localcolabfold linux setup script https://github.com/YoshitakaMo/localcolabfold/blob/main/install_colabbatch_linux.sh) the warning is not showing up.

The processes are still running on GPU and once tax is compiled there doesn't seem to be a drop in performance.

[cerabinowitz]

Would it be possible to provide a prebuilt docker image to run this codebase? It seems like many people are having issues correctly setting up all the dependencies on gpu