如何加速蛋白质组尺度的蛋白质结构预测

蛋白质很少作为单个单体孤立地发挥功能。大多数生物过程由蛋白质与其他蛋白质相互作用所支配，形成蛋白质复合物；在蛋白质结构层级中，这些复合物的结构被描述为四级结构表征。

这代表了比三级结构表征更高一层的复杂性；三级结构表征指单体的 3D 结构，自 AlphaFold2 出现以及 Protein Data Bank 建立以来，这类结构已为人们所熟知。

绝大多数复合物的结构信息仍然不可获得。尽管由 Google DeepMind 和 EMBL 的 European Bioinformatics Institute（EMBL-EBI）共同开发的 AlphaFold Protein Structure Database（AFDB）改变了人们获取单体蛋白质结构的方式，但在蛋白质组尺度上具备相互作用感知能力的结构生物学仍然是一个瓶颈，并面临独特挑战：

• 巨大的组合式相互作用空间
• 多序列比对（MSA）生成和蛋白质折叠的高计算成本
• 数百万复合物规模的推理扩展
• 置信度校准和基准测试
• 数据集一致性和生物学可解释性

在近期工作中，我们扩展了 AFDB，加入了由基于 AlphaFold-Multimer 的高通量流水线生成的大规模同源多聚体蛋白质复合物预测结果——这得益于 NVIDIA 加速计算。此外，我们还预测了异源多聚体复合物，以比较不同复合物预测方式的准确性。

具体而言，在这些数据集的预测中，我们利用了 MMseqs2-GPU 在 MSA 生成方面的内核级加速，以及 NVIDIA TensorRT 和 NVIDIA cuEquivariance 在基于深度学习的蛋白质折叠方面的加速。随后，我们通过最大化利用所有可用 GPU（包括横向扩展到多个集群），将工作负载映射到 HPC 规模的推理。

本文介绍了我们为提高蛋白质折叠吞吐量所采用的主要原则，从采用库和 SDK，到通过优化降低工作负载的计算复杂度。这些原则可以帮助你借鉴我们用于创建这一新数据集的技术，自行搭建类似的流水线。

因此，如果你是：

• 扩展结构预测流水线的计算生物学家
• 训练生成式蛋白质模型的 AI 研究人员
• 优化 GPU 工作负载的 HPC 工程师
• 构建结构资源的生物信息学团队

你将学习如何：

• 设计蛋白质组规模的复合物预测策略
• 将 MSA 生成与结构推断分离以提高效率
• 在 GPU 集群上扩展 AlphaFold-Multimer 工作流

先决条件

技术知识

• Python 和 shell 脚本编写
• 作为 HPC 工作负载调度器的 SLURM
• 基础结构生物学
• 熟悉 AlphaFold/ColabFold/OpenFold 或类似流程

基础设施

• 我们描述了在多 GPU、多节点 NVIDIA DGX H100 Superpod 集群上的扩展
• 该集群包括用于存储 MSA 和中间输出的高速存储

软件

• 访问 MMseqs2-GPU
• 熟悉 TensorRT
• 如果未使用集成了 cuEquivariance 的模型，则需要了解三角注意力和乘法运算

步骤/流程

1. 定义你想要计算的数据集

首先定义预测范围。由于预测蛋白质复合物可能会变成一个组合问题，因此了解哪些内容可能最有意义很有帮助。在某些情况下，如果你的蛋白质组足够小，那么进行全对全（二聚体）复合物预测可能是可行的；然而，如果你想预测大规模蛋白质组数据集，情况可能会发生变化。我们是这样决定开展工作的：同源多聚体复合物：我们选择了 AFDB 中所代表的所有蛋白质组，并按感知重要性对其排序（例如，与人类相关或常被访问的蛋白质组）。这使我们能够按特定顺序对蛋白质组进行计算排序，从而使执行更易于管理。

异聚体复合物：事情在这里可能会迅速变得复杂。对于我们的异聚体运行，我们决定重点关注源自若干参考蛋白质组以及 WHO 重要蛋白质组列表中所列蛋白质组的复合物。由于可由这些蛋白质组衍生出的复合物组合数量难以处理，因此在我们的运行中，我们重点关注同一蛋白质组内（不涉及跨蛋白质组复合物）、在 STRING 中具有“物理”相互作用证据的二聚体（由两个蛋白质组成的复合物）。由于我们追求覆盖率，因此决定考虑 STRING 中针对这些蛋白质组报告的所有相互作用，而不是进一步筛选。文献中的证据表明，按 STRING 评分 >700 进行筛选可以进一步减少输入数量，同时提高获得良好预测复合物的可能性。

2. 将 MSA 生成与结构预测解耦

MSA 生成和结构推断都需要大量计算，但二者的扩展方式不同，正如我们最近在一份白皮书中所介绍的。因此，我们将这些计算作为独立步骤处理，并实现了独立的 SLURM 流水线。一般而言，为了最佳利用节点，我们以这种方式设置 MSA 生成和结构预测。

MSA 生成

我们使用带有 MMseqs2-GPU 后端的 colabfold_search 生成 MSA。虽然 MMseqs2-GPU 原生支持在一个节点的多个 GPU 上扩展运行，但为了更便于进程管理，我们选择在一个节点上为每个 GPU 启动一个 MMseqs2-GPU 服务器进程。在 colabfold_search 中，GPU 仅用于 ungappedfilter 阶段，而不用于后续的比对阶段（后续阶段是多线程 CPU 进程）。因此，通过监控 colabfold_search 的输出，我们可以叠加 colabfold_search 调用，并在前一个调用不再使用 GPU 时启动下一个调用，以减少 GPU 空闲时间。尽管这种方法会超额占用 CPU 资源，但在实践中我们发现，在 DGX H100 节点上，通过三个错峰运行的 colabfold_search 进程，整体吞吐量最多可提升 25%，代价是单个输入块的处理速度变慢。

在确定合理的输入分块大小时，需要考虑两个因素。较小的分块大小会产生更多分块，这意味着更多的单进程开销，例如数据库加载；即使在快速存储上，每次加载也可能需要几分钟。（将数据库预先暂存到可用的最快存储上，例如节点内 SSD，也有助于提高吞吐量。）另一方面，较大的分块需要更长时间才能完成。在具有作业时间限制的 SLURM 集群上，这会导致更多未完成的分块。最佳平衡点取决于集群配置，但对于我们具有 4 小时墙钟时间限制的 DGX H100 节点，采用交错式 colabfold_search 方法时，300 条序列的分块大小似乎效果良好。

结构预测

为了提高结构预测吞吐量，我们既利用了通过 ColabFold 进行基于 JAX 的折叠时的数据处理优化，也利用了 NVIDIA 开发的加速工具，包括 TensorRT，以及用于基于 OpenFold 的折叠的 cuEquivariance。

深度学习推理参数

首先，我们选择了在准确性和速度之间取得良好平衡的推理参数。因此，所有深度学习推理流程（ColabFold 和 OpenFold）的蛋白质推理设置均采用：

• 权重：来自 AlphaFold Multimer（model_1_multimer_v3）的 1x 权重
• 四次循环（带提前停止）
• 不进行松弛
• MSA：通过 ColabFold-search（使用 MMseqs2-GPU）生成的冻结 MSA，如上所述

准确性验证

由于我们使用了不同的推理流程，因此采用一个经过整理的基准集进行了准确性验证，该基准集包含 125 个在 AlphaFold2 推出后发布的、经 X 射线解析的 PDB 同源二聚体，从而最大限度地降低信息泄漏的可能性。针对每种深度学习实现预测的复合物，均使用 DockQ 与实验参考结构进行比较；DockQ 通过天然接触比例（Fnat）、非天然接触比例（Fnonnat）、受体对齐后的界面 RMSD（iRMS）和配体 RMSD（LRMS）来评估界面准确性，并给出高、中、可接受或错误的标准 CAPRI 分类。

在整个 PDB 同源二聚体基准测试中，通过 TensorRT 和 cuEquivariance 加速的 OpenFold 复现了 ColabFold 的界面准确性，实现了相近比例的“高”评分预测，并获得了可比的平均 DockQ 分数。这表明，相较于 ColabFold 基线，加速实现保留了界面层面的结构准确性。

MSA 准备和序列打包

对于基于 ColabFold 的同源二聚体推理，可将长度相同的同源二聚体打包成一个批次进行处理，并按其 MSA 深度降序排序，从而实现更高吞吐量。这减少了 JAX 重新编译的次数，进而提高端到端吞吐量。不过，该技巧在处理异源二聚体时不起作用，因为各条链的长度不同。

对于 OpenFold，无论是同源二聚体还是异源二聚体，都不需要这种打包策略，因为该方法不需要重新编译。不过，鉴于序列长度与执行时间之间存在依赖关系，如果使用特定的 SLURM 运行时，将较长序列预留给单独作业可能是有益的。为进一步优化流程，在对当前查询执行推理步骤（GPU 受限）的同时，对下一个输入查询执行输入特征化（CPU 受限）。

此外，通过集成 NVIDIA cuEquivariance 库和 NVIDIA TensorRT SDK，OpenFold 的吞吐量得到了提升。这些模块化库和 SDK 可分别用于加速蛋白质结构 AI 和通用推理 AI 工作负载中常见的操作。我们此前曾介绍过如何利用 TensorRT 加速 OpenFold 推理。

3. 使用 SLURM 优化 GPU 利用率

如上一节所述，根据可用硬件的不同，你可以通过“打包”GPU 和节点来提高吞吐量。SLURM 是一款出色的编排器，我们将推理工作流划分到 SLURM 脚本中，以便：

• 每个节点打包多个预测任务
• 将 GPU 内存与序列长度相匹配
• 减少作业之间的空闲时间
• 将短序列队列与长序列队列分开

我们的工作负载被映射到一个 H100 DGX Superpod HPC 系统。因此，我们可以在多节点集群上的 NVIDIA H100 GPU 上部署推理，在单个节点上利用独占执行，并为每个 GPU 打包尽可能多的进程，以使 MSA 处理和深度学习推理的 GPU 利用率达到饱和。

实用提示：

• 按总残基长度对任务分组
• 监控 GPU 内存碎片化
• 使用异步 I/O 以避免磁盘瓶颈

4. 让全世界都能获得高质量预测

我们与 EMBL-EBI、首尔国立大学的 Steineggerlab 以及 Google DeepMind 合作，探索了复杂结构预测分析。我们强调，预测这些生物系统仍然具有挑战性。与蛋白质单体预测不同，在单体预测中，预测的局部距离差异测试（pLDDT）可以反映整体预测质量，从而产生数量适中的合理预测；而在复合物场景中，评估界面的合理性要困难得多。这与评估复合物需要全局和逐链置信度指标，以及界面处的局部置信度指标有关。简单来说，两个单体之间的界面是否合理，以及它是否被预测在正确的口袋中？鉴于可用数据非常有限，这些问题比关于单体可能性的更“局部”问题要难回答得多。因此，我们提供

入门

蛋白质组规模的四级结构预测需要的不仅仅是大规模运行 AlphaFold-Multimer。成功取决于：

• 证据驱动的交互选择
• 解耦并优化的计算工作流
• GPU 感知的作业编排
• 置信度校准与验证
• 数据集健康监测

通过结合 STRING 引导的选择、MMseqs2-GPU 加速以及由 NVIDIA H100 驱动的多聚体推理，本工作将 AFDB 扩展为一个统一的、具备相互作用感知能力的结构资源。

该基础设施支持：

• 界面处的变异解读
• 系统级结构生物学
• 药物靶点验证
• 生成式蛋白质设计基准测试

资源

• 在此阅读有关该项目的更多信息：https://research.nvidia.com/labs/dbr/assets/data/manuscripts/afdb.pdf
• 加速库和 SDK 可在此获取：MMseqs2-GPU NVIDIA cuEquivariance NVIDIA TensorRT
• 如果你希望轻松部署 MSA 搜索和蛋白质折叠，可以通过 NVIDIA 的 Inference Microservices（NIMs）获取加速推理流水线：MSA Search NIM OpenFold2 NIM
• 这项工作的预测结果可通过 https://alphafold.com 获取