ai-for-science-diffsbdd

DiffSBDD Ascend NPU Migration Skill

Prerequisites

Item	Requirement
Hardware	Ascend910 series (at least 1 card)
OS	openEuler / Ubuntu / KylinOS (aarch64 or x86_64)
CANN	≥ 8.2 (8.2.RC1 or 8.3.RC1 recommended)
Python	3.10
PyTorch	2.5.1
torch_npu	2.5.1

Migration Process

1. Clone the repository

git clone https://github.com/arneschneuing/DiffSBDD.git
cd DiffSBDD

2. Create a Conda environment

conda create -n diffSBDD-kernerlcat python=3.10 -y
conda activate diffSBDD-kernerlcat

3. Configure the CANN environment and Huawei mirror sources

source /home/Ascend/ascend-toolkit/set_env.sh
export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/

4. Install PyTorch + torch_npu

pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1
pip install torch_npu==2.5.1
pip install numpy==1.26.4
pip install pyyaml decorator attrs psutil absl-py cloudpickle ml-dtypes scipy tornado

Key: numpy must be 1.26.4; numpy 2.x is incompatible with the CANN TBE compiler.

5. Install project dependencies

pip install pytorch-lightning==1.8.4 wandb biopython==1.79 imageio==2.21.2 \
    pandas==2.2.2 seaborn==0.13.2 torchmetrics==1.4.2 tqdm==4.66.5
conda install -c conda-forge rdkit openbabel -y
# conda 安装后可能覆盖 numpy，需重新固定
pip install numpy==1.26.4 --force-reinstall
# 检查并移除 conda 残留的 numpy-base 2.x
conda remove numpy-base --force -y 2>/dev/null
pip install numpy==1.26.4 --force-reinstall

6. Compile torch_scatter from source

Mandatory requirement: must be compiled from source; using PyG precompiled wheels is prohibited.

source /home/Ascend/ascend-toolkit/set_env.sh
export FORCE_CUDA=0
export TORCH_CUDA_ARCH_LIST=''
git clone https://github.com/rusty1s/pytorch_scatter.git
cd pytorch_scatter
python setup.py bdist_wheel
pip install dist/*.whl
cd ..

7. Code adaptation

7.1 Inject transfer_to_npu (automatic migration)

Add at the top of the following 6 entry files:

import torch_npu
from torch_npu.contrib import transfer_to_npu

Files that need to be modified:

generate_ligands.py
test.py
inpaint.py
optimize.py
train.py
lightning_modules.py

transfer_to_npu will automatically redirect all torch.cuda.* calls to the NPU, including patterns such as device = 'cuda' if torch.cuda.is_available() else 'cpu' .

7.2 Add the missing `init.py`

touch analysis/__init__.py
touch analysis/SA_Score/__init__.py

8. Verification

8.1 Download the pretrained model

mkdir -p checkpoints
wget -P checkpoints/ https://zenodo.org/record/8183747/files/crossdocked_fullatom_cond.ckpt

8.2 De novo inference verification

source /home/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0
python generate_ligands.py checkpoints/crossdocked_fullatom_cond.ckpt \
    --pdbfile example/3rfm.pdb \
    --ref_ligand example/3rfm_B_CFF.sdf \
    --outfile example/3rfm_generated.sdf \
    --n_samples 2 --num_nodes_lig 44

Verification pass criteria:

The program exits normally (exit code 0)
The output SDF file is non-empty and correctly formatted

Special adaptation points

Adaptation item	Description
numpy version	Must be 1.26.4; after installing rdkit via conda it will be overwritten to 2.x and must be manually restored
numpy-base	conda may leave behind numpy-base 2.x, requiring `conda remove numpy-base --force`
torch_scatter	Must be compiled from source, set `FORCE_CUDA=0`
wandb	The original 0.13.1 is incompatible with numpy 1.26.4 and needs to be upgraded to the latest version
`__init__.py`	`analysis/` and `analysis/SA_Score/` missing `__init__.py`
fp64 precision	Ascend910 does not support double and automatically degrades to fp32; for diffusion models it is sufficient to compare distribution metrics
transfer_to_npu	Automatically handles all CUDA→NPU device mapping, with no need to manually change device logic

Companion scripts

Dependency and NPU runtime precheck:python scripts/validate_diffsbdd_env.py

References

DiffSBDD dependency and verification notes:references/dependency-notes.md