2024年11月11日の夜(日本時間)にAlphaFold3のコードが公開されました。GitHubリポジトリはこちら。
Dockerがある場合はそちらを使えば簡単に実行できますが、共用計算機環境向けのインストール方法を記述します。
はじめに
AlphaFold3のモデルパラメータは、Google DeepMind社に利用申請して認可された人のみが利用可能という規約になっています。この仕様のため、AlphaFold2のときのように何も準備なしに構造予測を利用するということはできないことになっています。しかしながら、ダウンロードしてきたモデルパラメータのディレクトリをmodels
という名前に設定して、使いたい計算機の上にアップロードすれば利用可能なので(※他の人がアクセスできるパーミッション設定にはしないようにしましょう)、スパコンの上で使うことも可能です。
申請はGoogle Formsから行います(当然英語)。また、Gmailアカウントが必要で、非営利目的で使用する誓約を行います。申請後、通常は2〜3営業日で認可され、ダウンロードリンクがメールで送られてきます。
システム要件
- GCC 11以上。9.4未満ではAlphaFold3に必要な一部のソフトウェア(PDB-REDOのlibcifpp)をビルドできません(ビルドできてしまえばあとは不要)。GCC 10でOKかは未確認。
- Python 3.11以上。3.10以下では動作しません。
ハードウェア要件
- RAM 64GB
- 高速にMSAを作成するために必要
- SSD 最低650GB以上の容量
- AlphaFold2のときより少なくて済みます。展開前でおよそ250GB, 展開後で640GBほどの空き容量を必要とします。
- CUDA 12.6が動作するNVIDIA製GPU。
- RTX30xxシリーズや40xxシリーズでも動作する。VRAMはもちろん大きい方が良い(RTX3090, RTX4090とか)。スパコンに搭載されているA100, H100でももちろん動作する。
インストール手順
gccコンパイラ, pythonバージョンの確認
GCCは11以上(9.4以上ならOK?)、Pythonは3.11以上が必要です。これが計算環境にある場合は省略してください。
Ubuntu 22.04でGCC 11コンパイラは以下で簡単に入ります(rootで実行)。
apt -y install gcc g++
次にPythonですが、3.10以下ではrun_alphafold.py
で
ImportError: cannot import name 'Self' from 'typing' (/usr/lib/python3.10/typing.py)
が出るため動作しません。Ubuntu 22.04ではデフォルトがpython3.10であり、システム/usr/bin/python3
のバージョンを変えてしまうと後々に思いもよらない不具合が出てしまうので、なんとかしてpython3.11をインストールする必要があります。なお、apt
を使うことでpython3.11.0rc1をインストールすることは一応可能で、これを用いてAlphaFold3をインストールすること自体は可能なのですが……これを使うとどうも構造推論開始時にSegmentation Faultで落ちてしまうようです。
AlphaFold3の公式Dockerfileではadd-apt-repository ppa:deadsnakes/ppa
とあり、このやり方でもpython3.11をインストールすることもできますが、これはシステムのpython3を変更してしまうものなので、よくわからず実行することはやめておいたほうが良いです(Docker内だから使える技ですね)。
なので、手動でPython 3.11をインストールするなりなんなりしてください。 以下にpython3.11を手動ビルド&インストールする手順を書きます。CentOS 7 or 8の方はこちらを参考にしてください。
# rootで実行
# Python 3.11.10をインストールする場所の指定
# ここでは/home/moriwaki/apps/python/3.11.10にインストールすることにします
INSTALL_PREFIX="/home/moriwaki/apps/python/3.11.10"
### 以下はそのまま入力 ###
# libssl1.1をインストールしておく必要があるかもしれない(不要かもしれないけど未確認)
wget http://archive.ubuntu.com/ubuntu/pool/main/o/openssl/libssl1.1_1.1.0g-2ubuntu4_amd64.deb
sudo dpkg -i libssl1.1_1.1.0g-2ubuntu4_amd64.deb
# python3ビルドのための依存関係をインストール
apt-get build-dep python3
apt-get install pkg-config
apt-get install build-essential gdb lcov pkg-config \
libbz2-dev libffi-dev libgdbm-dev libgdbm-compat-dev liblzma-dev \
libncurses5-dev libreadline6-dev libsqlite3-dev libssl-dev \
lzma lzma-dev tk-dev uuid-dev zlib1g-dev
# Python 3.11.10を手動インストール
wget https://www.python.org/ftp/python/3.11.10/Python-3.11.10.tar.xz
tar xvf Python-3.11.10.tar.xz
cd Python-3.11.10
export TCLTK_LIBS="-ltk8.6 -ltkstub8.6 -ltcl8.6"
./configure --prefix=${INSTALL_PREFIX} --enable-ipv6 \
--enable-loadable-sqlite-extensions --with-system-ffi --with-lto \
--enable-optimizations
make -j8
make altinstall
openssl 1.1がある状態でビルドすることでpipコマンドが使えるようになっているはずです。
cudaとnvidia driver
すでにNVIDIA driverとCUDA 12.6が入っている場合は、この手順は不要です。入っていない場合は以下の手順でインストールしてください。
現状では公式のalphafold3/dev-requirements.txt
においてnvidia-cublas-cu12==12.6.3.3
が指定されているので、これに対応するCUDA 12.6を入れておくのが無難です(もしかしたら12.3とかでも動くかもしれませんが未確認です)。
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
apt-get -y install nvidia-open
apt-get -y install cuda-toolkit-12-6
# 他に共存させたいcuda toolkitバージョンがあれば入れておく
# apt-get -y install cuda-toolkit-12-3 cuda-toolkit-12-4 cuda-toolkit-12-5
終わったらマシンを再起動し、nvidia-smi
でGPUの状態が表示されればOKです。
$ nvidia-smi [/mnt/database]
Fri Nov 15 14:56:40 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 565.57.01 Driver Version: 565.57.01 CUDA Version: 12.7 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA GeForce RTX 4090 On | 00000000:01:00.0 Off | Off |
| 0% 48C P8 22W / 450W | 2MiB / 24564MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| No running processes found |
+-----------------------------------------------------------------------------------------+
NVIDIA DriverとCUDAのインストール方法は、常に公式ページの手順に従ったほうが良いです。ネットの日本語記事は古い方法のままになっていることがあります。
HMMER 3.4のインストール
HMMERを使ったMSA作成を行うので、これをインストールします(AlphaFold2のときはhhblits
などでしたが変わったようです)。下の例では/home/moriwaki/apps/hmmer/3.4
にインストールしていますが、好きなところにインストールしてOKです。
HMMER_DIR="/home/moriwaki/apps/hmmer/3.4"
wget http://eddylab.org/software/hmmer/hmmer-3.4.tar.gz
tar zxvf hmmer-3.4.tar.gz
cd hmmer-3.4
./configure --prefix=${HMMER_DIR}
make -j8
make -j8 install
# easelの追加インストール
cd easel
make install
AlphaFold3のコード、配列・構造データベースの配置
パブリックな配列・構造データベースをダウンロードします。AlphaFold2のときと違い、DeepMindが動作に必要なファイルをまとめてくださっているので簡単にセットアップできます。
ただし、最新の圧縮&展開ソフトウェアZstandard (zstd)が必要です。ない場合は先にインストールしておきましょう。
# Ubuntuの場合
sudo apt -y install zstd
# RedHat系(CentOSなど)の場合
sudo yum -y install zstd
TIPS: zstdはCMakeがあれば手動ビルド&インストールも簡単にできると思う
次に以下のコマンドでダウンロードします。下の例ではAlphaFold3ディレクトリ内のpublic_databases
というディレクトリ内に作成し、そこに入れています。
これはDockerfileでのやり方を真似ています。
# ディレクトリの指定
APPDIR="/home/moriwaki/apps"
##### 以下コピペ #####
mkdir -p $APPDIR
cd $APPDIR
git clone https://github.com/google-deepmind/alphafold3.git
ALPHAFOLD3DIR="$APPDIR/alphafold3"
# クローンしてきたalphafold3ディレクトリに移動
cd ${ALPHAFOLD3DIR}
# 現在はpython3.6でも動作可能
python3 fetch_databases.py --download_destination="$ALPHAFOLD3DIR/public_databases"
public_databases
のサイズは展開後で641,428MB (640GB)でした。これくらいの容量を確保しておいてください。
/home/moriwaki/apps/alphafold3/public_databases
├── 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_2022_09_28_mmcif_files.tar
├── 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
配列・構造データベースディレクトリの場所
public_databases
ディレクトリはどこに配置しても構いませんが、必ずHDDやネットワークファイルドライブ上ではなくSSDの中に置いてください。これはMSA作成時に大量のファイルアクセスが発生するため、HDDやネットワークファイルドライブだと計算時間が大幅に遅くなるからです。models
ディレクトリはそこまでファイルサイズが大きくないので、HDDやネットワークファイルドライブ上でも問題ありません。
モデルパラメータのダウンロードと配置
メールで送られてきた学習済みパラメータはたった1.1GBくらいで、これもZstdでファイル展開する必要があります。
zstd -d af3.bin.zst
手元のmacOSなどでHomebrewを用いてbrew install zstd
でZstandardをインストールしておき、同コマンドで展開してからGPUつきLinux計算機上にアップロードしてもOKです。
展開し終わったら、中身af3.bin
を適当なmodels
ディレクトリの中に配置します。
models/
├── af3.bin
└── af3.bin.zst
例えばこれを/home/moriwaki/models
に配置した場合、後述のrun_alphafold.sh
の中で--model_dir=/home/moriwaki/models
と指定します。
AlphaFold3のpipインストール
先ほどインストールした、またはもともと存在するpython3.11を使って仮想環境を作成し、それを用いてpipで必要なモジュールと、AlphaFold3本体のインストールを行います。
# alphafold3ディレクトリの中に仮想環境を作ると仮定します。
cd ${ALPHAFOLD3DIR}
# 仮想環境の作成。先ほどインストールしたpython3.11を使う。絶対パスの方が安全。
/home/moriwaki/apps/python/3.11.10/bin/python3.11 -m venv .venv
# 仮想環境の有効化
. .venv/bin/activate
# python3.11が仮想環境のものになっているか確認
which python
# ${ALPHAFOLD3DIR}/.venv/bin/python3.11 のような表示になっていることを確認
# 必要なライブラリのインストール
python3.11 -m pip install absl-py==2.1.0 chex==0.1.87 dm-haiku==0.0.13 dm-tree==0.1.8 filelock==3.16.1 "jax[cuda12]==0.4.34" jax-cuda12-pjrt==0.4.34 jax-triton==0.2.0 jaxlib==0.4.34 jaxtyping==0.2.34 jmp==0.0.4 ml-dtypes==0.5.0 numpy==2.1.3 nvidia-cublas-cu12==12.6.3.3 nvidia-cuda-cupti-cu12==12.6.80 nvidia-cuda-nvcc-cu12==12.6.77 nvidia-cuda-runtime-cu12==12.6.77 nvidia-cudnn-cu12==9.5.1.17 nvidia-cufft-cu12==11.3.0.4 nvidia-cusolver-cu12==11.7.1.2 nvidia-cusparse-cu12==12.5.4.2 nvidia-nccl-cu12==2.23.4 nvidia-nvjitlink-cu12==12.6.77 opt-einsum==3.4.0 pillow==11.0.0 rdkit==2024.3.5 scipy==1.14.1 tabulate==0.9.0 toolz==1.0.0 tqdm==4.67.0 triton==3.1.0 typeguard==2.13.3 typing-extensions==4.12.2 zstandard==0.23.0
# AlphaFold3のインストール(必要なソフトウェアのコンパイル含む)
python3.11 -m pip install --no-deps .
# ....
# ....
# Successfully installed alphafold3-3.0.0
インストール後、.venv/bin/build_data
を実行してalphafold3/constants/converters/chemical_component_sets.pickle
というファイルを作成します。これはリガンドのccdファイルlibcifpp/components.cif
と構造式の対応表を作成しているものです(その仕様上、定期的にbuild_data
を再実行して更新してあげると、インプット作成時にユーザー側で使えるccdコードが増えてうれしいかもしれない)。
$ .venv/bin/build_data
# writing to /home/moriwaki/apps/alphafold3/.venv/lib/python3.11/site-packages/alphafold3/constants/converters/chemical_component_sets.pickle
# Done
最後に、python3.11 run_alphafold.py --help
でエラーなくヘルプメッセージが表示されればOKです。
$ python3.11 run_alphafold.py --help
AlphaFold 3 structure prediction script.
AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
...
...
--uniprot_cluster_annot_database_path: UniProt database path, used for protein paired MSA search.
(default: '${DB_DIR}/uniprot_all_2021_04.fa')
--uniref90_database_path: UniRef90 database path, used for MSA search. The MSA obtained by searching it is used to construct the
profile for template search.
(default: '${DB_DIR}/uniref90_2022_05.fa')
Try --helpfull to get a list of all flags.
システム側の実行スクリプトの準備
Dockerではなくスパコンなどの共用環境でAlphaFold3のサービスを提供したい場合、配列・構造データベースへのパスDB_DIR
やHMMER系ソフトウェアへの引数をあらかじめ設定しておいたほうが良いでしょう。そして残りの引数はユーザー側が指定するようにします。
ユーザー側の実行スクリプトの準備
環境変数HOME
を設定している場合、そのディレクトリの直下にあるmodels
とpublic_databases
ディレクトリを使える仕様になっています。つまりexport HOME=/home/moriwaki/apps
としている場合、/home/moriwaki/apps/models
と/home/moriwaki/apps/public_databases
が使われます。実行用シェルスクリプトにそれを設定しても良いでしょう。ただ、run_alphafold.py
の引数で指定したそれらのファイルパスを指定した場合はそちらが優先されます。
該当箇所のコードは以下の通りです。
_HOME_DIR = pathlib.Path(os.environ.get('HOME'))
DEFAULT_MODEL_DIR = _HOME_DIR / 'models'
DEFAULT_DB_DIR = _HOME_DIR / 'public_databases'
引数一覧
flags:
/alphafold3/run_alphafold.py:
--buckets: Strictly increasing order of token sizes for which to cache compilations. For any input with more tokens than the largest bucket size, a new bucket is created for exactly that number of tokens.
(default: '256,512,768,1024,1280,1536,2048,2560,3072,3584,4096,4608,5120')
(a comma separated list)
--db_dir: Path to the directory containing the databases.
(default: '/home/moriwaki/public_databases')
--flash_attention_implementation: <triton|cudnn|xla>: Flash attention implementation to use. 'triton' and 'cudnn' uses a Triton and cuDNN flash attention implementation, respectively. The Triton kernel is fastest and has been tested more thoroughly. The
Triton and cuDNN kernels require Ampere GPUs or later. 'xla' uses an XLA attention implementation (no flash attention) and is portable across GPU devices.
(default: 'triton')
--hmmalign_binary_path: Path to the Hmmalign binary.
--hmmbuild_binary_path: Path to the Hmmbuild binary.
--hmmsearch_binary_path: Path to the Hmmsearch binary.
--input_dir: Path to the directory containing input JSON files.
--jackhmmer_binary_path: Path to the Jackhmmer binary.
--jackhmmer_n_cpu: Number of CPUs to use for Jackhmmer. Default to min(cpu_count, 8). Going beyond 8 CPUs provides very little additional speedup.
(default: '8')
(an integer)
--jax_compilation_cache_dir: Path to a directory for the JAX compilation cache.
--json_path: Path to the input JSON file.
--mgnify_database_path: Mgnify database path, used for protein MSA search.
(default: '${DB_DIR}/mgy_clusters_2022_05.fa')
--model_dir: Path to the model to use for inference.
(default: '/home/moriwaki/models')
--nhmmer_binary_path: Path to the Nhmmer binary.
--nhmmer_n_cpu: Number of CPUs to use for Nhmmer. Default to min(cpu_count, 8). Going beyond 8 CPUs provides very little additional speedup.
(default: '8')
(an integer)
--ntrna_database_path: NT-RNA database path, used for RNA MSA search.
(default: '${DB_DIR}/nt_rna_2023_02_23_clust_seq_id_90_cov_80_rep_seq.fasta')
--output_dir: Path to a directory where the results will be saved.
--pdb_database_path: PDB database directory with mmCIF files path, used for template search.
(default: '${DB_DIR}/pdb_2022_09_28_mmcif_files.tar')
--rfam_database_path: Rfam database path, used for RNA MSA search.
(default: '${DB_DIR}/rfam_14_9_clust_seq_id_90_cov_80_rep_seq.fasta')
--rna_central_database_path: RNAcentral database path, used for RNA MSA search.
(default: '${DB_DIR}/rnacentral_active_seq_id_90_cov_80_linclust.fasta')
--[no]run_data_pipeline: Whether to run the data pipeline on the fold inputs.
(default: 'true')
--[no]run_inference: Whether to run inference on the fold inputs.
(default: 'true')
--seqres_database_path: PDB sequence database path, used for template search.
(default: '${DB_DIR}/pdb_seqres_2022_09_28.fasta')
--small_bfd_database_path: Small BFD database path, used for protein MSA search.
(default: '${DB_DIR}/bfd-first_non_consensus_sequences.fasta')
--uniprot_cluster_annot_database_path: UniProt database path, used for protein paired MSA search.
(default: '${DB_DIR}/uniprot_all_2021_04.fa')
--uniref90_database_path: UniRef90 database path, used for MSA search. The MSA obtained by searching it is used to construct the profile for template search.
(default: '${DB_DIR}/uniref90_2022_05.fa')
テストラン
AlphaFold3の論文のFig.3eに使われていた AziU3とリガンドの複合体構造(PDB ID: 7WUX)を予測してみます。
Kurosawa, S. et al. Molecular basis for enzymatic aziridine formation via sulfate elimination. J. Am. Chem. Soc. 144, 16164–16170 (2022)
>B4XYC0_AziU2
MTHVAETSAPTRSEPDTRVLTLPGTASAPEFRLIDIDGLLNNRATTDVRDLGSGRLNAWGNSFPAAELPAPGSLITVAGI
PFTWANAHARGDNIRCEGQVVDIPPGQYDWIYLLAASERRSEDTIWAHYDDGHADPLRVGISDFLDGTPAFGELSAFRTS
RMHYPHHVQEGLPTTMWLTRVGMPRHGVARSLRLPRSVAMHVFALTLRTAAAVRLAEGATT
>B4XYC1_AziU3
MTTTAPPVELWTRDLGSCLHGTLATALIRDGHDPVTVLGAPWEFRRRPGAWSSEEYFFFAEPDSLAGRLALYHPFESTWH
RSDGDGVDDLREALAAGVLPIAAVDNFHLPFRPAFHDVHAAHLLVVYRITETEVYVSDAQPPAFQGAIPLADFLASWGSL
NPPDDADVFFSASPSGRRWLRTRMTGPVPEPDRHWVGRVIRENVARYRQEPPADTQTGLPGLRRYLDELCALTPGTNAAS
EALSELYVISWNIQAQSGLHAEFLRAHSVKWRIPELAEAAAGVDAVAHGWTGVRMTGAHSRVWQRHRPAELRGHATALVR
RLEAALDLLELAADAVS
結合するリガンドはccdコードで"6OI"となっています。このccdコードが利用可能かどうかはPDBeChemで検索することができます。ない場合はSMILES表記で入力することになります。
これらを元にインプットとなるjsonファイルを作成します。
{
"name": "AziU2_AziU3_pred",
"modelSeeds": [1],
"sequences": [
{
"protein": {
"id": ["A", "B"],
"sequence": "MTHVAETSAPTRSEPDTRVLTLPGTASAPEFRLIDIDGLLNNRATTDVRDLGSGRLNAWGNSFPAAELPAPGSLITVAGIPFTWANAHARGDNIRCEGQVVDIPPGQYDWIYLLAASERRSEDTIWAHYDDGHADPLRVGISDFLDGTPAFGELSAFRTSRMHYPHHVQEGLPTTMWLTRVGMPRHGVARSLRLPRSVAMHVFALTLRTAAAVRLAEGATT"
}
},
{
"protein": {
"id": ["C", "D"],
"sequence": "MTTTAPPVELWTRDLGSCLHGTLATALIRDGHDPVTVLGAPWEFRRRPGAWSSEEYFFFAEPDSLAGRLALYHPFESTWHRSDGDGVDDLREALAAGVLPIAAVDNFHLPFRPAFHDVHAAHLLVVYRITETEVYVSDAQPPAFQGAIPLADFLASWGSLNPPDDADVFFSASPSGRRWLRTRMTGPVPEPDRHWVGRVIRENVARYRQEPPADTQTGLPGLRRYLDELCALTPGTNAASEALSELYVISWNIQAQSGLHAEFLRAHSVKWRIPELAEAAAGVDAVAHGWTGVRMTGAHSRVWQRHRPAELRGHATALVRRLEAALDLLELAADAVS"
}
},
{
"ligand": {
"id": ["E", "F"],
"ccdCodes": ["6OI"]
}
}
],
"dialect": "alphafold3",
"version": 1
}
以下はこれを実行するためのシェルスクリプトです。
#!/bin/bash
# cudaのロード(環境によっては不要)
module load cuda/12.6
### 各自の環境に合わせて変更してください #
ALPHAFOLD3DIR="/home/moriwaki/apps/alphafold3" # AlphaFold3のコードのディレクトリ
HMMER3_BINDIR="/home/moriwaki/apps/hmmer/3.4/bin" # HMMER3のバイナリディレクトリ
DB_DIR="/mnt/database/public_databases" # 配列・構造データベースのディレクトリ
MODEL_DIR="/home/moriwaki/models" # モデルパラメータのディレクトリ
##########################################
### activate alphafold3's virtual environment
. "${ALPHAFOLD3DIR}/.venv/bin/activate"
python3.11 ${ALPHAFOLD3DIR}/run_alphafold.py \
--jackhmmer_binary_path="${HMMER3_BINDIR}/jackhmmer" \
--nhmmer_binary_path="${HMMER3_BINDIR}/nhmmer" \
--hmmalign_binary_path="${HMMER3_BINDIR}/hmmalign" \
--hmmsearch_binary_path="${HMMER3_BINDIR}/hmmsearch" \
--hmmbuild_binary_path="${HMMER3_BINDIR}/hmmbuild" \
--db_dir="${DB_DIR}" \
--model_dir=${MODEL_DIR} \
--json_path="AziU2_AziU3_6OI.json" \
--output_dir="."
計算時間はMSA作成部分が10~20分、続く構造推論部分RTX3090で9分、RTX4090で5分ほどで終わりました。
予測結果の構造はpymolやUCSF Chimeraなどのソフトウェアで開いてください。
Nature記事の図とほぼ完全に一致。
この図を書き出すときに使ったpymolの設定値
set depth_cue, 0
set ray_trace_mode, 1
set ray_trace_color, 0x404040
set antialias, 4
set ambient, 0.7
set cartoon_oval_width, 0.3
set cartoon_oval_length, 1.3
set ray_trace_gain, 0
set specular, 0
インプットjsonファイルの作成法TIPS
(書きかけ)