Open1
CUDA12.6以外でのAlphaFold3インストール
YomiyamaAlphaFold3
2024年12月2日
日本語の表現に気になる点があったので修正
背景
AlphaFold3のコードが公開され自分たちの手元でAF3のコードを走らせることが出来るようになった。
AF3に要求されているCUDAバージョンが12.6であるが、国立遺伝学研究所のクラスター計算機のnvidia-driverが対応しているバージョンはCUDA12.4であるし、名古屋大学の不老だとCUDA12.2にしか対応していない。
つまりdeepmind公式のパッケージのインストール方法では、これらの計算機においてAF3を走らせることが出来ない。
そこで、今回はnvidia-driverのバージョンがCUDA12.6に対応していない計算機（CUDA12.4がのった計算機上）で走らせることが出来るかどうか検証を行った。
ではどこを変えたのか
まずインストール手順はdeepmindのレポジストリと、東京科学大学難治疾患研究所の森脇先生の
を参考にし、miniforgeの仮想環境を使ってAF3を動かす。
まずgcc 13.2とCUDA 12.4をenviroment moduleで読み込んでおく。
mamba create -n alphafold3 python=3.11 zstd hmmer -c conda-forge -c bioconda -y
mamba activate alphafold3
cd 任意のfolder
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
となっている。
ここで問題になるのが、
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-nvjitlink-cu12==12.6.77
nvidia-cusparse-cu12==12.5.4.2
になることが予想される。
そこで以下のパッケージを対応するCUDAのバージョンになるように
nvidia-cublas-cu12==12.4.2.65
nvidia-cuda-cupti-cu12==12.4.99
nvidia-cuda-nvcc-cu12==12.4.99
nvidia-cuda-runtime-cu12==12.4.99
nvidia-nvjitlink-cu12==12.4.99
#必要に応じて以下も（筆者の環境では書き換えなくても実行できた）
nvidia-cusparse-cu12==12.5.4.2
を書き換えたうえで、alphafold3のrequirements.txtのファイルも書き換える（書き換えないとこの後の手順でエラーが起きる）。
この手順を実行したうえで、
python3.11 -m pip install --no-deps .
${miniforgeのディレクトリ}/envs/alphafold3/bin/build_data
を実行すればインストールに成功する。