【CPU vs GPU】コンフォメーション生成/最適化の速度比較【nvMolKit】

創薬向けのOSSである「nvMolKit」をローカル環境で動かしてみたいと思います。

実行環境
nvMolKit について
環境構築
nvMolKit の操作
最後に

実行環境

OS：WSL2, Ubuntu24.04 (Windows 11 Pro)
CPU: Intel(R) Core(TM) Ultra 9 185H (WSL2プロセッサ数; 22コア)
GPU: NVIDIA GeForce RTX 4070 (8GB)
RAM: 32GB

nvMolKit について

nvMolKitは、2025年9月に公開されたケモ/マテリアルインフォマティクス向けのOSSです（NVIDIA社開発）。

nvMolKitは、一般的なRDKit分子操作を高速化するためのCUDAベースのPythonライブラリです。初期バージョン（0.0.1）では、以下の4つの操作をGPU計算により高速化することができます。

Morgan Fingerprints
Tanimoto and Cosine similarity
ETKDG conformer generation
MMFF optimization

今回はこのnvMolKitについて、実行環境の構築を行います。また、コンフォメーション生成/最適化を題材にして、CPUとGPUでどの程度計算速度に差がでるか確かめたいと思います。

環境構築

Githubの情報を参考にconda環境を構築します。

今回の環境構築では、ローカルマシンのライブラリを使用してコンパイルする必要があります。PC環境によって再現性が得られない可能性があることにご留意ください。

また、Dockerの利用環境が整っている場合は、condaよりもDockerでの環境構築をお勧めします（~~私はWSL2へのDockerのインストールをサボっているためconda環境でやりました~~）。

nvMolKitの環境構築（Dockerを利用しない場合）には、ローカルマシンのCUDA ToolKit (>=12.5)　が必要です。導入していない場合、条件に合うCUDA ToolKitをインストールしましょう。

※2025年9月15日現在、CUDA-12.8+cuDNN9.10.2で正常稼働を確認。pip installで入ってくるtorchのバージョンに合わせる必要あり？

CUDA Toolkitの導入が完了したら以下の手順でインストールを進めます。

sudo apt-get update
sudo apt-get install build-essential libeigen3-dev
sudo apt-get install libstdc++-12-dev libomp-15-dev
conda create --name nvmolkit_dev_py312 python=3.12.1
conda activate nvmolkit_dev_py312
conda install -c conda-forge rdkit=2024.09.3 rdkit-dev=2024.09.3 -y
conda install -c conda-forge libboost libboost-python libboost-devel libboost-headers libboost-python-devel -y
pip install torch torchvision torchaudio
conda install cmake=3.30 -y
export PATH=/usr/local/cuda-12.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:$LD_LIBRARY_PATH
export CC=/usr/bin/gcc
export CXX=/usr/bin/g++
cd (任意のディレクトリ)
git clone https://github.com/NVIDIA-Digital-Bio/nvMolKit.git
cd nvmolkit
CMAKE_BUILD_PARALLEL_LEVEL=16 pip -v install .　# 16は利用するCPUのスレッド数

以上で環境構築は完了です。

nvMolKit の操作

Githubに公開されているnotebookを参考に、CPU・GPUを利用した場合のコンフォメーション生成/最適化の計算速度を比較します。

方法としては、ベンチマークのために用意された50種の化合物について、①ETKDGv3を用いたコンフォメーション生成（1化合物あたり5コンフォメーション）と②MMFFを用いた最適化、この2つの操作に要する計算時間を比べます。

先に結果を記載します。単位は「conformers/second」で「1秒あたりに処理したコンフォメーション数」、つまりは値が大きいほど優れた処理能力を示します。

	CPU	GPU
コンフォメーション生成	2.6	40.3
コンフォメーション最適化	56.0	177.1
計	58.6	217.4

私の環境ではありますが、コンフォメーション生成では約16倍、コンフォメーション最適化では約4倍、nvMolKitでの処理速度が速い結果となりました。

特にコンフォメーション生成では効果が顕著です。3D-QSARなどはコンフォメーション生成が律速となるため、nvMolKitのメリットは非常に大きいと考えられます。

また、以下に利用したコードを記載します。

公式のnotebookにはnvMolKitを利用したGPU計算のコードが記載されているため、自前でRDkitのChem.AllChemモジュールを利用したCPU計算のコードを追加しました。CPU・GPUの違い以外はフェアになるように設定しているつもりです…。

最後に

創薬向けのOSSである「nvMolKit」をローカル環境で動かしてみました。

GPU計算で得られる恩恵はやはり大きいと感じます。現在は利用できる機能が少ないですが、今後のバージョンアップで追加されていくのが楽しみですね。