AlphaFold3¶
AlphaFold3 是由谷歌 DeepMind 和 Isomorphic Labs 团队开发的人工智能程序,于 2024 年 5 月发布,11 月开源供学术用途。它能准确预测蛋白质、DNA、RNA 等多种生物分子及其复合体的结构和相互作用,预测精度较 AlphaFold2 大幅提升,在药物研发等领域应用广泛。其架构引入 Pairformer 和扩散模块,采用跨蒸馏技术训练,减少了多重序列比对处理量,提高了计算效率和泛化能力,为生命科学研究提供了更强大的工具。
可用的版本¶
硬件 |
平台 |
版本 |
A100-40GB |
思源一号 |
v3.0.1 |
A800-80GB |
思源一号 |
v3.0.1 |
V100-32GB |
Pi2.0(待更新) |
v3.0.0 |
版本区别¶
思源一号的A800-80GB镜像为原镜像,未做修改,AlphaFold3可以在单张A800-80GB上运行,最大能够处理5120个tokens。
对于A100-40GB,根据官方文档,我们对原镜像进行了以下修改:
开启unified memory,启用统一内存允许程序在空间不足时将 GPU 内存溢出到主机内存,这可以防止OOM,修改 /app/alphafold/run_alphafold.py 文件如下:
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE']='false'
os.environ['XLA_CLIENT_MEM_FRACTION']='3.2'
os.environ['TF_FORCE_UNIFIED_MEMORY']='true'
调整 /app/alphafold/src/alphafold3/model/model_config.py 中 pair_transition_shard_spec 参数如下:
pair_transition_shard_spec: Sequence[_Shape2DType] = (
(2048, None),
(3072, 1024),
(None, 512),
)
A100-40GB上最大能够处理4352个tokens的输入,与A800-80GB相比,计算精度一致,但吞吐量更小。
使用前准备¶
新建运行文件夹,如
alphafold
在运行文件夹中创建输入文件夹
input
和输出文件夹output
在
input
文件夹中放置输入 JSON 文件,例如 fold_input.json,自定义输入 JSON 文件可参考官方文档 https://github.com/google-deepmind/alphafold3/blob/main/docs/input.md
蛋白质结构预测 JSON 文件示例¶
{
"name": "2PV7",
"sequences": [
{
"protein": {
"id": ["A", "B"],
"sequence": "GMRESYANENQFGFKTINSDIHKIVIVGGYGKLGGLFARYLRASGYPISILDREDWAVAESILANADVVIVSVPINLTLETIERLKPYLTENMLLADLTSVKREPLAKMLEVHTGAVLGLHPMFGADIASMAKQVVVRCDGRFPERYEWLLEQIQIWGAKIYQTNATEHDHNMTYIQALRHFSTFANGLHLSKQPINLANLLALSSPIYRLELAMIGRLFAQDAELYADIIMDKSENLAVIETLKQTYDEALTFFENNDRQGFIDAFHKVRDWFGDYSEQFLKESRQLLQQANDLKQG"
}
}
],
"modelSeeds": [1],
"dialect": "alphafold3",
"version": 1
}
在思源一号上运行 AlphaFold3¶
作业脚本 af3.slurm 示例如下:
A100-40GB¶
#!/bin/bash
#SBATCH --job-name=alphafold3
#SBATCH --partition=a100
#SBATCH -N 1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=16
#SBATCH --gres=gpu:1 # use 1 GPU
#SBATCH --output=%j.out
#SBATCH --error=%j.err
singularity exec \
--nv \
--bind $PWD/input:/root/af_input \
--bind $PWD/output:/root/af_output \
--bind /dssg/share/data/alphafold3/models:/root/models \
--bind /dssg/share/data/alphafold3/database:/root/public_databases \
/dssg/share/imgs/ai/alphafold/alphafold3-a100.sif \
/alphafold3_venv/bin/python /app/alphafold/run_alphafold.py \
--json_path=/root/af_input/fold_input.json \
--model_dir=/root/models \
--db_dir=/root/public_databases \
--output_dir=/root/af_output
A800-80GB¶
#!/bin/bash
#SBATCH --job-name=alphafold3
#SBATCH --partition=a800
#SBATCH -N 1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=16
#SBATCH --gres=gpu:1 # use 1 GPU
#SBATCH --output=%j.out
#SBATCH --error=%j.err
singularity exec \
--nv \
--bind $PWD/input:/root/af_input \
--bind $PWD/output:/root/af_output \
--bind /dssg/share/data/alphafold3/models:/root/models \
--bind /dssg/share/data/alphafold3/database:/root/public_databases \
/dssg/share/imgs/ai/alphafold/alphafold3-a800.sif \
/alphafold3_venv/bin/python /app/alphafold/run_alphafold.py \
--json_path=/root/af_input/fold_input.json \
--model_dir=/root/models \
--db_dir=/root/public_databases \
--output_dir=/root/af_output
使用 sbatch af3.slurm 语句提交作业。
运行结束后,计算结果保存在 alphafold/output/ 下,具体可参考 https://github.com/google-deepmind/alphafold3/blob/main/docs/output.md
分阶段运行AlphaFold3¶
AlphaFold3运行分为 data_pipeline 和 inference 两个阶段, data_pipeline 阶段主要利用CPU进行MSA/模板搜索, inference 阶段主要利用GPU进行模型推理。
将 data_pipeline 与 inference 分开,减少GPU机时浪费。
输出结果的json中会保存缓存 MSA / 模板搜索的结果,可以用于跨种子或其他特征变体(例如配体)的多个不同推理。
以A100为例:
data_pipeline阶段脚本示例(64c512g队列运行):
#!/bin/bash
#SBATCH --job-name=alphafold3
#SBATCH --partition=64c512g
#SBATCH -N 1
#SBATCH --ntasks-per-node=32
#SBATCH --output=%j.out
#SBATCH --error=%j.err
/usr/bin/time -v singularity exec \
--nv \
--bind $PWD/input:/root/af_input \
--bind $PWD/output:/root/af_output \
--bind /dssg/share/data/alphafold3/models:/root/models \
--bind /dssg/share/data/alphafold3/database:/root/public_databases \
/dssg/share/imgs/ai/alphafold/alphafold3-a100.sif \
/alphafold3_venv/bin/python /app/alphafold/run_alphafold.py \
--norun_inference \
--jackhmmer_n_cpu=$SLURM_NTASKS \
--nhmmer_n_cpu=$SLURM_NTASKS \
--input_dir=/root/af_input \
--model_dir=/root/models \
--db_dir=/root/public_databases \
--output_dir=/root/af_output
alphafold/output/ 文件夹下会保存 data_pipeline 阶段缓存的结果,以该目录下的json文件作为 inference 阶段的输入。本示例中生成的结果文件为 2PV7_data.json。
inference阶段脚本示例(a100队列运行):
#!/bin/bash
#SBATCH --job-name=alphafold3
#SBATCH --partition=a100
#SBATCH -N 1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=16
#SBATCH --gres=gpu:1 # use 1 GPU
#SBATCH --output=%j.out
#SBATCH --error=%j.err
singularity exec \
--nv \
--bind $PWD/input:/root/af_input \
--bind $PWD/output:/root/af_output \
--bind /dssg/share/data/alphafold3/models:/root/models \
--bind /dssg/share/data/alphafold3/database:/root/public_databases \
/dssg/share/imgs/ai/alphafold/alphafold3-a800.sif \
/alphafold3_venv/bin/python /app/alphafold/run_alphafold.py \
--norun_data_pipeline \
--json_path=/root/af_input/2PV7_data.json \
--model_dir=/root/models \
--db_dir=/root/public_databases \
--output_dir=/root/af_output
data_pipeline性能测试¶
思源平台 A100-40GB
版本 |
用时(s) |
v3.0.0 |
2475 |
v3.0.1 |
677 |
参考资料¶
AlphaFold3 GitHub: https://github.com/google-deepmind/alphafold3