生物分子结构预测如今常常在蛋白质组规模上运行，其目标是将整个工作列表高效地通过流水线处理。
NVIDIA BioNeMo 推理运行时（BioIR）有助于在 NVIDIA GPU 上加速受支持的生物分子结构预测模型，同时保持熟悉的 PyTorch 工作流。它使用优化的内核，并在适用的情况下使用 CUDA Graphs 来加速模型执行。对于大批量的独立输入，Ray 可以在单个节点中的每个 GPU 上运行一个完整的模型副本，以提高整体吞吐量。
BioIR 也已被用于真实的蛋白质组规模工作，包括 AlphaFold 数据库的近期扩展（AFDB），加速了跨越 4,777 个蛋白质组的蛋白质复合物结构生成，总计约 3,100 万个候选复合物，其中 181 万个作为高置信度预测发布。
你可以通过两种方式使用它（见下方图 1）：
- 端到端处理器让 InputRequest 依次经过解析、分词、特征生成、GPU 推理以及 PDB 或 mmCIF 写出。
- 直接 PyTorch 集成让你可以构建受支持的模型（torch.nn.Module），或在自定义代码中复用选定的模块。
本教程将带你了解 BioIR 的端到端处理器，从输入准备到 GPU 推理以及 PDB 或 mmCIF 输出，并展示如何跟踪每小时结构数和资源效率。
图 1. BioNeMo 推理运行时支持端到端处理器和直接 PyTorch 模块集成；内核和符合条件的模块优化在模型执行内适用；五阶段 Ray 流水线适用于端到端处理器工作流
先决条件
- 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。端到端处理器支持配体结构预测，但不支持配体亲和力预测。
from bionemo_ir.data.schemas import InputRequest, MSARecord, Polymer
request = InputRequest(
input_id="demo",
polymers=[
Polymer(
polymer_type="protein",
chain_id=["A"],
sequence="GSHMSL...",
msas=[MSARecord(path="msa.a3m", format="a3m")],
paired_msas=[],
templates=None,
)
],
)
rows = [{"record": request, "__record_id": request["input_id"]}]
将截断的序列和 MSA 路径替换为有效值。对于 Ray 测试，应从包含比副本数更多记录的真实工作列表构建行；不要重复同一行作为有效扩展的证据。
步骤 2. 使用串行处理器验证单个预测
BioIR 为端到端处理器工作流提供了两种执行器后端：串行和 Ray。串行后端针对单个输入按顺序运行每个阶段，在进入下一个输入之前完成整个工作流。这使得它有助于在使用 Ray 后端于步骤 3 中并发处理独立输入之前检查你的设置。
import json
from bionemo_ir.pipeline.processor.engine_proc import (
EngineProcessorConfig,
build_processor,
)
from bionemo_ir.pipeline.stages.configs import (
FeatureGeneratorStageConfig,
WriterStageConfig,
)
serial_config = EngineProcessorConfig(
model_source="boltz-2",
executor_backend=None, # None mean the serial processor
runtime_args={
"recycling_steps": 3,
"num_sampling_steps": 50,
"diffusion_samples": 1,
},
feature_generator_stage=FeatureGeneratorStageConfig(
init_context={"random_seed": 42},
),
writer_stage=WriterStageConfig(
output_path="output/serial",
format="cif",
),
engine_kwargs={"profile_inference": True},
)
serial_processor = build_processor(serial_config)
serial_outputs = serial_processor(rows)
for row in serial_outputs:
scores = json.loads(row["scores"])
print(row["output_path"])
print(row["model_inference_time"])
print(scores.get("confidence_score"))
model_inference_time
是 BioIR 的 CUDA 同步折叠模型前向测量。它不包括解析、分词、特征生成、后处理和写入。通过特征生成器的 init_context 设置种子
。scores
字段必须解码，因为它是一个 JSON 字符串。
步骤 3. 使用 Ray 副本扩展独立输入
可以按如下方式选择 Ray 后端，以使用默认副本布局：
from bionemo_ir.pipeline.processor.engine_proc import EngineProcessorConfig
ray_config = EngineProcessorConfig.create_default_replica_mode_config(
model_source="boltz-2",
output_dir="output/ray",
output_format="cif",
)
此配置在当前节点上每个可见 GPU 上放置一个完整的模型副本，并根据 torch.cuda.device_count() 确定 CPU 阶段的大小
。以下是一个显式控制四个 GPU 的替代配置：
import ray
from bionemo_ir.pipeline.processor.engine_proc import (
EngineProcessorConfig,
build_processor,
)
from bionemo_ir.pipeline.stages.configs import (
EngineStageConfig,
FeatureGeneratorStageConfig,
ParallelismMode,
ParserStageConfig,
TokenizerStageConfig,
WriterStageConfig,
)
ray_config = EngineProcessorConfig(
model_source="boltz-2",
executor_backend="ray", # selects the ray backend
parser_stage=ParserStageConfig(compute=4),
tokenizer_stage=TokenizerStageConfig(compute=4, num_cpus=2),
feature_generator_stage=FeatureGeneratorStageConfig(
compute=8,
num_cpus=4,
init_context={"random_seed": 42},
),
engine_stage=EngineStageConfig(
parallelism_mode=ParallelismMode.REPLICA,
compute=4,
num_gpus=1.0,
num_cpus=4,
),
writer_stage=WriterStageConfig(
compute=4,
output_path="output/ray",
format="cif",
),
)
ray_processor = build_processor(ray_config)
dataset = ray.data.from_items(rows)
ray_outputs = list(ray_processor(dataset).materialize().iter_rows())
容量规则是 engine_stage.compute × engine_stage.num_gpus ≤ visible GPUs
。这个四副本示例是单节点配置，假设有四块可见 GPU。本教程不涉及多节点部署。这里 Ray 创建四个引擎 actor，并为每个保留一块 GPU。每个 actor 加载完整模型。build_processor
在需要时初始化 Ray。实际吞吐量取决于输入分布、阶段平衡、存储、调度和故障，因此需要实测。
步骤 4. 平衡五个处理器阶段
在 ray 端到端处理器中，五个处理器阶段按以下依赖顺序消费输入：Parser → Tokenizer → Feature generator → Folding engine → Writer。为每个阶段配置其对应的 *StageConfig；步骤 3 的示例展示了相关字段。enabled
字段不是公开的跳过控制。每个阶段都暴露 compute
；相关阶段还暴露 num_cpus
、memor
y 和 batch_size
。Ray 引擎额外添加 max_concurrent_batches
、accelerator_type
和 num_gpus
。
要调优 Ray 流水线，从 EngineProcessorConfig.create_default_replica_mode_config(...)
开始。增加某个阶段的 compute
以添加 worker；使用 num_cpus
、memory
，对于引擎 actor 还使用 num_gpus
来设置资源预留。如果引擎等待输入，则添加 parser、tokenizer 或 feature worker。当 GPU 或对象存储内存导致失败时，降低并发度或分离大输入。
Ray 旨在将 CPU 阶段与推理重叠，但这是否能改善目标工作负载，取决于在给定硬件配置上给定输入的解析、特征生成和输出写入阶段的运行时成本。请参考 ScaleFold 中的图 4，了解 OpenFold 预处理时间的多样性。
步骤 5. 将每副本加速与流水线扩展分开
BioIR 在三个不同层面提供优化：
内核选择：支持的操作根据模型配置、GPU、数据类型和张量形状，选择兼容的 BioIR 自定义、cuEquivariance 或 PyTorch 回退实现。模块优化：在支持的情况下，单独的 optimize() 机制为兼容模块启用 CUDA Graph 捕获。流水线扩展：Ray 执行器将完整模型副本放置在 GPU 上，并在它们之间分发独立输入。
这些层面针对不同的瓶颈。内核和模块优化减少副本内的模型前向时间。Ray 可以通过将 CPU 阶段与 GPU folding 重叠，以及在独立 GPU 上为独立输入运行完整模型副本，来提高工作列表吞吐量。Ray 不会将单个模型前向传播拆分到多个 GPU 上。上方的图 1 区分了处理器路径和直接集成路径；Ray 扩展仅适用于处理器路径。
我们使用 BioNeMo Inference Runtime 进行的早期基准测试估计了以下模型前向加速：
图 2. BioNeMo Inference Runtime 在 H100 和 H200 GPU 上相对于 OSS torch.compile 的 GPU 同步模型前向延迟几何平均加速比
加速比测量使用 1 次预热运行（丢弃）和 1 次测量调用，覆盖 17 个输入，长度范围为 29–1,734 个残基。OpenFold3 和 Boltz2 OSS 基线使用了 torch.compile
，配置为 dynamic=None
、fullgraph=False
、recompile_limit=128
、accumulated_recompile_limit=256
、fail_on_recompile_limit_hit=True
。Boltz2 OSS 使用了 cuEq；OpenFold3 OSS 使用了 use_cuequivariance=True
和 use_deepspeed=True
。
这些结果量化的是单个模型副本内的加速。它们不测量解析、特征生成、输出写入、Ray 调度、多 GPU 吞吐量或完整工作列表的墙钟时间。下方的图 3 说明了为什么模型前向与端到端测量必须保持分开。
要确定额外 GPU 在实际部署中能带来什么，请在单节点上使用一个、两个和四个 Ray 副本测量同一代表性工作列表。步骤 6 定义了所需的指标和比较方法。
步骤 6. 对折叠阶段效率和端到端交付进行基准测试：AFDB – 案例研究
为了使这些测量更具体，我们在 1,000 个人类二聚体靶标上运行了匹配基准测试，这些靶标的合并序列长度低于 2,800 个残基，代表了大量独立的生物分子结构预测任务类似于最近添加到 AlphaFold 数据库中的数据集。
这个具有代表性的折叠阶段基准测试在 8xH100 GPU 上比较了 BioIR 加速的 Boltz-2 与 torch 编译的开源 Boltz-2 实现。两种实现使用了相同的靶标、分阶段 MSA、推理配方和 GPU 配置；吞吐量指标和其他结果特定于此配置，不应推广到所有 BioIR 支持的模型、数据集或硬件。
该工作流使用了三次循环、200 个采样步骤，每个靶标五个扩散样本。BioIR 完成了全部 1,000 个靶标，每分配 GPU 小时成功折叠 58.5K 个残基，而公共实现为 20.2K——按残基归一化的吞吐量提升了 2.90 倍；开源实现在 29 个靶标上内存不足。
图 3. 在匹配的 1,000 靶标人类二聚体基准测试中，BioIR 的按残基归一化折叠吞吐量比所测试的公共开源实现高 2.90 倍。吞吐量指标取决于输入数据集的组成
上方图 3 的左面板比较了 BioIR 与 torch 编译的开源实现的模型前向时间，并直接显示了 BioIR 提供的更低模型前向时间。图 3 的右面板比较了 BioIR 与开源实现提供的吞吐量，其中吞吐量是预测结构中残基总数按分配 GPU 小时归一化后的结果。图 3 的右面板显示了 BioIR 中内核级、模块级和流水线级优化带来的加速。图 3 的左面板显示了内核级和模块级实现带来的加速。
与 Boltz2 的加速类似，BIR 也能为其他生物分子共折叠模型（如 OpenFold2 和 OpenFold3）实现更快的推理。BIR 的早期版本曾为 NVIDIA 内部版本的 OpenFold2-MM 贡献了加速模块，从而实现了对 AFDB 中 3100 万个蛋白质复合物结构的大规模蛋白质结构预测。
我们使用 8 x H100 80GB HBM3 节点的额定功率等效值，将图 3 中 1000 个目标的匹配基准线性外推至一百万个可比目标（见下方图 4）。
据估计，BioIR 需要 11 MWh，而使用 8-GPU TDP（热设计功耗）等效值的公开实现需要 35 MWh；使用全节点最大功率等效值时，则分别为 21 MWh 与 64 MWh。这些仅为仅折叠的 IT 设备估算值，并非实际计量的能耗测量值，且不包括数据中心开销，例如电能利用效率（PUE）。
图 4。BioIR 使用配备 8x H100 80 GB GPU 的节点，在 TDP 下将折叠一百万个可比目标的额定功率等效估算能耗降低了 23 MWh，在全节点最大功率下降低了 43 MWh
该受控比较使用相同的输入和 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 Inference Runtime (BioIR)，并将其集成到您的大规模结构预测工作流中：http://github.com/NVIDIA-BioNeMo/BioNeMo-Inference-Runtime
如需通过智能体编排进一步加速药物发现工作流，请查看 NVIDIA BioNeMo Agent Toolkit (BAT)。