README.md 8.92 KB
Newer Older
Shenggan's avatar
Shenggan committed
1
2
![](/assets/fold.jpg)

shenggan's avatar
shenggan committed
3
# FastFold
Shenggan's avatar
Shenggan committed
4

Shenggan's avatar
Shenggan committed
5
[![](https://img.shields.io/badge/Paper-PDF-green?style=flat&logo=arXiv&logoColor=green)](https://arxiv.org/abs/2203.00854)
Shenggan's avatar
Shenggan committed
6
![](https://img.shields.io/badge/Made%20with-ColossalAI-blueviolet?style=flat)
Shenggan's avatar
Shenggan committed
7
8
![](https://img.shields.io/github/v/release/hpcaitech/FastFold)
[![GitHub license](https://img.shields.io/github/license/hpcaitech/FastFold)](https://github.com/hpcaitech/FastFold/blob/main/LICENSE)
Shenggan's avatar
Shenggan committed
9

shenggan's avatar
shenggan committed
10
Optimizing Protein Structure Prediction Model Training and Inference on GPU Clusters
Shenggan's avatar
Shenggan committed
11
12
13
14
15
16

FastFold provides a **high-performance implementation of Evoformer** with the following characteristics.

1. Excellent kernel performance on GPU platform
2. Supporting Dynamic Axial Parallelism(DAP)
    * Break the memory limit of single GPU and reduce the overall training time
Shenggan's avatar
Shenggan committed
17
    * DAP can significantly speed up inference and make ultra-long sequence inference possible
Shenggan's avatar
Shenggan committed
18
3. Ease of use
Shenggan's avatar
Shenggan committed
19
    * Huge performance gains with a few lines changes
Shenggan's avatar
Shenggan committed
20
    * You don't need to care about how the parallel part is implemented
21
22
4. Faster data processing, about 3x times faster on monomer, about 3Nx times faster on multimer with N sequence.
5. Great Reduction on GPU memory, able to inference sequence containing more than **10000** residues.
Shenggan's avatar
Shenggan committed
23
24
25

## Installation

zhuww's avatar
zhuww committed
26
To install FastFold, you will need:
LuGY's avatar
LuGY committed
27
+ Python 3.8 or 3.9.
28
+ [NVIDIA CUDA](https://developer.nvidia.com/cuda-downloads) 11.1 or above
zhuww's avatar
zhuww committed
29
+ PyTorch 1.12 or above 
30

LuGY's avatar
LuGY committed
31

32
33
34
35
36
For now, You can install FastFold:
### Using Conda (Recommended)

We highly recommend installing an Anaconda or Miniconda environment and install PyTorch with conda.
Lines below would create a new conda environment called "fastfold":
Shenggan's avatar
Shenggan committed
37

Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
38
39
40
```shell
git clone https://github.com/hpcaitech/FastFold
cd FastFold
Shenggan's avatar
Shenggan committed
41
conda env create --name=fastfold -f environment.yml
Shenggan's avatar
Shenggan committed
42
conda activate fastfold
43
python setup.py install
Shenggan's avatar
Shenggan committed
44
45
```

46
47
#### Advanced

zhuww's avatar
zhuww committed
48
To leverage the power of FastFold, we recommend you to install [Triton](https://github.com/openai/triton).
49
50

```bash
zhuww's avatar
zhuww committed
51
pip install triton==2.0.0.dev20221005
52
53
54
```


55
56
### Using PyPi
You can download FastFold with pre-built CUDA extensions.
Shenggan's avatar
Shenggan committed
57

58
59
Warning, only stable versions available.

Shenggan's avatar
Shenggan committed
60
```shell
61
62
63
64
65
66
67
68
69
70
71
pip install fastfold -f https://release.colossalai.org/fastfold
```

## Use Docker

### Build On Your Own
Run the following command to build a docker image from Dockerfile provided.

> Building FastFold from scratch requires GPU support, you need to use Nvidia Docker Runtime as the default when doing `docker build`. More details can be found [here](https://stackoverflow.com/questions/59691207/docker-build-with-nvidia-runtime).

```shell
72
73
cd FastFold
docker build -t fastfold ./docker
74
75
76
77
78
```

Run the following command to start the docker container in interactive mode.
```shell
docker run -ti --gpus all --rm --ipc=host fastfold bash
Shenggan's avatar
Shenggan committed
79
80
```

Shenggan's avatar
Shenggan committed
81
82
## Usage

83
You can use `Evoformer` as `nn.Module` in your project after `from fastfold.model.fastnn import Evoformer`:
Shenggan's avatar
Shenggan committed
84
85

```python
86
from fastfold.model.fastnn import Evoformer
Shenggan's avatar
Shenggan committed
87
88
89
evoformer_layer = Evoformer()
```

Shenggan's avatar
Shenggan committed
90
If you want to use Dynamic Axial Parallelism, add a line of initialize with `fastfold.distributed.init_dap`.
Shenggan's avatar
Shenggan committed
91
92
93
94
95
96
97

```python
from fastfold.distributed import init_dap

init_dap(args.dap_size)
```

Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
98
99
100
101
102
### Download the dataset
You can down the dataset used to train FastFold  by the script `download_all_data.sh`:

    ./scripts/download_all_data.sh data/

103
104
### Inference

Shenggan's avatar
Shenggan committed
105
You can use FastFold with `inject_fastnn`. This will replace the evoformer from OpenFold with the high performance evoformer from FastFold.
106
107

```python
108
from fastfold.utils import inject_fastnn
109
110
111
112

model = AlphaFold(config)
import_jax_weights_(model, args.param_path, version=args.model_name)

113
model = inject_fastnn(model)
114
115
116
117
118
```

For Dynamic Axial Parallelism, you can refer to `./inference.py`. Here is an example of 2 GPUs parallel inference:

```shell
119
python inference.py target.fasta data/pdb_mmcif/mmcif_files/ \
Shenggan's avatar
Shenggan committed
120
    --output_dir ./ \
121
    --gpus 2 \
122
123
124
125
126
    --uniref90_database_path data/uniref90/uniref90.fasta \
    --mgnify_database_path data/mgnify/mgy_clusters_2018_12.fa \
    --pdb70_database_path data/pdb70/pdb70 \
    --uniclust30_database_path data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
Shenggan's avatar
Shenggan committed
127
128
129
130
    --jackhmmer_binary_path `which jackhmmer` \
    --hhblits_binary_path `which hhblits` \
    --hhsearch_binary_path `which hhsearch` \
    --kalign_binary_path `which kalign`
131
```
LuGY's avatar
LuGY committed
132
133
134
135
136
137
or run the script `./inference.sh`, you can change the parameter in the script, especisally those data path.
```shell
./inference.sh
```

#### inference with data workflow
lclgy's avatar
lclgy committed
138
Alphafold's data pre-processing takes a lot of time, so we speed up the data pre-process by [ray](https://docs.ray.io/en/latest/workflows/concepts.html) workflow, which achieves a 3x times faster speed. To run the inference with ray workflow, you should install the package and add parameter `--enable_workflow` to cmdline or shell script `./inference.sh`
LuGY's avatar
LuGY committed
139
```shell
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
140
pip install ray==2.0.0 pyarrow
LuGY's avatar
LuGY committed
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
```
```shell
python inference.py target.fasta data/pdb_mmcif/mmcif_files/ \
    --output_dir ./ \
    --gpus 2 \
    --uniref90_database_path data/uniref90/uniref90.fasta \
    --mgnify_database_path data/mgnify/mgy_clusters_2018_12.fa \
    --pdb70_database_path data/pdb70/pdb70 \
    --uniclust30_database_path data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
    --jackhmmer_binary_path `which jackhmmer` \
    --hhblits_binary_path `which hhblits` \
    --hhsearch_binary_path `which hhsearch` \
    --kalign_binary_path `which kalign`  \
    --enable_workflow 
```

oahzxl's avatar
oahzxl committed
158
159
160
161
#### inference with lower memory usage
Alphafold's embedding presentations take up a lot of memory as the sequence length increases. To reduce memory usage, 
you should add parameter `--chunk_size [N]` and `--inplace` to cmdline or shell script `./inference.sh`. 
The smaller you set N, the less memory will be used, but it will affect the speed. We can inference 
162
163
164
a sequence of length 10000 in bf16 with 61GB memory on a Nvidia A100(80GB). For fp32, the max length is 8000.
> You need to set `PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:15000` to inference such an extreme long sequence.

oahzxl's avatar
oahzxl committed
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
```shell
python inference.py target.fasta data/pdb_mmcif/mmcif_files/ \
    --output_dir ./ \
    --gpus 2 \
    --uniref90_database_path data/uniref90/uniref90.fasta \
    --mgnify_database_path data/mgnify/mgy_clusters_2018_12.fa \
    --pdb70_database_path data/pdb70/pdb70 \
    --uniclust30_database_path data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
    --jackhmmer_binary_path `which jackhmmer` \
    --hhblits_binary_path `which hhblits` \
    --hhsearch_binary_path `which hhsearch` \
    --kalign_binary_path `which kalign`  \
    --chunk_size N \
    --inplace
```
181

182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
#### inference multimer sequence
Alphafold Multimer is supported. You can the following cmd or shell script `./inference_multimer.sh`.
Workflow and memory parameters mentioned above can also be used.
```shell
python inference.py target.fasta data/pdb_mmcif/mmcif_files/ \
    --output_dir ./ \
    --gpus 2 \
    --model_preset multimer \
    --uniref90_database_path data/uniref90/uniref90.fasta \
    --mgnify_database_path data/mgnify/mgy_clusters_2018_12.fa \
    --pdb70_database_path data/pdb70/pdb70 \
    --uniclust30_database_path data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
    --uniprot_database_path data/uniprot/uniprot_sprot.fasta \
    --pdb_seqres_database_path data/pdb_seqres/pdb_seqres.txt  \
    --param_path data/params/params_model_1_multimer.npz \
    --model_name model_1_multimer \
    --jackhmmer_binary_path `which jackhmmer` \
    --hhblits_binary_path `which hhblits` \
    --hhsearch_binary_path `which hhsearch` \
    --kalign_binary_path `which kalign`
```

Shenggan's avatar
Shenggan committed
205
206
207
208
209
210
211
212
213
## Performance Benchmark

We have included a performance benchmark script in `./benchmark`. You can benchmark the performance of Evoformer using different settings.

```shell
cd ./benchmark
torchrun --nproc_per_node=1 perf.py --msa-length 128 --res-length 256
```

Shenggan's avatar
Shenggan committed
214
215
216
217
218
219
220
Benchmark Dynamic Axial Parallelism with 2 GPUs:

```shell
cd ./benchmark
torchrun --nproc_per_node=2 perf.py --msa-length 128 --res-length 256 --dap-size 2
```

Shenggan's avatar
Shenggan committed
221
222
223
224
225
226
227
228
229
230
231
If you want to benchmark with [OpenFold](https://github.com/aqlaboratory/openfold), you need to install OpenFold first and benchmark with option `--openfold`:

```shell
torchrun --nproc_per_node=1 perf.py --msa-length 128 --res-length 256 --openfold
```

## Cite us

Cite this paper, if you use FastFold in your research publication.

```
Shenggan's avatar
Shenggan committed
232
233
234
235
236
237
238
239
@misc{cheng2022fastfold,
      title={FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours}, 
      author={Shenggan Cheng and Ruidong Wu and Zhongming Yu and Binrui Li and Xiwen Zhang and Jian Peng and Yang You},
      year={2022},
      eprint={2203.00854},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}
Shenggan's avatar
Shenggan committed
240
```