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

权利要求书

1.一种基于并行和张量分块的AlphaFold推理高性能实现方法，其特征在于，包括前处理部分和推理部分；

所述前处理部分包括如下步骤：

S11基于Ray的数据前处理并行策略：在单条氨基酸序列的前处理过程中，只有极少数数据库搜索之前存在依赖关系，使用Ray将不存在依赖关系的数据库搜索之间实现了并行；

S12对于多条氨基酸序列的蛋白质复合物结构预测：将每一条氨基酸序列作为一个工作流，每一个工作流下又分出多条工作流进行序列内前处理并行；

所述推理部分包括如下步骤：

S21优化的动态轴并行：在模型推理的过程中，主要由Embedding，Evoformer和Structure Module组成，将动态轴并行扩展到了AlphaFold模型的全部推理部分，均采用切分中间计算结果的思路，动态地在特征序列方向上进行数据划分，并插入适当的通信保证结果正确性；

S22自动张量分组：通过张量分组技术可以将张量从与计算无关的维度中拆开，依次进行计算，有效地降低计算时的峰值显存，通过自动张量分组策略，根据GPU的显存和输入序列长度调整是否分组以及分组大小，降低张量分组对推理速度的影响；

S23显存共享：通过显存共享技术减少显存占用，在传入参数时，使用列表进行张量传入，达到显存共享的功能，在推理中零成本的降低显存消耗。

2.根据权利要求1所述的一种基于并行和张量分块的AlphaFold推理高性能实现方法，其特征在于，所述推理部分使用动态轴并行将计算分散到多张GPU上。