AlphaFold3源码终于开放了!
六个月前,AlphaFold3横空出世震撼了整个学术界。AlphaFold的开发人也凭借它在上个月赢得了诺贝尔化学奖。
然而,这个诺奖级AI的「不开源」一直引起学界的不满。谷歌DeepMind只推出了一个免费研究平台「AlphaFold Server」,而且该服务有每日的次数限制。相比于开源的AlphaFold2来说,这种使用方式缺失了很多自由度。
好在它现在终于开源了!开源后,生化医药的科学家们可以在本地部署AlphaFold3,极大地缩短了新药、疫苗等研发进程。
现在,任何人都可以下载AlphaFold3软件代码并进行非商业使用,但目前只有学术背景的科学家可申请访问训练权重。
GitHub上的AlphaFold3开源项目代码目前已斩获1.8k星。
开源项目:https://github.com/google-deepmind/alphafold3
AlphaFold3的「效仿者」们
在过去的几个月中,不少公司都依靠AlphaFold3论文中的伪代码,争相发布了各自受到AlphaFold3启发的类似模型。
比如,获得OpenAI投资的AI生物初创Chai Discovery,就在9月发布了用于分子结构预测的新型多模态基础模型Chai-1,并附带了一份技术报告,比较了Chai-1与AlphaFold等模型的性能。
官网地址:https://www.chaidiscovery.com/
另一家位于美国旧金山的公司Ligo Biosciences则发布了一个无使用限制的AlphaFold3版本。但它尚未具备完整的功能,比如模拟药物和蛋白质以外分子的能力。
项目地址:https://github.com/Ligo-Biosciences/AlphaFold3
其他团队也正在开发没有使用限制的AlphaFold3版本:AlQuraishi希望在年底前推出一个名为OpenFold3的完全开源模型。这将使制药公司能够使用专有数据(例如结合不同药物的蛋白质结构)重新训练模型,从而有可能提高性能。
开源的重要性
过去一年里,许多公司发布了新的生物AI模型,这些公司对开放性采取了不同的态度。
威斯康星大学麦迪逊分校的计算生物学家Anthony Gitter对盈利性公司加入他的领域没有异议——只要他们在期刊和预印本服务器上分享工作时遵循科学界的标准。
「我和其他人希望盈利性公司们也分享关于如何进行预测的信息,并以我们可以审查的方式发布AI模型和代码,」Gitter补充道,「我的团队不会基于无法审查的工具进行构建和使用。」
DeepMind科学AI负责人Pushmeet Kohli表示,几种AlphaFold3复制品的出现表明,即使没有开源代码,该模型也是可复现的。
他补充说,未来他希望看到更多关于出版规范的讨论,因为这一领域越来越多地由学术界和企业研究人员共同参与。
此前,AlphaFold2的开源推动了其他科学家的大量创新。
例如,最近一次蛋白质设计竞赛的获胜者使用该AI工具设计出能够结合癌症靶标的新蛋白质。
AlphaFold项目的负责人Jumper最喜欢的一个AlphaFold2创新,是一个团队使用该工具识别出一种帮助精子附着在卵细胞上的关键蛋白。
Jumper迫不及待地想看到在分享AlphaFold3后出现这样的惊喜。
安装和运行
安装AlphaFold3需要一台运行Linux的机器;AlphaFold3不支持其他操作系统。
完整安装需要多达1TB的磁盘空间来存储基因数据库(建议使用SSD存储)以及一块具有计算能力8.0或更高的 NVIDIA GPU(具有更多内存的GPU可以预测更大的蛋白质结构)。
经过验证,单个NVIDIA A100 80 GB或NVIDIA H100 80 GB可以适配最多5120个token的输入。在NVIDIA A100和H100 GPU上的数值准确性也已被验证。
尤其是对于较长的目标,基因搜索阶段可能会消耗大量RAM——建议至少使用64GB的RAM运行。
配置步骤:
1. 在GCP上配置机器
2. 安装Docker
3. 为A100安装NVIDIA驱动程序
4. 获取基因数据库
5. 获取模型参数
6. 构建AlphaFold3 Docker容器或Singularity镜像
获取AlphaFold3源代码
通过git下载AlphaFold3的代码库:
复制
git clone https://github.com/google-deepmind/alphafold3.git
获取基因数据库
此步骤需要「curl」和「zstd」。
AlphaFold3需要多个基因(序列)蛋白质和RNA数据库来运行:
– BFD small
– MGnify
– PDB(mmCIF格式的结构)
– PDB seqres
– UniProt
– UniRef90
– NT
– RFam
– RNACentral
Python程序「fetch_databases.py」可以用来下载和设置所有这些数据库。
建议在「screen」或「tmux」会话中运行以下命令,因为下载和解压数据库需要一些时间。完整数据库的总下载大小约为252GB,解压后的总大小为630GB。
复制
cd alphafold3 # Navigate to the directory with cloned AlphaFold3 repository. python3 fetch_databases.py --download_destinatinotallow=<DATABASES_DIR>
该脚本从托管在GCS上的镜像下载数据库,所有版本与AlphaFold3论文中使用的相同。
脚本完成后,应该有以下目录结构:
复制
pdb_2022_09_28_mmcif_files.tar # ~200k PDB mmCIF files in this tar. bfd-first_non_consensus_sequences.fasta mgy_clusters_2022_05.fa nt_rna_2023_02_23_clust_seq_id_90_cov_80_rep_seq.fasta pdb_seqres_2022_09_28.fasta rfam_14_9_clust_seq_id_90_cov_80_rep_seq.fasta rnacentral_active_seq_id_90_cov_80_linclust.fasta uniprot_all_2021_04.fa uniref90_2022_05.fa
获取模型参数
访问AlphaFold3模型参数需要向Google DeepMind申请并获得授权。
数据管线
数据管线的运行时间(即基因序列搜索和模板搜索)可能会因输入的大小、找到的同源序列数量以及可用的硬件(磁盘速度尤其会影响基因搜索的速度)而显著变化。
如果想提高性能,建议提高磁盘速度(例如通过利用基于RAM的文件系统),或增加可用的CPU核心并增加并行处理。
此外,请注意,对于具有深度MSA的序列,Jackhmmer或Nhmmer可能需要超出推荐的64 GB RAM的大量内存。
模型推理
AlphaFold3论文的补充信息中的表8提供了在配置为运行在16个NVIDIA A100上时的AlphaFold3的无需编译的推理时间,每个设备具有40GB的内存。
相比之下,该存储库支持在单个NVIDIA A100上运行AlphaFold3,具有80GB内存,并在配置上进行了优化以最大化吞吐量。
下表中使用GPU秒(即使用16个A100时乘以16)比较了这两种设置的无需编译的推理时间。该存储库中的设置在所有token大小上效率更高(提高至少2倍),表明其适合高吞吐量应用。
硬件要求
AlphaFold3正式支持以下配置,并已对其进行了广泛的数值准确性和吞吐量效率测试:
– 1 NVIDIA A100(80GB)
– 1 NVIDIA H100(80GB)
通过以下配置更改,AlphaFold3可以在单个NVIDIA A100 (40GB) 上运行:
1. 启用统一内存。
2. 调整model_config.py中的pair_transition_shard_spec:
复制
pair_transition_shard_spec: Sequence[_Shape2DType] = ( (2048, None), (3072, 1024), (None, 512), )
虽然数值上准确,但由于可用内存较少,因此与NVIDIA A100 (80GB) 的设置相比,该配置的吞吐量会较低。
虽然也可以在单个NVIDIA V100上使用run_alphafold.py中的–flash_attention_implementatinotallow=xla来运行长度最多为1280 token的AlphaFold3,但此配置尚未经过数值准确性或吞吐量效率的测试,因此请谨慎操作。