From 0a4c75bf81c18ca47ec4508bf2100a2a39c2a37c Mon Sep 17 00:00:00 2001 From: Henry Ding Date: Tue, 23 Jan 2024 14:50:42 -0600 Subject: [PATCH 1/2] Modified scgpt/model/model.py & scgpt/model/multiomic_model.py for flash-attn2. Updated README.md with flash-attn2 installation and notes --- README.md | 14 ++++++++++---- pyproject.toml | 2 +- scgpt/model/model.py | 8 ++++---- scgpt/model/multiomic_model.py | 8 ++++---- 4 files changed, 19 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index f5fe45ef7..2340f38d8 100644 --- a/README.md +++ b/README.md @@ -30,11 +30,19 @@ scGPT works with Python >= 3.7.13 and R >=3.6.1. Please make sure you have the c scGPT is available on PyPI. To install scGPT, run the following command: ```bash -pip install scgpt "flash-attn<1.0.5" # optional, recommended +pip install scgpt ninja packaging && FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation # optional, recommended # As of 2023.09, pip install may not run with new versions of the google orbax package, if you encounter related issues, please use the following command instead: -# pip install scgpt "flash-attn<1.0.5" "orbax<0.1.8" +# pip install scgpt ninja packaging "orbax<0.1.8" && FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation ``` +**Note**: +The `flash-attn` dependency requires `CUDA >= 11.6, PyTorch >= 1.12, Linux`. +If you encounter any issues, please refer to the [flash-attn](https://github.com/HazyResearch/flash-attention/tree/main) repository for installation instructions. +For now, ~~May 2023, we recommend using CUDA 11.7 and flash-attn<1.0.5 due to various issues reported about installing new versions of flash-attn.~~ +we are using the latest flash-attn2, but please be aware that depends on the gpu architecture, +Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100) are tested and supported, however, +Turing GPUs (T4, RTX 2080) are not supported in flash-attn2 but in flash-attn 1.x + [Optional] We recommend using [wandb](https://wandb.ai/) for logging and visualization. ```bash @@ -49,8 +57,6 @@ $ cd scGPT $ poetry install ``` -**Note**: The `flash-attn` dependency usually requires specific GPU and CUDA version. If you encounter any issues, please refer to the [flash-attn](https://github.com/HazyResearch/flash-attention/tree/main) repository for installation instructions. For now, May 2023, we recommend using CUDA 11.7 and flash-attn<1.0.5 due to various issues reported about installing new versions of flash-attn. - ## Pretrained scGPT Model Zoo Here is the list of pretrained models. Please find the links for downloading the checkpoint folders. We recommend using the `whole-human` model for most applications by default. If your fine-tuning dataset shares similar cell type context with the training data of the organ-specific models, these models can usually demonstrate competitive performance as well. A paired vocabulary file mapping gene names to ids is provided in each checkpoint folder. If ENSEMBL ids are needed, please find the conversion at [gene_info.csv](https://github.com/bowang-lab/scGPT/files/13243634/gene_info.csv). diff --git a/pyproject.toml b/pyproject.toml index f402c050c..b6471bba6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ orbax = "<0.1.8" pytest = "^5.2" black = "^22.3.0" tensorflow = "^2.8.0" -flash-attn = "^1.0.1" +flash-attn = ">=2.4.2" torch-geometric = "^2.3.0" dcor = "~0.5.3" wandb = "^0.12.3" diff --git a/scgpt/model/model.py b/scgpt/model/model.py index ad0639bee..132ff3cb1 100644 --- a/scgpt/model/model.py +++ b/scgpt/model/model.py @@ -12,7 +12,7 @@ from tqdm import trange try: - from flash_attn.flash_attention import FlashMHA + from flash_attn.modules.mha import MHA flash_attn_available = True except ImportError: @@ -634,11 +634,11 @@ def __init__( ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() - self.self_attn = FlashMHA( + self.self_attn = MHA( embed_dim=d_model, num_heads=nhead, - batch_first=batch_first, - attention_dropout=dropout, + dropout=dropout, + use_flash_attn=True, **factory_kwargs, ) # Version compatibility workaround diff --git a/scgpt/model/multiomic_model.py b/scgpt/model/multiomic_model.py index 5128fac74..ab807e248 100644 --- a/scgpt/model/multiomic_model.py +++ b/scgpt/model/multiomic_model.py @@ -12,7 +12,7 @@ from tqdm import trange try: - from flash_attn.flash_attention import FlashMHA + from flash_attn.modules.mha import MHA except ImportError: import warnings @@ -668,11 +668,11 @@ def __init__( ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() - self.self_attn = FlashMHA( + self.self_attn = MHA( embed_dim=d_model, num_heads=nhead, - batch_first=batch_first, - attention_dropout=dropout, + dropout=dropout, + use_flash_attn=True, **factory_kwargs, ) # Implementation of Feedforward model From 5d29fb87905297b5c85f215fc6fab4d8886a4db4 Mon Sep 17 00:00:00 2001 From: Henry Ding Date: Fri, 2 Feb 2024 15:54:55 -0600 Subject: [PATCH 2/2] Updated flag::use_flash_attn to False in Multi-Head Attention because using padding_mask --- .DS_Store | Bin 0 -> 6148 bytes scgpt/model/model.py | 2 +- scgpt/model/multiomic_model.py | 2 +- 3 files changed, 2 insertions(+), 2 deletions(-) create mode 100644 .DS_Store diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..d5cffbd02e139bc25175f6d3c28c9b5197e6e497 GIT binary patch literal 6148 zcmeHKO>fgc5S?vP;t)VNfCL95OI)iE+Jc0*q-k=f1RQDv2SCA&qvF=}Mz%w!qDZ;J zZ{RO*=8y1SIKi9UZDl*<&V=+QfNYFv}+9oc{OqWgMK<{or*po()^K_hgd$Q8G>y zAqvMxdG{ttLOJWnX%c3N>*;{ga2mtbz4?4==W%D#-RW#EHr@HtEv|1rSu7gPod*w} z9So1-lSF>ex(j@ml-)FV247J*-_R#7O=6iGqL1^rU($%OiXh(+D+=P!m_9?g6Pi$r zdL>`nujJS1IqlOP4d^9#G$@ZQc*Vuhn;2DM)SOZ}r8Z{ZH62%2zbdR-V4cEKWFLV& z0>2HHSY=%1uJL}0GdAb|Yulwhy#Ths2+It_h`tqm5*w6ZhC|GkS29_va%Rp7VI>!- zRLpE3XR215QUrD56$UF+WT0v7@(P%v;WRH#wG69(Rp5#%z~2WK&e%1$)TllkDAW@G zSVOlm)YV^qpvxV=uEC{7^uUCs0yR~bCx$RD2chYRcTN6Mqo$KkGh-YxvoJ3dVICgB zRp}(U8f|G6unLqFSl7)K-~R)g;lI59%Ou;f3RnfMlmepG_xe3-$=t0go8!CIg@1sv pao$p+qM$I>v8wP@d>5_^F_$aAuEC{7^uX*N0VRVitOEa3fj@h@%f0{r literal 0 HcmV?d00001 diff --git a/scgpt/model/model.py b/scgpt/model/model.py index 132ff3cb1..b3550d050 100644 --- a/scgpt/model/model.py +++ b/scgpt/model/model.py @@ -638,7 +638,7 @@ def __init__( embed_dim=d_model, num_heads=nhead, dropout=dropout, - use_flash_attn=True, + use_flash_attn=False, **factory_kwargs, ) # Version compatibility workaround diff --git a/scgpt/model/multiomic_model.py b/scgpt/model/multiomic_model.py index ab807e248..aded09131 100644 --- a/scgpt/model/multiomic_model.py +++ b/scgpt/model/multiomic_model.py @@ -672,7 +672,7 @@ def __init__( embed_dim=d_model, num_heads=nhead, dropout=dropout, - use_flash_attn=True, + use_flash_attn=False, **factory_kwargs, ) # Implementation of Feedforward model