BioNeMo Inference Runtime

英伟达正式开源了 BioNeMo Inference Runtime (BioIR)，这是一个专为生物分子结构预测模型设计的 Python 推理运行时。它解决从单目标预测转向蛋白质组规模（Proteome-scale）工作负载时的性能瓶颈，通过三层优化架构，在保持纯 PyTorch 工作流的同时，显著提升了 Boltz-2、OpenFold2 和 OpenFold3 等模型的推理吞吐量。

BioIR 已成功应用于生产环境，为扩展后的 AlphaFold Database 提供了动力，生成了约 3,100 万个蛋白质复合物候选结构。

核心架构：三层加速策略

BioIR 并非通过导出静态图来牺牲灵活性，而是针对通用推理栈中未充分优化的操作（如 Pairformer、Evoformer、三角形操作等）进行深度定制：

1. 内核选择 (Kernel Selection)：
   • 根据模型配置、GPU 型号、数据类型和张量形状，动态选择最优实现：BioIR 自定义内核、cuEquivariance 或 PyTorch 原生回退。
2. 模块优化 (Module Optimization)：
   • 通过 optimize() 机制为兼容模块启用 CUDA Graph 捕获，大幅削减内核启动开销。
3. 管道扩展 (Pipeline Scaling)：
   • 利用 Ray 执行器在每个可见 GPU 上部署一个完整的模型副本。
   • 副本模式：将独立的工作列表分发给多个 GPU，而非将单个前向传递拆分到多卡。
   • CPU/GPU 重叠：CPU 阶段（解析、特征化、写入）与 GPU 折叠计算并行执行，最大化硬件利用率。

注意：目前尚不支持上下文并行（Context Parallelism），即单个超大蛋白质的推理仍需在单卡内存限制内完成。

基准测试：8×H100 上的性能飞跃

在针对 1,000 个人类二聚体（总残基数 < 2,800）的端到端基准测试中，BioIR 展现了显著优势：

注：测试环境为 8×H100 80GB，使用相同的 MSA、推理配方（3次循环、200采样步、5扩散样本）。

能耗估算

若线性外推至 100 万个目标：

• BioIR 预估能耗：11 MWh
• 基线预估能耗：35 MWh
• 节省：每百万目标可节省 23–43 MWh 的 IT 设备电力，对大规模蛋白质组项目而言意义重大。

部署与使用

• 开源可用：代码已在 GitHub 发布，提供包含预编译 CUBIN 的 Wheel 包。
• 极简依赖：
  • 需要 Python 3.12+、兼容的 NVIDIA GPU 及驱动。
  • 无需 nvcc、CUDA 源代码、CMake 或 CUDA Toolkit。
  • 需准备暂存的模型检查点和每条链的 A3M MSA 文件。
• 两种集成方式：
  1. 端到端处理器：输入 InputRequest，自动完成解析、特征生成、推理及 PDB/mmCIF 输出。
  2. 直接 PyTorch 集成：开发者可在自定义代码中重用经过优化的特定模块。

关键要点

1. 纯 PyTorch 体验：模型保持为标准的 torch.nn.Module，无需复杂的导出步骤，便于调试和迭代。
2. 吞吐量优先：通过 Ray 副本模式扩展独立任务队列，完美适配蛋白质组学中海量小任务的场景。
3. 生产级验证：已支撑 AlphaFold Database 的大规模扩展，证明了其在真实世界高负载下的可靠性。
4. 局限性：目前不支持配体亲和力预测，且模板搜索（HHsearch/HMMsearch）需外部运行后手动提供。