生物分子结构预测如今通常在整个蛋白质组规模上运行，其目标是高效地推动整个工作列表通过处理流程。
NVIDIA BioNeMo推理运行时（BioIR）有助于在NVIDIA GPU上加速受支持的生物分子结构预测模型，同时保持熟悉的PyTorch工作流程。它使用优化的内核，并在适用的情况下使用CUDA Graphs来加速模型执行。对于大批量的独立输入，Ray可以在单个节点上的每个GPU上运行一个完整的模型副本，以提高整体吞吐量。
BioIR还已用于实际的蛋白质组规模工作，包括最近扩展的AlphaFold数据库（AFDB），加速了跨4777个蛋白质组生成蛋白质复合物结构，总计约3100万个候选复合物，其中181万个作为高置信度预测被发布。
您可以通过两种方式使用它（见图1）：
端到端处理器通过解析、Token化、特征生成、GPU推理以及PDB或mmCIF写入来处理InputRequest。
直接PyTorch集成允许您构建受支持的模型（torch.nn.Module）或在自定义代码中重用选定的模块。
本教程演示了BioIR的端到端处理器，从输入准备到GPU推理和PDB或mmCIF输出，并展示了如何跟踪每小时结构数和资源效率。
先决条件
Python 3.12或更高版本
兼容的NVIDIA GPU和驱动程序，以及BioIR wheel包或受支持的开发环境
暂存的模型检查点（下面的示例中为Boltz-2）和所需的化学元数据
每个蛋白质链需要一个A3M MSA。对于具有多个非相同蛋白质链的输入，接受配对或非配对MSA
对于Ray吞吐量扩展，请在同一节点上使用多个可见GPU，以及比副本数量更多的独立记录
该wheel包中包含预编译的CUBIN，因此运行时使用不需要nvcc、CUDA源代码、CMake或CUDA工具包。
步骤1. 选择受支持的结构预测工作流程
下面，我们演示了BioIR中Boltz2的端到端工作流程。使用model_source="boltz-2"。蛋白质链需要MSA；对于具有多个非相同蛋白质链的输入，接受配对或非配对MSA。您可以选择自行提供模板，因为BioIR不运行HHsearch或HMMsearch。端到端处理器支持配体结构预测，但不支持配体亲和力预测。
（代码示例：构建InputRequest，包含蛋白质链、序列及MSA记录）
请用有效值替换截断的序列和MSA路径。对于Ray测试，请从真实的工作列表构建数据行，记录数量要多于副本数量；不要重复使用一行数据来证明有效的扩展性。
步骤2. 使用串行处理器验证单次预测
BioIR的端到端处理器工作流程有两种执行器后端：串行和Ray。串行后端为单个输入依次运行每个阶段，在完成整个工作流程后再处理下一个输入。这使其在使用Ray后端并发处理独立输入（步骤3）之前，非常适合用来检查您的设置。
（代码示例：配置串行处理器，设置回收步骤、采样步骤、扩散样本数等参数，并运行处理器获取输出）
model_inference_time是BioIR经CUDA同步的折叠模型前向计算测量值。它不包括解析、Token化、特征生成、后处理和写入。请通过特征生成器的init_context设置随机种子。scores字段必须解码，因为它是一个JSON字符串。
步骤3. 使用Ray副本扩展独立输入
可以按如下方式选择Ray后端，用于默认的副本布局：
（代码示例：使用create_default_replica_mode_config创建默认Ray配置）
此配置在当前节点上每个可见GPU上放置一个完整的模型副本，并根据torch.cuda.device_count()调整CPU阶段的大小。下面是一个显式控制四个GPU的替代配置：
（代码示例：显式配置解析器、Token化器、特征生成器、引擎和写入器阶段，指定计算资源、CPU数量、GPU数量等）
容量规则为：engine_stage.compute × engine_stage.num_gpus ≤ 可见GPU数量。这个四副本示例是一个单节点配置，假定有四个可见GPU。本教程不涉及多节点部署。这里Ray创建了四个引擎actor，每个都预留一个GPU。每个actor加载完整模型。build_processor会在需要时初始化Ray。实际吞吐量取决于输入分布、阶段平衡、存储、调度和故障情况，因此需要实际测量。
步骤4. 平衡五个处理器阶段
在Ray端到端处理器中，五个处理器阶段按以下依赖顺序处理输入：解析器→Token化器→特征生成器→折叠引擎→写入器。使用相应的*StageConfig配置每个阶段；步骤3的示例展示了相关字段。enabled字段不是公开的跳过控制项。每个阶段都公开compute；相关阶段还公开num_cpus、memory和batch_size。Ray引擎添加了max_concurrent_batches、accelerator_type和num_gpus。
要调整Ray流水线，请从EngineProcessorConfig.create_default_replica_mode_config(...)开始。增加某个阶段的compute以添加工作节点；使用num_cpus、memory以及（对于引擎actor）num_gpus来设置资源预留。如果引擎在等待输入，请添加解析器、Token化器或特征生成器工作节点。当GPU或对象存储内存导致故障时，请降低并发度或分离大型输入。
Ray旨在使CPU阶段与推理重叠，但这是否能改善目标工作负载，取决于在给定硬件配置下，给定输入的解析、特征生成和输出写入阶段的运行时间成本。请参阅ScaleFold的图4，了解OpenFold预处理时间的多样性。
步骤5. 区分单副本加速与流水线扩展
BioIR在三个不同层面提供优化：
内核选择：受支持的操作根据模型配置、GPU、数据类型和张量形状，选择兼容的BioIR自定义实现、cuEquivariance实现或PyTorch后备实现。
模块优化：在支持的情况下，独立的optimize()机制为兼容模块启用CUDA Graph捕获。
流水线扩展：Ray执行器将完整的模型副本放置在多个GPU上，并在它们之间分配独立的输入。
这些层面针对不同的瓶颈。内核和模块优化减少了单个副本内的模型前向时间。Ray可以通过将CPU阶段与GPU折叠计算重叠，以及在独立的GPU上运行完整模型副本处理独立输入，来提高工作列表的吞吐量。Ray不会将单次模型前向传递拆分到多个GPU上。上面的图1区分了处理器路径和直接集成路径；Ray扩展仅适用于处理器路径。
我们对BioNeMo推理运行时的早期基准测试估计了以下模型前向加速效果：
加速比是使用1次预热运行（丢弃）和1次测量调用，跨越29-1734个残基的17个输入进行测量的。OpenFold3和Boltz2开源基线使用了torch.compile，参数为dynamic=None、fullgraph=False、recompile_limit=128、accumulated_recompile_limit=256、fail_on_recompile_limit_hit=True。Boltz2开源版本使用了cuEq；OpenFold3开源版本使用了use_cuequivariance=True和use_deepspeed=True。
这些结果量化了单个模型副本内的加速效果。它们并未测量解析、特征生成、输出写入、Ray调度、多GPU吞吐量或完整工作列表的实际耗时。下面的图3展示了为什么模型前向测量和端到端测量必须保持独立。
要确定额外GPU在实际部署中能带来什么提升，请在单个节点上使用一个、两个和四个Ray副本，测量同一代表性工作列表。步骤6定义了所需的指标和比较方法。
步骤6. 基准测试折叠阶段效率和端到端交付：AFDB案例研究
为了使这些测量更加具体，我们在1000个人类二聚体目标（组合序列长度低于2800个残基）上进行了匹配基准测试，代表了一大批独立的生物分子结构预测任务，类似于AlphaFold数据库中最近添加的数据集。
这个代表性的折叠阶段基准测试在8xH100 GPU上，比较了BioIR加速的Boltz-2与torch编译的开源Boltz-2实现。两种实现使用了相同的目标、暂存的MSA、推理方案和GPU配置；吞吐量指标和其他结果是特定于此配置的，不应推广到所有BioIR支持的模型、数据集或硬件。
该工作流程使用了三次回收、200个采样步骤，每个目标五个扩散样本。BioIR完成了全部1000个目标，每分配的GPU小时交付了58.5K个成功折叠的残基，而公开实现为20.2K——在残基归一化吞吐量方面提升了2.90倍；开源实现在29个目标上出现了内存不足的情况。
上面图3的左侧面板比较了BioIR与torch编译的开源实现的模型前向时间，直接展示了BioIR实现的更低模型前向时间。图3的右侧面板比较了BioIR与开源实现所交付的吞吐量，其中吞吐量为预测结构中的残基总数除以分配的GPU小时数。图3的右侧面板展示了BioIR中内核级、模块级和流水线级优化所带来的加速效果。图3的左侧面板展示了内核级和模块级实现所带来的加速效果。
与Boltz2的加速效果类似，BIR还为其他生物分子共折叠模型（如OpenFold2和OpenFold3）提供了更快的推理速度。BIR的早期版本为NVIDIA内部版本的OpenFold2-MM贡献了加速模块，使得能够为拥有3100万个蛋白质复合物结构的AFDB进行大规模蛋白质结构预测。
我们将图3中1000个目标的匹配基准测试线性外推到100万个可比目标，使用8×H100 80GB HBM3节点的额定功率等效值（见下面的图4）。
据估计，使用8-GPU TDP（热设计功耗）等效值时，BioIR需要11兆瓦时，而公开实现需要35兆瓦时；使用整节点最大功率等效值时，两者分别为21兆瓦时和64兆瓦时。这些仅是针对IT设备的折叠估算值，而非实测能耗数据，且不包括数据中心开销，如电源使用效率（PUE）。
这项受控比较使用每种实现相同的输入和MSA来测量折叠吞吐量；不包括MSA生成、预处理CPU分配、存储、数据传输、重试和工程开销。请将端到端流水线性能指标（包括每小时完成的结构数、GPU和CPU利用率、GPU内存峰值、完成率、故障和重试）与模型前向指标区分开来报告。
故障排除
只有一个GPU处于活动状态：确认使用了Ray、REPLICA模式、多于一个副本、多个可见GPU，以及足够的独立记录数量。
处理器构建引发`ValueError`：检查compute * num_gpus是否超过了可见GPU数量。
蛋白质输入失败：检查所需的非配对A3M文件以及工作节点可见的路径。
GPU处于等待状态：在添加副本之前，先检查CPU阶段、CPU预留、排队情况以及Ray对象存储容量。
推理后数据行处于等待状态：检查写入器并发度和目标存储的吞吐量。
Ray无法放置actor：检查CPU、GPU、内存以及accelerator_type标签。
单条记录导致任务停止：默认的快速失败模式会引发FoldingPredictionError。仅在有意进行行级容错的情况下才设置should_continue_on_error=True，然后检查__inference_error__。
开始使用
探索BioNeMo推理运行时（BioIR），并将其集成到您大规模的结构预测工作流程中：http://github.com/NVIDIA-BioNeMo/BioNeMo-Inference-Runtime
要通过智能体编排进一步加速药物发现工作流程，请查看NVIDIA BioNeMo Agent工具包（BAT）。
有关最新的加速数据，请查阅API参考文档和支持矩阵。
Q&A
Q1：BioNeMo推理运行时（BioIR）是什么？
A：BioIR是NVIDIA推出的用于加速生物分子结构预测模型的运行时工具，它可以在GPU上加速支持的模型执行，同时保持PyTorch工作流程不变，适用于蛋白质组规模的结构预测任务。
Q2：BioIR相比开源实现能带来多大的性能提升？
A：在1000个人类二聚体目标的基准测试中，BioIR每GPU小时能交付58.5K个成功折叠残基，而公开实现只有20.2K，吞吐量提升了2.90倍，同时能耗也更低。
Q3：使用BioIR进行大规模结构预测需要什么条件？
A：需要Python 3.12及以上版本、兼容的NVIDIA GPU和驱动、模型检查点及化学元数据，每个蛋白质链还需要A3M格式的MSA文件。若要使用Ray进行吞吐量扩展，需要同一节点上有多个可见GPU。
