生物分子结构预测如今常常以蛋白质组规模运行，目标是把整份工作清单高效地推过流水线。NVIDIA BioNeMo Inference Runtime（BioIR）可以在 NVIDIA GPU 上加速受支持的结构预测模型，同时保留开发者熟悉的 PyTorch 工作流：它使用优化后的内核，并在适用场景下使用 CUDA Graphs 加速模型执行。对于大批量的独立输入，Ray 可以在单节点内的每块 GPU 上运行一个完整的模型副本，从而提高整体吞吐。
BioIR 已在真实的蛋白质组规模工作中得到应用，包括近期 AlphaFold 数据库（AFDB）的扩展：该工作加速了跨 4777 个蛋白质组的蛋白质复合物结构生成，候选复合物总量约 3100 万个，其中 181 万个作为高置信度预测对外发布。
在用法上，BioIR 提供两种方式。其一是端到端处理器，把一条 InputRequest 依次经过解析、分词、特征生成、GPU 推理以及 PDB 或 mmCIF 写出；其二是直接集成 PyTorch，开发者可以构建受支持的模型（torch.nn.Module），或在自定义代码中复用选定的模块。内核与符合条件的模块优化适用于模型执行环节，而五阶段的 Ray 流水线则适用于端到端处理器工作流。
使用前提包括：Python 3.12 或更高版本；一块兼容的 NVIDIA GPU 与驱动，以及 BioIR 的 wheel 包或受支持的开发环境；一份已暂存的模型检查点（文中示例使用 Boltz-2）与所需的化学元数据；每条蛋白质链需要一个 A3M 格式的 MSA，对于包含多条非同源蛋白链的输入，可使用配对或非配对的 MSA；若要用 Ray 做吞吐扩展，则需要同一节点上多块可见 GPU，且独立记录数量要多于副本数。
值得注意的是，该 wheel 包内含预编译的 CUBIN，因此运行时不需要 nvcc、CUDA 源码、CMake 或 CUDA 工具包。
在工作流选择上，文中演示了 BioIR 中 Boltz-2 的端到端流程，使用 model_source="boltz-2"：蛋白链需要 MSA，多链非同源输入可使用配对或非配对 MSA；模板可选自行提供，因为 BioIR 不运行 HHsearch 或 HMMsearch。端到端处理器支持配体结构预测，但不支持配体亲和力预测。
BioIR 的端到端处理器工作流有两种执行器后端：串行与 Ray。串行后端按顺序为单个输入依次运行各阶段，完成整个工作流后再处理下一个输入，因此适合在进入下一步之前验证环境配置；Ray 后端则用于并发处理相互独立的输入。此外，model_inference_time 是 BioIR 经 CUDA 同步的折叠模型前向测量值，它不包含解析、分词、特征生成、后处理与写出环节。
WeChat
Profile