README.md 9.39 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)
7
![](https://img.shields.io/badge/Habana-support-blue?style=flat&logo=intel&logoColor=blue)
Shenggan's avatar
Shenggan committed
8
9
![](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
10

11
12
13
14
15
16
17
## News :triangular_flag_on_post:
- [2023/01] Compatible with AlphaFold v2.3
- [2023/01] Added support for inference and training of AlphaFold on [Intel Habana](https://habana.ai/) platform. For usage instructions, see [here](#Inference-or-Training-on-Intel-Habana).

<br>

Optimizing Protein Structure Prediction Model Training and Inference on Heterogeneous Clusters
Shenggan's avatar
Shenggan committed
18
19
20
21
22
23

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
24
    * DAP can significantly speed up inference and make ultra-long sequence inference possible
Shenggan's avatar
Shenggan committed
25
3. Ease of use
Shenggan's avatar
Shenggan committed
26
    * Huge performance gains with a few lines changes
Shenggan's avatar
Shenggan committed
27
    * You don't need to care about how the parallel part is implemented
28
29
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
30
31
32

## Installation

oahzxl's avatar
oahzxl committed
33
To install FastFold, you will need:
LuGY's avatar
LuGY committed
34
+ Python 3.8 or 3.9.
oahzxl's avatar
oahzxl committed
35
36
+ [NVIDIA CUDA](https://developer.nvidia.com/cuda-downloads) 11.3 or above
+ PyTorch 1.12 or above 
LuGY's avatar
LuGY committed
37

38
39
40
41
42
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
43

Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
44
45
46
```shell
git clone https://github.com/hpcaitech/FastFold
cd FastFold
Shenggan's avatar
Shenggan committed
47
conda env create --name=fastfold -f environment.yml
Shenggan's avatar
Shenggan committed
48
conda activate fastfold
49
python setup.py install
Shenggan's avatar
Shenggan committed
50
51
```

52
53
#### Advanced

oahzxl's avatar
oahzxl committed
54
To leverage the power of FastFold, we recommend you to install [Triton](https://github.com/openai/triton).
55
56

```bash
oahzxl's avatar
oahzxl committed
57
pip install triton==2.0.0.dev20221005
58
59
60
```


61
62
### Using PyPi
You can download FastFold with pre-built CUDA extensions.
Shenggan's avatar
Shenggan committed
63

64
65
Warning, only stable versions available.

Shenggan's avatar
Shenggan committed
66
```shell
67
68
69
70
71
72
73
74
75
76
77
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
78
79
cd FastFold
docker build -t fastfold ./docker
80
81
82
83
84
```

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
85
86
```

Shenggan's avatar
Shenggan committed
87
88
## Usage

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

```python
92
from fastfold.model.fastnn import Evoformer
Shenggan's avatar
Shenggan committed
93
94
95
evoformer_layer = Evoformer()
```

Shenggan's avatar
Shenggan committed
96
If you want to use Dynamic Axial Parallelism, add a line of initialize with `fastfold.distributed.init_dap`.
Shenggan's avatar
Shenggan committed
97
98
99
100
101
102
103

```python
from fastfold.distributed import init_dap

init_dap(args.dap_size)
```

Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
104
105
106
107
108
### 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/

109
110
### Inference

Shenggan's avatar
Shenggan committed
111
You can use FastFold with `inject_fastnn`. This will replace the evoformer from OpenFold with the high performance evoformer from FastFold.
112
113

```python
114
from fastfold.utils import inject_fastnn
115
116
117
118

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

119
model = inject_fastnn(model)
120
121
122
123
124
```

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

```shell
125
python inference.py target.fasta data/pdb_mmcif/mmcif_files/ \
126
    --output_dir .outputs/ \
127
    --gpus 2 \
128
    --uniref90_database_path data/uniref90/uniref90.fasta \
129
    --mgnify_database_path data/mgnify/mgy_clusters_2022_05.fa \
130
    --pdb70_database_path data/pdb70/pdb70 \
131
    --uniref30_database_path data/uniref30/UniRef30_2021_03 \
132
    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
Shenggan's avatar
Shenggan committed
133
134
135
    --jackhmmer_binary_path `which jackhmmer` \
    --hhblits_binary_path `which hhblits` \
    --hhsearch_binary_path `which hhsearch` \
136
137
138
    --kalign_binary_path `which kalign` \
    --enable_workflow \
    --inplace
139
```
LuGY's avatar
LuGY committed
140
141
142
143
144
or run the script `./inference.sh`, you can change the parameter in the script, especisally those data path.
```shell
./inference.sh
```

145
146
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, we add parameter `--enable_workflow` by default.
To reduce memory usage of embedding presentations, we also add parameter `--inplace` to share memory by defaul.
LuGY's avatar
LuGY committed
147

oahzxl's avatar
oahzxl committed
148
149
#### inference with lower memory usage
Alphafold's embedding presentations take up a lot of memory as the sequence length increases. To reduce memory usage, 
150
you should add parameter `--chunk_size [N]` to cmdline or shell script `./inference.sh`. 
oahzxl's avatar
oahzxl committed
151
The smaller you set N, the less memory will be used, but it will affect the speed. We can inference 
152
153
154
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
155
156
```shell
python inference.py target.fasta data/pdb_mmcif/mmcif_files/ \
157
    --output_dir .outputs/ \
oahzxl's avatar
oahzxl committed
158
159
    --gpus 2 \
    --uniref90_database_path data/uniref90/uniref90.fasta \
160
    --mgnify_database_path data/mgnify/mgy_clusters_2022_05.fa \
oahzxl's avatar
oahzxl committed
161
    --pdb70_database_path data/pdb70/pdb70 \
162
    --uniref30_database_path data/uniref30/UniRef30_2021_03 \
oahzxl's avatar
oahzxl committed
163
164
165
166
167
    --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`  \
168
    --enable_workflow \
oahzxl's avatar
oahzxl committed
169
    --inplace
170
    --chunk_size N \
oahzxl's avatar
oahzxl committed
171
```
172

173
174
175
176
177
178
179
180
181
#### 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 \
182
    --mgnify_database_path data/mgnify/mgy_clusters_2022_05.fa \
183
    --pdb70_database_path data/pdb70/pdb70 \
184
    --uniref30_database_path data/uniref30/UniRef30_2021_03 \
185
    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
186
    --uniprot_database_path data/uniprot/uniprot.fasta \
187
188
189
190
191
192
193
194
195
    --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`
```

196
197
### Inference or Training on Intel Habana

198
To run AlphaFold inference or training on Intel Habana, you can follow the instructions in the [Installation Guide](https://docs.habana.ai/en/latest/Installation_Guide/) to set up your environment on Amazon EC2 DL1 instances or on-premise environments, and please use SynapseAI R1.7.1 to test as it was verified internally.
199
200
201
202

Once you have prepared your dataset and installed fastfold, you can use the following scripts:

```shell
203
cd fastfold/habana/fastnn/custom_op/; python setup.py build (this is for Gaudi, for Gaudi2 please use setup2.py) ; cd -
204
205
206
207
bash habana/inference.sh
bash habana/train.sh
```

Shenggan's avatar
Shenggan committed
208
209
210
211
212
213
214
215
216
## 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
217
218
219
220
221
222
223
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
224
225
226
227
228
229
230
231
232
233
234
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
235
236
237
238
239
240
241
242
@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
243
```
244
245
246
247

## Acknowledgments

We would like to extend our special thanks to the Intel Habana team for their support in providing us with technology and resources on the Habana platform.