Update to flash-attn2

[Ding3LI]

Update Comments

Modified scgpt/model/model.py & scgpt/model/multiomic_model.py for flash-attn2. Updated README.md with flash-attn2 installation and notes.

Updated Dependencies

flash-attn 1.x --> flash-att 2.x

Is Code Tested?

Yes. The updated code was tested on A100 GPU (Linux), and it works fine doing fine-tuning integration with the latest Flash Attention 2.

Note:
In the current version of code, scGPT uses FlashMHA. And the updated flash-attn2 uses a module named mpa.py. The details of newer MHA implementation can be found HERE starting at line 354 flash_attn/modules/mha.py

[NozomiMizore]

In ur latest changes, u have set the use_flash_attn=False in model.py. I have also modified the code to it. It seems like the model actually didn't call the flash-attn but call the standard mha. And the metric result of my "tutorial_annotation" is terrible. But if I don't set use_flash_attn to False as you did, the program would get some error about padding_mask. Do u have the same problem?

[moa4020]

Hey,

I've been having trouble installing this package for days and I thought I could benefit from compiling your version. Although the installation process didn't raise any errors, when importing the package, I receive the following error. All I did was create a venv with Python3 and install scGPT from a downloaded directory of your version.

pip install . ninja packaging && FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation

Did you encounter a similar issue while importing your version?

python -c "import scgpt; import flash_attn; print('Install successful.')"

A module that was compiled using NumPy 1.x cannot be run in
NumPy 2.0.0 as it may crash. To support both 1.x and 2.x
versions of NumPy, modules must be compiled with NumPy 2.0.
Some module may need to rebuild instead e.g. with 'pybind11>=2.12'.

If you are a user of the module, the easiest solution will be to
downgrade to 'numpy<2' or try to upgrade the affected module.
We expect that some modules will need time to support NumPy 2.

Traceback (most recent call last):  File "<string>", line 1, in <module>
  File "/home/fs01/moa4020/scGPT-update-flash_attn2/scgpt/__init__.py", line 18, in <module>
    from . import model, tokenizer, scbank, utils, tasks
  File "/home/fs01/moa4020/scGPT-update-flash_attn2/scgpt/model/__init__.py", line 1, in <module>
    from .model import (
  File "/home/fs01/moa4020/scGPT-update-flash_attn2/scgpt/model/model.py", line 5, in <module>
    import torch
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/__init__.py", line 1382, in <module>
    from .functional import *  # noqa: F403
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/functional.py", line 7, in <module>
    import torch.nn.functional as F
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/nn/__init__.py", line 1, in <module>
    from .modules import *  # noqa: F403
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/nn/modules/__init__.py", line 35, in <module>
    from .transformer import TransformerEncoder, TransformerDecoder, \
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/nn/modules/transformer.py", line 20, in <module>
    device: torch.device = torch.device(torch._C._get_default_device()),  # torch.device('cpu'),
/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/nn/modules/transformer.py:20: UserWarning: Failed to initialize NumPy: _ARRAY_API not found (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:84.)
  device: torch.device = torch.device(torch._C._get_default_device()),  # torch.device('cpu'),
/home/fs01/moa4020/scGPT-update-flash_attn2/scgpt/model/model.py:21: UserWarning: flash_attn is not installed
  warnings.warn("flash_attn is not installed")
/home/fs01/moa4020/scGPT-update-flash_attn2/scgpt/model/multiomic_model.py:19: UserWarning: flash_attn is not installed
  warnings.warn("flash_attn is not installed")
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/flash_attn/__init__.py", line 3, in <module>
    from flash_attn.flash_attn_interface import (
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 15, in <module>
    import flash_attn_2_cuda as flash_attn_gpu
ImportError: /lib64/libc.so.6: version `GLIBC_2.32' not found (required by /home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so)```

[Ding3LI]

@moa4020 Thanks for your comment. I believe I have encountered this issue before, but I cannot recall the details at the moment. Nevertheless, I consider it a valid approach to evaluate scGPT based on its original design without enabling the flash-attn module. For a quick reference on running the scGPT model, you may refer to this repo, which provides a complete pipeline for both inference and fine-tuning.
I also remain somewhat skeptical about evaluating flash-attn2 instead of the original flash-attn1. Although flash-attn1, which is included in the package list, has been deprecated, its multi-head attention (MHA) implementation appears to differ from that of flash-attn2. Therefore, I suggest using the original plain Torch MHA rather than the flash-attn2 MHA for early experiments.

[codycjy]

Hey,

I've been having trouble installing this package for days and I thought I could benefit from compiling your version. Although the installation process didn't raise any errors, when importing the package, I receive the following error. All I did was create a venv with Python3 and install scGPT from a downloaded directory of your version.

pip install . ninja packaging && FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation

Did you encounter a similar issue while importing your version?

python -c "import scgpt; import flash_attn; print('Install successful.')"

A module that was compiled using NumPy 1.x cannot be run in
NumPy 2.0.0 as it may crash. To support both 1.x and 2.x
versions of NumPy, modules must be compiled with NumPy 2.0.
Some module may need to rebuild instead e.g. with 'pybind11>=2.12'.

If you are a user of the module, the easiest solution will be to
downgrade to 'numpy<2' or try to upgrade the affected module.
We expect that some modules will need time to support NumPy 2.

Traceback (most recent call last):  File "<string>", line 1, in <module>
  File "/home/fs01/moa4020/scGPT-update-flash_attn2/scgpt/__init__.py", line 18, in <module>
    from . import model, tokenizer, scbank, utils, tasks
  File "/home/fs01/moa4020/scGPT-update-flash_attn2/scgpt/model/__init__.py", line 1, in <module>
    from .model import (
  File "/home/fs01/moa4020/scGPT-update-flash_attn2/scgpt/model/model.py", line 5, in <module>
    import torch
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/__init__.py", line 1382, in <module>
    from .functional import *  # noqa: F403
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/functional.py", line 7, in <module>
    import torch.nn.functional as F
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/nn/__init__.py", line 1, in <module>
    from .modules import *  # noqa: F403
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/nn/modules/__init__.py", line 35, in <module>
    from .transformer import TransformerEncoder, TransformerDecoder, \
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/nn/modules/transformer.py", line 20, in <module>
    device: torch.device = torch.device(torch._C._get_default_device()),  # torch.device('cpu'),
/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/torch/nn/modules/transformer.py:20: UserWarning: Failed to initialize NumPy: _ARRAY_API not found (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:84.)
  device: torch.device = torch.device(torch._C._get_default_device()),  # torch.device('cpu'),
/home/fs01/moa4020/scGPT-update-flash_attn2/scgpt/model/model.py:21: UserWarning: flash_attn is not installed
  warnings.warn("flash_attn is not installed")
/home/fs01/moa4020/scGPT-update-flash_attn2/scgpt/model/multiomic_model.py:19: UserWarning: flash_attn is not installed
  warnings.warn("flash_attn is not installed")
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/flash_attn/__init__.py", line 3, in <module>
    from flash_attn.flash_attn_interface import (
  File "/home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 15, in <module>
    import flash_attn_2_cuda as flash_attn_gpu
ImportError: /lib64/libc.so.6: version `GLIBC_2.32' not found (required by /home/fs01/moa4020/scgpt_venv/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so)```

Try using this method @moa4020

[lidanwu]

@Ding3LI If you're still interested in flash-attn2, I had a PR #351 that creates a FA1-style of FlashMHA using FA2 backend and verified to be working equivalently for scGPT model inference. Please feel free to try it out and let me know if it works on your ends too.