AI4S/Boltz-1
模型介绍文件和版本Pull Requests讨论分析

组件版本

hdk: 25.5.0
cann:8.3.RC1
python:3.11
torch：2.5.1
torch_npu：2.5.1

下载源码

git clone https://atomgit.com/AI4Science/Boltz-1.git
cd Boltz-1

环境准备

  • 1.创建虚拟环境
conda create -n Boltz-1 python=3.11 -y
conda activate Boltz-1
  • 2.安装依赖
# 临时使用华为镜像源安装 PyPI 包
export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple
export PIP_TRUSTED_HOST=repo.huaweicloud.com

# 安装项目基础依赖包 
pip install torch==2.5.1 torch_npu==2.5.1 pyyaml numpy decorator attrs psutil scipy
    1. 验证 PyTorch 与 torch_npu 安装 执行以下命令检查安装是否成功：
python3 -c "import torch;import torch_npu; a = torch.randn(3, 4).npu(); print(a + a);"

若输出类似以下信息，说明安装成功：

tensor([[-0.6066,  6.3385,  0.0379,  3.3356],
        [ 2.9243,  3.3134, -1.5465,  0.1916],
        [-2.1807,  0.2008, -1.1431,  2.1523]], device='npu:0')

若报错，排查顺序：

  • set_env.sh 是否已 source
  • decorator 等运行时依赖是否安装
  • CANN 与 torch_npu 版本是否匹配

安装boltz

pip install boltz

预测前准备

  • 导入torch_npu
vi /root/miniconda3/envs/Boltz-1/lib/python3.11/site-packages/boltz/main.py
import torch_npu
from torch_npu.contrib import transfer_to_npu
  • 修改lightning_fabric
vi /root/miniconda3/envs/Boltz-1/lib/python3.11/site-packages/lightning_fabric/accelerators/cuda.py
找到_is_ampere_or_later函数，注销原始代码，直接返回True，保存退出。
  • 修改boltz1.py
vi /root/miniconda3/envs/Boltz-1/lib/python3.11/site-packages/boltz/model/models/boltz1.py
# 找到268行附近，修改setup函数如下：
def setup(self, stage: str) -> None:
    """Set the model for training, validation and inference."""
    if stage == "predict" and not (
        torch.cuda.is_available()
        and False  # 直接让条件不成立，适配NPU环境
    ):
        self.use_kernels = False
  • 下载权重
mkdir ~/.boltz

# 下载huggingface上的所有文件，放到~/.boltz/下
https://huggingface.co/boltz-community/boltz-1/tree/main

预测

  • NPU 运行环境配置
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 指定可见 NPU 卡（0,1,2,3)
export ASCEND_RT_VISIBLE_DEVICES=0

# NPU 性能优化
export CPU_AFFINITY_CONF=1
export TASK_QUEUE_ENABLE=2
  • 执行
boltz predict ./examples/affinity.yaml --model boltz1 --devices 1 --out_dir ./boltz_output --num_workers 0

Docker镜像

已预置 Conda 运行环境、模型权重、Boltz-1 源码，开箱即用，可快速部署运行。

# 拉取镜像
sudo docker pull swr.cn-north-4.myhuaweicloud.com/ascend_ai4s/boltz1:v2
  • 项目源码路径：/root/Boltz-1/
  • 已配置好 Conda 虚拟环境
  • 已内置模型权重文件，无需额外下载