用途:把 BoltzGen 从 CUDA 迁移到昇腾 NPU,完成生成式蛋白设计与逆折叠流程的部署与复现。
主要能力:给出前置条件(Ascend 910、NPU 驱动、CANN Toolkit、Linux aarch64、conda、Python 3.11、PyTorch 2.1.0 与 torch_npu 2.1.0.post17 严格配套);列出需预先放到 ~/.cache/ 的权重文件(boltzgen1_diverse.ckpt、boltzgen1_adherence.ckpt、boltzgen1_ifold.ckpt、boltz2_conf_final.ckpt、mols.zip)及打包方法;提供 conda 环境、依赖安装、CANN 环境变量设置和 NPU 验证命令;通过 transfer_to_npu 劫持 CUDA API,新增 npu_kernels.py 以标准 PyTorch 算子实现 cuequivariance_torch 的 triangle attention 与 triangle multiplicative update,并在 triangular.py、primitives.py 中按设备类型分发;修补 get_device_capability、numpy/torch 类型比较,注释 CUDA 专属依赖。
使用方式:激活环境并设置环境变量后,克隆仓库、安装项目并运行,例如 pip install -e .;推理示例为 boltzgen run example/vanilla_protein/1g13prot.yaml --output workbench/test_run --protocol protein-anything --num_designs 100 --budget 5 --devices 1 --num_workers 0。可用 python scripts/check_boltzgen_assets.py --cache-dir ~/.cache --check-runtime 检查缓存与运行时。
适用场景:在昇腾 NPU 上离线或受控环境复现 BoltzGen 设计、逆折叠、折叠、分析与筛选全流程;单元测试预期 44 passed。
限制:NPU 下 --num_workers 必须为 0;transfer_to_npu 会禁用 TorchScript;torch/torch_npu/CANN 版本须严格匹配;use_kernels 可显式开关;具体硬件与版本以实际资料为准,其他参数未提供。