上一期文章讲到，幻方AI通过优化数据处理，采用特征预处理和特征裁切两种方式提高了Alphafold整体的训练性能。众所周知，幻方AI有很多并行训练加速神器，比如hfreduce，3FS，hfai.nn算子库等，它们是否能对Alphafold整体的训练进一步加速呢？本期文章将就这些问题进行试验。
hfreduce
之前的文章《幻方力量 | 模型并行训练工具：hfreduce》提到过，由于幻方AI架构的特点，Nvidia官方提供的NCCL工具并不能充分发挥萤火二号集群的通信带宽，因此幻方AI自研了hfreduce工具优化显卡间的通信，提高多机多卡的并行训练效率。使用方式非常简单，只需要将pytorch DDP替换为hfai DDP即可：
model = hfai.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])对于Alphafold模型，其本身的特点是前向传播路径比较长，而反向传播路径较短，具有比较鲜明的“参数规模不大，但计算量大的特点”。我们采用数据并行将数据分散到不同的显卡上进行并行加速，以加快训练效率。因此我们为Alphafold的训练引入了hfreduce，利用hfai.ddp工具进行数据并行，尝试进一步推高Alphafold在萤火二号集群上的计算效率。
然而在实践中我们发现，hfreduce给Alphafold训练带来的训练提升在绝对值上并不大。这是什么原因呢？在随后的实验部分中我们结合实验结果来分析这一问题。
模型仓库：https://github.com/HFAiLab/alphafold-optimized
实验设定
我们设计并进行了三组实验，以对比使用hfreduce后的加速效果。实验的设定如下表所示
| 序号 | GPU数 | DDP Backend | 迭代次数 | 耗时统计对象 | 特征裁切 |
|---|---|---|---|---|---|
| 1 | 1 | 无 | 200 | loss.backward() | 最大 |
| 2 | 128 | NCCL | 200 | loss.backward() | 最大 |
| 3 | 128 | hfreduce | 200 | loss.backward() + model.sync_grad() | 最大 |
我们希望通过多卡和单卡训练的耗时对比，来直接看到多卡通信带来的额外开销。因此，我们对3种训练方式测试了其200次迭代的平均backward耗时。
其中，“特征裁切：最大” 代表了对特征进行了尽可能大的裁切，目的是避免因特征处理导致多卡训练时部分GPU上的梯度同步长期处于等待状态，影响结果的准确性。使用hfreduce时由于梯度同步被单独拆出，因此需要将两个函数合并统计梯度计算+梯度同步的总时间。
此外，训练的基本参数还包括：
- 训练样本数量：73472
- Batch Size: 1
实验结果
根据上文中的实验设定，我们测试NCCL并行加速和hfreduce并行加速的完整训练结果如下表所示：
| NCCL | hfreduce | |
|---|---|---|
| 每个step平均时长（s） | 11.27 | 11.26 |
直接对比单次迭代的总耗时，从结果来看在Alphafold的训练场景下HFReduce和NCCL的不同选择并不会带来显著的效率差异。理论上来说，HFReduce比NCCL更加适合萤火二号的集群架构，一定能带来比DDP更快速的梯度同步。但在Alphafold的实验结果中，两次实验中每个Step的迭代速度几乎完全相同，这是为什么呢？
对于这种现象，最主要的原因可能是Alphafold独特的模型结构和蛋白质数据的特点使得模型虽然需要很大量的GPU运算，却只需要传输很少的梯度数据。Alphafold的模型参数大小只有93.2M，模型的参数量不到1亿，使用FP32单精度的情况下占用显存372.8MB。在前向传播时Alphafold会将模型中的主要部分（Evoformer）复制4次，这会带来一个很长的前向传播通路，使得时间复杂度大大增加。然而，最终模型需要更新的参数依然只有原先的数量，因此显得梯度同步的开销在整个迭代中显得微不足道了。
那么，既然梯度同步的耗时不显著，我们是否可以通过一些profile方法拆分出每个部分的耗时，进行定量优化呢？我们对backward部分进行更细致的测试，最终测试结果如下表所示：
| 序号 | GPU数 | DDP Backend | Backward平均耗时（s） | 通信开销（s） |
|---|---|---|---|---|
| 1 | 1 | 无 | 5.63 | - |
| 2 | 128 | NCCL | 6.03 | 0.4 |
| 3 | 128 | hfreduce | 5.91 | 0.28 |
这里，不考虑多卡时因数据处理进度不同产生的等待，那么多卡训练时的backward时间减去单卡训练的backward时间既可以认为是多机多卡带来的额外通信开销。
上表的结果中可以看到，在萤火二号上使用DDP训练时，选择hfreduce能够比NCCL节省30%的通信开销，这是一个非常可观的数字，说明hfreduce确实更适合在幻方萤火二号上进行多机多卡训练的加速。同时，90%以上的耗时在forward的计算上，特别是layer_norm和attention的算子计算开销比较大，这里需要进一步使用hfai的优化算子进行进一步提升。
hfai.nn
如前文《幻方萤火 | 性能卓越的深度学习算子 hfai.nn》所提到的，幻方AI针对Pytorch框架进行了深度优化，结合萤火二号的集群特点，对一些常用的AI算子重新设计研发，进一步提升了模型整体的训练效率。其使用非常简单，只需要在原有的任意模型训练代码中加入一行：
model = hfai.nn.to_hfai(model)hfai会自动扫描您的代码，替换为优化后的算子。hfai.nn提供了当前主流深度学习模型中常用的MultiHeadAttention，LayerNorm，LSTM等结构的优化算子，能够大幅加速模型中这些算子的运算。我们尝试在Alphafold中引入hfai.nn算子库，测试其能够对模型训练产生多大的增益效果。
Alphafold GPU开销分析
由于hfai.nn的算子加速取决于模型中各类型运算占总GPU开销的比重情况，因此我们首先尝试使用Pyorch Profiler工具对Alphafold模型使用标准torch算子时的GPU运算时间进行一些分析。在一次总耗时为12.6s的迭代中，主要开销情况如下：
| Name | CUDA Time | CUDA % |
|---|---|---|
| aten::native_layer_norm | 2.105s | 16.81% |
| void at::native::(anonymous namespace)::RowwiseMomen… | 1.866s | 14.90% |
| aten::mm | 1.716s | 13.70% |
| aten::native_layer_norm_backward | 1.642s | 13.11% |
| atten::add_ | 1.454s | 11.61% |
仅前5类运算占了Alphafold训练开销的70%，而其中LayerNorm与Attention中的矩阵运算又是最主要的耗时来源：仅LayerNorm运算就占了总耗时的30%左右。Alphafold相对一般的BERT等Transformer类模型更为复杂，使用了自己实现的Attention，因此在注意力运算上无法获得hfia.nn的加速。但同样占据耗时大头的LayerNorm却能够使用到hfai算子加速，因此可以预期hfai.nn能给alphafold带来较好的加速效果。
为了探究具体加速效果，我们首先尝试对Alphafold中LayerNorm的使用情况进行一些分析。一般来说输入Tensor的形状往往会对算子加速的效果产生较大的影响，因此在这里我们首先尝试对输入模型中的LayerNorm层的不同的Tensor形状进行一些理论上的分析。在一个典型Transformer类模型中，LayerNorm的输入形状往往是[BatchSize, SeqLen, EmbDim]。在Alphafold中由于其模型结构与简单的Transformer模型有较大差异，输入的张量形状会更为特殊一些。模型中LayerNorm层的不同Input Shape对应的出现频率可见下表：
| Input Shape | Run% | Relative Perf% | Torch Run Time(s) | Hfai Run Time(s) |
|---|---|---|---|---|
| [1,256,256,128] | 59.96% | 364.07% | 0.003 | 0.0008 |
| [1,132,256,256] | 13.2% | 267.02% | 0.0017 | 0.0007 |
| [1,1,256,256,64] | 8.03% | 379.80% | 0.0028 | 0.0007 |
| [1,128,256,256] | 7.92% | 276.52% | 0.0017 | 0.0006 |
| [1,256,132,256] | 4.4% | 276.13% | 0.0017 | 0.0006 |
| [1,256,128,256] | 2.64% | 272.92% | 0.0017 | 0.0006 |
| [1,256,384] | 1.87% | 80.50% | 0.0002 | 0.0003 |
| [1,1024,256,64] | 1.32% | 471.77% | 0.0111 | 0.0024 |
| [1,256,1024,64] | 0.44% | 474.66% | 0.0111 | 0.0024 |
| [1,4,256,256,64] | 0.11% | 474.41% | 0.0111 | 0.0024 |
| [1,256,256] | 0.11% | 93.81% | 0.0003 | 0.0003 |
可见在大多数输入时，hfai.nn提供的LayerNorm算子都能取得相比torch.nn数倍的大幅性能提升，只在极少数形状下性能可能会稍差于torch。另外根据上表中不同输入Tensor形状的占比可以计算得，使用hfai.nn在Alphafold中可以在LayerNorm的GPU开销上取得总计350%的加速。由于Alphafold中的LayerNorm使用频率很高，且在模型的不同部分都有使用，因此可以预计对模型整体也能有较明显的加速效果。
实际训练加速结果
在理论分析了hfai.nn能给Alphafold带来的加速幅度后，我们也希望能了解在实际训练时Alphafold能够从hfai.nn中获得的收益。因此我们在与Alphafold预训练时相同设定的真实场景下使用128卡进行了并行训练，分别测试了使用torch算子和hfai算子进行Alphafold训练时的训练耗时情况，结果如下表：
| 算子 | GPU数 | 迭代次数 | 平均时长 |
|---|---|---|---|
| torch.nn | 128 | 200 | 11.26s |
| hfai.nn | 128 | 200 | 8.64s |
相比于使用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将持续加大研发，提供更多的优化方案和工具。