2024年11月11日の夜(日本時間)にAlphaFold3のコードが公開されました。GitHubリポジトリはこちら。
Dockerがある場合はそちらを使えば簡単に実行できますが、共用計算機環境向けのインストール方法を記述します。
更新情報
- 2024-12-27 スパコン不老での使い方、インプットjsonについてのTIPS、予測結果の解釈を追加。
- 2024-12-23 TSUBAME 4.0で
module load alphafold3_database/202411
の仕様の追加。 - 2024-12-06 TSUBAME 4.0での実行例を追加。
- 2024-11-25 11月19日にデータベースの構築方法が変更されていたので修正
はじめに
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でのやり方を真似ています。
※ 2024年11月25日: 公式のやり方が変更されていたので修正
# ディレクトリの指定
APPDIR="/home/moriwaki/apps"
# データベースディレクトリの指定
DB_DIR="/mnt/database/public_databases"
##### 以下コピペ #####
mkdir -p $APPDIR
cd $APPDIR
git clone https://github.com/google-deepmind/alphafold3.git
ALPHAFOLD3DIR="$APPDIR/alphafold3"
# クローンしてきたalphafold3ディレクトリに移動
cd ${ALPHAFOLD3DIR}
# 11月19日以降はpython3を使ったデータベース構築を行わなくなりました。
# ./fetch_databases.shを使ってください。
# python3 fetch_databases.py --download_destination="$DB_DIR/public_databases"
./fetch_databases.sh $DB_DIR
public_databases
のサイズは展開後で641,428MB (640GB)でした。これくらいの容量を確保しておいてください。
mmcif_files/ # この直下に200,000以上のmmcifファイルが存在する
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
配列・構造データベースディレクトリの場所
public_databases
ディレクトリはどこに配置しても構いませんが、必ずHDDやネットワークファイルドライブ上ではなくSSDの中に置いてください。これはMSA作成時に大量のファイルアクセスが発生するため、HDDやネットワークファイルドライブだと計算時間が大幅に遅くなるからです。models
ディレクトリはそこまでファイルサイズが大きくないので、HDDやネットワークファイルドライブ上でも問題ありません。
また、Dockerを使う場合、public_databases
ディレクトリはalphafold3リポジトリのディレクトリの外に配置してください。そうでないと、Dockerが大きなイメージを作成しようとすることになり、ビルドがとても遅くなります。
モデルパラメータのダウンロードと配置
メールで送られてきた学習済みパラメータはたった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 python3.11
# ${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.
ユーザー側の実行スクリプトの準備
環境変数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}/mmcif_files')
--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に使われていた AziU2, 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
TSUBAME 4.0でのAlphaFold3の利用方法
2024-12-23: TSUBAME 4.0内にAlphaFold3の動作に必要なDBが配列・構造データベース(約650GB)がTSUBAME 4.0の全ユーザーがアクセス可能な領域に内蔵されました。module load alphafold3_database/202411
を実行することで環境変数ALPHAFOLD_DATA_DIR
にデータベースへのパスがロードされます。
ビルド作業は多くのコアを利用するため、TSUBAME4.0の利用の手引きの5.2.2. インタラクティブジョブ(または5.3. インタラクティブジョブ専用キュー)のやり方にしたがって、ログインノードではない計算ノードを使ってビルドしてください。
AlphaFold 3.0.0から現在のGit最新コミットの間に様々な仕様変更が起きているため、TSUBAME 4.0では、3.0.1以降の安定バージョンがリリースされるまで、管理者側でのビルド済みAlphaFold3の提供を待っている状態です。このため、早めにAlphaFold3を利用したい方は、以下の手順に従って自身のディレクトリ内にインストールする必要があります。
/path/to/your/apps_directory
は各自のディレクトリに置き換えてください。
APPS_DIRECTORY="/path/to/your/apps_directory"
# インストールディレクトリの作成
mkdir -p ${APPS_DIRECTORY}
cd ${APPS_DIRECTORY}
# Python 3.11をMiniforgeでインストール
wget -q -P . https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh
bash ./Miniforge3-Linux-x86_64.sh -b -p "${APPS_DIRECTORY}/miniforge"
rm Miniforge3-Linux-x86_64.sh
CONDAROOT="${APPS_DIRECTORY}/miniforge"
source "${CONDAROOT}/etc/profile.d/mamba.sh"
source "${CONDAROOT}/etc/profile.d/conda.sh"
mamba create --name alphafold3 python=3.11 -y
mamba activate alphafold3
# Zlibが必要
mamba install zlib -n alphafold3 -y
# CUDA 12.6をnvidiaからインストール
mamba install cuda -c nvidia -n alphafold3 -y
# HMMER 3.4をbiocondaからインストール
mamba install bioconda::hmmer -n alphafold3 -y
# hmmalignが使えることを確認
${APPS_DIRECTORY}/miniforge/envs/alphafold3/bin/hmmalign -h
# AlphaFold3
cd ${APPS_DIRECTORY}
git clone https://github.com/google-deepmind/alphafold3.git
cd alphafold3
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 .
# build_dataは途中でプロセスがkillされることがあるのでそのときはもう1回実行
${APPS_DIRECTORY}/miniforge/envs/alphafold3/bin/build_data
TSUBAME 4.0のログインノードは1ユーザーあたりの同時起動プロセス数(?)に制限があるので、インタラクティブキューを使ってビルドしてください。
実行スクリプト例はこちら。モデルパラメータは各自でアップロードして配置してください。
#!/bin/sh
#$ -cwd
#$ -l gpu_1=1
#$ -l h_rt=0:40:00
echo `hostname`
### TSUBAME 4.0に内蔵されているAlphaFold3用データベースを呼び出す
### 環境変数ALPHAFOLD_DATA_DIRに/gs/bs/datasets/alphafold/3/$ver/dbがロードされる
module load alphafold3_database/202411
# qsub時に-v INPUT=fooとすることでインプットのJSON"foo.json"について予測できる
if test "$INPUT" = ""; then
echo "INPUT is not set" 1>&2 ; exit
fi
### activate alphafold3's virtual environment
APPS_DIRECTORY="/path/to/your/apps_directory"
ALPHAFOLD3DIR="${APPS_DIRECTORY}/alphafold3"
CONDAROOT="${APPS_DIRECTORY}/miniforge"
source "${CONDAROOT}/etc/profile.d/conda.sh"
source "${CONDAROOT}/etc/profile.d/mamba.sh"
mamba activate alphafold3
HMMER3_BINDIR="${APPS_DIRECTORY}/miniforge/envs/alphafold3/bin"
### 重みパラメータファイルaf3.zstは自身でGoogle DeepMindに申し込んでダウンロードしたものを
### MODEL_DIRに置く必要がある
MODEL_DIR="/path/to/models"
export XLA_PYTHON_CLIENT_PREALLOCATE=true
export XLA_CLIENT_MEM_FRACTION=0.95
python3.11 ${ALPHAFOLD3DIR}/run_alphafold.py \
--run_data_pipeline=True \
--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="${ALPHAFOLD_DATA_DIR}" \
--json_path=${INPUT}.json \
--model_dir="${MODEL_DIR}" \
--output_dir="."
--run_data_pipeline
はMSAサーチとテンプレートサーチを行うかどうかを規定します。デフォルトではTrue
ですが、すでにインプットのjsonファイルにMSAやテンプレートが含まれている場合、False
に指定することで、その検索を省略できます。
TSUBAME 4.0のGPUのNVIDIA H100 SXM5(VRAM 94GB)では4000残基程度まで予測可能らしいので、1GPUで十分なことが多く、# -l gpu_1=1
を指定しています。計算時間も、短くて20分程度、長くて1〜2時間前後で終わることがほとんどですので、# -l h_rt=0:40:00
で40分制限にしています。
また、TSUBAME 4.0ではqsub
時に-v VARIABLE=foo
を追加することで環境変数の値を指定させることができます。例えばfoo.json
, bar.json
について構造予測したいとき、このrun_alphafold3.sh
の中身を変更せずに、qsub -v INPUT=foo run_alphafold3.sh
, qsub -v INPUT=bar run_alphafold3.sh
とすることでそれぞれの構造予測を実行する事が可能です。
名古屋大学・不老でのAlphaFold3の利用方法
こちらのType-IIのCXノードにGPUが搭載されています。GPUはVolta世代のV100 32GBです。このためか、AlphaFold3が本来想定している動作環境であるAmpere世代以降のGPUでないため、bf16が使えないこともあり、依存パッケージを変更する必要があります。またV100 32GBでは1520トークン(=残基)くらいまでしか予測できず、それ以上ではRESOURCE_EXHAUSTED: Out of memory
となり、予測できません。なお、RTX4090 24GB(bf16対応)では2000残基程度まで予測可能です。
APPS_DIRECTORY="/path/to/your/apps_directory"
# インストールディレクトリの作成
mkdir -p ${APPS_DIRECTORY}
cd ${APPS_DIRECTORY}
# Python 3.11をMiniforgeでインストール
wget -q -P . https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh
bash ./Miniforge3-Linux-x86_64.sh -b -p "${APPS_DIRECTORY}/miniforge"
rm Miniforge3-Linux-x86_64.sh
CONDAROOT="${APPS_DIRECTORY}/miniforge"
source "${CONDAROOT}/etc/profile.d/mamba.sh"
source "${CONDAROOT}/etc/profile.d/conda.sh"
mamba create --name alphafold3 python=3.11 -y
mamba activate alphafold3
# Zlibが必要
mamba install zlib -n alphafold3 -y
# CUDA 12.6をnvidiaからインストール
mamba install cuda -c nvidia -n alphafold3 -y
# HMMER 3.4をbiocondaからインストール
mamba install bioconda::hmmer -n alphafold3 -y
# hmmalignが使えることを確認
${APPS_DIRECTORY}/miniforge/envs/alphafold3/bin/hmmalign -h
# AlphaFold3
cd ${APPS_DIRECTORY}
git clone https://github.com/google-deepmind/alphafold3.git
cd alphafold3
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.0.50 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.2 \
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 --no-cache-dir
# AlphaFold3のビルド&インストール
python3.11 -m pip install --no-deps .
# build_dataは途中でプロセスがkillされることがあるのでそのときはもう1回実行
${APPS_DIRECTORY}/miniforge/envs/alphafold3/bin/build_data
実行シェルスクリプト例はこちら
#!/bin/bash
#PJM -L rscgrp=cxgfs-share
#PJM -L elapse=02:00:00
#PJM -j
#PJM -S
## SET YOUR DIRECTORY PATH ##
APPS_DIRECTORY="/path/to/your/apps_directory"
MODEL_DIR="/path/to/your/models"
#############################
test $PJM_O_WORKDIR && cd $PJM_O_WORKDIR
CONDAROOT="${APPS_DIRECTORY}/miniforge"
source "${CONDAROOT}/etc/profile.d/mamba.sh"
source "${CONDAROOT}/etc/profile.d/conda.sh"
mamba activate alphafold3
ALPHAFOLD3DIR="${APPS_DIRECTORY}/alphafold3"
HMMER3_BINDIR="${CONDAROOT}/envs/alphafold3/bin"
DB_DIR="/beegfs/share/alphafold/af3"
export XLA_PYTHON_CLIENT_PREALLOCATE=true
export XLA_CLIENT_MEM_FRACTION=0.95
# 名古屋大学Flow-cxのGPUにあるV100などのCUDA Capability 7.x GPUsに対しては以下の設定が必要不可欠。
export XLA_FLAGS="--xla_disable_hlo_passes=custom-kernel-fusion-rewriter"
if test "$INPUT" = ""; then
echo "INPUT is not set" 1>&2 ; exit
fi
python3.11 ${ALPHAFOLD3DIR}/run_alphafold_mod.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}" \
--flash_attention_implementation="xla" \
--json_path=${INPUT}.json \
--model_dir=${MODEL_DIR} \
--output_dir="."
このrun_alphafold3.sh
を用いて、foo.json
に対して構造予測を行いたいときは
INPUT=foo ; pjsub -x INPUT=${INPUT} -N ${INPUT} run_alphafold3.sh
で予測が実行できます。
- GPU V100の場合、現状のAlphaFold3では
export XLA_FLAGS="--xla_disable_hlo_passes=custom-kernel-fusion-rewriter"
と--flash_attention_implementation="xla"
の設定追記が必須となっています。
予測可能限界
GPUの予測可能限界を超えて予測させようとすると以下のようなRESOURCE_EXHAUSTED: Out of memory
エラーが発生します。残基数を減らして予測してください。
Traceback (most recent call last):
File "/data/group1/z00000a/apps/alphafold3/run_alphafold.py", line 759, in <module>
app.run(main)
File "/data/group1/z00000a/apps/miniforge/envs/alphafold3/lib/python3.11/site-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/data/group1/z00000a/apps/miniforge/envs/alphafold3/lib/python3.11/site-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
^^^^^^^^^^
File "/data/group1/z00000a/apps/alphafold3/run_alphafold.py", line 742, in main
process_fold_input(
File "/data/group1/z00000a/apps/alphafold3/run_alphafold.py", line 599, in process_fold_input
all_inference_results = predict_structure(
^^^^^^^^^^^^^^^^^^
File "/data/group1/z00000a/apps/alphafold3/run_alphafold.py", line 392, in predict_structure
result = model_runner.run_inference(example, rng_key)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/group1/z00000a/apps/alphafold3/run_alphafold.py", line 309, in run_inference
result = self._model(rng_key, featurised_example)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17196646400 bytes.
インプットjsonファイルの作成法TIPS
ディレクトリ内にあるすべてのjsonファイルをまとめて予測させる
run_alphafold.py
に与える引数で、入力のインプットJSONファイルを指定する--json_path
というものがありますが、これは--input_dir=<ディレクトリ名>
という引数に置き換えることが可能です。これを使うと、指定したディレクトリ内に存在するすべてのjsonファイルをまとめて予測させることができます。
化合物を変えた予測を高速に実行させる
AlphaFold3の計算の大部分は、入力のタンパク質(またはRNA)配列についてのMultiple Sequence Alignment (MSA)の取得部分に時間を費やしています。こちらがほぼ9割で、構造を予測する部分は残りの1割程度です。そのため、予測したいタンパク質の配列が変わらない場合は、MSAを再取得する必要が本来なく、計算時間を大幅に短縮させることができます。
具体的には、以下のような手順で実行することが可能です。
まず、予測したいタンパク質配列が入ったjsonファイル(例:foo.json
)を作成します。そしてMSAを取得するために、一度その配列予測を実行します。このとき、--oupput_dir
で指定していた出力ディレクトリの中にはfoo_data.json
というファイルが生成されています。これは、foo.json
の上から、AlphaFold3が取得してきたMSAやテンプレート情報が追記されているファイルとなっています。つまり、このファイルを再利用すれば、タンパク質についてのMSA, テンプレート情報を再取得する必要がなく、構造予測の部分だけを再実行することが可能です。
以下はfoo_data.json
の例です。この例はタンパク質ホモ六量体の予測なので、"protein"のidが6つ、"ligand"のidも6つとなっていますが、ここは予測対象に応じて適切に変えてください。特に、"unpairedMsa"
, "pairedMsa"
, "templates"
の部分がMSA+テンプレート取得の結果が埋め込まれた部分であり、これらの情報を再利用することで次回以降の計算をスキップすることができます。そして、目的である"ligand"
の情報を変更することで、新たな化合物についての予測を行うことが可能です。
{
"dialect": "alphafold3",
"version": 1,
"name": "1bjp_6_pred",
"sequences": [
{
"protein": {
"id": [
"A",
"B",
"C",
"D",
"E",
"F"
],
"sequence": "PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASKVRR",
"modifications": [],
"unpairedMsa": ">101\nPIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASKVRR\n>UniRef100_A0A7J3FD47\t90\t0.500\t8.638E-19\t0\t59\t62\t1\t60\t63\nPVVTIELWEGRTPEQKRELVRAVSSAISRVLGCPEEAVHVILHEVPKANWGIGGRLASEL--\n>UniRef100_A0A7C4BER2\t87\t0.406\t1.096E-17\t0\t58\t62\t1\t59\t65\nPVVTIEMWEGRTPEQKKALVEAVTSAVAGAIGCPPEAVEVIIHEVPKVNWGIGGQIASE---\n>UniRef100_A0A5B8RYW4\t86\t0.451\t1.506E-17\t0\t61\t62\t1\t62\t63\nPFVQIHMLEGRTPEQKKAVIEKVTQALVQAVGVPASAVRVLIQEVPKEHWGIGGVSARELGR\n>UniRef100_UPI00140DD609\t85\t0.440\t2.842E-17\t0\t58\t62\t1\t59\t62\nPIIQVQMLKGRSPELKKQLISEITDTISRTLGSPPEAVRVILTEVPEENWGVGGVPINE---\n>UniRef100_A0A1B7KRL2\t85\t0.467\t3.905E-17\t0\t61\t62\t1\t62\t67\nPFIQVYLLEGRSPEKKEALIYELTETVRKVLGAPRENVRVMIQEIPLEHWGIAGESVKKRRK\n",
"pairedMsa": "",
"templates": []
}
},
{
"ligand": {
"id": ["G","H","I","J","K","L"],
"ccdCodes": ["OXP"]
}
}
],
"modelSeeds": [
1
],
"bondedAtomPairs": null,
"userCCD": null
}
このファイルをコピーして、新たなjsonファイル(例:bar.json
)を作成します。そして、化合物("ligand"
)情報を入力する"ccdCodes":
または"smiles":
の値を変更します。
{
"ligand": {
"id": ["G","H","I","J","K","L"],
"smiles": "CC(=O)OC1C[NH+]2CCC1CC2"
}
}
最後に、run_alphafold.py
を実行する際に、--json_path
でbar.json
を指定し、--run_data_pipeline
をTrue
からFalse
に変更して実行します。これにより、MSAの再取得を行わずに、構造予測の部分だけを再実行することが可能です。
予測結果の解釈
予測構造の解釈については、羊土社から出ているAlphaFold時代の構造バイオインフォマティクス実践ガイドを読んでください……。ってこれだとただのダイレクトマーケティングなので、以下にいくつか無料で読める資料を紹介します。
得られた予測構造はアウトプットディレクトリ内に**_model.cif
という拡張子.cif
のファイル形式で出力されます。これをPyMOLやUCSF Chimera Xを用いて映し出してください。
まずは、予測構造の局所的な信頼性を表すpLDDT値を可視化する必要があります。これは以前私が書いた記事AlphaFold予測構造についてのpLDDTカラーリングコマンドを使うと便利です。基本的には青いほど信頼性が高く、赤いほど信頼性が低いです。赤いところは信用しないでOKです。
color 0x0053D6, b < 100
color 0x65CBF3, b < 90
color 0xFFDB13, b < 70
color 0xFF7D45, b < 50
# 必要であれば原子ごとの色分け(#を外して使う)
# util.cnc
次に、全体構造の相対的な位置関係の信頼性を表すPredicted Aligned Error (PAE)について表示したいところですが、こちらはAlphaFold3のアウトプット情報ではそのまま画像として出てきません。これを表示させるためには私が作成したalphafold3_toolsに含まれるpaeplot
コマンドを使うと便利です。使い方はリンクを読んでください。PAEの値が0に近いところは信頼性が高く、32に近いところは信頼性が低いです。
上海老師氏がこれらの解釈と無料まとめ記事を書いてくれています。
- タンパク質立体構造予測プログラムAlphaFold2の非専門家向け活用法 第1回「予測の良し悪しの判断」
- AlphaFold2の非専門家向け活用法 第1.5回
- AlphaFold2の非専門家向け活用法 第2回「予測構造を元にいろいろ考察してみる(実践編)」
- AlphaFold2の非専門家向け活用法 第3回「予測構造を活用したタンパク質発現や構造解析」
- AlphaFold2の非専門家向け活用法 第4回前半「AlphaFoldが現状カバーできていない構造情報」
- AlphaFold2の非専門家向け活用法 第4回後半「AlphaFoldが現状カバーできていない構造情報」
- AlphaFold2の非専門家向け活用法 第5回(最終回)「最近の進展と今後の展望」