Ascend-SACT/Boltz2
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

一键部署脚本:Boltz2 在昇腾 NPU 上的完整部署流程

关键词:【模型-Boltz-2】【场景-推理】【问题-模型部署适配】 适配环境:昇腾 NPU(Atlas 800T A3) 代码路径:/home/tz/Boltz2 (模型代码下载到的路径,实际配置需要切换到用户自己路径) 权重路径:/weights/Boltz2(模型权重下载到的路径,实际配置需要切换到用户自己路径)

Boltz2 部署操作手册(昇腾 NPU)

本文档用于在昇腾 NPU 环境下完成 Boltz2 模型的容器化部署、依赖安装与推理验证。


一、前置条件

项目要求
NPUAscend Atlas 800T A3(至少 1 卡可用)
CANN8.5.0
Python3.11
Docker20.10+
依赖文件Boltz2.zip、requirements_pypi.txt
模型权重已放置 /weights/Boltz2

二、模型介绍

Boltz-2 是 MIT 与 Recursion 联合开发的生物分子基础模型,支持蛋白质/复合物 3D 结构预测与蛋白-配体结合亲和力预测,在虚拟筛选、靶点验证、分子设计场景广泛使用。


三、宿主机操作步骤

1. 克隆目录

# 进入目标工作目录
cd /home/tz
# 克隆官方仓库代码
git clone https://github.com/jwohlwend/boltz
# 进入克隆后的项目目录
cd boltz
# 下载补丁包到当前目录
wget https://ai.gitcode.com/Ascend-SACT/Boltz2/blob/main/boltz_custom.patch
# 应用补丁
git apply boltz_custom.patch

2. 一键创建并启动容器

docker rm -f Boltz2 && docker run -itd --net=host \
    --name Boltz2 \
    --shm-size=1g \
    --privileged \
    --device /dev/davinci0 \
    --device /dev/davinci1 \
    --device /dev/davinci2 \
    --device /dev/davinci3 \
    --device /dev/davinci4 \
    --device /dev/davinci5 \
    --device /dev/davinci6 \
    --device /dev/davinci7 \
    --device /dev/davinci8 \
    --device /dev/davinci9 \
    --device /dev/davinci10 \
    --device /dev/davinci11 \
    --device /dev/davinci12 \
    --device /dev/davinci13 \
    --device /dev/davinci14 \
    --device /dev/davinci15 \
    --device /dev/davinci_manager \
    --device /dev/devmm_svm \
    --device /dev/hisi_hdc \
    -v /usr/local/dcmi:/usr/local/dcmi \
    -v /usr/local/sbin/npu-smi:/usr/local/bin/npu-smi \
    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
    -v /etc/ascend_install.info:/etc/ascend_install.info \
    -v /weights/Boltz2:/weights/Boltz2 \
    -v /home/tz/Boltz2:/home/tz/Boltz2 \
harbor.stonewise.cn/kubeflow/ascend-cann-mol:8.5.0-a3-ubuntu22.04-py3.11

2. 一键创建并启动容器

docker exec -it Boltz2 bash

四、容器内环境配置

1. 加载 CANN 环境

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0

2. 安装 Python 依赖

cd /home/tz/Boltz2
pip install --upgrade pip
pip install torch==2.5.1 torch_npu==2.5.1
pip install -r requirements_pypi.txt
pip install "click>=8.2.1"

3. 安装 Boltz2 包

cd /home/tz/Boltz2/boltz
pip install --no-build-isolation -e .

4. 验证安装

boltz --help

五、SSL 证书问题修复

cd /home/tz/Boltz2/boltz/src/boltz/data/msa

sed -i 's/    auth=auth,$/    auth=auth,\n                    verify=False,/' mmseqs2.py
sed -i 's/timeout=6.02, headers=headers, auth=auth$/timeout=6.02, headers=headers, auth=auth, verify=False/' mmseqs2.py
sed -i 's/timeout=6.02, headers=headers, auth=auth$/timeout=6.02, headers=headers, auth=auth, verify=False/' mmseqs2.py

六、模型权重下载(容器内执行)

权重已统一存放至 /weights/Boltz2,若目录为空可执行以下命令下载:

mkdir -p /weights/Boltz2
cd /weights/Boltz2

wget https://huggingface.co/boltz-community/boltz-2/resolve/main/boltz2_conf.ckpt
wget https://huggingface.co/boltz-community/boltz-2/resolve/main/boltz2_aff.ckpt
wget https://huggingface.co/boltz-community/boltz-1/resolve/main/ccd.pkl
wget https://huggingface.co/boltz-community/boltz-2/resolve/main/mols.tar

tar -xf mols.tar
rm -f mols.tar

ls -lh /weights/Boltz2/

七、单配体推理测试(核心验证)

cd /home/tz/Boltz2/boltz

boltz predict examples/MCL1_single/lig_26.yaml \
    --use_msa_server \
    --out_dir /home/tz/Boltz2/npu_test \
    --sampling_steps 50 \
    --recycling_steps 3 \
    --diffusion_samples 1 \
    --seed 42 \
    --num_workers 0 \
    --accelerator npu \
    --override \
    --cache /weights/Boltz2

查看结果:

ls /home/tz/Boltz2/npu_test/boltz_results_lig_26/predictions/lig_26/

八、批量推理测试

boltz predict examples/MCL1_single \
    --use_msa_server \
    --out_dir /home/tz/Boltz2/npu_MCL1_test \
    --sampling_steps 50 \
    --recycling_steps 3 \
    --diffusion_samples 1 \
    --seed 42 \
    --num_workers 0 \
    --accelerator npu \
    --override \
    --cache /weights/Boltz2

九、输出结果说明

输出目录:

  • 单例:/home/tz/Boltz2/npu_test

  • 批量:/home/tz/Boltz2/npu_MCL1_test

包含文件:

  • *_model_0.cif:结构文件

  • plddt_*.npz:置信度

  • pae_*.npz、pde_*.npz:结构评价

  • affinity_*.json:亲和力预测结果


十、常见问题

1. NPU 不可用

运行以下命令排查:

npu-smi info
python -c "import torch; print(torch.npu.is_available())"

2. 切换 NPU 卡号

运行以下命令切换:

export ASCEND_RT_VISIBLE_DEVICES=1

3. 依赖冲突

忽略 warning,不影响推理运行。