用途:将 Boltz2 蛋白质结构预测模型从 CUDA 迁移到华为昇腾 NPU,并在 Ascend 910/910B/910C 上完成端到端推理复现。
主要能力:提供权重与数据文件清单、CANN 与 torch/torch_npu 环境搭建步骤、源码迁移改动方案(transfer_to_npu 入口映射、npu_adapter 设备无关封装、npu_kernels 等价实现、Lightning NPU accelerator 注册、CLI 与 Trainer 适配)。
使用方式:先确认 npu-smi info 正常、CANN Toolkit 8.2.RC1,创建 conda 环境并安装 torch==2.5.1 与 torch_npu==2.5.1,拉起 set_env.sh;将 boltz2_conf.ckpt、boltz2_aff.ckpt、ccd.pkl、mols.tar 等放入 ~/.boltz/;克隆源码后按文档修改 main.py、boltz2.py、boltz1.py、triangular_mult.py、triangular_attention/primitives.py,并新增 npu_adapter.py、npu_accelerator.py、npu_kernels.py,再 pip install -e .;最后执行 boltz predict examples/prot_no_msa.yaml --devices 1 --out_dir ./boltz_output_e2e --num_workers 0。
适用场景:在昇腾 NPU 上复现 Boltz2 推理,产出 CIF、pLDDT、PAE、PDE 与 confidence 文件。
限制:文档针对特定已验证版本组合,其他 CANN/torch_npu 版本可能不兼容;NPU 上建议 --num_workers 0 避免 DataLoader 崩溃;CUDA-only 依赖缺失时走 fallback 路径,性能与数值行为可能与 CUDA 存在差异。其它功能与收费信息暂未提供。