模型高性能优化 | 如何实现一行代码 30% Alphafold 提速

实验｜Daylight       算力支持｜幻方AIHPC 

上一期关于 Alphafold 训练提速的文章讲到，幻方 AI 通过优化数据处理，采用特征预处理和特征裁切两种方式通过降低 CPU 开销，避免数据加载阻塞训练的方式提高了 Alphafold 整体的训练性能。

本期，我们将尝试用另外一种方式，对这个“AI 顶流明星模型”进行二次提速。这次我们的实验思路是用上幻方 AI 的并行训练加速神器之二 hfreduce 以及 hfai.nn 算子加速，它们是否能对 Alphafold 整体的训练进一步加速呢？本期文章将通过实验给出答案。扫描文末二维码，即可申请试用幻方的AI 训练平台。

---

hfreduce

之前的文章《幻方萤火 | 模型并行训练工具 hfreduce》提到过，由于幻方 AI 集群的架构特点， Nvidia 官方提供的 NCCL 工具并不能充分发挥萤火二号超算集群的通信带宽，因此幻方 AI 自研了 hfreduce 工具优化显卡间的通信，提高多机多卡的并行训练效率。使用 hfreduce 也只需要简单地将 pytorch DDP 替换为 hfai DDP 即可：

model = hfai.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

对于 Alphafold 模型，其本身的特点是前向传播路径比较长，而反向传播路径较短，具有比较鲜明的“参数规模不大，但计算量大的特点”。我们采用数据并行将数据分散到不同的显卡上进行并行加速，以加快训练效率。因此我们为 Alphafold 的训练引入了 hfreduce，利用 hfai.ddp 工具进行数据并行，尝试进一步推高 Alphafold 在萤火二号超算集群上的计算效率。

然而在实践中我们发现，hfreduce 给 Alphafold 训练带来的训练提升在绝对值上并不大。这是什么原因呢？在随后的实验部分中我们结合实验结果来分析这一问题。  

1►实验设定

我们设计并进行了三组实验，以对比使用 hfreduce 后的加速效果。实验的设定如下表所示：

我们希望通过多卡和单卡训练的耗时对比，来直接看到多卡通信带来的额外开销。因此，我们对 3 种训练方式测试了其 200 次迭代的平均 backward 耗时。

其中，“特征裁切：最大”代表了对特征进行了尽可能大的裁切，目的是避免因特征处理导致多卡训练时部分 GPU 上的梯度同步长期处于等待状态，影响结果的准确性。使用 hfreduce 时由于梯度同步被单独拆出，因此需要将两个函数合并统计梯度计算+梯度同步的总时间。

此外，训练的基本参数还包括：

• 训练样本数量：73472

• Batch Size: 1

2►实验结果

根据上文中的实验设定，我们测试 NCCL 并行加速和 hfreduce 并行加速的完整训练结果如下表所示：

直接对比单次迭代的总耗时，从结果来看在 Alphafold 的训练场景下 hfreduce 和 NCCL 的不同选择并不会带来显著的效率差异。理论上来说，hfreduce 比 NCCL 更加适合萤火二号的集群架构，一定能带来比 DDP 更快速的梯度同步。

 但在 Alphafold 的实验结果中，两次实验中每个 step 的迭代速度几乎完全相同。接下来我们将在理论上对这种实验结果和预期的差异做出一些分析。

对于这种现象，最主要的原因是 Alphafold 独特的模型结构和蛋白质数据的特点使得模型虽然需要很大量的 GPU 运算，却只需要传输很少的梯度数据。Alphafold 的模型参数只有 93.2M，参数量不到 1 亿，使用 FP32 单精度的情况下占用显存 372.8MB。在前向传播时 Alphafold 会将模型中的主要部分(Evoformer)复制 4 次，这会带来一个很长的前向传播通路，使得时间复杂度大大增加。然而，最终模型需要更新的参数依然只有原先的数量，因此显得梯度同步的开销在整个迭代中显得微不足道了。

那么，既然梯度同步的耗时不显著，我们是否可以通过一些 profile 方法拆分出每个部分的耗时，进行定量分析优化呢？我们对 backward 部分进行更细致的测试，最终测试结果如下表所示：

这里，不考虑多卡时因数据处理进度不同产生的等待，那么多卡训练时的 backward 时间减去单卡训练的 backward 时间既可以认为是多机多卡带来的额外通信开销。

上表的结果中可以看到，在萤火二号上使用 DDP 训练时，选择 hfreduce 能够比 NCCL 节省 30% 的通信开销，这是一个非常可观的数字，说明 hfreduce 确实更适合在幻方萤火二号超算平台上进行多机多卡训练的加速。同时，90% 以上的耗时在 forward 的计算上，特别是 LayerNorm 和 Attention 的算子计算开销比较大，这里需要进一步使用 hfai 的优化算子进行进一步提升。

hfai.nn

上一期的文章《幻方萤火 | 性能卓越的深度学习算子 hfai.nn》与大家分享了：幻方 AI 针对 Pytorch 框架进行了深度优化，结合萤火二号的集群特点，对一些常用的 AI 算子重新设计研发，进一步提升了模型整体的训练效率。其使用非常简单，只需要在原有的任意模型训练代码中加入一行：

model = hfai.nn.to_hfai(model)

hfai 会自动扫描您的代码，替换为优化后的算子。hfai.nn 提供了对当前主流深度学习模型中常用的 MultiHeadAttention，LayerNorm，LSTM 等结构中算子的深度优化，能够大幅加速模型中这些算子的运算。我们尝试在 Alphafold 中引入 hfai.nn 算子库，测试其能够对模型训练产生多大的增益效果。

1►Alphafold GPU 开销分析

由于 hfai.nn 的算子加速取决于模型中各类型运算占总 GPU 开销的比重情况，因此我们首先尝试使用 Pyorch Profiler 工具对 Alphafold 模型使用标准 torch 算子时的 GPU 运算时间进行一些分析。在一次总耗时为 12.6s 的迭代中，主要开销情况如下：

仅前 5 类运算占了 Alphafold 训练开销的 70%，而其中 LayerNorm 与 Attention 中的矩阵运算又是最主要的耗时来源：仅 LayerNorm 运算就占了总耗时的 30% 左右。Alphafold 相对一般的 BERT 等 Transformer 类模型更为复杂，使用了自己实现的 Attention，因此在注意力运算上无法获得 hfai.nn 的加速。但同样占据耗时大头的 LayerNorm 却能够使用到 hfai 算子加速，因此可以预期 hfai.nn 能给 alphafold 带来较好的加速效果。

为了探究具体加速效果，我们首先尝试对 Alphafold 中 LayerNorm 的使用情况进行一些分析。一般来说输入 Tensor 的形状往往会对算子加速的效果产生较大的影响，因此在这里我们首先尝试对输入模型中的 LayerNorm 层的不同的 Tensor 形状进行一些理论上的分析。在一个典型 Transformer 类模型中， LayerNorm 的输入形状往往是 [BatchSize, SeqLen, EmbDim]。在 Alphafold 中由于其模型结构与简单的 Transformer 模型有较大差异，输入的张量形状会更为特殊一些。模型中 LayerNorm 层的不同 Input Shape 对应的出现频率可见下表：

可见在大多数输入时，hfai.nn 提供的 LayerNorm 算子都能取得相比 torch.nn 数倍的大幅性能提升，只在极少数形状下性能可能会稍差于 torch。另外根据上表中不同输入 Tensor 形状的占比可以计算得，使用hfai.nn 在 Alphafold 中可以在 LayerNorm 的 GPU 开销上取得总计 350% 的加速。由于 Alphafold 中的 LayerNorm 使用频率很高，且在模型的不同部分都有使用，因此可以预计对模型整体也能有较明显的加速效果。

2►实际训练加速结果

在理论分析了 hfai.nn 能给 Alphafold 带来的加速幅度后，我们也希望能了解在实际训练时 Alphafold 能够从 hfai.nn 中获得的收益。因此我们在与 Alphafold 预训练时相同设定的真实场景下使用 128 卡进行了并行训练，分别测试了使用 torch 算子和 hfai 算子进行 Alphafold 训练时的训练耗时情况，结果如下表：

相比于使用 torch.nn 算子进行训练，使用 hfai.nn 时模型整体训练中单次迭代的平均时长能够从 11.26 秒减少到 8.64 秒。可见在 Alphafold 训练时只需要添加一行代码引入 hfai.nn，就足足能够获得 30% 的训练性能提升。值得注意的是，由于 Alphafold 中自定义了 Attention 实现没有使用标准的 nn.MultiHeadAttention，在这里能够获得的 hfai.nn 算子加速效果其实只来自于 LayerNorm 运算减少的耗时。由此可见在常见的标准 Transformer 类模型中，使用 hfai.nn 将更加容易获得大幅的性能提升。

总 结

通过上述两个实验，我们验证了 hfreduce 和 hfai.nn 算子对 Alphafold 模型的加速效果。他们分别表现出了不同的加速能力，说明在模型优化的过程中我们需要根据模型的结构特性选择合适的加速方案。这里，幻方 AI 将持续加大研发，提供更多的优化方案和工具。

---

幻方 AI 团队， 紧跟  AI  研究的前沿浪潮，致力于用领先算力助力 AI 实验落地与价值创造，降低超大算力需求的使用门槛，欢迎各方数据科学家与开发者们一同共建。

点击下方链接，幻方AI BLOG更多干货奉上

幻方 | 技术博客