数据中心/云端

使用 QUELO-G 和 CUDA 图进行量子力学增强型药物研发

在药物研发中,人们经常使用基于所谓经典力场的方法,并认为这些方法很有用。然而,人们也广泛认识到,力场模型中缺少一些重要的物理特性,导致应用受限。

例如,在比较两种具有不同形式电荷的分子时,力场模型无法提供准确的预测,因为模型中未考虑蛋白质极化。此外,由于缺乏形成和打破化学键的能力,这些模型不适用于共价药物分子。

许多人一致认为,基于量子力学 (管理微观世界的基本物理学) 的模拟是解决这些问题的方法,但量子力学模拟被认为过于昂贵和耗时,不实用。

QSimulate 近期宣布推出 QUELO-G 平台,该平台能够实现基于量子力学的自由能微扰(FEP)模拟,并实现出色的吞吐量。每个 GPU 卡每天的吞吐量超过 100 纳秒,因此这些模拟可以在几小时内完成。基于量子力学的模拟的直接应用将改变计算机辅助药物研发,特别是在识别新的药物分子方面。

本文将介绍 QSimulate 如何通过开发创新算法来应对挑战,并将其应用到利用 NVIDIA 提供的最新 GPU 硬件的软件中,例如 NVIDIA A100NVIDIA H100。这些紧密优化的软件利用了最近在 NVIDIA Omniverse 中提供的一些功能,包括 CUDA 工具包 中的条件图节点和 CUDA 图,这是在 版本 12.3 中引入,并在 版本 12.4 中进一步完善。

量子力学的吞吐量挑战

药物研发模拟的一个重要方面是,每个能量和力评估都必须在几毫秒内完成。这一要求源于这样一个事实,即由于蛋白质是灵活的,并且随着时间的推移会改变其形状,因此有必要通过分子动力学 (MD) 模拟在多个纳秒内对热力学集合进行采样,以准确预测特性。例如,药物分子与蛋白质标的结合亲和力。

如果目标是实现每天 100 纳秒的动力学吞吐量(标准时间步长为 2 飞秒),则每个时间步长必须在不到 2 毫秒的墙面时间内完成。经典力场模拟已被证明能够实现此吞吐量。例如,请参阅以下相关文章:GROMACSNAMD 等经典力场模拟软件的相关文章。

通常认为,即使使用现代 NVIDIA GPU 硬件,也很难甚至不可能将量子力学模拟加速到毫秒级步进机制。这也是 GPU 上大多数现有量子力学模拟实现专注于加速大型计算(例如 集群计算大规模密度泛函理论计算),通常从几小时到几分钟不等,或者从几天到几小时不等。如需了解详情,请参阅 GPU 加速的量子化学和分子动力学 会议。

虽然这些模型有用,但并不能直接用于药物研发的自由能计算,因为吞吐量降低了几个数量级。从技术角度讲,量子力学的吞吐量挑战源于量子力学本质上是一个优化问题,必须以迭代方式解决大量参数。

这势必会引入复杂的控制逻辑并限制并发程度。单纯地实施量子力学的复杂逻辑会导致设备到主机的频繁通信,从而降低 GPU 的利用率。

QUELO-G 通过将 CUDA 图形与紧密结合的量子力学方法相结合,克服了使用 GFN-xTB混合量子力学/分子力学 (QM/MM) 方案 时所面临的挑战。

使用 CUDA 图形和条件节点的技术解决方案

作为迭代量子力学算法的示例,请考虑 Krylov 子空间算法。该方法通过构建一个由试验向量组成的子空间来迭代求解,其中新的试验向量是从上一次迭代的最佳解生成的。该过程会进行迭代,直到满足收标准。图 1 显示了算法的原理图表示。

Two side-by-side flow charts comparing the standard and CUDA Graph-based algorithms.
图 1. 使用 CPU 上的控制代码在 GPU 上实现迭代求解器的标准方法 (左),与使用条件图形节点功能的基于 CUDA 图形的方法 (右) 相比

实现此算法的标准方法是让主机负责循环和条件分支,同时将所有其他部分卸载到设备。试验向量 (以及子空间) 可以保留在设备内存中。但是,残差误差是标量值,需要在主机上执行的控制逻辑的每次迭代中从设备传输到主机。

即使是标量变量,设备到主机通信的延迟也绝非可以忽略不计。由于底层算法的结构,这种情况会反复发生,因此在推动实现毫秒级每步吞吐量的同时,它对基于量子力学的动力学模拟的整体性能具有重要意义。

QUELO-G 中的软件实现使用最近在 CUDA Graphs (自 CUDA 12.3 起) 中引入的条件图节点。使用此功能,量子力学算法中的整个迭代过程可以映射到 CUDA Graphs,在其中在设备中执行循环和条件分支 (图 2)。这消除了在每次迭代中从设备到主机通信标量的需求。这种方法不仅可以显著提高性能,而且由于更好的抽象化,还简化了代码结构。

CUDA Graph for quantum mechanics optimization.
图 2.CUDA Graph 表示 QUELO-G 中用于量子力学波函数优化的算法。每列都对应于条件节点的主体

此外,通过使用 CUDA Graphs 提取代码,该软件将自动受益于针对现有 NVIDIA A100 和 NVIDIA H100 GPU 以及未来 NVIDIA GPU 的 CUDA Graphs 运行时的进一步开发和优化。

基于量子力学的自由能微扰模拟性能

通过采用上述实施策略,可以实现每天超过 100 纳秒的 QM/MM 动态吞吐量。由于 QSimulate 的生产平台旨在出于商业原因最大限度地降低每次模拟的成本,通常在 GPU 上执行七个并发动力学模拟,以充分发挥 多实例 GPU (MIG) 的优势。图 3 编译了在 A100 和 H100 GPU 上测量的蛋白质配体系统 (由单元细胞中大约 25 K 个经典原子组成) 的吞吐量。

在 A100 中,由 74 个原子 (一种典型的小分子药物) 组成的量子力学区域的 QM/MM MD 模拟吞吐量已超过每天 100 纳秒。使用一个 H100 时,QM/MM MD 和 FEP 的吞吐量分别为 120 和 90 纳秒。当量子力学区域由 200 个原子组成时,A100 和 H100 之间的速度加快了50%以上。请注意,未来的创新和软件优化可以进一步改善这些时间。

NVIDIA A100 GPU acceleration of QM/MM FEP simulations against those on a CPU core.
NVIDIA H100 GPU acceleration of QM/MM FEP simulations against those on a CPU core.
图 3.在 NVIDIA A100 和 H100 GPU 上测量的 QM/MM MD (绿色) 和 FEP (蓝色) 模拟的吞吐量。使用 MIG 执行七个并发动力学模拟

我们将这些结果与之前由 QSimulate 团队与 NVIDIA Omniverse 合作开发的优化 CPU 实现进行比较,即 RIKEN 的 Sugita 小组。图 4 显示了结果。图 4 中的四列对应图 3 中的四个量子力学区域。值得注意的是,CPU 和 GPU 实现中的基础算法相同,包括截止、值和参数。由于 CPU 代码与 GPU 实现紧密优化,因此图 4 表示不同硬件架构之间的合理比较。

CPU 计时基准测试使用 Intel Xeon 8375C (Ice Lake) 的四个 CPU 核心进行测量,并为简便起见,将结果转换为一个 CPU 核心的结果。请注意,与此计时基准测试一样,QSimulate 的面向 CPU 硬件的生产平台使用四个 CPU 核心,主要原因是难以将 CPU 模拟扩展到更多线程和 CPU 核心。

GPU acceleration of QM/MM FEP simulations against those on a CPU core.
图 4.与 CPU 内核上的 QM/MM FEP 模拟相比,QM/MM FEP 模拟的 GPU 加速

即使是由 74 个量子力学原子组成的最小量子力学区域,我们也可以观察到,对于 A100 和 H100 卡,GPU 卡和 CPU 核心上的吞吐量比分别约为 140 和 170.随着量子力学区域的扩大,该比率变得更加显著。对于 200 个量子力学原子,使用 H100 卡的吞吐量等于使用 265 个 CPU 核心的吞吐量。

AI 时代的模拟

随着 AI 正在改变社会,基于物理的模拟比以往任何时候都更加重要。复杂工作流程(如药物研发中的工作流程)的数字化转型需要通过准确的模拟方法快速可靠地评估 AI 预测,然后再向 AI 模型提供更多数据。本文总结了 QSimulate 的进步(在毫秒级时间分步机制中首次基于量子力学的模拟),展示了 QSimulate 的领域知识和软件工程如何与 NVIDIA 硬件和软件结合使用,包括最新的 CUDA 工具包 和软件堆栈,以帮助实现下一个数字发现范式的飞跃。

欲了解更多信息并开始使用 GPU 加速的量子力学模拟进行药物研发,请访问 QUELO-G 产品页面,获取详细信息。

 

Tags