如何加速蛋白质组规模的蛋白质结构预测

蛋白质很少像单个单体那样独立发挥作用。大多数生物过程由与其他蛋白质相互作用的蛋白质控制，形成蛋白质配合物，其结构在蛋白质结构的层次结构中被描述为四级表征。

这代表了从三级表征 (单体的 3D 结构) 到上一级的复杂性，自 AlphaFold2 出现和 蛋白质数据库 创建以来，这种结构就一直为人所知。

绝大多数配合物的结构信息仍然不可用。虽然 Google DeepMind 和 EMBL 的欧洲生物信息学研究所 (EMBL-EBI) 联合开发的 AlphaFold 蛋白质结构数据库 (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 脚本
• SLURM 作为 HPC 工作负载调度程序
• 基础结构生物学
• 熟悉 AlphaFold/ ColabFold/ OpenFold 或类似工作流

基础架构

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

软件

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

过程/ 步骤

1. 定义要计算的数据集

首先定义预测范围。由于预测蛋白质复合体可能会成为一个组合问题，因此了解最有趣的蛋白质复合体非常有用。在某些情况下，如果您的蛋白质组足够小，则 all-against-all (二聚体) 复杂预测可能易于实现；但是，如果您想预测蛋白质组的大型数据集，这种预测可能会发生变化。

这里是我们的决定：

单体复合体： 我们选择了 AFDB 中表示的所有蛋白质组，并根据感知的重要性 (例如，人类关注的或常见访问的蛋白质组) 对其进行排序。这让我们能够按特定顺序对蛋白质组进行排序，使执行更易于管理。

异构配合物：在这种情况下，事物会变得复杂、快速。在异构运行中，我们决定专注于源自 WHO 重要蛋白质组列表中的多个参考蛋白质组和蛋白质组的复合体。由于可以从这些蛋白质组中衍生出大量复杂的复合体组合，因此在我们的运行中，我们专注于在 STRING 中具有“物理”相互作用证据的同一蛋白质组 (无蛋白质组间复合体) 内的二聚体 (两种蛋白质的复合体) 。在寻求覆盖范围时，我们决定考虑在字符串中报告的所有蛋白质组相互作用，而不是进一步过滤。文献中的证据表明，对 STRING 分数 >700 进行过滤可以进一步减少输入数量，同时增加准确预测复合体的可能性。

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

MSA 生成和结构推理都是计算密集型任务，但规模不同，正如我们最近在 白皮书 中介绍的那样。因此，我们将这些计算视为单独的步骤，并实施了单独的 SLURM 工作流。一般来说，为了优化节点的使用，我们以这种方式设置了 MSA 生成和结构预测。

生成 MSA

我们使用 colabfold_search 和 MMseqs2-GPU 后端生成了 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 节点而言，300 个序列的块大小似乎非常适合惊人的 colabfold_search 方法。

结构预测

为了提高结构预测吞吐量，我们通过 ColabFold 优化了基于 JAX 的折叠数据处理，并利用了 NVIDIA 开发的加速工具，包括 TensorRT，和用于 OpenFold-based folding 的 cuEquivariance。

深度学习推理参数

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

• 权重：1 个 AlphaFold Multimer (model_1_multimer_v3) 权重
• 四次回收 (提前终止)
• 无松
• MSA：如上所述，通过 ColabFold-search (使用 MMseqs2-GPU) 生成的冻结 MSA

准确性验证

当我们使用不同的推理流程时，我们使用在推出 AlphaFold2 后发布的一套精选基准测试集 (包含 125 个 X 射线分辨率的 PDB 同构二聚体) 进行了准确性验证，从而最大限度地减少了信息泄露的可能性。

使用 DockQ 将每个深度学习实现的预测复合体与实验参考结构进行比较，DockQ 通过受体比对 (LRMS) 后的原生接触分数 (Fnat) 、非原生接触分数 (Fnonnat) 、界面 RMSD (iRMS) 和配体 RMSD 来评估接口准确性，并分配高、中、可接受或 inco 的标准 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 推理。

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

MSA 生成和结构推理都是计算密集型任务，但规模不同，正如我们最近在 白皮书 中介绍的那样。因此，我们将这些计算视为单独的步骤，并实施了单独的 SLURM 工作流。一般来说，为了优化节点的使用，我们以这种方式设置了 MSA 生成和结构预测。

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 数据库提供了一组高置信度结构，从而首次实现了对蛋白质复合体的探索。我们打算进一步完善我们的方法，扩大 AlphaFold 数据库中可用蛋白质复合物的范围。

开始使用

蛋白质组尺度四级结构预测需要的不仅仅是大规模运行 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
  • NVIDIATensorRT
• 如果您希望轻松部署 MSA 搜索和蛋白质折叠，可以通过 NVIDIA 的推理微服务 (NIM) 获得加速推理流程：
  • MSA 搜索 NIM
  • OpenFold2 NIM
• 这项研究的预测结果可通过 https://alphafold.com 获取。