前两期文章展示了幻方AI对Alphafold的优化，采用特征预处理和特征裁切两种方式提高了Alphafold数据处理性能，通过并行训练加速神器进一步提高模型的训练速度，将Alphafold深度融合进幻方AI的集群特点中，发挥最大的计算效能。
那么从整体上看，在幻方萤火二号上训练Alphafold还有哪些需要我们注意的，以及未来同类型深度学习模型该如何优化？就这些话题，本期文章将和大家聊聊幻方AI的一些思考。
模型仓库：https://github.com/HFAiLab/alphafold-optimized
总体优化效果
首先，我们看一下在进行了前两期优化后，最终Alphafold在幻方萤火二号上的训练效果。我们使用16个节点共128张A100显卡，对将近7万条蛋白质序列进行了比较完整的Alphafold训练，得到了如下的实验结果：
| 平均GPU利用率 | 90.14% |
|---|---|
| 平均每Epoch时间（min） | 84.15 |
| 平均每次迭代时间 （s） | 11.32 |
| 每Epoch开始时数据加载时间（s） | 197 |
| 总迭代数 | 35680 |
| 总训练时长（h） | 112 |
训练时的GPU利用率曲线如下：
训练时的GPU利用率大多数时候在90%左右。可见经过上述的多种优化后，Alphafold 已经能够在幻方AI平台上高效进行训练，而有中间偶尔出现的利用率骤降则是训练开始时的数据加载耗时。
另外，从CPU利用率图可以看出，整个训练过程中CPU的利用率长期保持在大于50%（35/64核）的情况，说明Alphafold训练时确实对CPU算力有非常大的需求。
多卡训练中的一些踩坑
在优化Alphafold模型的过程中，特别是将单机单卡拓展至多机多卡时，我们遇到了不少的问题。
1. OpenFold代码问题
Deepmind没有开源Alphafold训练代码，现在github主流的Alphafold训练代码是由其他三方实现的，比如OpenFold（https://github.com/aqlaboratory/openfold），其通过PyTorch比较完整地复现了Deepmind开源的Alphafold 2（https://github.com/deepmind/alphafold）。
OpenFold的复现使用了比较复杂的代码构造模型结构，其中不同蛋白质序列数据在训练时可能会有不完全一样的前向传播路径，这导致在多卡训练时，一些机器上有一部分参数没有参与梯度计算，而Pytorch DDP提供的find_unused_parameters方法并不能正确的找到这部分参数，因此会时不时报错，无法正常进行多卡训练。
经我们一段时间问题排查发现，是Loss Function部分中存在部分不参与计算的Loss，需要将其手动注释掉后方能正常进行多卡训练。
2. 环境问题
另外一个比较容易踩坑的就是环境问题。幻方AI平台内置了PyTorch等环境，适配了常见的深度学习机器学习。而Alphafold本身作为一个深度学习领域的新方向，所依赖的环境必然与常规环境不同。
因为幻方AI训练平台采用的是计算存储分离的方式，环境统一配置，不托管镜像，所以Alphafold的运行环境需要一步步重头安装，这里不可避免地会发生与原有环境版本冲突，需要补充系统库等一系列问题。
因此，我们通过 haienv 完全创建一个全新的虚拟环境，并安装了 Alphafold 的依赖项。然而，在进行多机多卡训练时会出现NCCL Error的问题。经过排查后发现：由于 haienv 创建了全新的环境，使用了与原先不同的全新Python bin文件，而新建的Python并没有获得IPC Lock，因此在多机多卡训练时NCCL会出现问题。需要在sudo权限下使用setcap指令setcap cap_ipc_lock+iep /venv/openfold38_0/bin/python来解决问题。
最后，我们将Alphafold所依赖的环境单独列出来，供用户调用。
局限与展望
虽然经过我们的优化后模型本身训练的效率已经很高，GPU利用率达到90%以上，但我们仍然对模型性能做出了一定的取舍，并且模型内部仍然存在不少优化的空间。一方面来说我们已经完成的优化都基于不修改模型内部结构的前提，因此对模型内不同部分的差异造成的额外开销难以处理；另一方面来说训练时显存的利用还不够理想：A100 40GB的显存只能使用到21GB而有接近一半的空余，但batch_size只能设置为1，这造成了大量浪费。
从以上两个角度入手，后续Alphafold模型还可以进行的优化有以下三个方向：
- 结构简化：Alphafold 2本身虽然达到了很好的蛋白质预测准确率，但其复杂的模型结构和极高的资源占用使得复现和部署模型推理都有很高的门槛。因此，仍然可以考虑对模型内部进行细致的Profile，寻求对模型结构进行一定程度的简化。
- 算子优化：幻方AI在算子优化，尤其是Transformer类模型的算子优化方面有很多的积累，而Alphafold作为一个同样基于Transformer结构的模型，还没有用到这些算子优化技术。因此，在这里也有进一步优化的空间。
- 模型并行：Alphafold目前训练时的显存占用率仍然不高，单卡batch_size也只能设置为1。可以考虑尝试在Alphafold上应用更多的模型并行方式，例如和Pipeline结合来增大训练时的单卡batch_size，提升显存使用效率。
欢迎各方数据科学家与开发者们一同共研共建。