GPU 加速药物研发与基因组学分析

作者：Baidu AICLOUD

2023-04-24
北京
本文字数：3334 字
阅读完需：约 11 分钟

本文整理自 2022 年 12 月的智算峰会 · 生命科学分论坛上的同名分享。

大家好，我是来自 NVIDIA GPU 解决方案团队的翟健。此次分享介绍了 NVIDIA 如何利用 GPU 和加速软件推动 AI 驱动的药物研发与基因组学分析，包括 NVIDIA 与百度智能云在赋能药物研发、基因测序分析方面的合作。

1. NVIDIA SDK 平台：CLARA

CLARA 是 NVIDIA 在医疗方向的 SDK 平台，它借助 NVIDIA 基础软硬件平台，为医疗用户提供了医疗设备、制药、影像、基因以及智慧医院等五个方面的能力，针对具体的应用场景提供完备的加速能力。

其中的 CLARA Discovery 是 CLARA 平台下面向药物研发场景的重要解决方案，可以基于 NVIDIA 的 AI 和 HPC 能力来辅助加速药物研发的工作流程。

目前该方案已经纳入到了百度智能云赋能药物研发的解决方案中。

2. 使用 AI/HPC 来加速药物研发的工作流程

医疗领域的前期研发过程当中会涵盖如下几个过程：靶点发现、化学分子生成、蛋白质性质结构预测、药物分子与蛋白质进行对接打分，构建自由能 FEP 等。

整个制药行业涉及到两个学科，结构生物学与计算化学。这两个学科都涉及到传统的 HPC 与新兴的 AI 方法，而且 AI 的方法相较于 HPC 的方法会有比较明显的速度上的优势。

目前 AI 正在颠覆整个药物研发的过程，包括在靶点发现方面、虚拟筛选以及分子生成、结构预测，甚至在临床上应用的自然语言模型，都是 AI 加速的体现。

Transformer 模型正在逐渐地应用在药物研发领域中。下图左边的四张小图来自于 MegaMolBART 与 ProtTrans 一系列突破性的论文，表明基于 Transformer 的预训练模型可以有效地加速分子生成和蛋白质的结构预测。

下图右上角展示的是 Transformer 模型的 Encoder-Decoder 的架构形式，它们可以用在诸如右下方的几个典型领域，包括：小分子模型的生成、反应序列的预测、蛋白质结构预测、生物医学领域的 NLP 以及图像分析等。

基于在大语言模型上的经验，NVIDIA 推出了 BioNeMo 的解决方案，它是一款可以在云端进行训练和部署的服务框架，主要面向有大语言模型需求的药物研发人员。

此外，它基于 NVIDIA 的 GPU 硬件也做了很好的优化工作，并且提供了多种预训练模型，支持云原生的服务，极大地加速了药物研发的工作流程。

下图展示的是 BioNeMo 的逻辑架构，最底层的是 NVIDIA 加速的计算平台。第二层是 NVIDIA 的大语言模型平台 NeMo Megatron，是 BioNeMo 的快速训练和部署的重要保证。

BioNeMo 提供了多种预训练的 Transformer 类的模型，分别针对化学分子生成、蛋白质结构序列预测、DNA 的 embedding 等几个方面。最终这些都服务于我们最上层中药物研发领域里几个典型应用场景。

BioNeMo 的三个典型特点是：支持分子、蛋白质和核苷酸的 SMILES 和 FASTA 的表征；含有多种预训练模型，像我们刚才提到的 MegaMolBART 等；可以在云端部署相应的可视化界面的服务。

这里跟大家分享一个案例，Vyasa 是一个面向药物研发的 AI 解决方案提供商，他们在方案中整合了 BioNeMo 中的 MegaMolBART 模型，从而实现了终端用户在本地和云端都可以利用 GPU 对药物分子生成过程进行加速。

除了在分子生成、蛋白质结构性质预测方面，NVIDIA 提供了很好的加速优化。在虚拟筛选和仿真的过程当中，NVIDIA 也协同了众多的开源社区、高校，加速了传统 HPC 领域当中的对接、分子动力学和量子化学里的常用软件。

这张图中展示的就是我们在三个领域当中常用的一些软件。这些软件 NVIDIA 都提供了相应的 GPU 加速版本，大家可以在 NVIDIA 的 NGC 平台进行下载：https://catalog.ngc.nvidia.com 。

下面展示的是分子动力学软件 Gromacs 的 GPU 版本的性能评测。可以看到，在 NVIDIA 的 A100 和 V100 GPU 上，Gromacs 都同比 CPU 获得了极好的加速。

类似的，这一页展示的是量子化学软件 VASP CPU-GPU 的性能对比。NVIDIA 的 V100 和 A100 GPU 同样都获得了极高的加速。

3. GPU 加速基因组学分析管线

接下来让我们介绍一下 NVIDIA GPU 在加速基因组学方面的方案。

在 22 年春季的 NVIDIA GTC 大会上，来自于斯坦福大学的团队介绍了他们如何基于 NVIDIA GPU 打破了基因测序的世界纪录。他们将人类的基因测序缩短到了 7 小时 18 分钟，完成了湿实验和在计算机上的数据分析等过程。

而在基因测序方面，一般包括如下三个环节：第一，通过测序仪得到数据之后进行的一级分析过程，完成四分类任务；第二，在计算机上完成一致性对比处理、变体识别等二级分析过程；第三，最后的三级分析则是对大量的数据进行处理。

这三个过程目前都是可以用 GPU 实现加速的，分别可以通过 GPU 加速的 TensorFlow、PyTorch，以及 TensorRT 进行一级分析加速。通过 CLARA Parabricks 对二级分析进行加速。利用 RAPIDS、MONAI 等可以加速三级分析。

4. 利用 GPU 加速高通量、高精度的 DNA 和 RNA 测序分析

下面让我们来介绍一下二级分析的软件 CLARA Parabricks。

CLARA Parabricks 是一款利用 GPU 加速高通量、高精度的 DNA 和 RNA 测序分析工具，主要用于人类基因组学分析、癌症基因筛查、RNA 测序分析等。目前其中含有 60 多个工具模块，包括基因数据的一致性比对、金标准处理和质量把控、高精度变体识别等范畴类的多种常用工具。

下图展示的是 CLARA Parabricks 目前支持的工具模块的部分内容，基本上涵盖了主流的基因测序二级分析中的大部分工具。

使用 CLARA Parabricks 可以实现对典型的应用的加速，它是针对金标准的种系、体细胞和 RNA 的加速工具。而且目前使用 Ampere 架构的 GPU 可以实现 80 倍的加速，精度方面也能够保证，且具有比较灵活的工作流程。

那么谈到这里就要说一下 CLARA Parabricks 的几大特点了。因为它是模块化的工具，所以可以通过各种倾向性的组合，可以灵活选择 CLARA Parabricks 当中的各种模块。

同时刚才提到它的工作流程灵活，是因为它对主流的基因组学分析中的 workflow 管理器都支持，包括 WDL，nextflow 等。此外，它还对 Google 的 DeepVariant 1.4 版本的变体识别工具也做了很好的支持。

接下来我们来看看 Parabricks 的 benchmark。可以看到，列举的几个模块在 GPU 上，同比 CPU 都具有非常好的加速效果。

而且在 A100 上可以最高实现 80 倍的加速。

这一页展示的就是分别在 2 张、4 张和 8 张 V100 上同比 CPU 实现加速的效果。同时，在癌症的基因筛查中，端到端的流程也可以在 GPU 上实现比较明显的加速。

2022 年春季 GTC 大会上，NVIDIA 发布了 Hopper 架构的 GPU。该架构的 GPU 对于动态编程做了很好的加速，单就动态编程这一特性，Hopper 架构就比上一代的 Ampere 架构有了 7 倍的理论加速。而在基因组学分析中，动态编程技术是需要被频繁使用到的。所以，Hopper 架构的 GPU 可以给基因测序带来重大的收益。

这里再跟大家分享一个案例，就是英国的 BioBank 这家公司，他们要处理 50 万个外显子，这些在 CPU 上需要 1 个小时才能得到结果，而 GPU 将这个过程仅仅缩短到了 5 分钟，成本下降了 60%。

这里展示了 CLARA Parabricks 的资料参考页，感兴趣的听众可以登陆 CLARA Parabricks 的网页了解更多的内容。

同时，CLARA Parabricks 已经可以在百度智能云上使用了，在云上的 GPU 最佳实践专栏中还提供了 CLARA Parabricks 的使用方式和详尽介绍，大家可以访问：https://cloud.baidu.com/doc/GPU/s/pl6vzliqu?track=infoQ03 。