重新思考内存与通信成本以实现高效大语言模型训练
机器学习
2023-10-31 v2 人工智能
摘要
近期，人们提出了多种用于大语言模型训练的分布式策略。然而，这些方法为内存消耗与通信成本之间的权衡提供的解决方案有限。在本文中，我们重新思考内存消耗与通信成本对大语言模型训练速度的影响，并提出一种内存-通信均衡策略集——部分冗余优化器（PaRO）。PaRO 提供全面的选项，通过细粒度分片策略以轻微的内存冗余降低组间通信的量与频率，从而在各类训练场景中提升训练效率。此外，我们提出一种层次化重叠环（HO-Ring）通信拓扑，以增强大语言模型训练中节点间或跨交换机的通信效率。我们的实验表明，与 SOTA 方法相比，PaRO 将训练吞吐量显著提升了 1.19x–2.50x，并实现了近线性可扩展性。HO-Ring 算法相比传统 Ring 算法将通信效率提升了 36.5%。
引用
@article{arxiv.2310.06003,
title = {Rethinking Memory and Communication Cost for Efficient Large Language Model Training},
author = {Chan Wu and Hanxiao Zhang and Lin Ju and Jinjing Huang and Youshao Xiao and Zhaoxin Huan and Siyuan Li and Fanzhuang Meng and Lei Liang and Xiaolu Zhang and Jun Zhou},
journal= {arXiv preprint arXiv:2310.06003},
year = {2023}
}