太极AngelPTM MoE组件性能优化策略——Part2

太极AngelPTM MoE组件性能优化策略-Part2

​| 导言 近期鹅厂“混元”多模态大模型在全球最具权威的MSR-VTT，MSVD，LSMDC，DiDeMo和ActivityNet五大跨模态视频检索数据集榜单中，先后拿下第一名的成绩，实现了该领域的大满贯[1]。随后“混元”NLP大模型在CLUE榜单（中文语言理解评测集合）的分类榜中，获得了第一的成绩[2,3]，不久又刷新了CLUE总榜及阅读理解榜、KgCLUE的最好成绩，实现了跨模态检索和自然语言理解的双佳绩！

“混元”AI大模型基于腾讯太极机器学习平台进行研发，借助GPU算力，实现快速的算法迭代和模型训练，本文将介绍腾讯太极机器学习平台助力“混元”AI大模型的训练框架——AngelPTM。

太极AngelPTM系列文章分Part 1和Part 2两部分，Part 1介绍本文的背景、太极AngelPTM概览和太极AngelPTM训练框架优化策略中的数据并行和模型并行部分，Part 2介绍太极AngelPTM训练框架优化策略中的MoE组件性能优化策略和混合并行策略，实验效果以及总结与展望。

MoE组件性能优化策略

T5-MoE模型在DeepSpeed框架下各时间占比

在常规网络环境下，评测当前主流的框架DeepSpeed在2*8 A100机器运行T5-MoE并统计各部分时间占比。结果显示AllToAll的时间开销（约55%）在当前集群下训练MoE模型的主要挑战。AngelPTM在MoE组件中提出了拓扑感知的层次化AllToAll来优化单次的AllToAll通信的时间，提出了MoE重计算策略来减少训练时的AllToAll通信调用次数，提出专家计算流水来重叠AllToAll通信与计算以隐藏AllToAll通信时间。此外，AngelPTM在MoE组件还实现了高效的Gating算子以及专家计算算子，提高了GPU计算利用率。       

      图11 T5-MoE各部分时间示意图

拓扑感知的层次化AllToAll

NCCL并未实现专用的AllToAll通信算子，而是并行化Send/Recv操作。如图所示，每个GPU有长度为4的数据需要进行AllToAll操作。GPU0会启动4个点对点通信（P2P），每个P2P通信的消息大小为1。对于N个节点，每个节点G张GPU，每个GPU数据量为B的情形下，点对点通信量为B/(N*G)。       

     图12 AllToAll通信工作原理

高效执行AllToAll通信在常规网络环境主要面临着两点挑战：

1.常规网络环境由层次网络组成，机内GPU通过8*NvLINK相互连接，带宽为300GB/s，机间由一张网卡的RoCE网络连接，带宽为25GB/s。现有的AllToAll实现中，机内8张GPU中有6张需要跨PCIe访问网卡进行机器间通信，无法发挥机内NvLink的性能优势。

2.在大规模集群中，由于点对点通信量太小，无法充分利用跨机的RoCE网络带宽。

针对以上两个挑战，AngelPTM设计了拓扑感知的层次化AllToAll, 实现了混合网络通信优化技术（NvLink和NIC）和消息聚合技术[18]。

混合网络通信优化技术是指在单个节点内，将各GPU的数据通过NvLink传输到距离网卡（NIC）最近的GPU上（Gather），有效的利用了机器内的高速NvLink带宽。消息聚合技术是指将目标机器相同的多个Message合并为1个Message，增大单次的通信量。图13表示了的层次AllToAll的具体流程，Gather/Scatter和All2All inter-node组成了混合网络通信优化技术，Layout Transform在保持原有的All2All语义下实现了消息聚合技术。通过以上优化，层次AllToAll利用了混合的网络链路，并将网卡间的信息传输量从原有的B/(NG)提升到了BG/N，将消息大小提升了G*G倍。

  图13 层次AllToAll操作示意图     

在4*8 A100的环境下评测了优化前后的AllToAll技术在不同Batch Size下的运行时间，其中sequence length = 512, embedding dimension = 1024。实验结果如图所示，通过引入层次AllToAll相关技术，取得了平均约2.75倍的加速。在Batch Size=16的时候，最大取得了约3.13倍的加速。

图14 层次AllToAll优化结果

Gate算子优化

AngelPTM设计并实现了多个高效的 GPU kernel，为特定于 MoE 的计算提供算子。例如，对于“将发往同一个expert的token映射到连续的存储空间”这一操作，对DeepSpeed MoE里相关函数进行了等价变化，将计算的时间复杂度从 O(N^3) 降低到 O(N^2)，从而显着提高了数据分发效率。同时，AngelPTM利用NVRTC动态编译技术对不同shape的MoE kernel进行运行时编译。

表3 Gate算子优化效果