模型高性能优化 | 两步数据预处理让Alphafold 提速一倍

实验｜Daylight       算力支持｜幻方萤火 

若说2021年人工智能学术界最令人振奋的成果，那么Alphafold可谓当之无愧。Alphafold2在CASP14 蛋白质预测挑战上取得了远远超出同类模型的准确率，并首次将蛋白质结构预测的精度提高到了原子级别——已经接近了实验测量的水准。

幻方AI团队在Alphafold2推出不久就在萤火二号成功将Alphafold2训练运行了起来，详情如我们上一期文章。然而，由于Alphafold2本身模型的复杂性和开源代码的缺陷，并且幻方AI的训练平台也比较特殊，使得一开始复现资源利用效率非常低下，训练成本不可接受。

---

为了提高Alphafold的训练效率，帮助科研人员与开发者们降低研究门槛，幻方AI团队对Alphafold2模型进行了大量优化，大幅提高了模型训练的效率。我们的工作可以被概括为以下三点： 

1. 重构代码，解决开源版本中无法多机多卡训练的问题，利用数据并行大幅提高Alphafold2训练效率，降低复现门槛；

2. 优化数据处理逻辑，大幅增加GPU的利用率，从40%左右提升到了90%以上；

3. 利用hfai高性能工具改造模型，将Alphafold2深度融合进幻方萤火二号中，大幅提高训练性能。

本次文章先和大家聊聊，如何优化蛋白质数据的处理，以获得高效的训练性能。

---

输入优化

Alphafold的输入特征可以分为这样三类： 

1. 蛋白质序列本身 

2. Multi-Sequence Alignment(MSA)多序列对齐特征 

3. Template蛋白质模板特征

它们在不同的蛋白质序列长度之间差异非常大，处理所需时间可以从几秒到上千秒不等，这会极大的影响模型训练的效率。例如在一次20分钟的训练过程中，单个长样本的处理就可能阻塞整体训练10几分钟。而在多机多卡大批次场景下，每个批次内有部分数据的处理时间极长几乎成了必然事件。

如上图所示，在训练Alphafold的初始阶段，输入的特征是通过序列数据库搜索，生物学工具（JackHmmer/HHSearch/...）计算而得出，需要依赖大量的CPU计算能力支持。在整个Alphafold模型训练的过程中，CPU算力往往成为了阻碍Alphafold训练性能提升的重要瓶颈。

针对上述问题，幻方AI采用特征预处理与特征裁切两种方案对输入进行优化。简单来说，特