aqlaboratory
Openfold not compatible with newest Flash attention 2.0.4
Hi,
I upgraded to flash attention 2.0.4, which breaks imports in the OpenFold package.
Will there be an upgrade to make OpenFold compatible?
Thanks
Here the traceback:
Traceback (most recent call last): File "/usr/src/code//subprojects/moppo/moppo_esm_stab_train.py", line 13, in <module> from rlif.tools.advantages import deepstabp_advantage_fn, ESMFoldAdvantage, ESMFoldAdvantageConfig File "/usr/src/code/rlif/tools/advantages.py", line 12, in <module> from rlif.folding_models.esmfold import ESMFoldModel File "/usr/src/code/rlif/folding_models/esmfold.py", line 2, in <module> from openfold.utils.feats import atom14_to_atom37 File "/root/miniconda3/envs/piai/lib/python3.11/site-packages/openfold/__init__.py", line 1, in <module> from . import model File "/root/miniconda3/envs/piai/lib/python3.11/site-packages/openfold/model/__init__.py", line 11, in <module> _modules = [(m, importlib.import_module("." + m, __name__)) for m in __all__] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/piai/lib/python3.11/site-packages/openfold/model/__init__.py", line 11, in <listcomp> _modules = [(m, importlib.import_module("." + m, __name__)) for m in __all__] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/piai/lib/python3.11/importlib/__init__.py", line 126, in import_module return _bootstrap._gcd_import(name[level:], package, level) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/piai/lib/python3.11/site-packages/openfold/model/triangular_multiplicative_update.py", line 22, in <module> from openfold.model.primitives import Linear, LayerNorm File "/root/miniconda3/envs/piai/lib/python3.11/site-packages/openfold/model/primitives.py", line 28, in <module> from flash_attn.flash_attention import FlashAttention ModuleNotFoundError: No module named 'flash_attn.flash_attention'
And in addition to that: will there be version compatible with pytorch lightning 2.0? That causes a couple downstream import issues too.
Hi, these upgrades are on our to-do list and will be included in a future release.
