一种基于并行和张量分块的AlphaFold推理高性能实现方法技术

本发明专利技术公开了一种基于并行和张量分块的AlphaFold推理高性能实现方法，包括前处理部分和推理部分；所述前处理部分包括如下步骤：基于Ray的数据前处理并行策略；对于多条氨基酸序列的蛋白质复合物结构预测；所述推理部分包括如下步骤：优化的动态轴并行；自动张量分组；显存共享。本发明专利技术前处理部分：单条氨基酸序列的前处理速度为约原方法前处理速度的3倍；推理部分：能够更好地发挥GPU平台的性能，提高推理速度；可以有效地降低计算所需的显存，从而支持更长的输入氨基酸序列。而支持更长的输入氨基酸序列。而支持更长的输入氨基酸序列。

【技术实现步骤摘要】
一种基于并行和张量分块的AlphaFold推理高性能实现方法


[0001]本专利技术涉及人工智能
，尤其涉及一种基于并行和张量分块的AlphaFold推理高性能实现方法。

技术介绍

[0002]蛋白质是构成生命的基础，蛋白质的结构能够帮助我们理解其功能。由于使用传统的实验方法解出一个蛋白的结构可能需要花费数月甚至数年的时间，目前人类掌握的蛋白质结构仅有10万种左右，这在数以十亿计的的已知氨基酸序列中只占一小部分。
[0003]随着深度学习的发展，研究者们在积极探索使用深度学习进行蛋白质结构推理，其中的佼佼者是由DeepMind提出的AlphaFold。目前AlphaFold的推理结果已经达到了原子精度。但由于推理速度较慢且对于GPU显存要求较高，AlphaFold在工业上落地还并不成熟。
[0004]AlphaFold是由DeepMind公司针对蛋白质结构预测任务提出的深度学习模型。2018年，DeepMind首次提出了基于卷积神经网络的AlphaFold；在2020年，受Transformer模型启发，DeepMind提出了基于Transformer的AlphaFold，该模型在CASP14(第十四届全球蛋白质结构预测比赛)中摘得桂冠，并将预测误差降低至原子量级。使用AlphaFold进行蛋白质结构预测的步骤主要分为前处理和模型推理两部分：通过前处理得到输入氨基酸序列的多序列比对特征和配对表征，然后将这两组特征作为输入进行模型推理。
[0005]模型并行：为了应对超大模型对软、硬件提出的巨大挑战，研究人员提出了流水并行和张量并行，把大模型的参数分布到多个设备上进行计算。其中流水并行将模型按层结构进行分段，每个处理器从前一个处理器获得输入，再将输出结果传给下一个处理器。张量并行，如英伟达的Megatron
‑
LM系统，通过将模型层内的权重参数按行(横切并行)或按列(竖切并行)切分到不同处理器上，利用分块矩阵乘法，将一个运算分布到多个处理器上同时进行。张量并行需要引入通信从而获得正确的结果：竖切并行需要将不同的列进行跨处理器结果的拼接，横切并行需要进行跨处理器结果的加法。动态轴并行技术是根据AlphaFold的计算特征创新性开发的并行技术。不同于传统的张量并行，动态轴并行选择在AlphaFold的特征的序列方向上进行数据划分，并在并行模块间插入适当的通信。因为AlphaFold模型模块的输入由两个序列维度，动态轴并行技术在动态轴方向上进行切分，让计算方向在处理器上是连续的，所以模块内部不需要通信，只需要在计算方向发生变化的时候引入通信。动态轴并行对比张量并行，支持了AlphaFold中的更多计算模块，并需要更少的通信量。
[0006]AlphaFold进行蛋白质结构预测的步骤主要分为前处理和模型推理两部分：
[0007]1.前处理：
[0008]a.速度慢，时间成本高；在前处理中，AlphaFold使用Jackhmmer、HHblits、HHsearch等软件在Uniref90、BFD、PDB70等蛋白质数据库中搜索输入氨基酸序列的多序列比对特征和配对表征。由于这些数据库的量级均达到了数十GB甚至TB级，前处理过程耗时往往很长。尤其是在多条氨基酸序列的蛋白质复合物结构预测上，要对每条氨基酸序列单
独进行前处理，前处理所需时间与氨基酸序列的数量成正比。
[0009]2.模型推理：
[0010]a.速度慢，时间和经济成本高；因为AlphaFold使用了注意力机制，并且因为输入特征拥有两个序列维度，所以需要沿着两个方向分别进行注意力计算，从而导致计算复杂度较高。并且现有的推理方案没有针对AlphaFold的性能特征进行优化，使得其在GPU平台上无法发挥出硬件性能。从而使得AlphaFold实际推理过程时间和经济成本过高。
[0011]b.内存占用大，无法进行长序列推理；AlphaFold输入包含多序列比对特征和配对表征，这两种表征都包含两个序列维度，在进行注意力计算的同时会产生大量的中间张量。随着序列长度的增加，中间显存的大小呈指数型增长。在序列较长时，中间显存就是达到几十GB甚至上百GB。而目前最先进的GPU也只有数十GB的显存，导致长序列推理十分困难。

技术实现思路

[0012]基于
技术介绍
存在的技术问题，本专利技术提出了一种基于并行和张量分块的AlphaFold推理高性能实现方法。
[0013]本专利技术提出的一种基于并行和张量分块的AlphaFold推理高性能实现方法，包括前处理部分和推理部分；
[0014]所述前处理部分包括如下步骤：
[0015]S11基于Ray的数据前处理并行策略：在单条氨基酸序列的前处理过程中，只有极少数数据库搜索之前存在依赖关系，使用Ray将不存在依赖关系的数据库搜索之间实现了并行；
[0016]S12对于多条氨基酸序列的蛋白质复合物结构预测：将每一条氨基酸序列作为一个工作流，每一个工作流下又分出多条工作流进行序列内前处理并行；
[0017]所述推理部分包括如下步骤：
[0018]S21优化的动态轴并行：在模型推理的过程中，主要由Embedding，Evoformer和Structure Module组成，将动态轴并行扩展到了AlphaFold模型的全部推理部分，均采用切分中间计算结果的思路，动态地在特征序列方向上进行数据划分，并插入适当的通信保证结果正确性；
[0019]S22自动张量分组：通过张量分组技术可以将张量从与计算无关的维度中拆开，依次进行计算，有效地降低计算时的峰值显存，通过自动张量分组策略，根据GPU的显存和输入序列长度调整是否分组以及分组大小，降低张量分组对推理速度的影响；
[0020]S23显存共享：通过显存共享技术减少显存占用，在传入参数时，使用列表进行张量传入，达到显存共享的功能，在推理中零成本的降低显存消耗。
[0021]优选的，所述推理部分使用动态轴并行将计算分散到多张GPU上。
[0022]本专利技术中，所述一种基于并行和张量分块的AlphaFold推理高性能实现方法，具有如下有益效果：
[0023]1.前处理部分：
[0024]a.单条氨基酸序列的前处理速度为约原方法前处理速度的3倍。
[0025]b.多条氨基酸序列的前处理时间取决于其中最长序列的前处理时间，假设共有n条氨基酸序列作为蛋白质复合物结构预测的输入，平均情况下本专利技术的前处理速度约为原
方法前处理速度的3*n倍。
[0026]2.推理部分：
[0027]a.使用动态轴并行将计算分散到多张GPU上，能够更好地发挥GPU平台的性能，提高推理速度。
[0028]b.使用张量分块技术和显存共享技术，可以有效地降低计算所需的显存，从而支持更长的输入氨基酸序列。
附图说明
[0029]图1为本专利技术提出的一种基于并行和张量分块的AlphaFold推理高性能实现方法的流程示意图；
[00本文档来自技高网...

【技术保护点】

【技术特征摘要】
1.一种基于并行和张量分块的AlphaFold推理高性能实现方法，其特征在于，包括前处理部分和推理部分；所述前处理部分包括如下步骤：S11基于Ray的数据前处理并行策略：在单条氨基酸序列的前处理过程中，只有极少数数据库搜索之前存在依赖关系，使用Ray将不存在依赖关系的数据库搜索之间实现了并行；S12对于多条氨基酸序列的蛋白质复合物结构预测：将每一条氨基酸序列作为一个工作流，每一个工作流下又分出多条工作流进行序列内前处理并行；所述推理部分包括如下步骤：S21优化的动态轴并行：在模型推理的过程中，主要由Embedding，Evoformer和Structure Module组成，将动态轴并行扩展到了AlphaFold模型的全...

【专利技术属性】
技术研发人员：程盛淦，路广阳，赵轩磊，马千里，方佳瑞，李升桂，卞正达，李永彬，柳泓鑫，麦思琪，刘育良，黄海晨，娄宇轩，吴俊铭，陈巍文，
申请(专利权)人：北京潞晨科技有限公司，
类型：发明
国别省市：