arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2407.04480cs.LGmath.OC
收藏

LoCo:面向大规模模型训练的低比特通信适配器

LoCo: Low-Bit Communication Adaptor for Large-scale Model Training

National University of Singapore(新加坡国立大学);Singapore Management University(新加坡管理大学);Sea AI Lab(Sea AI实验室)

Xingyu Xie, Zhijie Lin, Kim-Chuan Toh, Pan Zhou

AI总结:

针对低比特梯度通信压缩导致训练质量下降的问题,提出LoCo适配器,通过历史补偿误差移动平均估计当前压缩误差并补偿梯度,兼容通用优化器与分片策略,理论证明不影响收敛速度,实验提升训练速度14%-40%且无性能损失。

摘要:

为高效训练大规模模型,低比特梯度通信将本地GPU节点的全精度梯度压缩为低精度梯度,以提升GPU节点间的梯度同步效率。但压缩信息丢失常导致训练质量下降。为此,我们提出低比特通信适配器(LoCo),在压缩前于本地GPU节点补偿梯度,确保同步高效且不影响训练质量。具体而言,LoCo设计历史补偿误差的移动平均以稳定估计当前压缩误差,再用该估计补偿当前梯度压缩,实现低损失压缩。该机制使其兼容Adam等通用优化器及FSDP等分片策略。理论分析表明,将LoCo集成到Adam、SGD等全精度优化器中不会降低其在非凸问题上的收敛速度。实验结果显示,在Megatron-LM和PyTorch的FSDP等大规模模型训练框架中,LoCo显著提升通信效率,例如在LLAMAs和MoE等大语言模型上,在不降低性能的情况下将Adam的训练速度提升14%至40%。

英文摘要:

To efficiently train large-scale models, low-bit gradient communication compresses full-precision gradients on local GPU nodes into low-precision ones for higher gradient synchronization efficiency among GPU nodes. However, it often degrades training quality due to compression information loss. To address this, we propose the Low-bit Communication Adaptor (LoCo), which compensates gradients on local GPU nodes before compression, ensuring efficient synchronization without compromising training quality. Specifically, LoCo designs a moving average of historical compensation errors to stably estimate concurrent compression error and then adopts it to compensate for the concurrent gradient compression, yielding a less lossless compression. This mechanism allows it to be compatible with general optimizers like Adam and sharding strategies like FSDP. Theoretical analysis shows that integrating LoCo into full-precision optimizers like Adam and SGD does not impair their convergence speed on nonconvex problems. Experimental results show that across large-scale model training frameworks like Megatron-LM and PyTorch's FSDP, LoCo significantly improves communication efficiency, e.g., improving Adam's training speed by 14% to 40% without performance degradation on large language models like LLAMAs and MoE.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑