BioNeMo Inference Runtime 如何提升大规模生物分子结构预测吞吐

BioNeMo 的关键进步不只是更快内核，而是完整流水线的资源经济学

NVIDIA 的 BioNeMo Inference Runtime 将 GPU 专用内核、可选 CUDA Graph 优化和基于 Ray 的模型副本结合起来，面向独立的结构预测任务提升工作列表吞吐。真正需要验证的不是单次 forward 是否更快，而是预处理、显存、调度和写盘是否能让整条流水线持续受益。

Token4AI Cloud2026-09-12约 5 分钟阅读

THE BRIEF

博客要点

NVIDIA 将 BioNeMo Inference Runtime 定义为兼容 PyTorch 的生物分子结构预测运行时。端到端处理器覆盖解析、分词、特征生成、GPU 推理以及 PDB 或 mmCIF 写出；另一条路径允许直接集成受支持的 PyTorch 模块。Ray 执行器针对相互独立的输入，在每张 GPU 上放置完整模型副本，而不是把一次 forward 拆到多张 GPU。原文在 8 张 H100 上测试了 1,000 个长度合计低于 2,800 个残基的人类二聚体目标，采用 3 次 recycle、200 次扩散采样和每个目标 5 个 diffusion samples。该配置下，NVIDIA 报告 BioIR 达到每 GPU 小时 58.5K 个成功折叠残基，测试的 torch.compile 开源实现为 20.2K；后者有 29 个目标因显存不足失败。上述数字只对应给定工作负载。能源数据则是仅针对折叠阶段的额定功率等效估算，不是完整设施的电表读数。

来源公告日期 · 2026-09-10

这项工作的核心不是把结构预测当成孤立的 GPU 函数，而是把工作列表当成一条需要持续供料和排空的流水线。BioIR 的端到端路径包含 Parser、Tokenizer、Feature Generator、Folding Engine 和 Writer。结构预测输入通常来自不规则文件，输出又要落成可复用的结构文件。因此，即使模型 forward 变快，如果 A3M 读取、特征生成或写盘跟不上，GPU 仍会等待。把这些阶段纳入同一执行模型，才有机会优化真正的交付吞吐。

三层优化分别对应不同瓶颈。内核选择会根据模型配置、GPU、数据类型和张量形状，在 BioIR、cuEquivariance 或 PyTorch fallback 之间选择。对兼容模块进行 CUDA Graph 捕获，可以减少稳定执行路径中的启动和调度开销。Ray 副本则是另一种优化：每个 actor 加载完整模型，处理独立记录。这是面向任务队列的数据并行，不是把一个超大复合物的单次推理做模型并行。

这个区别也限定了 Ray 的适用范围。只有在独立记录足够多，并且 CPU 准备与输出写入能持续供给 GPU 时，副本数量才可能转化为吞吐。它不能让一次 forward 自动跨 GPU 分摊显存或计算。每增加一个副本，都要在另一张 GPU 上加载完整模型，因而并发能力与显存容量之间存在直接交换。长度差异很大的复合物还会形成长尾任务，使部分副本等待最慢输入，降低平均利用率。

原文的残基每 GPU 小时指标适合描述混合长度任务中的有效产出，但不能单独代表系统质量。它应与端到端墙钟时间、每阶段排队时间、峰值显存、失败率和重试成本一起解读。残基归一化吞吐还能避免只统计任务数带来的偏差，不过它仍然受数据集组成影响。文中报告的 2.90 倍提升只适用于指定的 1,000 个目标、MSA、推理配方和 8 张 H100，不能直接外推到其他模型、长度分布或硬件。

开源实现出现 29 个显存失败，说明吞吐同时也是可靠性指标。若系统只在容易样本上速度很高，却频繁丢失长样本，实际每天交付的结构数量未必更高。但 BioIR 的可靠性也不能只由厂商案例证明。文档所述默认行为可能在单条记录失败时快速终止整个任务；开启逐行继续后，又会改变完成率的定义。生产统计必须区分尝试、成功、失败、重试和跳过，不能只看已完成任务数。

能源结论需要同样谨慎。NVIDIA 将基准结果外推到一百万个相似目标，并分别给出额定功率等效和整机最大功率情景。这些是折叠计算的估算，排除了 MSA 生成、CPU 预处理、存储、数据传输、重试、工程开销以及 PUE。较低的单位能耗可能来自更短的 forward 时间和更少的失败任务，但这只是合理假设，必须用 GPU 或节点级功率遥测，以及端到端成功产出进行验证。

可复现实验应先冻结工作列表、序列、MSA、模板、随机种子、recycle 次数、采样设置、精度和输出格式。先用 serial 处理器验证结构正确性，再在同一节点上使用 1、2、4、8 个 Ray 副本，并记录每个阶段的队列。指标至少包括每小时成功结构数、每 GPU 小时成功残基数、GPU 与 CPU 利用率、峰值显存、对象存储压力、p50 与尾延迟、写出带宽、失败类别和重试次数。还应分别测试短、长、同质和异质复合物，只有确认瓶颈后，才能决定增加 CPU worker、降低并发、调整批大小或增加模型副本。