AlphaFold3重磅開源,諾獎級AI顛覆世界!GitHub斬獲1.8k星,本地即可部署

新智元報導  

編輯:靜音

【新智元導讀】六個月的爭議後,諾獎級AI AlphaFold3開源了。這個在蛋白質結構預測領域掀起波瀾的AI——期待它的開源推動更多科學家的大量創新。文後附有安裝和運行步驟詳解哦!

AlphaFold3源碼終於開放了!

六個月前,AlphaFold3橫空出世震撼了整個學術界。AlphaFold的開發人也憑藉它在上個月贏得了盧保化學獎。

然而,這個諾獎級AI的「不開源」一直引起學界的不滿。GoogleDeepMind只推出了一個免費研究平台「AlphaFold Server」,而且該服務有每日的次數限制。相比於開源的AlphaFold2來說,這種使用方式缺失了很多自由度。

好在它現在終於開源了!開源後,生化醫藥的科學家們可以在本地部署AlphaFold3,極大地縮短了新藥、疫苗等研發進程。

現在,任何人都可以下載AlphaFold3軟件代碼並進行非商業使用,但目前只有學術背景的科學家可申請訪問訓練權重。

GitHub上的AlphaFold3開源項目代碼目前已斬獲1.8k星。

開源項目:https://github.com/google-deepmind/alphafold3開源項目:https://github.com/google-deepmind/alphafold3

AlphaFold3的「效仿者」們

在過去的幾個月中,不少公司都依靠AlphaFold3論文中的偽代碼,爭相發佈了各自受到AlphaFold3啟發的類似模型。

比如,獲得OpenAI投資的AI生物初創Chai Discovery,就在9月發佈了用於分子結構預測的新型多模態基礎模型Chai-1,並附帶了一份技術報告,比較了Chai-1與AlphaFold等模型的性能。

官網地址:https://www.chaidiscovery.com/官網地址:https://www.chaidiscovery.com/

另一家位於美國舊金山的公司Ligo Biosciences則發佈了一個無使用限制的AlphaFold3版本。但它尚未具備完整的功能,比如模擬藥物和蛋白質以外分子的能力。

項目地址:https://github.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_destination=<DATABASES_DIR>

該腳本從託管在GCS上的鏡像下載數據庫,所有版本與AlphaFold3論文中使用的相同。

腳本完成後,應該有以下目錄結構:

pdb_2022_09_28_mmcif_files.tar  # ~200k PDB mmCIF files in this tar.bfd-first_non_consensus_sequences.fastamgy_clusters_2022_05.fant_rna_2023_02_23_clust_seq_id_90_cov_80_rep_seq.fastapdb_seqres_2022_09_28.fastarfam_14_9_clust_seq_id_90_cov_80_rep_seq.fastarnacentral_active_seq_id_90_cov_80_linclust.fastauniprot_all_2021_04.fauniref90_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_implementation=xla來運行長度最多為1280 token的AlphaFold3,但此配置尚未經過數值準確性或吞吐量效率的測試,因此請謹慎操作。

參考資料:

https://www.nature.com/articles/d41586-024-03708-4