[论文解读] 2.5-dimensional distributed model training.

Boxiang Wang, Qifan Xu+2 more|arXiv (Cornell University)|May 30, 2021

Topic Modeling参考文献 9被引用 8

一句话总结

本文提出 SUMMA2.5-LM，一种用于大规模语言模型的 2.5 维分布式张量并行方法，通过结合 SUMMA 和 2.5D 矩阵乘法的原理，降低了通信开销。与 2D [8,8,1] 相比，其在弱缩放下的效率提高了 1.45 倍，优于 2.5D [4,4,4] 配置，显著提升了大规模集群上的缩放效率。

ABSTRACT

Data parallelism does a good job in speeding up the training. However, when it comes to the case when the memory of a single device can not host a whole model, data parallelism would not have the chance to do anything. Another option is to split the model by operator, or horizontally. Megatron-LM introduced a 1-Dimensional distributed method to use GPUs to speed up the training process. Optimus is a 2D solution for distributed tensor parallelism. However, these methods have a high communication overhead and a low scaling efficiency on large-scale computing clusters. To solve this problem, we investigate the 2.5-Dimensional distributed tensor parallelism.Introduced by Solomonik et al., 2.5-Dimensional Matrix Multiplication developed an effective method to perform multiple Cannon's algorithm at the same time to increase the efficiency. With many restrictions of Cannon's Algorithm and a huge amount of shift operation, we need to invent a new method of 2.5-dimensional matrix multiplication to enhance the performance. Absorbing the essence from both SUMMA and 2.5-Dimensional Matrix Multiplication, we introduced SUMMA2.5-LM for language models to overcome the abundance of unnecessary transmission loss result from the increasing size of language model parallelism. Compared to previous 1D and 2D model parallelization of language models, our SUMMA2.5-LM managed to reduce the transmission cost on each layer, which could get a 1.45X efficiency according to our weak scaling result between 2.5-D [4,4,4] arrangement and 2-D [8,8,1] arrangement.

研究动机与目标

• 解决现有 1D 和 2D 模型并行方法在大规模语言模型中通信开销高、缩放效率低的问题。
• 在模型规模超过单个设备内存容量时，降低模型并行中的传输成本。
• 设计一种新型的 2.5D 矩阵乘法技术，在保持高计算效率的同时最小化不必要的数据移动。
• 通过融合 SUMMA 和 2.5D 矩阵乘法的混合通信模式，提升大规模计算集群上的缩放效率。

提出的方法

• 将 2.5D 矩阵乘法的原理应用于模型并行，支持多个类似 Cannon 的算法同时执行。
• 提出一种新型 2.5D 矩阵乘法方法，相比传统 Cannon 算法，减少了移位操作和通信瓶颈。
• 将 SUMMA（可扩展通信）与 2.5D 矩阵乘法（负载均衡和减少数据移动）的优势整合到统一框架中。
• 采用跨设备的 3D 张量切片策略，2.5D 配置为 [4,4,4]，2D 配置为 [8,8,1]，以优化数据分布并减少冗余传输。
• 设计一种通信优化流水线，在前向和反向传播过程中最小化空闲时间并平衡 GPU 间负载。
• 采用混合通信模式，通过复用层间中间结果，减少冗余数据传输。

实验结果

研究问题

• RQ1 2.5D 张量并行相比 1D 和 2D 方法，如何降低大规模语言模型训练中的通信开销？
• RQ2 最小化移位操作和冗余数据传输对分布式模型训练中缩放效率的影响是什么？
• RQ3 结合 SUMMA 和 2.5D 矩阵乘法的混合方法是否能在通信成本和加速比方面超越现有的 2D 和 1D 并行策略？
• RQ4 在弱缩放效率方面，2.5D [4,4,4] 配置与 2D [8,8,1] 配置相比表现如何？
• RQ5 在大规模语言模型训练中应用所提出的 SUMMA2.5-LM 方法，可实现的性能提升是多少？

主要发现

• 所提出的 SUMMA2.5-LM 方法通过优化模型并行中的数据移动模式，降低了每层的传输成本。
• 在弱缩放实验中，2.5D [4,4,4] 配置的效率比 2D [8,8,1] 配置高出 1.45 倍。
• 通过最小化不必要的数据传输，SUMMA2.5-LM 显著提升了大规模计算集群上的缩放效率。
• 该方法有效缓解了传统 1D 和 2D 模型并行方法固有的通信瓶颈。
• SUMMA2.5-LM 的混合设计有效平衡了计算负载，并减少了设备间的空闲时间。
• 新型 2.5D 矩阵乘法技术相比 Cannon 算法，减少了移位操作和通信开销。