[发明专利]一种基于并行和张量分块的AlphaFold推理高性能实现方法

说明书

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

技术领域

本发明涉及人工智能技术领域，尤其涉及一种基于并行和张量分块的AlphaFold推理高性能实现方法。

背景技术

蛋白质是构成生命的基础，蛋白质的结构能够帮助我们理解其功能。由于使用传统的实验方法解出一个蛋白的结构可能需要花费数月甚至数年的时间，目前人类掌握的蛋白质结构仅有10万种左右，这在数以十亿计的的已知氨基酸序列中只占一小部分。

随着深度学习的发展，研究者们在积极探索使用深度学习进行蛋白质结构推理，其中的佼佼者是由DeepMind提出的AlphaFold。目前AlphaFold的推理结果已经达到了原子精度。但由于推理速度较慢且对于GPU显存要求较高，AlphaFold在工业上落地还并不成熟。

AlphaFold是由DeepMind公司针对蛋白质结构预测任务提出的深度学习模型。2018年，DeepMind首次提出了基于卷积神经网络的AlphaFold；在2020年，受Transformer模型启发，DeepMind提出了基于Transformer的AlphaFold，该模型在CASP14(第十四届全球蛋白质结构预测比赛)中摘得桂冠，并将预测误差降低至原子量级。使用AlphaFold进行蛋白质结构预测的步骤主要分为前处理和模型推理两部分：通过前处理得到输入氨基酸序列的多序列比对特征和配对表征，然后将这两组特征作为输入进行模型推理。

模型并行：为了应对超大模型对软、硬件提出的巨大挑战，研究人员提出了流水并行和张量并行，把大模型的参数分布到多个设备上进行计算。其中流水并行将模型按层结构进行分段，每个处理器从前一个处理器获得输入，再将输出结果传给下一个处理器。张量并行，如英伟达的Megatron-LM系统，通过将模型层内的权重参数按行(横切并行)或按列(竖切并行)切分到不同处理器上，利用分块矩阵乘法，将一个运算分布到多个处理器上同时进行。张量并行需要引入通信从而获得正确的结果：竖切并行需要将不同的列进行跨处理器结果的拼接，横切并行需要进行跨处理器结果的加法。动态轴并行技术是根据AlphaFold的计算特征创新性开发的并行技术。不同于传统的张量并行，动态轴并行选择在AlphaFold的特征的序列方向上进行数据划分，并在并行模块间插入适当的通信。因为AlphaFold模型模块的输入由两个序列维度，动态轴并行技术在动态轴方向上进行切分，让计算方向在处理器上是连续的，所以模块内部不需要通信，只需要在计算方向发生变化的时候引入通信。动态轴并行对比张量并行，支持了AlphaFold中的更多计算模块，并需要更少的通信量。

AlphaFold进行蛋白质结构预测的步骤主要分为前处理和模型推理两部分：

1.前处理：

a.速度慢，时间成本高；在前处理中，AlphaFold使用Jackhmmer、HHblits、HHsearch等软件在Uniref90、BFD、PDB70等蛋白质数据库中搜索输入氨基酸序列的多序列比对特征和配对表征。由于这些数据库的量级均达到了数十GB甚至TB级，前处理过程耗时往往很长。尤其是在多条氨基酸序列的蛋白质复合物结构预测上，要对每条氨基酸序列单独进行前处理，前处理所需时间与氨基酸序列的数量成正比。

2.模型推理：