AlphaFold3源碼終于開(kāi)放了！

六個(gè)月前，AlphaFold3橫空出世震撼了整個(gè)學(xué)術(shù)界。AlphaFold的開(kāi)發(fā)人也憑借它在上個(gè)月贏得了諾貝爾化學(xué)獎(jiǎng)。

然而，這個(gè)諾獎(jiǎng)級(jí)AI的「不開(kāi)源」一直引起學(xué)界的不滿。谷歌DeepMind只推出了一個(gè)免費(fèi)研究平臺(tái)「AlphaFold Server」，而且該服務(wù)有每日的次數(shù)限制。相比于開(kāi)源的AlphaFold2來(lái)說(shuō)，這種使用方式缺失了很多自由度。

好在它現(xiàn)在終于開(kāi)源了！開(kāi)源后，生化醫(yī)藥的科學(xué)家們可以在本地部署AlphaFold3，極大地縮短了新藥、疫苗等研發(fā)進(jìn)程。

現(xiàn)在，任何人都可以下載AlphaFold3軟件代碼并進(jìn)行非商業(yè)使用，但目前只有學(xué)術(shù)背景的科學(xué)家可申請(qǐng)?jiān)L問(wèn)訓(xùn)練權(quán)重。

GitHub上的AlphaFold3開(kāi)源項(xiàng)目代碼目前已斬獲1.8k星。

開(kāi)源項(xiàng)目：https://github.com/google-deepmind/alphafold3

AlphaFold3的「效仿者」們

在過(guò)去的幾個(gè)月中，不少公司都依靠AlphaFold3論文中的偽代碼，爭(zhēng)相發(fā)布了各自受到AlphaFold3啟發(fā)的類似模型。

比如，獲得OpenAI投資的AI生物初創(chuàng)Chai Discovery，就在9月發(fā)布了用于分子結(jié)構(gòu)預(yù)測(cè)的新型多模態(tài)基礎(chǔ)模型Chai-1，并附帶了一份技術(shù)報(bào)告，比較了Chai-1與AlphaFold等模型的性能。

官網(wǎng)地址：https://www.chaidiscovery.com/

另一家位于美國(guó)舊金山的公司Ligo Biosciences則發(fā)布了一個(gè)無(wú)使用限制的AlphaFold3版本。但它尚未具備完整的功能，比如模擬藥物和蛋白質(zhì)以外分子的能力。

項(xiàng)目地址：https://github.com/Ligo-Biosciences/AlphaFold3

其他團(tuán)隊(duì)也正在開(kāi)發(fā)沒(méi)有使用限制的AlphaFold3版本：AlQuraishi希望在年底前推出一個(gè)名為OpenFold3的完全開(kāi)源模型。這將使制藥公司能夠使用專有數(shù)據(jù)（例如結(jié)合不同藥物的蛋白質(zhì)結(jié)構(gòu)）重新訓(xùn)練模型，從而有可能提高性能。

開(kāi)源的重要性

過(guò)去一年里，許多公司發(fā)布了新的生物AI模型，這些公司對(duì)開(kāi)放性采取了不同的態(tài)度。

威斯康星大學(xué)麥迪遜分校的計(jì)算生物學(xué)家Anthony Gitter對(duì)盈利性公司加入他的領(lǐng)域沒(méi)有異議——只要他們?cè)谄诳皖A(yù)印本服務(wù)器上分享工作時(shí)遵循科學(xué)界的標(biāo)準(zhǔn)。

「我和其他人希望盈利性公司們也分享關(guān)于如何進(jìn)行預(yù)測(cè)的信息，并以我們可以審查的方式發(fā)布AI模型和代碼，」Gitter補(bǔ)充道，「我的團(tuán)隊(duì)不會(huì)基于無(wú)法審查的工具進(jìn)行構(gòu)建和使用?！?/span>

DeepMind科學(xué)AI負(fù)責(zé)人Pushmeet Kohli表示，幾種AlphaFold3復(fù)制品的出現(xiàn)表明，即使沒(méi)有開(kāi)源代碼，該模型也是可復(fù)現(xiàn)的。

他補(bǔ)充說(shuō)，未來(lái)他希望看到更多關(guān)于出版規(guī)范的討論，因?yàn)檫@一領(lǐng)域越來(lái)越多地由學(xué)術(shù)界和企業(yè)研究人員共同參與。

此前，AlphaFold2的開(kāi)源推動(dòng)了其他科學(xué)家的大量創(chuàng)新。

例如，最近一次蛋白質(zhì)設(shè)計(jì)競(jìng)賽的獲勝者使用該AI工具設(shè)計(jì)出能夠結(jié)合癌癥靶標(biāo)的新蛋白質(zhì)。

AlphaFold項(xiàng)目的負(fù)責(zé)人Jumper最喜歡的一個(gè)AlphaFold2創(chuàng)新，是一個(gè)團(tuán)隊(duì)使用該工具識(shí)別出一種幫助精子附著在卵細(xì)胞上的關(guān)鍵蛋白。

Jumper迫不及待地想看到在分享AlphaFold3后出現(xiàn)這樣的驚喜。

安裝和運(yùn)行

安裝AlphaFold3需要一臺(tái)運(yùn)行Linux的機(jī)器；AlphaFold3不支持其他操作系統(tǒng)。

完整安裝需要多達(dá)1TB的磁盤(pán)空間來(lái)存儲(chǔ)基因數(shù)據(jù)庫(kù)（建議使用SSD存儲(chǔ)）以及一塊具有計(jì)算能力8.0或更高的 NVIDIA GPU（具有更多內(nèi)存的GPU可以預(yù)測(cè)更大的蛋白質(zhì)結(jié)構(gòu)）。

經(jīng)過(guò)驗(yàn)證，單個(gè)NVIDIA A100 80 GB或NVIDIA H100 80 GB可以適配最多5120個(gè)token的輸入。在NVIDIA A100和H100 GPU上的數(shù)值準(zhǔn)確性也已被驗(yàn)證。

尤其是對(duì)于較長(zhǎng)的目標(biāo)，基因搜索階段可能會(huì)消耗大量RAM——建議至少使用64GB的RAM運(yùn)行。

配置步驟：

1. 在GCP上配置機(jī)器

2. 安裝Docker

3. 為A100安裝NVIDIA驅(qū)動(dòng)程序

4. 獲取基因數(shù)據(jù)庫(kù)

5. 獲取模型參數(shù)

6. 構(gòu)建AlphaFold3 Docker容器或Singularity鏡像

獲取AlphaFold3源代碼

通過(guò)git下載AlphaFold3的代碼庫(kù)：

git clone https://github.com/google-deepmind/alphafold3.git

獲取基因數(shù)據(jù)庫(kù)

此步驟需要「curl」和「zstd」。

AlphaFold3需要多個(gè)基因（序列）蛋白質(zhì)和RNA數(shù)據(jù)庫(kù)來(lái)運(yùn)行：

- BFD small

- MGnify

- PDB（mmCIF格式的結(jié)構(gòu)）

- PDB seqres

- UniProt

- UniRef90

- NT

- RFam

- RNACentral

Python程序「fetch_databases.py」可以用來(lái)下載和設(shè)置所有這些數(shù)據(jù)庫(kù)。

建議在「screen」或「tmux」會(huì)話中運(yùn)行以下命令，因?yàn)橄螺d和解壓數(shù)據(jù)庫(kù)需要一些時(shí)間。完整數(shù)據(jù)庫(kù)的總下載大小約為252GB，解壓后的總大小為630GB。

cd alphafold3  # Navigate to the directory with cloned AlphaFold3 repository.
python3 fetch_databases.py --download_destinatinotallow=<DATABASES_DIR>

該腳本從托管在GCS上的鏡像下載數(shù)據(jù)庫(kù)，所有版本與AlphaFold3論文中使用的相同。

腳本完成后，應(yīng)該有以下目錄結(jié)構(gòu)：

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

獲取模型參數(shù)

訪問(wèn)AlphaFold3模型參數(shù)需要向Google DeepMind申請(qǐng)并獲得授權(quán)。

數(shù)據(jù)管線

數(shù)據(jù)管線的運(yùn)行時(shí)間（即基因序列搜索和模板搜索）可能會(huì)因輸入的大小、找到的同源序列數(shù)量以及可用的硬件（磁盤(pán)速度尤其會(huì)影響基因搜索的速度）而顯著變化。

如果想提高性能，建議提高磁盤(pán)速度（例如通過(guò)利用基于RAM的文件系統(tǒng)），或增加可用的CPU核心并增加并行處理。

此外，請(qǐng)注意，對(duì)于具有深度MSA的序列，Jackhmmer或Nhmmer可能需要超出推薦的64 GB RAM的大量?jī)?nèi)存。

模型推理

AlphaFold3論文的補(bǔ)充信息中的表8提供了在配置為運(yùn)行在16個(gè)NVIDIA A100上時(shí)的AlphaFold3的無(wú)需編譯的推理時(shí)間，每個(gè)設(shè)備具有40GB的內(nèi)存。

相比之下，該存儲(chǔ)庫(kù)支持在單個(gè)NVIDIA A100上運(yùn)行AlphaFold3，具有80GB內(nèi)存，并在配置上進(jìn)行了優(yōu)化以最大化吞吐量。

下表中使用GPU秒（即使用16個(gè)A100時(shí)乘以16）比較了這兩種設(shè)置的無(wú)需編譯的推理時(shí)間。該存儲(chǔ)庫(kù)中的設(shè)置在所有token大小上效率更高（提高至少2倍），表明其適合高吞吐量應(yīng)用。

硬件要求

AlphaFold3正式支持以下配置，并已對(duì)其進(jìn)行了廣泛的數(shù)值準(zhǔn)確性和吞吐量效率測(cè)試：

- 1 NVIDIA A100（80GB）

- 1 NVIDIA H100（80GB）

通過(guò)以下配置更改，AlphaFold3可以在單個(gè)NVIDIA A100 (40GB) 上運(yùn)行：

1. 啟用統(tǒng)一內(nèi)存。

2. 調(diào)整model_config.py中的pair_transition_shard_spec：

pair_transition_shard_spec: Sequence[_Shape2DType] = (
      (2048, None),
      (3072, 1024),
      (None, 512),
  )

雖然數(shù)值上準(zhǔn)確，但由于可用內(nèi)存較少，因此與NVIDIA A100 (80GB) 的設(shè)置相比，該配置的吞吐量會(huì)較低。

雖然也可以在單個(gè)NVIDIA V100上使用run_alphafold.py中的--flash_attention_implementatinotallow=xla來(lái)運(yùn)行長(zhǎng)度最多為1280 token的AlphaFold3，但此配置尚未經(jīng)過(guò)數(shù)值準(zhǔn)確性或吞吐量效率的測(cè)試，因此請(qǐng)謹(jǐn)慎操作。