gpu4pyscf を RTX 5000 シリーズで動かす

概要

gpu4pyscf は、PySCF を GPU で高速化するためのプラグインである。
特に、4中心2電子積分や density fitting における 3中心積分を GPU で高速化しており、SCF や DFT 計算を大幅に高速化できる。行列演算やテンソル演算には CUDA ベースの CuPy が用いられている。

ベンチマークでは、32 コア CPU を用いた Q-Chem と比較して、条件によっては数十倍程度の高速化が報告されている。gaussian と比較すると 30 倍くらい早い。

また、gpu4pyscf を Blackwell 世代（RTX 5000 シリーズ）で利用する場合は、ソースからのビルドが必要になる。この記事は、ソースビルドの作業記録である。

参考文献：https://wires.onlinelibrary.wiley.com/doi/full/10.1002/wcms.70008

検証環境

AlmaLinux 8.10
cuda 12.8
python 3.11
nvidia driver 590.48.01
RTX 5080
Ryzen 9 9950x

インストール方法

まずは、gpu4pyscf をインストールする。その後 pyscf をインストールする。
ここでのポイントは、-DCUDA_ARCHITECTURES="-DCUDA_ARCHITECTURES="86-real;89-real;120-real"というオプションをつけることである。120-real というオプションだけだと Blackwell 世代しか動かなくなってしまうので、自分の使用する計算機クラスターで使用されている GPU 全ての世代を指定する必要がある。
ビルドには、結構時間がかかる。

git clone https://github.com/pyscf/gpu4pyscf.git
cd gpu4pyscf
cmake -S gpu4pyscf/lib -B build/temp.gpu4pyscf -DCUDA_ARCHITECTURES="86-real;89-real;120-real"
cmake --build build/temp.gpu4pyscf -j8
CURRENT_PATH=`pwd`
export PYTHONPATH="${PYTHONPATH}:${CURRENT_PATH}"

ここで、仮想環境の設定を行う。
個人で使用する場合は conda、クラスター計算機で module システムを利用している際は venv を使うと良い。

仮想環境設定後に、cupy をインストールする。cuda のバージョンに敏感のため、cuda 13 系は使用しない方が良い。

venvでの仮想環境の構築

python3 -m venv /home/apps/gpu4pyscf/all/venv
source /home/apps/gpu4pyscf/all/venv/bin/activate
python -m pip install -U pip setuptools wheel
python -m pip install "pyscf[dispersion]" cupy-cuda12x geometric

module ファイルは、以下のように設定した。

conflict gpu4pyscf
prereq cuda/12.8
prereq python3.11

set base        /home/apps/gpu4pyscf/all
set CUDA_HOME   /usr/local/cuda
setenv GPU4PYSCF_HOME $base

# 仮想環境を venv で設定
set root $base/venv
setenv VIRTUAL_ENV $root
setenv PYTHONNOUSERSITE 1
prepend-path PATH $root/bin

prepend-path PATH $CUDA_HOME/bin
prepend-path LD_LIBRARY_PATH $CUDA_HOME/lib64
prepend-path LD_LIBRARY_PATH $CUDA_HOME/targets/x86_64-linux/lib
# PYTHONPATH という環境変数は、gpu4pyscfで使用するので必要
append-path PYTHONPATH $base